- L5 (Stability & Reliability) |
- Online: long-term stability test + reliability test |
- Long-term stability and reliability testing for high-priority models Low-priority models: function and doc testing |
+ L5 (Stability & Reliability & selected Perf) |
+ Online: long-term stability + reliability; coverage test |
+ Long-term stability and reliability testing for key models Non-critical scenario performance and accuracy testing for key models Other models: function and doc testing |
/ |
- Invalid-parameter validation for high-priority interfaces |
slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion |
Depends on reality |
@@ -178,8 +170,8 @@ For per-level test authoring (markers, examples), see [Test Writing Guide](./tes
Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues.
-1. ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
-2. ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
+1. ***PR Checklist ([`.github/PULL_REQUEST_TEMPLATE.md`](https://github.com/vllm-project/vllm-omni/blob/main/.github/PULL_REQUEST_TEMPLATE.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
+2. ***CI Failure Explanation ([CI Failures](./failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
## Notes
diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md
index 397f5e83c..469e69cac 100644
--- a/docs/contributing/ci/test_writing_guide.md
+++ b/docs/contributing/ci/test_writing_guide.md
@@ -17,47 +17,43 @@ Summary:
- **Reliability:** `tests/dfx/reliability/`
- Do **not** add new top-level `tests/` directories unrelated to a component (or to `e2e` / `dfx` / `helpers` / `examples` / `buildkite`)
-Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [stage configs](../../configuration/stage_configs.md)).
+Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [pipeline and deploy configs](../../configuration/stage_configs.md)).
## Markers for Tests
By adding markers before test functions, tests can later be executed uniformly by simply declaring the corresponding marker type.
### Current Markers
+
Defined in `pyproject.toml`:
-| Marker | Description |
-| ------------------ | --------------------------------------------------------- |
-| `core_model` | L1&L2 tests (run in each PR) |
-| `advanced_model` | L3 tests (run on each merge to main) |
-| `full_model` | L4 tests (run nightly) |
-| `diffusion` | Diffusion model tests |
-| `omni` | Omni multimodal model tests |
-| `tts` | TTS model tests |
-| `cache` | Cache backend tests |
-| `parallel` | Parallelism/distributed tests |
-| `cpu` | Tests that run on CPU |
-| `gpu` | Tests that run on GPU * |
-| `cuda` | Tests that run on CUDA * |
-| `rocm` | Tests that run on AMD/ROCm * |
-| `xpu` | Tests that run on Intel XPU * |
-| `npu` | Tests that run on NPU/Ascend * |
-| `H100` | Tests that require H100 GPU * |
-| `L4` | Tests that require L4 GPU * |
-| `MI325` | Tests that require MI325 GPU (AMD/ROCm) * |
-| `A2` | Tests that require A2 NPU * |
-| `A3` | Tests that require A3 NPU * |
-| `distributed_cuda` | Tests that require multi cards on CUDA platform * |
-| `distributed_rocm` | Tests that require multi cards on ROCm platform * |
-| `distributed_npu` | Tests that require multi cards on NPU platform * |
-| `skipif_cuda` | Skip if the num of CUDA cards is less than the required * |
-| `skipif_rocm` | Skip if the num of ROCm cards is less than the required * |
-| `skipif_npu` | Skip if the num of NPU cards is less than the required * |
-| `slow` | Slow tests (may skip in quick CI) |
+| Marker | Description |
+| ------------------ | ------------------------------------------------------------------------------------------------------- |
+| `core_model` | L1&L2 tests (run in each PR) |
+| `advanced_model` | L3 tests (run on each merge to main) |
+| `full_model` | L4 tests (run nightly) |
+| `diffusion` | Diffusion model tests |
+| `omni` | Omni multimodal model tests |
+| `tts` | TTS model tests |
+| `cache` | Cache backend tests |
+| `parallel` | Parallelism/distributed tests |
+| `cpu` | Tests that run on CPU |
+| `gpu` | Tests that run on GPU * |
+| `cuda` | Tests that run on CUDA * |
+| `rocm` | Tests that run on AMD/ROCm * |
+| `xpu` | Tests that run on Intel XPU * |
+| `npu` | Tests that run on NPU/Ascend * |
+| `H100` | Tests that require H100 GPU * |
+| `L4` | Tests that require L4 GPU * |
+| `MI325` | Tests that require MI325 GPU (AMD/ROCm) * |
+| `A2` | Tests that require A2 NPU * |
+| `A3` | Tests that require A3 NPU * |
+| `cards_{n}` | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` * |
+| `slow` | Slow tests (may skip in quick CI) |
| `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
-| `local_model` | Tests requiring local / non-HF-hub model weights |
+| `local_model` | Tests requiring local / non-HF-hub model weights |
-\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
+\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). Do **not** write `pytest.mark.H100` / `pytest.mark.L4` (or other SKU names) by hand: `check-mark` rejects that so `cards_{n}` cannot be skipped.
#### Example usage for markers
@@ -80,28 +76,26 @@ def test_video_to_audio()
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:
1. **Applies platform and resource markers**
- Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
- ```
+ Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`).
+ ```python
@pytest.mark.cuda
@pytest.mark.L4
```
-2. **Handles multi-card (distributed) scenarios**
- For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
- ```
- @pytest.mark.distributed_cuda(num_cards=num_cards)
- ```
+2. **Handles card-count filtering**
+ Adds a filterable marker `cards_{n}` (including `cards_1` when `num_cards` is omitted / 1). Use this with the SKU and platform, e.g. `-m "H100 and cards_4 and cuda"` vs `-m "H100 and cards_2 and cuda"`. Any multi-card CUDA test is `-m "cuda and not cards_1"`.
+
+ When `res` lists more than one platform **and** `num_cards` differs by platform, `@hardware_test` collects **one variant per platform** so those `cards_{n}` marks do not stack. `-m "H100 and cards_4 and cuda"` then selects only the CUDA/H100 4-card variant, not a ROCm 2-card sibling. The same count on every platform stays a single item.
3. **Supports flexible card requirements**
Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
4. **Integrates resource validation**
- On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
- Support for `skipif_rocm` and `skipif_npu` will be implemented later.
-
+ On CUDA and XPU, when `num_cards > 1`, attaches pytest's built-in `skipif` if the machine has fewer devices than requested. This is not a named `skipif_*` marker and is not used for `-m` filtering.
5. **Works with pytest filtering**
- Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
+ Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "H100 and cards_2"`, `-m "L4 and not cards_1"`).
##### Example usage for decorator
+
- Single call for multiple platforms:
```python
@hardware_test(
@@ -116,13 +110,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
num_cards=2,
)
```
-- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3)
+- `res` must be a dict. SKU values come from `pyproject.toml` markers tagged `[hardware-resource]` plus the platform key (`[cuda]`, `[rocm]`, `[xpu]`, `[npu]`, `[musa]`).
- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
-- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
- Filtering examples:
- - CUDA only: `pytest -m "distributed_cuda and L4"`
- - ROCm only: `pytest -m "distributed_rocm and MI325"`
- - NPU only: `pytest -m "distributed_npu"`
+ - CUDA H100 4-card: `pytest -m "H100 and cards_4 and cuda"`
+ - CUDA H100 2-card: `pytest -m "H100 and cards_2 and cuda"`
+ - CUDA multi-card (any count): `pytest -m "L4 and not cards_1"`
+ - ROCm multi-card: `pytest -m "MI325 and not cards_1"`
+ - NPU multi-card: `pytest -m "npu and not cards_1"`
#### Function: `hardware_marks`
@@ -178,37 +173,30 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
[tool.pytest.ini_options]
markers = [
# ... existing markers ...
- "tpu: Tests that require TPU device",
- "TPU_V3: Tests that require TPU v3 hardware",
- "distributed_tpu: Tests that require multiple TPU devices",
+ "tpu: [hardware-platform] Tests that require TPU device",
+ # Add [gpu] next to [hardware-platform] if this platform should also get pytest.mark.gpu.
+ "TPU_V3: [hardware-resource] [tpu] Tests that require TPU v3 hardware",
+ "cards_1: [hardware-cards] Tests that require 1 accelerator card",
+ "cards_2: [hardware-cards] Tests that require 2 accelerator cards",
]
```
+ `check-mark` treats `[hardware-platform]` names (including `cpu`/`gpu`) as allowed `pytest.mark.` values. `hardware_test(res=...)` still rejects `cpu`/`gpu` as `res` keys.
2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`:
```python
# In vllm-omni/tests/helpers/mark.py
- def tpu_marks(*, res: str, num_cards: int):
- test_platform = pytest.mark.tpu
- if res == "TPU_V3":
- test_resource = pytest.mark.TPU_V3
- else:
- raise ValueError(
- f"Invalid TPU resource type: {res}. Supported: TPU_V3")
-
- if num_cards == 1:
- return [test_platform, test_resource]
- else:
- test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
- # Optionally: add skipif_tpu when implemented
- return [test_platform, test_resource, test_distributed]
+ def _tpu_marks(*, res: str, num_cards: int):
+ _require_sku("tpu", res)
+ return [pytest.mark.tpu, getattr(pytest.mark, res), _cards_mark(num_cards)]
```
-3. **Update `hardware_marks` to recognize your new platform**:
- In the relevant place (see the `hardware_marks` implementation), add:
+3. **Register the builder** in `_marks_for_platform`:
```python
- if platform == "tpu":
- marks = tpu_marks(res=resource, num_cards=cards)
+ builders = {
+ # ... existing platforms ...
+ "tpu": _tpu_marks,
+ }
```
- (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
+ (`hardware_test` / `hardware_marks` both go through `_marks_for_platform`.)
4. **(Recommended) Add a test using your new markers**:
```python
@hardware_test(
@@ -220,12 +208,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
```
**Summary**:
-- Add pytest markers for your new platform/resources
-- Implement a marker function (`xxx_marks`)
-- Plug into `hardware_marks`
-- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
-See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
+- Add pytest markers for your new platform/resources
+- Implement a private marker function (`_xxx_marks`)
+- Register it in `_marks_for_platform`
+- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers and `cards_{n}` filtering.
+
+See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`_cuda_marks`, `_rocm_marks`, `_npu_marks`).
## Test case style
@@ -237,47 +226,46 @@ L1 and L2 level testing form the foundation of the quality assurance system. L1
L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
-
-
#### 1.2 Testing Content and Scope
-- ***L1 (Unit & Logic Testing)***:
-- - ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`.
- - ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack.
- - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback.
-- ***L2 (Basic End-to-End Testing)***:
-- - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference).
- - ***Focus***: Uses `dummy` weights (via deploy YAML patching at `core_model`) or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, and **cheap payload checks** at `--run-level core_model` (see below). Also includes some unit tests that require launching independent service instances.
- - ***L2 response validation (`core_model`)***: Implemented in `tests/helpers/assertions.py` and invoked by `OpenAIClientHandler` based on `--run-level`. At L2 we require **request success** plus minimal output sanity—not full accuracy:
- - **Speech / TTS** (`assert_audio_speech_response`): decoded audio must be present and non-empty (or exceed `min_audio_bytes` when set in `request_config`); `response_format` must match the returned content-type (e.g. `wav`, `pcm`). Whisper transcript similarity, PCM HNR, and preset-voice gender checks run only at L3+.
- - **Diffusion** (`assert_diffusion_response`): at least one non-empty image, video, or audio artifact. Resolution/frame-count parity and other deep checks run only at L3+.
- - **Omni multimodal** (`assert_omni_response`): L2 asserts successful completion; keyword, transcript, and cross-modal similarity checks run only at L3+.
- - ***Characteristic***: Requires ***GPU*** resources to perform model computations.
+- ***L1 (Unit & Logic Testing)***:
+ - ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`.
+ - ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack.
+ - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback.
+- ***L2 (Basic End-to-End Testing)***:
+ - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference).
+ - ***Focus***: Uses `dummy` weights (via deploy YAML patching at `core_model`) or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, and **cheap payload checks** at `--run-level core_model` (see below). Also includes some unit tests that require launching independent service instances.
+ - ***L2 response validation (`core_model`)***: Implemented in `tests/helpers/assertions.py` and invoked by `OpenAIClientHandler` based on `--run-level`. At L2 we require **request success** plus minimal output sanity—not full accuracy:
+ - **Speech / TTS** (`assert_audio_speech_response`): decoded audio must be present and non-empty (or exceed `min_audio_bytes` when set in `request_config`); `response_format` must match the returned content-type (e.g. `wav`, `pcm`). Whisper transcript similarity, PCM HNR, and preset-voice gender checks run only at L3+.
+ - **Diffusion** (`assert_diffusion_response`): at least one non-empty image, video, or audio artifact. Resolution/frame-count parity and other deep checks run only at L3+.
+ - **Omni multimodal** (`assert_omni_response`): L2 asserts successful completion; keyword, transcript, and cross-modal similarity checks run only at L3+.
+ - ***Characteristic***: Requires ***GPU*** resources to perform model computations.
#### 1.3 Test Directory and Execution Files
A clear directory structure is key to managing test cases efficiently. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules.
-- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component.
-- ***L2 (basic model E2E)***:
- - Online serving: `tests/e2e/online_serving/test_{model_name}.py`
- - Offline inference: `tests/e2e/offline_inference/test_{model_name}.py`
-- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features//` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`).
+- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component.
+- ***L2 (basic model E2E)***:
+ - Online serving: `tests/e2e/online_serving/test_{model_name}.py`
+ - Offline inference: `tests/e2e/offline_inference/test_{model_name}.py`
+- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features//` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`).
#### 1.4 Execution Method and Example
-- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
-- ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](ci_settings.md#step-filtering) prefixes; non-E2E groups still always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci).
-- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
-- ***Script Example***:
+- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
+- ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](ci_settings.md#step-filtering) prefixes; non-E2E groups still always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci).
+- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
+- ***Script Example***:
L1 Test Examples
Examples from `tests/model_executor/models/qwen2_5_omni/test_audio_length.py`
+
```python
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
@@ -331,6 +319,7 @@ def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames)
if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats:
assert target_mel_len <= int(max_mel_frames)
```
+
@@ -338,7 +327,7 @@ def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames)
You can refer to Test Examples in the L3 section to see example test cases that incorporate both L2 and L3 testing logic.
-- - ***Run Command***:
+- ***Run Command***:
`pytest -s -v /tests/e2e/online_serving/test_{model_name}.py`
`pytest -s -v -m 'core_model and cpu' --run-level=core_model`
@@ -351,36 +340,34 @@ L3 level testing executes after code is merged into the main branch. Its core pu
. It acts as the "quality gatekeeper" for the main branch, ensuring that no merge breaks the core capabilities of the model service. Testing needs to provide clear conclusions within a relatively short time (<30min), balancing test depth with feedback speed.
-
-
#### 2.2 Testing Content and Scope
-- ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc.
-- ***Core Verification***:
-- 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results.
- 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues.
- 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds.
-- ***L3 response validation (`advanced_model`)***: At `--run-level advanced_model`, `tests/helpers/assertions.py` adds semantic checks on top of L2 payload gates (see L1 & L2 above):
- - **Speech / TTS** (`assert_audio_speech_response`): Whisper transcript of returned audio vs `request_config["input"]` (cosine similarity > 0.9 when input text is set); optional `min_audio_bytes` floor; PCM harmonic-to-noise ratio when `response_format` is `pcm`; **preset voice gender** via `_assert_preset_voice_gender_from_audio` when `voice` matches a known preset in `_PRESET_VOICE_GENDER_MAP` (pitch/F0-based classifier; skipped for unknown voices or `pcm` output).
- - **Omni multimodal** (`assert_omni_response`): non-empty text/audio outputs; `key_words` in transcript or text; text–audio similarity / containment; preset **`speaker`** gender check (same helper as TTS).
- - **Diffusion** (`assert_*_diffusion_response`): image/video dimension and frame-count parity with request parameters where configured.
- - **Accuracy suites**: pixel/video similarity and metric baselines under `/tests/e2e/accuracy/` (separate from inline helper assertions).
+- ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc.
+- ***Core Verification***:
+ 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results.
+ 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues.
+ 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds.
+- ***L3 response validation (`advanced_model`)***: At `--run-level advanced_model`, `tests/helpers/assertions.py` adds semantic checks on top of L2 payload gates (see L1 & L2 above):
+ - **Speech / TTS** (`assert_audio_speech_response`): Whisper transcript of returned audio vs `request_config["input"]` (cosine similarity > 0.9 when input text is set); optional `min_audio_bytes` floor; PCM harmonic-to-noise ratio when `response_format` is `pcm`; **preset voice gender** via `_assert_preset_voice_gender_from_audio` when `voice` matches a known preset in `_PRESET_VOICE_GENDER_MAP` (pitch/F0-based classifier; skipped for unknown voices or `pcm` output).
+ - **Omni multimodal** (`assert_omni_response`): non-empty text/audio outputs; `key_words` in transcript or text; text–audio similarity / containment; preset **`speaker`** gender check (same helper as TTS).
+ - **Diffusion** (`assert_*_diffusion_response`): image/video dimension and frame-count parity with request parameters where configured.
+ - **Accuracy suites**: pixel/video similarity and metric baselines under `/tests/e2e/accuracy/` (separate from inline helper assertions).
#### 2.3 Test Directory and Execution Files
-- ***Functional / model E2E***:
- - Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py`
- - Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py`
- - Accuracy / similarity: `tests/e2e/accuracy/`
- - (`_expansion.py` holds broader scenario coverage than the L2 smoke file.)
-- ***Feature-level integration***: `tests/e2e/features//` (not under a new top-level `tests/` directory).
+- ***Functional / model E2E***:
+ - Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py`
+ - Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py`
+ - Accuracy / similarity: `tests/e2e/accuracy/`
+ - (`_expansion.py` holds broader scenario coverage than the L2 smoke file.)
+- ***Feature-level integration***: `tests/e2e/features//` (not under a new top-level `tests/` directory).
#### 2.4 Execution Method and Example
-- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge).
-- ***Diff-aware step skipping***: Same [`source_file_dependencies`](ci_settings.md#step-filtering) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci).
-- ***Execution Environment***: ***GPU*** servers.
-- ***Script Example***:
+- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge).
+- ***Diff-aware step skipping***: Same [`source_file_dependencies`](ci_settings.md#step-filtering) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [CI Settings — Diff-aware CI](ci_settings.md#diff-aware-ci).
+- ***Execution Environment***: ***GPU*** servers.
+- ***Script Example***:
???+ example "Test Examples"
@@ -496,9 +483,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](test_system_overview.md#test-helper-environment-variables)).
-- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
+- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
-- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
+- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
### L4 Level Testing - Full Functionality, Performance, and Documentation Testing
@@ -506,27 +493,25 @@ L3 level testing executes after code is merged into the main branch. Its core pu
L4 level testing is a comprehensive quality audit before a version release. It expands upon L3, executing ***full*** functional scenarios, conducting systematic ***performance stress tests***, and simultaneously verifying the correctness of accompanying ***example documentation***. Its purpose is to perform deep validation of the system during off-peak nighttime hours, providing quality trend reports for daytime development and data support for release decisions.
-
-
#### 3.2 Testing Content and Scope
-- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
-- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below.
-- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
+- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
+- ***Performance Testing***: JSON configs under `tests/dfx/perf/tests/` drive Omni, TTS, and diffusion serving benchmarks via `run_benchmark.py` or `run_diffusion_benchmark.py`, measuring throughput, latency, and related metrics under configured workloads.
+- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
#### 3.3 Test Directory and Execution Files
-- ***Functional Testing***: Same directories as L3.
-- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests.
-- ***Documentation Example Tests***:
- - `tests/examples/online_serving/test_{model_name}.py`
- - `tests/examples/offline_inference/test_{model_name}.py`
+- ***Functional Testing***: Same directories as L3.
+- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests.
+- ***Documentation Example Tests***:
+ - `tests/examples/online_serving/test_{model_name}.py`
+ - `tests/examples/offline_inference/test_{model_name}.py`
#### 3.4 Execution Method and Example
-- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
-- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
-- ***Script Example***:
+- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
+- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
+- ***Script Example***:
??? example "Test Examples: Documentation Example Tests"
@@ -540,7 +525,7 @@ L4 level testing is a comprehensive quality audit before a version release. It e
--8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
-- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
+- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
### L5 Level Testing - Stability and Reliability Testing
@@ -548,41 +533,43 @@ L4 level testing is a comprehensive quality audit before a version release. It e
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
-
-
#### 4.2 Testing Content and Scope
-- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
-- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
+- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
+- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See Section 4.3 for file-level responsibilities and Section 4.4 Test Examples (Reliability) for CI markers (`slow`, `hardware_test`, POSIX-only kill).
#### 4.3 Test Directory and Execution Files
-- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
-- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
- - `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
- - `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
- - `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
- - `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
- - `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
- - `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
- - `README.md` — Short local run commands for this directory.
+- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
+- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
+ - `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
+ - `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
+ - `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
+ - `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
+ - `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
+ - `README.md` — Short local run commands for this directory.
#### 4.4 Execution Method and Example
-- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
-- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
-- ***Script Example***:
+- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
+- ***Run Command***:
+ - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_.py` alongside a matching JSON config)
+ - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`).
+- ***Script Example***:
+
Test Examples
-When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape:
+##### Stability
+
+When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). Pair the JSON with `tests/dfx/stability/scripts/test_stability_.py`. The following illustrates the Qwen3-Omni shape:
```json
{
"test_name": "test_qwen3_omni_stability",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
- "stage_config_name": "qwen3_omni.yaml"
+ "stage_config_name": "qwen3_omni_moe.yaml"
},
"benchmark_params": [
{
@@ -601,112 +588,61 @@ When you want to add L5-level stability test cases, add or extend the appropriat
}
```
-##### Parameter Explanation
+###### Parameter Explanation
***Overview***
| Field | Required | Description |
| ---------------- | -------- | --------------------------------------------------------------------------- |
| test_name | Yes | Unique identifier for the stability test case |
-| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) |
+| server_params | Yes | Server-side configuration parameters (model, deploy configuration, etc.) |
| benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) |
-##### server_params Configuration
+###### server_params Configuration
-###### Basic Parameters
+Basic parameters:
| Parameter | Required | Example | Description |
| ----------------- | -------- | ---------------------------------- | ----------------------------------- |
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
-| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
+| stage_config_name | Yes | "qwen3_omni_moe.yaml" | Deploy configuration file name |
-###### Dynamic Configuration (update/delete)
-
-Supports incremental modifications based on the basic configuration:
+Dynamic configuration (`update` / `delete`) supports incremental modifications based on the basic configuration:
| Operation | Description |
| --------- | ------------------------------------ |
| update | Update or add configuration items |
| delete | Delete specified configuration items |
-***Example***:
-You can refer to Test Examples in L4 §3.4
+***Example***: You can refer to Test Examples in L4 §3.4.
-##### benchmark_params Configuration
+###### benchmark_params Configuration
For stability testing, the key parameters are:
-- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
-- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
- - `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
- - `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
-- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
+- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
+- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
+ - `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
+ - `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
+- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
All other optional parameters follow the same rules as in L4 §3.4.
-
+##### Reliability
-
- Reliability test suite (tests/dfx/reliability)
+Add pytest modules under `tests/dfx/reliability/` (for example `test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`). Reuse `helpers.py` and `conftest.py` rather than duplicating OOM / process-kill / raw HTTP plumbing.
-##### Purpose and relationship to stability
-
-Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load).
-
-##### Directory layout
-
-| Path | Responsibility |
-| ---- | -------------- |
-| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
-| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
-| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
-| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
-| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
-| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
-| `README.md` | Minimal run / collect examples. |
-
-##### Parametrization and markers
+###### Parametrization and markers
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
-- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100).
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
-##### CI trigger
-
-Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example:
-
-| Buildkite step | Test file | CI hardware |
-| -------------- | --------- | ----------- |
-| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
-| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
-| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
-| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
-
-```bash
-pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
-pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
-pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
-```
-
-##### Local commands
-
-```bash
-pytest --collect-only tests/dfx/reliability
-pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
-```
-
-##### Adding a new model suite
+###### Adding a new model suite
1. Add `test_reliability_.py` under `tests/dfx/reliability/`.
-2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files).
-3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML.
+2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e deploy-config directory (same pattern as existing files).
+3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs deploy YAML.
4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/cuda/test-weekly.yml`** when the suite should run in weekly L5.
-
-- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_.py` alongside a matching JSON config)
- - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models)
diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md
index 4698caf9c..42b9a53d7 100644
--- a/docs/contributing/model/adding_omni_model.md
+++ b/docs/contributing/model/adding_omni_model.md
@@ -9,7 +9,7 @@ This guide walks through the process of adding a new multi-stage model to vLLM-O
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Key Components](#key-components)
5. [Model Registration](#model-registration)
-6. [Stage Configuration](#stage-configuration)
+6. [Pipeline and Deploy Configuration](#pipeline-and-deploy-configuration)
7. [Stage Input Processors](#stage-input-processors)
8. [Testing](#testing)
9. [Adding a Model Recipe](#adding-a-model-recipe)
@@ -316,18 +316,28 @@ _OMNI_MODELS = {
The registry uses lazy loading, so the model class is imported only when needed.
-## Stage Configuration
+## Pipeline and Deploy Configuration
-Create a YAML configuration file in `vllm_omni/deploy/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml).
+Define fixed topology in `vllm_omni/model_executor/models//pipeline.py`
+with `PipelineConfig` and `StagePipelineConfig`, then register it in
+`vllm_omni/config/pipeline_registry.py`. For a complete example, see the
+[Qwen3-Omni pipeline](gh-file:vllm_omni/model_executor/models/qwen3_omni/pipeline.py).
-### Key Configuration Fields
+### Pipeline topology fields
-- **`model_stage`**: Which stage to run ("thinker", "talker", "code2wav", etc.)
-- **`model_arch`**: The model architecture name (must match registry)
-- **`engine_input_source`**: List of stage IDs that provide input to this stage
-- **`custom_process_input_func`**: Function to process inputs from previous stages
-- **`final_output`**: Whether this stage produces the final output (True/False)
-- **`final_output_type`**: Type of final output ("text", "audio", "image", etc.)
+- **`model_stage`**: Logical stage name ("thinker", "talker", "code2wav", etc.)
+- **`execution_type`**: Runtime family (`LLM_AR`, `LLM_GENERATION`, or `DIFFUSION`)
+- **`model_arch`**: Stage architecture name, when different from the pipeline default
+- **`input_sources`**: Upstream stage IDs that provide input to this stage
+- **`custom_process_input_func`**: Function that processes inputs from previous stages
+- **`final_output`**: Whether this stage produces user-facing output
+- **`final_output_type`**: Output modality ("text", "audio", "image", etc.)
+- **`owns_tokenizer`**: Whether this stage owns the pipeline tokenizer
+
+Create the matching runtime defaults in `vllm_omni/deploy/.yaml` and set
+`PipelineConfig.default_deploy_config_name` to that file. The deploy YAML owns
+GPU placement, memory sizing, parallelism, connectors, and sampling defaults;
+it must not redefine topology fields.
## Stage Input Processors
@@ -374,7 +384,7 @@ The stage transition process follows these steps:
2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
-3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
+3. **Input Processing**: The stage input processor registered by the pipeline (implemented under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
```python
def process_engine_inputs(
self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None
@@ -577,7 +587,7 @@ Adding a new model to vLLM-Omni involves:
1. **Create model directory structure** with stage implementations
2. **Implement unified model class** that orchestrates stages
3. **Register model** in `registry.py`
-4. **Create stage configuration** YAML file
+4. **Define pipeline topology and deploy defaults** in `pipeline.py` and `vllm_omni/deploy/`
5. **Implement stage input processors** for stage transitions
6. **Write tests** to verify functionality
7. **Add model recipe** to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository (see [Adding a Model Recipe](#adding-a-model-recipe) section)
@@ -590,7 +600,8 @@ For a complete reference implementation, see:
- **Thinker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py`
- **Talker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py`
- **Code2Wav**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`
-- **Stage config**: `vllm_omni/deploy/qwen3_omni_moe.yaml`
+- **Pipeline topology**: `vllm_omni/model_executor/models/qwen3_omni/pipeline.py`
+- **Deploy config**: `vllm_omni/deploy/qwen3_omni_moe.yaml`
- **Input processors**: `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`
- **Registry**: `vllm_omni/model_executor/models/registry.py`
- **Testing**: `vllm_omni/tests/e2e/offline_inference/test_qwen3_omni.py`
@@ -598,4 +609,4 @@ For a complete reference implementation, see:
For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Supported Models](../../models/supported_models.md)
-- [Stage Configuration Guide](../../configuration/stage_configs.md)
+- [Pipeline and Deploy Configuration Guide](../../configuration/stage_configs.md)
diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md
index 45680ee5d..264287306 100644
--- a/docs/contributing/model/adding_tts_model.md
+++ b/docs/contributing/model/adding_tts_model.md
@@ -20,7 +20,7 @@ pattern, refer to MOSS-TTS-Nano.
4. [Step-by-Step Implementation](#step-by-step-implementation)
5. [Key Components](#key-components)
6. [Model Registration](#model-registration)
-7. [Stage Configuration](#stage-configuration)
+7. [Pipeline and Deploy Configuration](#pipeline-and-deploy-configuration)
8. [Stage Input Processors](#stage-input-processors)
9. [Online Serving Integration](#online-serving-integration)
10. [Single-Stage Models](#single-stage-models)
@@ -78,7 +78,7 @@ Profile before optimizing.
new TTS model must pass all three levels:
| Layer | Catches | Tool |
-|-------|---------|------|
+| ------- | --------- | ------ |
| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
| Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts |
@@ -94,10 +94,10 @@ the symptom is crosstalk or truncation under load, nothing in single-request tes
vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
and can be placed on different devices. Qwen3-TTS has two stages:
-| Stage | Name | Input | Output |
-|-------|------|-------|--------|
-| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
-| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
+| Stage | Name | Input | Output |
+| ----- | ------------------- | --------------- | ------------------------ |
+| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
+| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
Each stage is a separate model class configured independently via YAML. The two stages
are connected by the `async_chunk` framework, which enables inter-stage streaming for
@@ -170,7 +170,7 @@ and Qwen3-Omni).
When adding a new TTS model, create the following structure:
-```
+```text
vllm_omni/model_executor/models/
your_model_name/
__init__.py
@@ -201,7 +201,7 @@ it would not be picked up — the hub README is the documented surface.
**Qwen3-TTS reference files:**
| File | Purpose |
-|------|---------|
+| ------ | --------- |
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
@@ -427,112 +427,93 @@ _OMNI_MODELS = {
The registry uses lazy loading - model classes are only imported when needed.
-## Stage Configuration
+## Pipeline and Deploy Configuration
-Each stage has a `worker_type` that determines how it is scheduled:
+Define fixed topology in `models//pipeline.py`. The execution type
+selects the scheduler and worker family:
-- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
-- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
+- `StageExecutionType.LLM_AR` uses the autoregressive runtime with PagedAttention.
+- `StageExecutionType.LLM_GENERATION` uses the generation runtime for stages such as audio decoders.
-Key configuration fields:
+Key topology fields belong to `StagePipelineConfig`:
| Field | Description |
-|-------|-------------|
-| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
-| `model_arch` | Architecture name, must match `registry.py` |
-| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
-| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
-| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
-| `final_output` | Whether this stage produces the final user-facing output |
-| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
+| ------- | ------------- |
+| `model_stage` | Logical stage name (`ar_stage`, `decoder`, etc.). |
+| `execution_type` | Runtime family for this stage. |
+| `input_sources` | Upstream stage IDs that provide input. |
+| `model_arch` | Stage-specific architecture name, when it differs from the pipeline default. |
+| `engine_output_type` | Output representation such as `latent` or `audio`. |
+| `custom_process_next_stage_input_func` | Processor for full-payload stage handoff. |
+| `async_chunk_process_next_stage_input_func` | Processor for streaming chunk handoff. |
+| `final_output` / `final_output_type` | Whether this stage produces user-facing output and its modality. |
+| `owns_tokenizer` | Whether this stage supplies the pipeline tokenizer. |
-!!! note
- New in-tree models should define frozen topology in `models//pipeline.py`,
- register it in `vllm_omni/config/pipeline_registry.py`, and put deployment
- knobs in `vllm_omni/deploy/.yaml`. Use `--deploy-config` to load a
- custom deployment YAML.
+For example, a two-stage TTS topology can be defined as:
-### Batch mode
+```python
+from vllm_omni.config.stage_config import (
+ PipelineConfig,
+ StageExecutionType,
+ StagePipelineConfig,
+)
-```yaml
-# stage_configs/your_model_name.yaml
+_PROC = "vllm_omni.model_executor.stage_input_processors.your_model_name"
-stage_args:
- - stage_id: 0
- stage_type: llm
- runtime:
- devices: "0"
- engine_args:
- model_stage: ar_stage
- max_num_seqs: 64
- model_arch: YourTTSModelForConditionalGeneration
- worker_type: ar
- scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
- engine_output_type: latent
- default_sampling_params:
- temperature: 0.9
- top_k: 50
- max_tokens: 2048
-
- - stage_id: 1
- stage_type: llm
- runtime:
- devices: "0"
- engine_args:
- model_stage: decoder
- model_arch: YourTTSModelForConditionalGeneration
- worker_type: generation
- scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
- engine_output_type: audio
- engine_input_source: [0]
- final_output: true
- final_output_type: audio
+YOUR_TTS_PIPELINE = PipelineConfig(
+ model_type="your_model_name",
+ default_deploy_config_name="your_model_name.yaml",
+ model_arch="YourTTSModelForConditionalGeneration",
+ stages=(
+ StagePipelineConfig(
+ stage_id=0,
+ model_stage="ar_stage",
+ execution_type=StageExecutionType.LLM_AR,
+ input_sources=(),
+ owns_tokenizer=True,
+ engine_output_type="latent",
+ custom_process_next_stage_input_func=f"{_PROC}.ar2decoder",
+ async_chunk_process_next_stage_input_func=f"{_PROC}.ar2decoder_async_chunk",
+ ),
+ StagePipelineConfig(
+ stage_id=1,
+ model_stage="decoder",
+ execution_type=StageExecutionType.LLM_GENERATION,
+ input_sources=(0,),
+ model_arch="YourTTSDecoder",
+ engine_output_type="audio",
+ final_output=True,
+ final_output_type="audio",
+ ),
+ ),
+)
```
-### Streaming mode (async_chunk)
-
-Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
-on Stage 0 to define how intermediate outputs are chunked and forwarded:
+Register the pipeline in `vllm_omni/config/pipeline_registry.py`. Runtime
+placement and sizing belong in `vllm_omni/deploy/your_model_name.yaml`:
```yaml
-# stage_configs/your_model_name_async_chunk.yaml
-
async_chunk: true
-stage_args:
+stages:
- stage_id: 0
- stage_type: llm
- runtime:
- devices: "0"
- engine_args:
- model_stage: ar_stage
- max_num_seqs: 64
- model_arch: YourTTSModelForConditionalGeneration
- worker_type: ar
- scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
- engine_output_type: latent
- custom_process_next_stage_input_func: >
- vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
+ devices: "0"
+ max_num_seqs: 64
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 2048
- stage_id: 1
- stage_type: llm
- runtime:
- devices: "0"
- engine_args:
- model_stage: decoder
- model_arch: YourTTSModelForConditionalGeneration
- worker_type: generation
- scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
- engine_output_type: audio
- engine_input_source: [0]
- final_output: true
- final_output_type: audio
+ devices: "0"
+ max_num_seqs: 1
```
+Set `async_chunk: false` for end-to-end handoff. The registered pipeline
+selects the full-payload or async chunk processor; processor paths and topology
+must not be moved into the deploy YAML. Use `--deploy-config` to load a custom
+deployment overlay.
+
## Stage Input Processors
Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
@@ -547,9 +528,9 @@ Understanding what's available in stage outputs:
- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
- Each `RequestOutput` has:
- - `token_ids` - generated token IDs
- - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
- - `prompt_token_ids` - original prompt token IDs
+ - `token_ids` - generated token IDs
+ - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
+ - `prompt_token_ids` - original prompt token IDs
### Batch mode (non-streaming)
@@ -629,6 +610,7 @@ def ar2decoder_async_chunk(
```
Key points:
+
- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
- The final (possibly partial) chunk is flushed when `is_finished` is true
@@ -683,6 +665,7 @@ Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`:
```python
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""YourModel serving adapter."""
from typing import TYPE_CHECKING
@@ -821,7 +804,7 @@ directly from the model's own generator.
### Directory structure
-```
+```text
vllm_omni/model_executor/models/your_model_name/
__init__.py
modeling_your_model_name.py # unified class: load_weights + forward + streaming
@@ -832,32 +815,40 @@ vllm_omni/deploy/your_model_name.yaml
No stage input processor is needed.
-### Stage config
+### Pipeline and deploy config
-Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the
-top-level `async_chunk: false` are required — omitting them causes silent
-misclassification in the serving layer. Set `max_num_seqs` to at least 4 for
-concurrent production use.
+Use a single `LLM_AR` stage that owns the tokenizer and produces final audio:
+
+```python
+YOUR_TTS_PIPELINE = PipelineConfig(
+ model_type="your_model_name",
+ default_deploy_config_name="your_model_name.yaml",
+ model_arch="YourModelForCausalLM",
+ stages=(
+ StagePipelineConfig(
+ stage_id=0,
+ model_stage="your_model_stage_key",
+ execution_type=StageExecutionType.LLM_AR,
+ input_sources=(),
+ owns_tokenizer=True,
+ engine_output_type="audio",
+ final_output=True,
+ final_output_type="audio",
+ ),
+ ),
+)
+```
+
+The deploy YAML controls runtime behavior. Set `max_num_seqs` to at least 4
+for concurrent production use:
```yaml
-# stage_configs/your_model_name.yaml
async_chunk: false
-stage_args:
+stages:
- stage_id: 0
- stage_type: llm
- is_comprehension: true # required for serving_speech.py dispatch
- runtime:
- devices: "0"
- engine_args:
- model_stage: your_model_stage_key
- model_arch: YourModelForCausalLM
- worker_type: ar
- scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
- engine_output_type: audio
- max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps
- final_output: true
- final_output_type: audio
+ devices: "0"
+ max_num_seqs: 4
```
### Generator-based streaming pattern
@@ -961,10 +952,17 @@ is correct. Stage the modified files and commit again — do not revert the chan
Common failures and fixes:
| Check | Cause | Fix |
-|-------|-------|-----|
+| ------- | ------- | ----- |
| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call |
| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block |
| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit |
+| `check-spdx-header` | Missing header, or copyright still says `vLLM project` | Use `Copyright contributors to the vLLM-Omni project`; restage the rewrite |
+| `check-forbidden-imports` | Stdlib `re`/`base64`, pickle, Hugging Face Hub API, or direct Triton/TileLang | `import regex as re` and `pybase64`; do not grow the allowlist without review |
+| `check-torch-cuda-call` | New `torch.cuda.*` call site | Use `current_omni_platform`; do not grow `ALLOWED_FILES` without review |
+| `check-tts-adapter-migration` | New `_tts_model_type` branch in `serving_speech.py` | Put logic in `tts_adapters/`; lower `MAX_MODEL_TYPE_BRANCHES` when removing branches |
+| `check-mark` | Test file missing level/hardware mark, or direct `pytest.mark.H100` | Add `core_model`/… plus `cpu`/`cuda`/`hardware_test(`; SKU marks only via helpers |
+| `shellcheck` | Native `shellcheck` missing, or script warning | Install via apt/dnf/brew (or `shellcheck.exe` on Windows). See [Linting](../README.md#linting) |
+| `mypy-3.10` / `markdownlint-cli2` / `check-buildkite` | Types, docs markdown, or Buildkite YAML | See [Linting](../README.md#linting) |
### DCO sign-off
@@ -1019,7 +1017,7 @@ Adding a TTS model to vLLM-Omni involves:
### Qwen3-TTS Reference Files
| File | Purpose |
-|------|---------|
+| ------ | --------- |
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
@@ -1031,4 +1029,4 @@ For more information, see:
- [Architecture Overview](../../design/architecture_overview.md)
- [Async Chunk Design](../../design/feature/async_chunk.md)
-- [Stage Configuration Guide](../../configuration/stage_configs.md)
+- [Pipeline and Deploy Configuration Guide](../../configuration/stage_configs.md)
diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md
index b57de441e..62fff29c5 100644
--- a/docs/contributing/profiling.md
+++ b/docs/contributing/profiling.md
@@ -322,6 +322,15 @@ Each 1-second window records:
On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima).
+`loop_idle` and `loop_active` count loop iterations, so their absolute scale is
+tied to which orchestration loop is running. Under the default poll loop an idle
+orchestrator records roughly one iteration per millisecond. Under the
+event-driven loop (`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, see
+[Speech API](../serving/speech_api.md#orchestration-loop-experimental)) an idle
+orchestrator wakes only once per 0.5 s reconcile timeout, while a busy one still
+records one iteration per routed output. Window counts and the `loop_active_pct`
+summary are therefore not comparable across the two modes.
+
### Relationship to other diagnostics
This monitor is intentionally separate from the existing profiling tools:
diff --git a/docs/design/ar_diffusion_pipeline_capability.md b/docs/design/ar_diffusion_pipeline_capability.md
index c9e05c8e8..3c1b4b9cd 100644
--- a/docs/design/ar_diffusion_pipeline_capability.md
+++ b/docs/design/ar_diffusion_pipeline_capability.md
@@ -4,6 +4,9 @@
> `experimental/ar_diffusion` contract. It follows the ownership and lifecycle
> direction in [#5137](https://github.com/vllm-project/vllm-omni/issues/5137),
> but it is not yet part of the stable module-design hierarchy.
+>
+> The end-to-end request, identity, failure, and memory contracts are documented
+> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).
`ARDiffusionEngine` is an opt-in diffusion backend for pipelines that keep
autoregressive attention KV across requests. The engine only selects
@@ -51,14 +54,13 @@ time, and its inherited batch/step entry points also fail explicitly. Supporting
those paths requires a batch-aware state-binding contract rather than silently
bypassing `bind_ar_diffusion_state()`.
-The current implementation retains one session per model instance. A request
-for a different `session_id` closes the resident session through the same LRU
-lifecycle used by explicit close and failure cleanup before allocating the new
-one. Self-attention and cross-attention memory are therefore sized for one
-resident session even if the pipeline declares a larger `session_capacity`.
-This prioritizes rollout throughput and prevents block-pool exhaustion before a
-count-based eviction can run. Multi-session residency can be added later with a
-memory-backed capacity calculation or allocation-pressure eviction.
+The implementation may retain multiple sessions per model instance. It selects
+the largest feasible capacity up to the pipeline's `session_capacity`, based on
+self-attention pages, cross-attention allocations, scratch, and declared
+model-owned state. One session is admitted whenever it fits actual free device
+memory; `gpu_memory_fraction` controls additional residency. Requests beyond
+that capacity evict the least-recently-used session through the same lifecycle
+hook used by explicit close and failure cleanup.
Scratch storage is also capability-driven. `frames_per_block` reserves space for
an uncommitted current video block, while `max_scratch_tokens_per_branch`
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index aa6f72770..d7938b62c 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -335,8 +335,8 @@ The single resolution boundary is implemented by
`StageConfigFactory.create_from_model()` and
`VllmOmniConfig.from_pipeline_config()` in
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
-The legacy `stage_args` YAML path remains only for models that have not yet
-migrated to `PipelineConfig` and `DeployConfig`.
+The legacy `stage_args` YAML path has been removed. Model topology now resolves
+through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
In the typed path, each stage configuration derives from
`BaseVllmOmniStageConfig` and is specialized as
@@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details.
stage contract.
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
resolves per-pipeline or per-component quantization configurations, while
- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
+ [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
streams diffusion blocks from host memory within the existing parallel
topology.
diff --git a/docs/design/feature/async_chunk.md b/docs/design/feature/async_chunk.md
index 57b4209b8..068a5ac1b 100644
--- a/docs/design/feature/async_chunk.md
+++ b/docs/design/feature/async_chunk.md
@@ -211,38 +211,39 @@ In sequential mode, each stage must wait for the previous stage to complete enti
## Configuration
-Enable async_chunk in stage configuration YAML:
+Enable async chunk mode in the deploy YAML:
```yaml
async_chunk: true
-stage_args:
- - stage_id: 0
- engine_args:
- custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
- - stage_id: 1
- engine_args:
- custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
```
-### Stage Configuration
+The registered `PipelineConfig` owns the async handoff processor functions.
+For Qwen3-Omni these are `thinker2talker_async_chunk` and
+`talker2code2wav_async_chunk`; deploy YAML only selects the mode and runtime
+sizing.
+
+### Configuration ownership
- `async_chunk: bool`: Enable/disable async chunk mode
-- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
-- `stage_connector_config: dict`: Connector configuration
-- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
+- `async_chunk_process_next_stage_input_func: str`: Pipeline-owned chunk processor path
+- `execution_type: StageExecutionType`: Pipeline-owned runtime family
+- `connectors: dict`: Deploy-owned connector definitions
- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
-
### Connector Configuration
```yaml
connectors:
- - from_stage: 0
- to_stage: 1
- spec:
- name: SharedMemoryConnector
- extra:
- stage_id: 0
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+
+stages:
+ - stage_id: 0
+ output_connectors:
+ to_stage_1: connector_of_shared_memory
+ - stage_id: 1
+ input_connectors:
+ from_stage_0: connector_of_shared_memory
```
### Code2Wav Batch Configuration
@@ -250,13 +251,10 @@ connectors:
For optimal performance with async_chunk, the code2wav stage should be configured with batching:
```yaml
-stage_args:
+stages:
- stage_id: 2 # code2wav stage
- runtime:
- devices: "1"
- engine_args:
- model_stage: code2wav
- max_num_seqs: 64 # Enables batched audio generation
+ devices: "1"
+ max_num_seqs: 64 # Enables batched audio generation
```
## Related Files
diff --git a/docs/design/feature/async_diffusion_output.md b/docs/design/feature/async_diffusion_output.md
index a81ed6811..733478883 100644
--- a/docs/design/feature/async_diffusion_output.md
+++ b/docs/design/feature/async_diffusion_output.md
@@ -117,7 +117,22 @@ After (async D2H):
### Batch Split
-When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
+When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
+
+### Reliability, Lifecycle & Timeout Behavior
+
+1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
+ When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
+
+2. **Timeout Diagnostics (`_async_output_timeout()`)**:
+ The engine waits for `wait_output_ready()` with a timeout of `_async_output_timeout()` — 600.0s by default, overridable per run with the `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` environment variable (a non-numeric or non-positive value is ignored with a warning and the default applies, since this resolves on the request path). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter. The packing itself takes milliseconds, but the wait is queued behind the step's GPU work, so its wall-clock duration tracks step time; the bound is generous because a hung engine is surfaced by the worker monitor and `check_health()`, not by this timeout.
+
+3. **Queue Serialization & Atomic Future Resolution**:
+ - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
+ - `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
+
+4. **Storage-Aware IPC Packing**:
+ `pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.
## Model Coverage
@@ -131,6 +146,7 @@ All models running in request-mode (`step_execution=False`, the default) automat
|---|---|---|
| **HunyuanImage-3.0** | Image | `False` |
| **Qwen-Image** | Image | `True` |
+| **LTX-2.3** | Video / Audio | `False` |
Other models with `step_execution=False` are also supported but not yet verified.
diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md
new file mode 100644
index 000000000..41cbaf3ff
--- /dev/null
+++ b/docs/design/feature/attention_backend_selection.md
@@ -0,0 +1,106 @@
+# Diffusion Attention Backend Selection
+
+This document defines the selection and extension contract for diffusion
+attention backends. User-facing backend choices, installation, and tuning are
+in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
+
+## Scope
+
+The contract applies to DiT and other diffusion attention layers. It is
+separate from vLLM's autoregressive attention selector.
+
+The selection path has four responsibilities:
+
+1. normalize user configuration into `AttentionConfig` and `AttentionSpec`;
+2. resolve a spec for an attention role;
+3. ask the active platform to validate an explicit backend or choose a
+ hardware default; and
+4. load the selected `AttentionBackend` class from the registry path.
+
+## Resolution contract
+
+`get_attn_backend_for_role()` resolves configuration in this order:
+
+1. exact `per_role[role]`;
+2. category `per_role[role_category]`;
+3. `default`;
+4. platform default.
+
+An explicit resolution returns both the backend class and its
+`AttentionSpec`. A platform-default resolution returns the class and `None`.
+Layers must therefore treat the spec as optional and must not infer that a
+missing spec means SDPA.
+
+Class resolution is cached by backend name, head size, and whether TRTLLM may
+be selected automatically. Logging is cached separately by role and source so
+the startup record identifies why each role received its backend.
+
+## Registry and platform boundary
+
+`DiffusionAttentionBackendEnum` maps stable configuration names to default
+qualified class paths. `register_diffusion_backend()` may replace a path at
+runtime without changing the public enum value.
+
+The active `OmniPlatform` owns compatibility policy through
+`get_diffusion_attn_backend_cls()`. It must:
+
+- validate explicit selections and fail with an actionable error when the
+ requested kernel cannot run;
+- choose only an available, compatible default when no backend is explicit;
+- account for head size and other platform-visible constraints; and
+- return a qualified class path implementing `AttentionBackend`.
+
+The selector must not duplicate device capability or package-availability
+policy that belongs to the platform.
+
+## Typed backend options
+
+Backend-specific settings remain typed fields on `AttentionSpec` rather than
+unstructured keyword dictionaries:
+
+- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
+ validation;
+- `skip_softmax` is consumed by TRTLLM; and
+- `block_sparse` is consumed by block-sparse backends such as RainFusion; and
+- `fastvideo_vsa_topk` is consumed by FastVideo VSA.
+
+A backend reads only the fields it owns and rejects incompatible values. New
+options should be added to a shared typed spec only when more than one backend
+shares their semantics; otherwise add a dedicated typed block.
+
+## Model integration contract
+
+Each diffusion attention call declares a stable role and, when appropriate, a
+broader category. Model-specific roles permit precise overrides; categories
+preserve common `self` or `cross` policies. A model also decides whether its
+path is compatible with automatic TRTLLM selection, because only the model
+knows whether masking and packing satisfy the kernel contract.
+
+## Adding or changing a backend
+
+An implementation change is complete only when it:
+
+1. implements the `AttentionBackend` interface and exposes a stable enum name;
+2. adds platform validation and default routing where appropriate;
+3. defines and validates any typed configuration;
+4. tests explicit selection, default selection, and incompatible paths;
+5. documents installation, fallback behavior, and quality implications in the
+ corresponding user guide; and
+6. updates the overview matrix without moving algorithm details back into the
+ landing page.
+
+## FastVideo VSA model metadata
+
+FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
+the original latent video grid. Wan integrations therefore attach the
+post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
+backend validates that `T * H * W` equals the sequence length, derives the
+runtime block count, and selects the configured top-k key/value blocks for
+each query block.
+
+Checkpoint-specific compensation remains model-owned. Wan layers pass an
+optional learned `gate_compress` projection through attention metadata when
+the checkpoint contains those weights. Checkpoints without the projection do
+not allocate or execute it. When top-k selects every block, native checkpoints
+route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
+to retain checkpoint semantics.
diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md
index 053360494..db8f738c1 100644
--- a/docs/design/feature/cfg_parallel.md
+++ b/docs/design/feature/cfg_parallel.md
@@ -63,7 +63,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
### N-Branch CFG (3+ branches)
-Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
+Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches.
`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
diff --git a/docs/design/feature/disaggregated_inference.md b/docs/design/feature/disaggregated_inference.md
index c31086c0d..ba8c8f9fc 100644
--- a/docs/design/feature/disaggregated_inference.md
+++ b/docs/design/feature/disaggregated_inference.md
@@ -59,19 +59,18 @@ This `metadata` must be passed through the control plane so `get()` can locate t
## Configuration Model
-Define connectors in runtime:
+Define connectors at the top level of the deploy YAML:
```yaml
-runtime:
- connectors:
- connector_of_shared_memory:
- name: SharedMemoryConnector
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
```
Wire stages to connectors:
```yaml
-stage_args:
+stages:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_shared_memory
diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md
deleted file mode 100644
index da735d60b..000000000
--- a/docs/design/feature/distributed_layerwise_offload.md
+++ /dev/null
@@ -1,156 +0,0 @@
-# Distributed Layerwise Offload
-
-This document describes distributed layerwise offload (DLO) for diffusion
-models. DLO keeps only a small number of DiT blocks on the accelerator and
-streams the remaining blocks from host memory. The distributed backend can
-either shard those host-side weights across an existing parallel group or keep
-the standard loader's rank-local weights and avoid an additional collective.
-
-For user-facing commands, see the
-[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
-and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
-
-## Status
-
-DLO is implemented for multi-device diffusion execution. The default
-AllGather path is the primary path for DP and SP deployments. The
-`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
-for standard-loader sharding, workstation bring-up, and systems where an
-additional DLO collective is undesirable, but it does not reduce host weight
-storage across ranks.
-
-The compatibility matrix below describes the current implementation. The
-unit-level guards are covered, but not every parallelism combination has a
-full model-and-hardware end-to-end test.
-
-## Design
-
-### DLO consumes the existing parallel topology
-
-DLO does not create a new DP, TP, or SP topology. It reads the configured
-`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
-standard distributed groups have been initialized.
-
-The DLO weight-sharding group is selected as follows:
-
-1. Use the existing DP group when `data_parallel_size > 1`.
-2. When DP is one and SP is greater than one, use the SP group.
-3. Otherwise, run rank-locally without a DLO process group.
-
-TP is deliberately not used as DLO's AllGather group. HSDP has its own
-parameter-sharding lifecycle and is not allowed to be sharded a second time by
-DLO's AllGather path.
-
-### AllGather path
-
-With the default `dlo_use_allgather=True`, each rank stores approximately
-`1 / group_size` of each streamable block in pinned host memory. The next
-block's shard is copied to a device buffer and reconstructed with
-`all_gather_into_tensor` on a communication stream while the current block is
-executing.
-
-```text
-Compute: [Block N] [Block N+1] [Block N+2]
-H2D: [shard N+1] [shard N+2]
-AllGather: [full N+1] [full N+2]
-Buffers: [current slot] [prefetch slot] [current slot]
-```
-
-
-
-The backend uses two shared device buffers, so accelerator weight residency is
-bounded by the largest streamed blocks rather than the complete model.
-
-When DP is greater than one, the engine can process one request per DP rank in
-the same denoising wave. Because AllGather is a collective, all participating
-requests must take the same execution path at every denoising step.
-
-### Rank-local path without DLO AllGather
-
-With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one. The regular model loader remains responsible for preparing each rank's
-weights, including TP-local tensors or HSDP-managed parameters. DLO then
-streams those rank-local tensors block by block using H2D copies only.
-
-This mode means:
-
-- DP still provides independent replicas, but DLO does not shard weights
- across DP ranks.
-- SP still performs its normal activation/attention collectives, but DLO does
- not shard weights across SP ranks.
-- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
- DLO's additional weight AllGather is disabled.
-- Pure DP deployments keep a full host-side model copy per rank, subject to
- shared OS page-cache behavior.
-- The scheduler does not require a synchronized DP request wave for DLO.
-
-This path is intentionally not implemented by reusing the AllGather mmap
-loader. It relies on the standard loader so model-specific TP/HSDP transforms
-remain intact.
-
-## Parallelism compatibility
-
-| Parallelism | DLO + AllGather | DLO without AllGather |
-|---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
-| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
-| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
-
-### Combined dimensions
-
-- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
- than one; SP continues to use its own sequence-parallel group. If DP is one,
- the SP group becomes DLO's sharding group in AllGather mode.
-- **DP + TP/SP without AllGather:** standard model loading defines the
- rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
- collective.
-- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
- DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
- materialization and synchronization.
-- **HSDP + DP or TP:** rejected independently by the diffusion parallel
- configuration.
-
-## Request and loading constraints
-
-AllGather DP multi-concurrency requires:
-
-- explicit `num_inference_steps`;
-- the same `num_inference_steps` for all requests in a wave; and
-- identical request arguments that affect the collective execution path.
-
-The no-AllGather path does not impose these DLO-specific synchronized-wave
-requirements.
-
-The mmap loader is used only by the supported DLO+AllGather path when the
-model has a compatible checkpoint layout. Online quantization is incompatible
-with that sharded mmap path; use `--dlo-no-use-allgather` or disable online
-quantization. Models that do not meet the mmap requirements use the regular
-loader path.
-
-## Validation coverage
-
-Current source-level validation includes:
-
-- HSDP + DLO + AllGather rejection;
-- HSDP + DLO without AllGather acceptance at configuration level;
-- TP rejection in the DLO+AllGather mmap path;
-- resident-layer requests requiring no-AllGather;
-- DP request-wave validation for denoising-step compatibility;
-- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
- tests.
-
-The highest-value missing coverage is end-to-end numerical comparison against
-ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
-on the target CUDA/NCCL or CANN/HCCL hardware.
-
-## Recommendations
-
-- Use **DP + DLO AllGather** for the supported throughput and host-memory
- scaling path.
-- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
- not the goal.
-- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with
- the expectation of higher host-memory use and lower validation confidence.
-- Prefer **HSDP alone** for production HSDP deployments until the combined
- HSDP + DLO no-AllGather path has broader end-to-end coverage.
diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md
new file mode 100644
index 000000000..a9bb1b0b8
--- /dev/null
+++ b/docs/design/feature/host_weight_runtime.md
@@ -0,0 +1,386 @@
+# Host Weight Runtime
+
+The Host Weight Runtime is a loader-facing feature for reusing immutable,
+runtime-ready model weights across workers on one host. It lets a model
+integration replace repeated checkpoint loading and transformation with an
+exact lookup of a previously published host representation, while keeping the
+canonical model source authoritative.
+
+This feature is shared infrastructure. Diffusion, autoregressive, speech, and
+future model integrations may depend on the top-level
+`vllm_omni.host_weight_runtime` package without depending on one another's
+loader or execution stack. The internal component and filesystem contracts are
+specified in the [Host Weight Runtime module design](../module/host_weight_runtime.md).
+
+## Status
+
+The implementation provides contracts and a CPU local-filesystem store. The
+diffusion consumer additionally defines typed,
+representation-independent final-layout identity/restoration mechanics plus a
+concrete BF16-with-preserved-FP32 policy for MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. The opt-in no-AllGather DLO integration
+selects, publishes, restores, and transfers these artifacts.
+
+V1 includes:
+
+- exact runtime-weight identity and immutable manifests;
+- coordinated local lookup and one-producer publication;
+- descriptor-backed safetensors mmap leases;
+- preferred and required resolution policy;
+- explicit, separately reported post-load publication;
+- validation, deny, quarantine, cleanup, and capacity controls; and
+- typed reports for every terminal resolution outcome.
+
+V1 does not include:
+
+- default-on activation or consumers outside no-AllGather DLO;
+- online FP8, quantized, merged-adaptation, or additional model producers;
+- HWR interaction with DLO AllGather;
+- a remote artifact provider or cross-node coordination;
+- automatic eviction; or
+- a change to DLO collective or execution behavior.
+
+## Motivation and use cases
+
+Model loading can create the same final host representation repeatedly. This
+is especially expensive when loading performs checkpoint decoding, tensor
+renaming, TP slicing, quantization, packing, or scale construction before GPU
+transfer. Independent workers may then retain private CPU copies of identical
+runtime weights.
+
+The Host Weight Runtime is useful when:
+
+- multiple same-node workers request semantically identical final weights;
+- constructing those weights is materially more expensive than validating and
+ mapping a local artifact;
+- the runtime layout can be identified exactly and reproduced deterministically;
+- the workers can share immutable file-backed pages through the OS page cache;
+ and
+- canonical loading remains available when policy permits fallback.
+
+Typical consumers include same-node diffusion replicas, DLO host-weight
+sources, and future transformed or quantized model loaders. Request routing and
+replica orchestration remain outside this feature.
+
+This feature is not zero-copy GPU execution. A lease exposes stable CPU tensor
+views and mapped ranges. A separate transport still decides whether to use
+registered mmap, private pinned staging, synchronous copies, or asynchronous
+H2D transfer.
+
+## Resolution behavior
+
+The loader resolves the immutable canonical source and computes the exact
+requested identity before asking the runtime for a representation.
+
+```mermaid
+flowchart TD
+ L["Loader resolves canonical source"] --> I["Compute exact WeightArtifactIdentity"]
+ I --> M{"Runtime mode"}
+ M -->|"disabled"| C["Canonical loader"]
+ M -->|"preferred or required"| A{"Validated local artifact?"}
+ A -->|"yes"| R["Return HostWeightLease"]
+ A -->|"no"| X["Remote exact artifact (future)"]
+ X --> P{"Registered producer available?"}
+ P -->|"yes"| B["Build, validate, and publish atomically"]
+ B --> R
+ P -->|"no or recoverable failure"| F{"Policy permits fallback?"}
+ F -->|"preferred"| C
+ F -->|"required"| E["Fail startup"]
+ P -->|"nonretryable failure"| E
+ C --> W{"Post-load publication enabled?"}
+ W -->|"yes"| B2["Publish final model through POST_LOAD_ONLY producer"]
+ B2 --> CL["Close validated publication lease"]
+ CL --> R2["Return separate publication report"]
+ R2 --> D
+ W -->|"no"| D["Keep canonical model"]
+```
+
+The modes are:
+
+| Mode | Observable behavior |
+| --- | --- |
+| `disabled` | Use canonical loading directly without probing storage, identity, credentials, or topology. |
+| `preferred` | Return an exact lease when available; otherwise use canonical fallback only for a miss, invalid cache entry, or typed retryable failure. |
+| `required` | Fail when an exact lease cannot be acquired. |
+
+During pre-load resolution, an unsupported capability, semantic identity
+collision, producer failure, or publication failure is nonretryable and remains
+visible even in preferred mode. A post-load publication failure is likewise
+visible in its own report, but cannot revise the canonical-fallback outcome.
+Storage policy must never disguise a semantic or configuration error as a cache
+miss.
+
+## Loader and restoration sequence
+
+The model integration owns the end-to-end loading transaction. The runtime
+does not construct a model or invoke the canonical loader.
+
+```mermaid
+sequenceDiagram
+ participant L as Loader adapter
+ participant R as HostWeightRuntime
+ participant S as HostWeightStore
+ participant P as WeightProducer
+ participant X as WeightRestorer
+ participant T as GPU transport
+
+ L->>L: Resolve canonical revision and exact identity
+ L->>R: resolve(identity, producer)
+ R->>S: lookup(exact identity)
+ alt validated hit
+ S-->>R: HostWeightLease
+ R-->>L: lease and LOCAL_HIT report
+ L->>X: plan_restore(model, lease)
+ X-->>L: validation-only restore plan
+ L->>X: commit() once
+ else miss with allowed producer
+ R->>S: get_or_build(identity, producer)
+ S->>P: produce(store-scoped writer)
+ P-->>S: final-layout tensors and metadata
+ S-->>R: validated HostWeightLease
+ R-->>L: lease and LOCAL_PRODUCTION report
+ L->>X: plan_restore(model, lease)
+ X-->>L: validation-only restore plan
+ L->>X: commit() once
+ else policy permits canonical fallback
+ R-->>L: CANONICAL_FALLBACK
+ L->>L: Run canonical loader
+ opt explicit post-load publication enabled
+ L->>R: publish_after_load(identity, POST_LOAD_ONLY producer)
+ R->>S: get_or_build(identity, producer)
+ S-->>R: validated HostWeightLease or typed failure
+ R->>R: close publication lease
+ R-->>L: separate publication report
+ end
+ else required or nonretryable failure
+ R-->>L: FAILED report
+ end
+ L->>T: consume final CPU tensors and any lease
+ T-->>L: transfer teardown complete
+ L->>L: close lease when one was acquired
+```
+
+`plan_restore()` must not mutate the model or lease. `commit() -> None` is the
+sole one-shot model mutation. If planning fails, canonical fallback may reuse
+the untouched model. If commit begins and fails, the partially hydrated model
+must be discarded and canonical fallback must construct a fresh model.
+
+`publish_after_load()` is synchronous in V1 and accepts only a
+`POST_LOAD_ONLY` producer. Its report is separate from the terminal resolution,
+so a publication failure cannot rewrite a successful canonical fallback. A
+successful publication closes the store-returned lease inside the runtime and
+warms only future startups. It does not restore, rebind, or otherwise mutate the
+canonically loaded model serving the current startup.
+`allow_local_build` gates producers during pre-load resolution, while
+`allow_post_load_publish` independently gates this explicit post-load path.
+
+## Exact representation identity
+
+The store performs exact lookup. It never silently converts one representation
+into another or substitutes the first backing that responds.
+
+Identity includes:
+
+- immutable model revision and source fingerprint;
+- model component and ownership boundary;
+- representation name, dtype, and format metadata;
+- final tensor layout and semantic parallel coordinates;
+- static adaptation identity; and
+- producer implementation, manifest, and restorer schema versions.
+
+The requested representation is selected by the loader. The store chooses
+where to obtain that exact representation: validated local artifact, future
+remote materialization, or a registered producer.
+
+### Parallelism
+
+Parallel coordinates are included only when they change bytes, shape, layout,
+or component ownership.
+
+| Dimension | Identity and sharing rule |
+| --- | --- |
+| DP | Exclude DP rank for replicated weights so same-node replicas can acquire the same artifact. Include a coordinate only if DP changes weight ownership. |
+| TP | Include TP size and rank when each rank owns different tensor slices or packed layouts. Each distinct TP shard is a separate exact artifact. |
+| SP | Include SP size and backend when they change the weight layout. Exclude SP rank when all SP ranks consume identical bytes. |
+| PP | Encode component ownership or stage-local layout when PP changes which weights belong to the consumer. |
+| EP | Encode expert ownership and layout whenever ranks own different expert weights. |
+
+The runtime does not infer these rules from process-group topology. The model
+adapter is responsible for constructing a correct identity.
+
+### Quantization and adaptations
+
+Runtime BF16, FP8, packed quantized weights, and future formats are distinct
+representations. Conversion is performed by an explicitly versioned producer,
+never by cache coercion. A producer must publish the final layout consumed by
+its matching restorer.
+
+Dynamic LoRA overlays are not part of a reusable base-weight artifact. A
+statically merged adapter is cacheable only as a separate identity containing
+the adapter fingerprint and merge semantics.
+
+### Initial diffusion final-layout contract
+
+The shared diffusion contract covers complete final-layout DiT parameters and
+persistent buffers. Text encoders, VAEs, non-persistent derived state, and other
+pipeline components remain outside the artifact. One explicit representation
+policy selects allowed dtypes, tensor roles, physical layout identity, producer
+ABI, manifest schema, and restoration schema.
+
+The contract is intentionally separate from loader activation:
+
+- `FinalLayoutRequest` contains typed loader identity/configuration fingerprints,
+ TP coordinate, and conservative SP semantics. It has no open metadata bag,
+ DP coordinate, SP rank, device identity, DLO transfer mode, registration
+ policy, or store path.
+- `FinalLayoutArtifactSpec` binds one `WeightRepresentation` and runtime-layout
+ name to explicit producer/restorer schemas and a canonical, versioned
+ implementation ABI descriptor. Compatibility never depends on reflective
+ source inspection.
+- `PreparedWeightSource` snapshots immutable revisions or exact local file
+ content plus a typed checkpoint-adapter identity before ordinary
+ materialization. Source replacement before or during production fails
+ publication. A hash-looking symlink basename is trusted only for an explicit
+ Hugging Face Hub source whose repository ID and
+ `models--.../snapshots/ -> blobs/` topology validate; every
+ local or otherwise unverified symlink target is content-hashed.
+- the tensor ownership digest records exact runtime names, kinds, shapes,
+ semantic roles, dtypes, and strides from a CPU or meta model skeleton;
+- `FinalLayoutTensorRestorer` accepts only an exact lease identity, validates
+ complete policy-defined coverage without mutation, and returns a one-shot
+ commit plan;
+- each model declares one dtype-neutral `FinalLayoutModelContract` with an
+ explicit implementation version and a post-commit validator; and
+- `FinalLayoutBF16Producer` accepts only the matching identity context and a
+ finalized CPU model. It is `POST_LOAD_ONLY` and `SINGLE_PROCESS` per exact TP
+ coordinate. Its BF16 policy preserves model-declared FP32 parameters and
+ buffers, revalidates MiniMax H3 mixed-precision invariants, and revalidates
+ FLUX.2-klein's two block stacks, packed QKV mapping, and BF16 base layout.
+
+Other representations reuse source identity, typed parallel identity, tensor
+ownership, and exact restoration only when their policy proves those semantics.
+For example, runtime FP8 needs a separate policy/producer for generated scales,
+quantization metadata, and Cutlass physical layouts; it is not enabled by
+changing a dtype string on the BF16 producer.
+
+This stage makes no startup, sharing, or DLO performance claim. A following
+consumer PR owns disabled/preferred/required precedence, mixed-component loader
+transactions, warm-hit restoration, and transactional lease handoff. A TP2
+prewarm deployment will require a matching TP2 producer cohort to populate both
+TP-coordinate identities even though the store coordinates each artifact
+independently.
+
+## Host sharing and GPU transport
+
+Every process receives its own virtual mappings and `HostWeightLease`, but
+workers mapping the same local artifact can share physical file-backed pages
+through the OS page cache. This avoids independent private tensor allocations;
+it does not guarantee that process PSS reports exactly one model copy.
+
+The ownership boundary is:
+
+- the store owns immutable artifact files, publication, and lifecycle;
+- the kernel owns page-cache residency and placement;
+- the lease owns process-local mappings and the shared artifact lock; and
+- transport owns page registration, page locking, private staging, H2D copies,
+ streams, device buffers, and lease release ordering.
+
+DLO AllGather and no-AllGather remain transport and execution choices. A DLO
+integration may consume a lease, but the Host Weight Runtime does not choose a
+parallel collective or orchestrate DP requests. See the
+[DLO feature design](offloader/distributed_layerwise_offload.md).
+
+## Locality and NUMA policy
+
+The V1 backend is node-local. It accepts an allowlist of known local disk
+filesystems plus tmpfs/ramfs, records the detected kernel filesystem type, and
+rejects known remote or unknown filesystems. NFS, CIFS, Lustre, Ceph, and other
+cross-node mounts cannot silently satisfy the local backend because their
+page-cache and advisory-lock behavior does not provide the required node-local
+contract.
+
+A NUMA-specific root creates a separate storage domain but does not by itself
+place page-cache pages on that NUMA node. A topology-aware integration must
+also pin workers, establish local first touch or prefault, perform transport
+registration in the same domain, and collect residency evidence.
+
+Tmpfs artifacts consume host memory and may consume swap. They must be
+accounted as memory-backed storage rather than ordinary disk capacity.
+
+## Concurrency and timeouts
+
+One process per exact identity owns a build; other workers wait and then acquire
+leases for the published artifact. Publication is invisible until all payloads
+and metadata are validated, hashed, fsynced, and atomically renamed.
+
+`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not
+cancel synchronous validation, a producer that has already started, or atomic
+publication. A hung in-process producer therefore blocks its owning process and
+must be handled by external process supervision. Enforceable producer
+cancellation requires a future process-isolated producer contract.
+
+Leases keep a shared artifact lock for their lifetime. A forked child may close
+its inherited descriptors but cannot unlock the parent process's lease. Cleanup
+uses a nonblocking exclusive lock and reports an active lease instead of
+removing live mappings.
+
+## Failure and fallback contract
+
+| Condition | Preferred mode | Required mode |
+| --- | --- | --- |
+| Exact local hit | Return lease | Return lease |
+| Local miss | Try later backing or canonical fallback | Try later backing or fail |
+| Invalid or denied artifact | Quarantine/rebuild when possible; otherwise canonical fallback | Quarantine/rebuild or fail |
+| Retryable lock, domain, or capacity failure | Canonical fallback | Fail |
+| Unsupported producer or backend | Fail | Fail |
+| Identity collision | Fail | Fail |
+| Pre-load producer or atomic publication failure | Fail | Fail |
+| Post-load publication failure after canonical fallback | Keep canonical model; report publication failure | Not reached through required resolution |
+| Restore planning failure | Canonical fallback may reuse untouched model | Fail |
+| Restore commit failure | Discard model; fallback requires a fresh instance | Fail |
+
+Cache failures should normally reduce performance rather than availability,
+but only when the typed failure explicitly permits fallback. Canonical source,
+authentication, configuration, and semantic errors remain visible.
+
+## Consumer integration requirements
+
+A consumer PR must:
+
+1. Resolve an immutable canonical model source before cache lookup.
+2. Construct the exact final representation identity, including relevant
+ parallel, quantization, and adaptation semantics.
+3. Register a deterministic producer for that identity or support lookup-only
+ operation.
+4. Provide a validation-only restorer with a one-shot commit.
+5. Keep GPU transport outside the producer, store, and restorer contracts.
+6. Handle preferred fallback and required failure without reusing a model after
+ a failed restore commit.
+7. Retain the lease until every transport operation that may access mapped
+ memory has completed.
+8. Emit and test the terminal resolution report and any separate post-load
+ publication report.
+
+Consumer validation must prove:
+
+- output parity with canonical loading;
+- exact warm-hit work avoidance;
+- correct TP/SP/PP/EP ownership and layout identity;
+- correct BF16, FP8, quantized, or adaptation semantics;
+- shared-backing evidence when host-memory savings are claimed;
+- clean unmapping, unregistration, and artifact-lock release; and
+- startup, host-memory, and transport effects for the claimed deployment.
+
+## Rollout
+
+The rollout is intentionally staged:
+
+1. Land the neutral contracts and local filesystem store.
+2. Add one model-specific producer/restorer integration with parity evidence.
+3. Connect eligible DLO or other transport paths without moving transport
+ ownership into the runtime.
+4. Add additional final-layout and quantized producers independently.
+5. Introduce remote materialization only through the same local lease contract.
+
+The architectural decisions and deferred work are tracked in
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414).
diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md
new file mode 100644
index 000000000..f071196b1
--- /dev/null
+++ b/docs/design/feature/offloader/README.md
@@ -0,0 +1,80 @@
+# CPU Offloading
+
+This document defines the shared architecture for diffusion CPU offloading.
+The three strategies have separate user and design pages:
+
+| Strategy | User guide | Design contract |
+| --- | --- | --- |
+| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
+| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
+| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
+
+## Strategy selection
+
+`OffloadConfig.from_od_config()` converts diffusion configuration into one
+`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
+
+1. `DISTRIBUTED_LAYERWISE`;
+2. `LAYERWISE`;
+3. `MODEL_LEVEL`; and
+4. `NONE`.
+
+The factory logs when a higher-priority strategy disables another. It derives
+parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
+a separate offload group size.
+
+`get_offload_backend()` then validates platform offload support, resolves the
+device, and creates exactly one backend. Returning `None` means offloading is
+disabled or unsupported and must not leave partially installed hooks.
+
+## Shared lifecycle
+
+Every backend implements `OffloadBackend`:
+
+- `enable(pipeline)` discovers modules, establishes initial residency, and
+ installs hooks;
+- `disable()` removes owned hooks and resources; and
+- `is_enabled()` reports lifecycle state.
+
+`disable()` does not promise to restore the pipeline's original device
+placement. The caller owns any subsequent rematerialization.
+
+Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
+must use distinct hook names and remove only hooks they own.
+
+## Discovery and topology
+
+Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
+
+- `_dit_modules`;
+- `_encoder_modules`;
+- `_vae_modules`; and
+- `_resident_modules`.
+
+Dotted paths are supported. Legacy pipelines may use the fallback scan of
+well-known attribute names, but new integrations should declare components
+explicitly.
+
+Ordinary layerwise offload obtains topology from declared
+`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
+compatibility path). Distributed layerwise offload may instead use a pipeline
+`OffloadPlan`, then falls back to declared block attributes and its documented
+compatibility discovery. Discovery metadata describes structure only; the
+backend remains responsible for transfer, synchronization, and storage
+ownership.
+
+## Cross-strategy invariants
+
+- At most one strategy owns offload hooks for a pipeline.
+- A parameter has one authoritative host representation while offloaded.
+- Device storage is not freed until dependent compute or transfer work is
+ complete.
+- Non-persistent and model-specific buffers remain correct after movement or
+ checkpoint rematerialization.
+- Unsupported parallel or loading combinations fail before hooks mutate the
+ model.
+- Platform streams, events, synchronization, and cache management go through
+ the vLLM-Omni platform abstraction.
+
+The diffusion [Offloader module design](../../module/diffusion/offloader.md)
+describes how these feature contracts fit into the larger diffusion runtime.
diff --git a/docs/design/feature/offloader/cpu_offload.md b/docs/design/feature/offloader/cpu_offload.md
new file mode 100644
index 000000000..2318611cc
--- /dev/null
+++ b/docs/design/feature/offloader/cpu_offload.md
@@ -0,0 +1,56 @@
+# Model-Level Offload
+
+Model-level offload enforces mutual exclusion between pipeline component
+groups. For commands and model support, see the
+[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
+
+## Generic hook path
+
+`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
+makes encoders and VAEs device resident for initialization, then installs
+`SequentialOffloadHook` instances:
+
+- each DiT hook offloads all encoders and other DiTs before loading itself;
+- each encoder hook offloads every DiT before loading itself; and
+- VAEs are not part of the mutual-exclusion hook set.
+
+The hook moves parameters and buffers without recursively calling
+`module.to()`. This avoids recursion through modules that retain references to
+other transformer blocks. Host tensors are pinned when configured; DTensors
+are not pinned through the ordinary tensor path.
+
+Before forward returns control to the model, the hook synchronizes the active
+platform so compute cannot race the transfer. XPU uses blocking device-to-host
+copies because its allocator cannot safely evict storage while an asynchronous
+copy is pending.
+
+## Split-model delegation
+
+If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
+delegates instead of installing generic encoder/DiT hooks. This supports models
+whose mutually exclusive components are nested inside one transformer, such
+as a reasoner and generator.
+
+The delegated pipeline owns its internal contexts, but must preserve the same
+contract:
+
+- only the executing component is device resident;
+- enable and disable are idempotent;
+- transfers use the supplied device, pinning, and HSDP settings; and
+- disabling removes all model-local offload hooks or contexts.
+
+The delegation is currently duck typed. If another independent implementation
+is added, formalize the enable/disable pair as a protocol rather than expanding
+attribute-name heuristics.
+
+## Extension invariants
+
+- New ordinary pipelines declare both DiT and encoder groups through
+ `SupportsComponentDiscovery`.
+- Models with multiple DiTs offload the inactive DiTs as well as encoders.
+- Empty or parameterless modules do not trigger transfers.
+- Hook removal does not move modules; callers must not assume original
+ residency is restored.
+
+Shared selection and lifecycle behavior is defined in the
+[CPU Offloading design](README.md).
diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md
new file mode 100644
index 000000000..2b5f929ed
--- /dev/null
+++ b/docs/design/feature/offloader/distributed_layerwise_offload.md
@@ -0,0 +1,555 @@
+# Distributed Layerwise Offload
+
+This document describes distributed layerwise offload (DLO) for diffusion
+models. DLO keeps only a small number of DiT blocks on the accelerator and
+streams the remaining blocks from host memory. The distributed backend can
+either shard those host-side weights across an existing parallel group or keep
+complete rank-local block sources and avoid an additional collective.
+
+For user-facing commands, see the
+[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
+and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
+
+## Feature compatibility
+
+Host-storage optimization and runtime compatibility are separate decisions.
+When direct checkpoint mmap is unavailable, DLO can still use tensors produced
+by the ordinary loader. "Compatibility path" below means that fallback is
+implemented but has less end-to-end coverage than the primary path.
+
+Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsupported.
+
+| Feature | DLO + AllGather | DLO without AllGather |
+| --- | --- | --- |
+| **DP** | ✅ Primary path; host weights are sharded across the DP group. | ✅ Each DP rank streams complete rank-local blocks. |
+| **SP** | ✅ When DP=1, DLO uses the SP group for weight sharding. | ✅ SP remains active without a DLO weight collective. |
+| **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
+| **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
+| **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Online INT8 linears** | ✅ Ordinary loader finalizes the int8 weight (contiguous or transposed view) and fp32 scale before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Online MXFP8 linears** | ✅ Ordinary loader finalizes the fp8 weight (contiguous) and e8m0 block scale (contiguous or transposed view) before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
+| **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
+| **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
+
+See [Parallelism compatibility](#parallelism-compatibility) and
+[Request and loading constraints](#request-and-loading-constraints) for the
+detailed contracts and validation boundaries.
+
+## Status
+
+DLO is implemented for multi-device diffusion execution. The default
+AllGather path is the primary path for DP and SP deployments. The
+`--dlo-no-use-allgather` path streams complete blocks independently and adds no
+DLO weight collective.
+
+Host storage is selected separately from the transfer protocol. The loader can
+produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
+otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
+the same node can share immutable checkpoint pages when direct mmap is
+selected, while the ordinary-loader fallback still keeps a private runtime
+copy per process.
+
+The Phase A shared-mmap support boundary is TP1. TP greater than one is an
+ordinary-loader compatibility path: DLO can consume the resulting TP-local
+tensors, but those configurations do not use checkpoint mmap and must not be
+used to claim shared-mmap host-memory savings.
+
+The compatibility matrix below describes the current implementation. The
+unit-level guards are covered, but not every parallelism combination has a
+full model-and-hardware end-to-end test.
+
+## Design
+
+### DLO consumes the existing parallel topology
+
+DLO does not create a new DP, TP, or SP topology. It reads the configured
+`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
+standard distributed groups have been initialized.
+
+The DLO weight-sharding group is selected as follows:
+
+1. Use the existing DP group when `data_parallel_size > 1`.
+2. When DP is one and SP is greater than one, use the SP group.
+3. Otherwise, run rank-locally without a DLO process group.
+
+TP is deliberately not used as DLO's AllGather group. HSDP has its own
+parameter-sharding lifecycle and is not allowed to be sharded a second time by
+DLO's AllGather path.
+
+### The loader owns host-weight planning
+
+Before it decides whether ordinary weight materialization can be skipped, the
+diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
+accepted only when preflight proves all of the following:
+
+- every required DiT parameter and persistent buffer has exactly one source;
+- runtime names, checkpoint keys, shapes, and dtypes match;
+- the runtime topology is TP1 without HSDP or online quantization; and
+- every custom loader operation is represented by a loader-owned checkpoint
+ adapter.
+
+The exact plan object is handed to DLO. The backend does not rescan checkpoint
+files, repeat the capability decision, or reconstruct names from its block
+topology. If preflight fails, the loader materializes weights normally and DLO
+consumes those runtime tensors.
+
+The plan owns only dedicated DiT component sources. If a pipeline also exposes
+ordinary sources for a text encoder or another non-DiT component, the loader
+still consumes those sources and includes their loaded names in its strict
+coverage check. Only the source prefixes covered by the plan skip ordinary
+materialization. A source that mixes DiT and non-DiT weights fails closed to
+the complete ordinary-loader path because it cannot be skipped safely as a
+unit.
+
+This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
+flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
+transforms. Model-specific direct-layout knowledge, when required, lives in a
+checkpoint adapter beside the ordinary loader.
+
+### AllGather path
+
+With the default `dlo_use_allgather=True`, each rank stores approximately
+`1 / group_size` of each streamable block in pinned host memory. The next
+block's shard is copied to a device buffer and reconstructed with
+`all_gather_into_tensor` on a communication stream while the current block is
+executing.
+
+```text
+Compute: [Block N] [Block N+1] [Block N+2]
+H2D: [shard N+1] [shard N+2]
+AllGather: [full N+1] [full N+2]
+Buffers: [current slot] [prefetch slot] [current slot]
+```
+
+
+
+The backend uses two shared device buffers, so accelerator weight residency is
+bounded by the largest streamed blocks rather than the complete model.
+
+When direct checkpoint mmap is selected, the checkpoint mappings are only the
+source used to prepare each rank's persistent shard. They can be closed after
+shard preparation. Across the AllGather group, those private shards total
+approximately one runtime model copy.
+
+An effective DLO group size of one performs no collective, even when
+`dlo_use_allgather=True`; it follows the rank-local transfer path described
+below.
+
+When DP is greater than one, the engine can process one request per DP rank in
+the same denoising wave. Because AllGather is a collective, all participating
+requests must take the same execution path at every denoising step.
+
+### Rank-local path without DLO AllGather
+
+With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
+one and streams complete blocks using H2D copies only. The host backing may be
+either a loader-approved checkpoint mapping or ordinary runtime tensors.
+
+For direct mmap, each process retains immutable safetensors views and uses two
+bounded pinned host staging slots. Processes on the same node that map the same
+files share physical checkpoint pages through the OS page cache. This removes
+the persistent private full-model copy per pure-DP process, but each process
+still packs and transfers every complete block. Sharing is node-local; each
+node has its own page cache.
+
+When direct mmap preflight fails, the regular model loader remains responsible
+for preparing each rank's weights, including TP-local tensors or HSDP-managed
+parameters. In that fallback, each pure-DP process keeps a private full runtime
+copy.
+
+This mode means:
+
+- DP still provides independent replicas, but DLO does not shard weights
+ across DP ranks.
+- SP still performs its normal activation/attention collectives, but DLO does
+ not shard weights across SP ranks.
+- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
+ DLO's additional weight AllGather is disabled.
+- Pure DP deployments share one checkpoint-backed copy per node when direct
+ mmap is selected; the ordinary-loader fallback keeps one private runtime
+ copy per rank.
+- The scheduler does not require a synchronized DP request wave for DLO.
+
+### Component allocator-cache retention
+
+MiniMax-H3 DLO keeps one bounded PyTorch allocator cache across its staged
+encoder, DiT, and VAE component boundaries so allocations released by one
+component can be reused by the next. The DiT prefetch path and its two shared
+device buffers are unchanged.
+
+After a component is offloaded, cached-but-unallocated memory is retained only
+while it is at most 25% of device capacity and at least 5% of device capacity
+remains physically free. Crossing either bound releases the allocator cache.
+Missing allocator telemetry also releases it conservatively. Component or
+staging failure forces a release, and an out-of-memory allocation gets one
+retry after release. This is a component-local policy rather than a global
+`empty_cache` override, so the executor's unconditional shutdown cleanup is
+preserved.
+
+Retaining allocator blocks trades higher reserved and peak device memory for
+lower lifecycle latency. The benefit is device- and workload-dependent, and a
+smaller or externally loaded device may cross a bound and return to the
+ordinary per-component cache-release behavior. Current performance validation
+covers MiniMax-H3 no-AllGather DLO on CUDA. Other MiniMax-H3 DLO transfer
+topologies and platforms receive the same bounded policy but are not claimed
+performance targets; other DLO models remain unchanged unless they explicitly
+adopt it.
+
+### Final-layout Host Weight Runtime consumer
+
+> **Status:** PR2 landed in
+> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds
+> registered direct H2D plus generic HWR publication/source-digest
+> optimizations without changing AllGather behavior.
+
+Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies
+only to no-AllGather DLO. Enable it with
+`--host-weight-runtime-mode preferred` (or `required`) and
+`--host-weight-runtime-root `.
+
+The modes express operator fallback policy, not different artifact formats:
+
+- `preferred` consumes an exact hit, but on a miss it allows canonical loading
+ followed by post-load publication. It is the normal population path.
+- `required` consumes the same exact artifact but fails startup on a miss or
+ unusable artifact. It never invokes canonical DiT fallback or post-load
+ publication and therefore cannot populate an empty store.
+
+PR2 has no separate prewarm command. Operators populate one matching producer
+cohort per node-local storage domain in `preferred` mode, then restart the same
+model revision and parallel layout with `required`. TP coordinates have distinct
+identities, while equivalent DP replicas share them.
+
+The current producer/consumer boundary is the model-declared final-layout BF16
+contract (MiniMax H3 and `black-forest-labs/FLUX.2-klein-4B`). The FLUX.2-klein
+contract covers both transformer block stacks, constructor-stable packed QKV
+mapping state, BF16 parameters, and any persistent loader-owned `beta`/`eps`
+buffers. The shared HWR machinery supports ordinary-loader final layouts for
+TP1 and TP2 rank identities plus SP layout identities. FLUX.2-klein evidence
+in this change is TP1-only; its TP2/SP layouts rely on the existing
+per-coordinate identity mechanics and remain unmeasured. Online quantization,
+HSDP, LoRA/adapted weights, and non-default load formats remain ineligible.
+HWR mode
+`disabled`, DLO-disabled, and DLO AllGather configurations stop before source
+identity or store construction and retain the existing checkpoint-mmap or
+ordinary-loader path.
+
+Eligibility is decided before constructing HWR, resolving canonical sources,
+hashing identity inputs, probing a filesystem, or emitting an HWR observer
+event:
+
+```text
+DLO disabled / HWR disabled / DLO AllGather
+ -> zero final-layout HWR interaction
+ -> preserve current checkpoint-mmap or ordinary-loader behavior
+```
+
+When final-layout HWR is explicitly enabled for no-AllGather DLO, runtime
+outcomes are authoritative:
+
+- `LOCAL_HIT`: plan and commit an exact restoration, then transfer the lease
+ transactionally to DLO.
+- `CANONICAL_FALLBACK`: bypass checkpoint mmap, canonically materialize and
+ finalize the DiT, then call `publish_after_load()` for a future startup. The
+ current startup retains its canonical tensors.
+- `FAILED`: fail startup. Preferred mode does not reinterpret nonretryable
+ identity, configuration, or compatibility failures as misses.
+- Required mode cannot bootstrap an empty store through a `POST_LOAD_ONLY`
+ producer.
+
+Checkpoint mmap remains an unchanged control path whenever final-layout HWR is
+not selected.
+
+Local source identity and publication avoid repeated full-payload passes while
+preserving the same semantic contract:
+
+- Canonical files without a trusted immutable Hub blob identity are hashed in
+ parallel. Their content IDs are cached per storage domain under the HWR root
+ and reused only when the path, inode, size, timestamps, symlink target, and
+ cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or
+ cache-I/O failure falls back to hashing the canonical file directly.
+- Producers that emit payload bytes in canonical storage-key order may declare
+ `ordered=True`. The filesystem writer then computes file and tensor SHA256
+ values during the write and overlaps payload `fsync` with later producer
+ work. Producers without that guarantee retain unordered writes and use a
+ parallel readback checksum fallback.
+- Manifest and `READY` publication still wait for every payload `fsync` and
+ checksum to complete. These optimizations do not weaken artifact validation
+ or change lease ownership.
+
+#### Registered mmap transport
+
+After DLO takes an HWR lease, the no-AllGather backend may register the lease's
+complete immutable mapped ranges under the existing `pin_cpu_memory` policy.
+Registration is transport state: it does not change artifact identity, the
+store, tensor ownership, or H2D payload. On success, each tensor view copies
+directly into the existing rotating HBM block buffers and the two private host
+staging slots are not allocated.
+
+`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling
+over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned
+memory policy, unsupported platform/capability, over-budget mapping, or fully
+rolled-back registration error selects the existing two-slot staging path. A
+partial registration that cannot be rolled back aborts startup because closing
+the lease would unmap memory still owned by the platform.
+
+Direct checkpoint mmap remains unchanged and continues to use staging. It may
+require loader-owned per-block transforms, while the HWR artifact already
+contains final runtime bytes. DLO AllGather never receives an HWR final-layout
+lease and therefore never enters this registration path.
+
+#### Pre-service transaction boundary
+
+The startup transaction does not end at restore commit. A lease-backed model is
+disposable until backend setup and initial prefetch complete:
+
+```text
+UNRESOLVED
+ -> LEASE_OWNED_BY_LOADER
+ -> RESTORE_PLANNED
+ -> COMMIT_STARTED
+ -> CARRIER_OWNED
+ -> BACKEND_OWNED
+ -> BACKEND_READY
+ -> IN_SERVICE
+```
+
+Failure before `COMMIT_STARTED` closes the loader-owned lease. Preferred mode
+may canonically load using the untouched skeleton; required mode fails.
+
+Failure from `COMMIT_STARTED` through `BACKEND_READY` must:
+
+1. synchronize or otherwise quiesce partial backend and initial-prefetch work;
+2. release hook, staging, restore-plan, and model references;
+3. close the lease through its current owner;
+4. discard the restored model; and
+5. in preferred mode, construct a fresh canonical model while bypassing HWR
+ lookup, HWR publication, and checkpoint mmap for that recovery attempt.
+
+Required mode fails instead of constructing the fresh fallback. Once the model
+enters service, failures follow normal runtime handling rather than startup
+fallback.
+
+#### Finalization phases
+
+Cold canonical loading separates byte-changing work from shared runtime-state
+finalization:
+
+```text
+Cold-only byte-changing:
+ casting, reordering, packing, quantization, generated scales,
+ model-specific weight transforms, and calibration that mutates
+ parameters or persistent buffers
+
+Shared cold/warm non-byte:
+ validation, hook installation, eval state, bookkeeping,
+ and non-persistent runtime state
+```
+
+The warm path must snapshot restored tensor bytes and backing pointers before
+shared finalization and prove both are unchanged afterward.
+
+#### Lease ownership
+
+Lease transfer is single-owner and single-take:
+
+```text
+loader owns
+ -> restore plan borrows
+ -> commit succeeds
+ -> runner carrier takes once
+ -> backend takes once before asynchronous work
+ -> backend drains pending H2D work
+ -> backend releases hooks, staging, and model references
+ -> backend unregisters every mapped range
+ -> backend closes lease
+```
+
+The runner carrier is process-local, rejects serialization, and rejects a
+duplicate `take()`. If the carrier still owns the lease, runner cleanup closes
+it. Once the backend takes ownership, only backend abort or teardown may drain
+work and close it.
+
+The implementation keeps the final-layout lease through backend setup and
+initial prefetch. Successful registration copies directly from immutable HWR
+views; otherwise the backend uses the existing two bounded rank-local staging
+slots. Transactional cleanup drains device work, removes source references,
+unregisters mappings, and only then closes the lease. An unregistration failure
+retains both registration and lease for retry/process teardown. Preferred mode
+then uses the runner's fresh canonical retry; required mode propagates the
+failure.
+
+#### PR2 and PR3 promotion gates
+
+- Warm hit performs zero ordinary DiT materialization and zero producer calls.
+- Shared warm finalization changes neither restored bytes nor backing pointers.
+- Preferred `FAILED` outcomes do not silently fall back.
+- Planning failure leaves the skeleton untouched.
+- Any pre-service failure after commit begins discards the restored model;
+ preferred mode uses a fresh canonical model and required mode fails.
+- Mixed components load normally and no required tensor remains on `meta`.
+- Disabled mode and AllGather emit zero final-layout HWR interaction.
+- The lease carrier rejects duplicate take and serialization.
+- Backend setup or prefetch failure drains asynchronous work before lease close.
+- Registration is all-or-nothing; rollback/unregistration failure never closes
+ a mapping still owned by the platform.
+- Registered HWR transport bypasses host staging while preserving output and
+ H2D payload; unsupported registration retains the bounded staging path.
+- Compatible checkpoint mmap remains an unchanged benchmark and control path.
+- Prewarm uses one matching producer cohort per TP/SP topology and storage
+ domain.
+- Benchmarks compare ordinary loading, cold publication, and warm staging with
+ output parity, startup latency, aggregate PSS, page-cache state, HBM, and H2D
+ payload.
+
+## Parallelism compatibility
+
+| Parallelism | DLO + AllGather | DLO without AllGather |
+| --- | --- | --- |
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
+| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
+| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
+
+### Combined dimensions
+
+- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
+ than one; SP continues to use its own sequence-parallel group. If DP is one,
+ the SP group becomes DLO's sharding group in AllGather mode.
+- **DP + TP/SP without AllGather:** standard model loading defines the
+ rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
+ collective. When the model declares the final-layout contract, HWR keys the
+ reusable artifact by TP rank/size and SP layout.
+- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
+ DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
+ materialization and synchronization.
+- **HSDP + DP or TP:** rejected independently by the diffusion parallel
+ configuration.
+
+## Request and loading constraints
+
+AllGather DP multi-concurrency requires:
+
+- explicit `num_inference_steps`;
+- the same `num_inference_steps` for all requests in a wave; and
+- identical request arguments that affect the collective execution path.
+
+The no-AllGather path does not impose these DLO-specific synchronized-wave
+requirements.
+
+Direct checkpoint mmap can back either transfer path. It is currently limited
+to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
+ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears
+can use DLO AllGather after the ordinary loader finalizes their runtime
+weights and scales; DLO then shards and reconstructs those tensors with their
+recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
+views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
+fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
+finalized tensor, all covered by the same physical-order packing as online
+FP8). Other online methods must use `--dlo-no-use-allgather` or disable
+online quantization until their runtime layouts are validated.
+
+The Host Weight Runtime representation, publication, and no-AllGather consumer
+contracts are merged; see
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414),
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and
+[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered
+direct H2D is an optional transport layer over that merged lease contract.
+
+## Validation coverage
+
+Current source-level validation includes:
+
+- HSDP + DLO + AllGather rejection;
+- HSDP + DLO without AllGather acceptance at configuration level;
+- loader preflight fallback for TP, HSDP, online quantization, unknown custom
+ loaders, missing keys, and shape/dtype mismatches;
+- ordinary-loader fallback for per-tensor online FP8, online INT8, and online
+ MXFP8 linears followed by DLO sharding of finalized weights and scales;
+- exact loader-to-backend plan transfer and ordinary-loader fallback;
+- ordered publication hashing, overlapped durability, unordered parallel
+ checksum fallback, and node-local source-digest reuse/invalidation;
+- rank-local mmap source retention, bounded two-slot staging, and adapter
+ transforms without parameter-side flags;
+- read-only registration preflight, direct source-to-device copies, safe
+ fallback, partial rollback, retryable unregistration, and lease ordering;
+- resident-layer requests requiring no-AllGather;
+- DP request-wave validation for denoising-step compatibility;
+- bounded component allocator-cache retention, conservative/forced release,
+ OOM retry, and immutable encoder non-block staging;
+- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
+ tests.
+
+### B300 parallel-topology smoke matrix
+
+A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
+CUDNN attention backend, 256x256 output, two denoising steps, and
+`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
+VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
+or shared-mmap host-memory savings.
+
+| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
+| --- | ---: | ---: | ---: | ---: |
+| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
+| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
+| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
+| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
+
+Within each topology, the AllGather and no-AllGather video and audio outputs
+were byte-identical. All four runs completed without an `ERROR` or traceback
+and released their device allocations. For DP4xTP1, no-AllGather direct mmap
+reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
+125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
+ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
+AllGather, because DP replicas did not share checkpoint-backed runtime
+weights. This is a functional and memory smoke test, not a production-quality
+performance or output-quality benchmark.
+
+### Host-memory measurement
+
+A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
+ordinary-loader fallback with direct mmap. Both runs used
+DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
+256x256 four-second request. The ordinary-loader workers were sampled after
+initialization. The mmap workers were sampled after one completed request, so
+the checkpoint working set had been faulted into the page cache; this is the
+more conservative point for mmap.
+
+The values below come from `/proc//smaps_rollup` and include the whole
+worker, not only the DiT. The stable rank-to-rank difference comes from other
+pipeline components, so each worker should be compared with the same worker in
+the other storage mode.
+
+| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
+| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
+| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
+
+The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
+31.20 GiB `Pss_File`, which is the proportional charge expected when the same
+resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
+from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
+1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
+because it counts a shared physical page in every process that maps it; summed
+PSS is the appropriate node-memory comparison.
+
+The highest-value missing coverage is broader end-to-end numerical and
+lifecycle comparison against ordinary layerwise offload for DP+SP,
+HSDP+SP+no-AllGather, and TP greater than one across additional models and
+target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
+change the Phase A direct-mmap TP1 support boundary.
+
+## Recommendations
+
+- Use **DP + DLO AllGather** for the supported throughput and host-memory
+ scaling path.
+- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
+ not the goal.
+- Use **no-AllGather** when independent replica execution is required. TP1
+ direct-mmap deployments can share checkpoint pages per node; other layouts
+ retain the ordinary loader's private host memory behavior and are outside
+ the Phase A shared-mmap support scope.
+- Prefer **HSDP alone** for production HSDP deployments until the combined
+ HSDP + DLO no-AllGather path has broader end-to-end coverage.
diff --git a/docs/design/feature/offloader/layerwise_offload.md b/docs/design/feature/offloader/layerwise_offload.md
new file mode 100644
index 000000000..3fff501da
--- /dev/null
+++ b/docs/design/feature/offloader/layerwise_offload.md
@@ -0,0 +1,57 @@
+# Layerwise Offload
+
+Layerwise offload streams transformer blocks through one accelerator. For
+commands and model support, see the
+[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
+
+## Hook ring
+
+`LayerWiseOffloadBackend` discovers every streamable block and registers one
+`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
+host representation of its registered block and a reference to the next block
+it must prefetch; the final hook prefetches the first block for the next
+denoising iteration.
+
+During initialization, parameters and buffers are grouped by dtype and copied
+into contiguous pinned CPU tensors. Original tensor storage is replaced with
+an empty or meta placeholder while identity and module structure are retained.
+Metadata records names, offsets, shapes, and dtypes for rematerialization.
+
+## Forward lifecycle
+
+For each block:
+
+1. the previous hook's completion event ensures this block is materialized;
+2. the copy stream asynchronously allocates and fills the next block's device
+ tensors;
+3. parameter and buffer storage is rebound to views of those tensors;
+4. the current block computes on the normal stream; and
+5. its post-forward hook replaces device storage with placeholders after all
+ dependent work is safe.
+
+The copy stream waits on the compute stream before reusing storage, and the
+compute stream waits on the per-block prefetch event before forward. These
+dependencies, not incidental global synchronization, define correctness.
+
+## Topology contract
+
+A DiT declares ordered block-container attributes using
+`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
+be an executable `nn.Module`; ordering must match forward execution. Multiple
+containers are concatenated in execution order.
+
+Non-block DiT modules, encoders, VAEs, and declared resident modules are not
+streamed by this backend.
+
+## Invariants and limitations
+
+- Each offloaded tensor has one pinned host backing store.
+- Parameter identity is preserved when storage is rebound.
+- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
+ implementations must not assume identical blocks.
+- The ring is initialized before the first forward so block zero is available.
+- This backend is single-device; distributed weight reconstruction belongs to
+ [Distributed Layerwise Offload](distributed_layerwise_offload.md).
+
+Shared selection and lifecycle behavior is defined in the
+[CPU Offloading design](README.md).
diff --git a/docs/design/feature/omni_connectors/mooncake_store_connector.md b/docs/design/feature/omni_connectors/mooncake_store_connector.md
index 2ff116859..ed8e1ef49 100644
--- a/docs/design/feature/omni_connectors/mooncake_store_connector.md
+++ b/docs/design/feature/omni_connectors/mooncake_store_connector.md
@@ -32,26 +32,25 @@ mooncake_master \
## Configuration
-Define the connector in runtime:
+Define the connector at the top level of the deploy YAML:
```yaml
-runtime:
- connectors:
- connector_of_mooncake:
- name: MooncakeStoreConnector
- extra:
- host: "127.0.0.1"
- metadata_server: "http://:8080/metadata"
- master: ":50051"
- segment: 512000000
- localbuf: 64000000
- proto: "tcp"
+connectors:
+ connector_of_mooncake:
+ name: MooncakeStoreConnector
+ extra:
+ host: "127.0.0.1"
+ metadata_server: "http://:8080/metadata"
+ master: ":50051"
+ segment: 512000000
+ localbuf: 64000000
+ proto: "tcp"
```
Wire stages to the connector:
```yaml
-stage_args:
+stages:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_mooncake
diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md
index 306a0620b..e23ea897b 100644
--- a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md
+++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md
@@ -21,26 +21,25 @@ InfiniBand/RoCE NICs).
## Configuration
-Define the connector in runtime:
+Define the connector at the top level of the deploy YAML:
```yaml
-runtime:
- connectors:
- rdma_connector:
- name: MooncakeTransferEngineConnector
- extra:
- host: "auto" # Auto-detect local RDMA IP
- zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
- protocol: "rdma" # "rdma" or "tcp"
- device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
- memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU
- memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
+connectors:
+ rdma_connector:
+ name: MooncakeTransferEngineConnector
+ extra:
+ host: "auto" # Auto-detect local RDMA IP
+ zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
+ protocol: "rdma" # "rdma" or "tcp"
+ device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
+ memory_pool_size: 4294967296 # 4 GB (CPU); use 2147483648 (2 GB) for GPU
+ memory_pool_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
```
Wire stages to the connector:
```yaml
-stage_args:
+stages:
- stage_id: 0
output_connectors:
to_stage_1: rdma_connector
@@ -415,7 +414,7 @@ A receiver connector:
- allocates receive buffers from its own pool
- requests metadata or transfer service from the sender
-The role is not inferred dynamically. It is injected by the stage configuration layer:
+The role is not inferred dynamically. It is injected from the deploy connector wiring:
- incoming edge for a stage -> `role="receiver"`
- outgoing edge for a stage -> `role="sender"`
diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md
index 62d77cd6d..1c0272340 100644
--- a/docs/design/feature/omni_connectors/shared_memory_connector.md
+++ b/docs/design/feature/omni_connectors/shared_memory_connector.md
@@ -13,10 +13,9 @@ segment name is returned in metadata.
## Configuration
```yaml
-runtime:
- connectors:
- connector_of_shared_memory:
- name: SharedMemoryConnector
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
```
## Notes
diff --git a/docs/design/feature/omni_connectors/yuanrong_connector.md b/docs/design/feature/omni_connectors/yuanrong_connector.md
index 12127ba3e..828feb99f 100644
--- a/docs/design/feature/omni_connectors/yuanrong_connector.md
+++ b/docs/design/feature/omni_connectors/yuanrong_connector.md
@@ -65,23 +65,22 @@ dscli stop --worker_address "${WORKER_IP}:31501"
## Configuration
-Define the connector in runtime:
+Define the connector at the top level of the deploy YAML:
```yaml
-runtime:
- connectors:
- connector_of_yuanrong:
- name: YuanrongConnector
- extra:
- host: "127.0.0.1"
- port: 31501
- get_sub_timeout_ms: 1000
+connectors:
+ connector_of_yuanrong:
+ name: YuanrongConnector
+ extra:
+ host: "127.0.0.1"
+ port: 31501
+ get_sub_timeout_ms: 1000
```
Wire stages to the connector:
```yaml
-stage_args:
+stages:
- stage_id: 0
output_connectors:
to_stage_1: connector_of_yuanrong
diff --git a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
index 21d5f1841..c98a34896 100644
--- a/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
+++ b/docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
@@ -94,26 +94,23 @@ For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
memory pool:
```yaml
-runtime:
- enabled: true
- defaults:
+connectors:
+ yuanrong_te_connector:
+ name: YuanrongTransferEngineConnector
+ extra:
+ host: "auto"
+ zmq_port: 50051
+ rpc_port: "auto"
+ protocol: "ascend"
+ device_name: "auto"
+ memory_pool_size: 1073741824
+ memory_pool_device: "npu"
+
+edges:
+ - from: 0
+ to: 1
window_size: -1
max_inflight: 1
- connectors:
- yuanrong_te_connector:
- name: YuanrongTransferEngineConnector
- extra:
- host: "auto"
- zmq_port: 50051
- rpc_port: "auto"
- protocol: "ascend"
- device_name: "auto"
- memory_pool_size: 1073741824
- memory_pool_device: "npu"
- edges:
- - from: 0
- to: 1
- window_size: -1
```
Important fields:
@@ -141,31 +138,27 @@ For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
both stage `omni_kv_config` blocks:
```yaml
-stage_args:
+stages:
- stage_id: 0
- runtime:
- devices: "0,1,2,3"
- engine_args:
- tensor_parallel_size: 4
- omni_kv_config:
- need_send_cache: true
- rank_mapping:
- from_tp: 4
- to_tp: 4
+ devices: "0,1,2,3"
+ tensor_parallel_size: 4
+ omni_kv_config:
+ need_send_cache: true
+ rank_mapping:
+ from_tp: 4
+ to_tp: 4
output_connectors:
to_stage_1: yuanrong_te_connector
- stage_id: 1
- runtime:
- devices: "4,5,6,7"
- engine_args:
- parallel_config:
- tensor_parallel_size: 4
- omni_kv_config:
- need_recv_cache: true
- rank_mapping:
- from_tp: 4
- to_tp: 4
+ devices: "4,5,6,7"
+ parallel_config:
+ tensor_parallel_size: 4
+ omni_kv_config:
+ need_recv_cache: true
+ rank_mapping:
+ from_tp: 4
+ to_tp: 4
input_connectors:
from_stage_0: yuanrong_te_connector
```
diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md
index ebad8b691..a8fe4fc50 100644
--- a/docs/design/feature/prefix_caching.md
+++ b/docs/design/feature/prefix_caching.md
@@ -41,7 +41,7 @@ With this in mind, consider the set of blocks in a 2D layout, where the row repr
### Example
!!! note "Note 3"
- Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. It can be enabled/disabled per-stage via the `enable_prefix_caching` parameter in the model's stage config.
+ Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config.
The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
diff --git a/docs/design/feature/quantization.md b/docs/design/feature/quantization.md
index 37dd5e26b..84c9641b3 100644
--- a/docs/design/feature/quantization.md
+++ b/docs/design/feature/quantization.md
@@ -76,6 +76,12 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
ModelOpt checkpoint metadata is detected from its method, producer metadata, or
`quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.
+Unless the input is a per-component dictionary, the resolved configuration is
+passed unchanged to every quantization-aware component constructed by the
+pipeline. It affects only layers supported by that component's quantization
+method; it does not rewrite arbitrary `torch.nn` modules. A per-component
+dictionary narrows this global scope through runtime layer-prefix routing.
+
The factory uses Omni-specific overrides for methods whose diffusion behavior
or platform dispatch is not provided by the vLLM registry. The current override
set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -154,9 +160,11 @@ quantizable vLLM layers. Each method configuration is responsible for:
Model code must keep precision-sensitive or unsupported modules unquantized.
The shared `safe_quant_config()` helper is used by model integrations for
normalization and modulation layers where a generic FP8 configuration would be
-unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
-are not implicitly quantized just because the main transformer is quantized;
-they require explicit component routing and method support.
+unsafe. Encoders exposed through the quantization factory receive a global
+configuration just like the main transformer, but only their supported vLLM
+quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
+schedulers, and tokenizers are not rewritten automatically. A component map can
+select a narrower scope, but it cannot add quantization support to a layer.
Weight quantization and quantized KV cache are separate paths. The former is
resolved through `quantization_config`; runtime attention/KV-cache quantization
@@ -182,12 +190,15 @@ configured parallelism strategy:
- `WeightsMapper` transformations must be applied before matching ignored-layer
lists or component prefixes.
-Distributed layerwise offload has an additional boundary: its sharded mmap
-AllGather path is not compatible with online quantization because the path
-expects already materialized checkpoint weights. Use the rank-local
-`--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
-combining these features. See the [DLO design document](distributed_layerwise_offload.md)
-for the full compatibility matrix.
+Distributed layerwise offload has an additional boundary: direct checkpoint
+mmap cannot prepare online-quantized tensors because it bypasses the ordinary
+loader that creates their runtime weights and scales. Per-tensor online FP8
+linears can still use DLO AllGather through the ordinary loader: each rank first
+finalizes the FP8 tensors, then DLO shards those runtime tensors for
+reconstruction during prefetch. Other online methods remain limited to
+`--dlo-no-use-allgather` until their packing and scale layouts are validated.
+See the [DLO design document](offloader/distributed_layerwise_offload.md) for
+the full compatibility matrix.
## Adding a quantization backend
diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md
new file mode 100644
index 000000000..5963befe7
--- /dev/null
+++ b/docs/design/feature/realtime_ar_diffusion.md
@@ -0,0 +1,213 @@
+# Realtime AR-Diffusion sessions
+
+> **Status:** experimental. This document describes the internal execution
+> contract implemented under `vllm_omni.experimental.ar_diffusion`. It does not
+> define a public HTTP or WebSocket schema.
+
+## Scope
+
+Realtime AR-Diffusion serves world models that generate one autoregressive
+latent block per request while retaining model state across requests. The
+generic runtime owns identity, event ordering, session lifecycle, routing, and
+KV allocation. A model adapter owns model-specific controls and conditioning.
+
+LingBot World v2 is the first single-KV-branch integration. DreamZero uses the
+same runner capability with two logical CFG branches.
+
+## Identities
+
+Four identifiers have distinct lifetimes:
+
+| Identifier | Lifetime | Owner | Contract |
+| --- | --- | --- | --- |
+| `session_id` | One persistent world | Session manager | Selects runner KV and model-owned state. |
+| `event_id` | One control/prompt update | Transport adapter | Unique and monotonically increasing within a session, including across reset. |
+| `chunk_index` | One committed AR block position | Session | Contiguous from zero; reset starts again at zero. |
+| `request_id` | One chunk execution | Session | Correlates the tick with its returned metadata. |
+
+`AsyncOmni` keeps its normal UUID-suffixed engine routing ID. The tick
+`request_id` remains inside the immutable AR-Diffusion snapshot and the
+standard output metadata, while caller-visible `OmniRequestOutput.request_id`
+continues to use AsyncOmni's existing external-ID contract. Protocol identity
+therefore does not require a private entrypoint override.
+
+## Request lifecycle
+
+1. `ARDiffusionSession.accept_event()` validates and queues an immutable
+ `ARDiffusionSessionEvent`. Prompt and opaque control tracks may be updated in
+ the same event.
+2. `ARDiffusionSession.next_chunk()` serializes chunk execution and calls
+ `_build_tick_locked()`. The resulting `ARDiffusionTickRequest` is a snapshot
+ of the current prompt, reduced controls, ordered event IDs, and chunk index.
+3. `ARDiffusionOmniTickConsumer.execute_tick()` copies the stage sampling
+ parameters, stores the typed tick in
+ `sampling_params.extra_args["ar_diffusion_tick"]`, constructs the normal Omni
+ prompt through its `prompt_provider`, and invokes `AsyncOmni.generate()`.
+4. `AsyncOmni` allocates its normal unique engine routing ID and submits one
+ `OmniDiffusionRequest` through the orchestrator.
+5. `ARDiffusionModelRunner._request_session()` parses the typed tick.
+ `_get_or_create_session()` obtains runner-owned paged KV, and
+ `bind_ar_diffusion_state()` exposes it to the pipeline only for the duration
+ of `execute_model()`.
+6. The model pipeline validates the tick and generates exactly one latent
+ block. LingBot interprets camera controls in
+ `LingBotCameraControlReducer` and
+ `LingBotWorldCausalDMDPipeline.forward()`; the generic runtime does not parse
+ camera or action schemas.
+7. The pipeline returns one `DiffusionOutput` with the generated payload and an
+ `ARDiffusionChunkMetadata` mapping under `metadata["ar_diffusion"]`.
+ `AsyncOmni` exposes it as
+ `OmniRequestOutput.multimodal_output["metadata"]["ar_diffusion"]`.
+8. The consumer parses the standard envelope. The session commits its reducer
+ state, prompt, controls, event IDs, and next chunk index only when returned
+ metadata exactly equals the submitted tick snapshot.
+
+The metadata envelope is:
+
+```json
+{
+ "metadata": {
+ "ar_diffusion": {
+ "session_id": "world-7",
+ "request_id": "ar-world-7-3-...",
+ "chunk_index": 3,
+ "applied_event_ids": [10, 11]
+ }
+ }
+}
+```
+
+## Transaction and failure semantics
+
+Accepted events are not acknowledged as applied until the model output and
+metadata are validated. Reducer `prepare()` is speculative; reducer `commit()`
+and the session snapshot form one logical commit. A model, metadata, or reducer
+failure leaves the chunk index and queued events uncommitted and transitions the
+session to `FAILED`.
+
+After a tick failure, worker state is closed because a forward may have
+partially modified or already released KV. The same chunk is not retried in
+place. A caller must explicitly reset the session, starting again at chunk
+zero, or close it.
+
+Explicit close follows this state machine:
+
+```text
+ACTIVE or FAILED
+ |
+ v
+ CLOSING ---- worker cleanup succeeds ----> CLOSED
+ |
+ +------- worker cleanup fails --------> CLEANUP_FAILED
+ |
+ +-- close retry --> CLOSING
+```
+
+`CLEANUP_FAILED` is a local tombstone. The session manager retains it, rejects
+creation of the same `session_id`, and rejects reset. Only an idempotent close
+retry that succeeds removes the session. Transport disconnect uses the same
+path.
+
+## KV ownership and capacity
+
+`ARDiffusionKVCacheSpec` is the pipeline-owned immutable geometry contract. It
+declares:
+
+- TP-local self-attention geometry, AR block size, sink, and recent window;
+- logical KV branches and their worker-local slots;
+- fixed-length cross-attention caches;
+- scratch tokens required by an uncommitted forward;
+- persistent model-owned CUDA bytes per resident session; and
+- the requested resident-session capacity.
+
+The runner computes:
+
+```text
+required(capacity) =
+ scratch
+ + managed_self_attention_pages(capacity)
+ + capacity * cross_attention_bytes_per_session
+ + capacity * model_owned_state_bytes_per_session
+```
+
+`gpu_memory_fraction` is a soft expansion budget:
+
+```text
+configured_budget = available_device_bytes * gpu_memory_fraction
+effective_budget = max(configured_budget, required(1))
+```
+
+If `required(1)` exceeds actual available device memory, initialization fails.
+Otherwise the runner always admits one viable session and uses the configured
+fraction to select additional resident sessions up to the model-declared cap.
+All model-owned reservations are deducted before the paged self-attention pool
+is allocated.
+
+LingBot reports its persistent full-horizon image-condition tensor. DreamZero
+reports the measured 603 MiB per-session Wan VAE causal-convolution state.
+
+## Concurrency and routing limits
+
+The current runner executes one request at a time:
+
+- `max_num_seqs` must be one;
+- request batching and diffusion step execution are rejected;
+- one AR block is generated per engine request; and
+- an AR-Diffusion stage must have exactly one replica.
+
+The single-replica restriction is intentional. Multi-replica support requires
+session-affine routing so every tick, reset, and close reaches the worker that
+owns the session. Random or round-robin routing is not correct.
+
+Multiple resident sessions may share one runner. Capacity selection accounts
+for all persistent pools, and the runner performs LRU eviction through the same
+model lifecycle hook used by explicit close.
+
+## Model boundary
+
+The generic runtime transports controls as opaque
+`ARDiffusionControlInput(track, schema, data)` values. Model adapters validate
+and reduce those values. This keeps target/velocity, key state, pose,
+intrinsics, and other model-specific conventions out of the session protocol.
+
+For LingBot World v2:
+
+- `lingbot_world/actions.py` parses key-state and trajectory controls;
+- `lingbot_world/camera.py` loads or constructs camera trajectories and Plücker embeddings;
+- `lingbot_world/transformer.py` implements checkpoint-compatible causal attention; and
+- `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and
+ produces one block plus the standard metadata envelope.
+
+## Non-goals
+
+This contract does not currently provide:
+
+- a public HTTP or WebSocket request schema (the structured-interaction
+ frontend is tracked separately in
+ [#5527](https://github.com/vllm-project/vllm-omni/pull/5527));
+- camera/action semantics shared by every world model;
+- session migration or replication across workers;
+- retry of an ambiguous partially executed chunk;
+- cross-stage KV transfer; or
+- stateful streaming VAE decode.
+
+Those features can be layered on top without adding LingBot-specific identity
+or lifecycle fields to the generic runtime.
+
+## Required regression coverage
+
+CPU contract tests cover:
+
+- separation of tick correlation IDs from internal engine routing IDs;
+- event ordering, metadata equality, and reducer fail-closed commit;
+- close/disconnect failure tombstones and cleanup retry;
+- capacity one under the default fraction for shipped LingBot geometry;
+- capacity expansion under a tuned fraction;
+- DreamZero model-owned state accounting;
+- runner session reuse, reset, close, LRU eviction, and forward cleanup; and
+- LingBot model registration, imports, camera/action reduction, and one-block
+ output metadata.
+
+GPU validation must additionally exercise real model loading, action input,
+prompt switching, rolling-window replacement, reset, close, finite output
+latents, and exact metadata/request identity.
diff --git a/docs/design/feature/skip_softmax.md b/docs/design/feature/skip_softmax.md
index 8e6375ee8..0527c7eeb 100644
--- a/docs/design/feature/skip_softmax.md
+++ b/docs/design/feature/skip_softmax.md
@@ -2,8 +2,10 @@
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
and how to pick an operating point — is in
-[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
-the algorithm.
+[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
+The shared selector contract is documented in
+[Diffusion Attention Backend Selection](attention_backend_selection.md). This
+page explains the algorithm.
## Motivation
diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md
index c465449f8..c6cd19987 100644
--- a/docs/design/feature/vae_parallel.md
+++ b/docs/design/feature/vae_parallel.md
@@ -22,59 +22,60 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and *
**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
-+ Split into multiple spatial tiles
+- Split into multiple spatial tiles
-+ Distributed across multiple ranks
+- Distributed across multiple ranks
-+ Encoded/Decoded in parallel
+- Encoded/Decoded in parallel
-+ Merged to reconstruct the final output
+- Merged to reconstruct the final output
This approach:
-+ Distributes computation across multiple devices
+- Distributes computation across multiple devices
-+ Reduces peak memory usage per device
+- Reduces peak memory usage per device
-+ Accelerates encoding/decoding latency
+- Accelerates encoding/decoding latency
### When to Use Encode vs Decode Parallel
-| Operation | Use Case | Example |
-|-----------|----------|---------|
-| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
-| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
+| Operation | Use Case | Example |
+| ------------------- | ---------------------------- | --------------------------------- |
+| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
+| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
### Architecture
+
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
The executor is model-agnostic and accepts three function parameters:
-+ split – Partition the latent into tiles
+- split – Partition the latent into tiles
-+ exec – Decode a single tile
+- exec – Decode a single tile
-+ merge – Combine decoded tiles into the final output
+- merge – Combine decoded tiles into the final output
#### Execution Flow
-+ Call split(z) to generate a list of TileTask and a GridSpec
+- Call split(z) to generate a list of TileTask and a GridSpec
-+ Dispatch tasks across ranks using workload-based balancing
+- Dispatch tasks across ranks using workload-based balancing
-+ Each rank executes exec(task) on its assigned tiles
+- Each rank executes exec(task) on its assigned tiles
-+ Gather decoded tile results to rank 0
+- Gather decoded tile results to rank 0
-+ Rank 0 performs merge(...)
+- Rank 0 performs merge(...)
-+ (Optional) Broadcast final result to all ranks
+- (Optional) Broadcast final result to all ranks
This design separates:
-+ Distributed execution logic
+- Distributed execution logic
-+ Model-specific tiling and merging logic
+- Model-specific tiling and merging logic
#### Why split / exec / merge is necessary?
@@ -82,23 +83,23 @@ The latent tensor cannot be arbitrarily partitioned.
During decoding:
-+ Each output pixel may depend on neighboring pixels
+- Each output pixel may depend on neighboring pixels
-+ The receptive field is model-dependent
+- The receptive field is model-dependent
Therefore:
-+ Tiles must include overlap
+- Tiles must include overlap
-+ Merge must perform blending to avoid seams
+- Merge must perform blending to avoid seams
## Step-by-Step Implementation (Decode)
### Step 1: Implement DistributedAutoencoderKLQwenImage
+
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
-
-```
+```python
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
@classmethod
def from_pretrained(cls, *args: Any, **kwargs: Any):
@@ -106,28 +107,35 @@ class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMi
model.init_distributed()
return model
```
+
**Key points**:
-+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
-+ Call init_distributed() after loading weights
+
+- Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
+- Call init_distributed() after loading weights
### Step 2: Implement split/exec/merge
+
Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
-```
+
+```python
class GridSpec:
split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols)
tile_spec: dict = field(default_factory=dict) # Metadata required for merging
output_dtype: torch.dtype | None = None # Final output dtype
```
-```
+
+```python
class TileTask:
tile_id: int # task id
grid_coord: tuple[int, ...] # Tile position in grid
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
workload: int | float = 1 # Used for load balancing (e.g., tile area)
```
+
And tiled base split/exec/merge as follow:
-```
+
+```python
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
# mostly copy from AutoencoderKL
_, _, num_frames, height, width = z.shape
@@ -210,11 +218,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
```
### Step 3: Override tiled_decode
+
We need to override tiled_decode, the main logic is:
-+ check distributed is enabled
-+ select split/exec/merge
-+ Invoke self.distributed_executor.execute to decode
-```
+
+- check distributed is enabled
+- select split/exec/merge
+- Invoke self.distributed_executor.execute to decode
+
+```python
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
if not self.is_distributed_enabled():
return super().tiled_decode(z, return_dict=return_dict)
@@ -230,11 +241,14 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
return DecoderOutput(sample=result)
```
+
`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
### Step 4: Modify Pipeline
+
Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
-```
+
+```diff
class YourModelPipeline(nn.Module):
def __init__(
self,
@@ -259,8 +273,8 @@ For models that require VAE encoding (e.g., Image-to-Video), you can also parall
Similar to decode, split the input tensor into tiles. Key considerations:
-+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
-+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
+- **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
+- **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
```python
def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
@@ -406,23 +420,28 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
**Key differences from decode parallel:**
| Aspect | Decode Parallel | Encode Parallel |
-|--------|-----------------|-----------------|
+| -------- | ----------------- | ----------------- |
| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
## Testing
-Verify numerical consistency between:
-+ vae_patch_parallel_size = 1
-+ vae_patch_parallel_size = N
+Verify numerical consistency between:
+
+- vae_patch_parallel_size = 1
+- vae_patch_parallel_size = N
Example:
+
+```python
torch.allclose(output_1, output_n, atol=1e-5)
+```
Testing requirements:
-+ Fix random seed
-+ Use identical tiling strategy
+
+- Fix random seed
+- Use identical tiling strategy
```python
m = Omni(
@@ -434,6 +453,7 @@ m = Omni(
),
)
```
+
When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
## Reference Implementations
@@ -441,7 +461,7 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
Complete examples in the codebase:
| Model | Path | Decode Parallel | Encode Parallel |
-|-------|------|-----------------|-----------------|
+| ------- | ------ | ----------------- | ----------------- |
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
@@ -456,7 +476,7 @@ The tile-parallel executor above assigns independent spatial **tiles** to ranks.
### How it differs from tile parallel
| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) |
-|--------|--------------------------|-----------------------------------------------|
+| -------- | -------------------------- | ----------------------------------------------- |
| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W |
| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) |
| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) |
@@ -481,7 +501,7 @@ serve.py (--vae-parallel-mode) / OmniEngineArgs
### Notes
- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance.
-- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`).
+- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `cards_2`).
---
diff --git a/docs/design/index.md b/docs/design/index.md
index d81ccb806..066d257d1 100644
--- a/docs/design/index.md
+++ b/docs/design/index.md
@@ -18,10 +18,12 @@ implementation contract; it is not, by itself, a general support claim.
### Runtime and stage execution
- [Disaggregated Inference](feature/disaggregated_inference.md)
+- [Host Weight Runtime](feature/host_weight_runtime.md)
- [Async Chunk](feature/async_chunk.md)
- [Async Diffusion Output](feature/async_diffusion_output.md)
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
+- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
### Communication
@@ -53,15 +55,22 @@ implementation contract; it is not, by itself, a general support claim.
#### Attention optimization
The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
-guide lists all selectable backends and their platform defaults. Skip-Softmax
-is the backend-related optimization with a standalone feature-design contract:
+guides list selectable backends, platform defaults, installation, and tuning.
+The design contracts separate selection mechanics from backend algorithms:
+- [Attention Backend Selection](feature/attention_backend_selection.md)
- [Skip-Softmax](feature/skip_softmax.md)
+#### CPU offloading
+
+- [Overview and Shared Contracts](feature/offloader/README.md)
+- [Model-Level Offload](feature/offloader/cpu_offload.md)
+- [Layerwise Offload](feature/offloader/layerwise_offload.md)
+- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
+
- [Cache-DiT](feature/cache_dit.md)
- [TeaCache](feature/teacache.md)
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
-- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
## Infrastructure and Performance
@@ -80,14 +89,15 @@ is the backend-related optimization with a standalone feature-design contract:
- [Model Integration](module/model_integration.md)
- [Autoregressive Runtime](module/ar_runtime.md)
- Diffusion
- - [Overview](module/diffusion/index.md)
- - [Runtime](module/diffusion/diffusion_runtime.md)
- - [Model Integration](module/diffusion/diffusion_model_integration.md)
- - [Continuous Batching](module/diffusion/continuous_batching.md)
- - [Parallelism](module/diffusion/parallelism.md)
- - [Offloader](module/diffusion/offloader.md)
+ - [Overview](module/diffusion/index.md)
+ - [Runtime](module/diffusion/diffusion_runtime.md)
+ - [Model Integration](module/diffusion/diffusion_model_integration.md)
+ - [Continuous Batching](module/diffusion/continuous_batching.md)
+ - [Parallelism](module/diffusion/parallelism.md)
+ - [Offloader](module/diffusion/offloader.md)
- [Execution Platforms](module/execution_platforms.md)
- [Cache Management](module/cache_management.md)
+- [Host Weight Runtime](module/host_weight_runtime.md)
- [Quantization](module/quantization.md)
- [Observability](module/observability.md)
- [Profiling](module/profiling.md)
diff --git a/docs/design/metrics.md b/docs/design/metrics.md
index 7ee35813f..85eb7ac38 100644
--- a/docs/design/metrics.md
+++ b/docs/design/metrics.md
@@ -6,7 +6,7 @@ This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
-- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
+- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss, plus image and diffusion workload, queueing, execution, and memory signals.
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
@@ -20,7 +20,7 @@ vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` col
### The Problem
-vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm_omni:` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
@@ -38,7 +38,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
|
+--------+--------+--------+--------+--------+
| |
- vllm:omni_* vllm:*
+ vllm_omni:* vllm:*
collectors collectors
| |
+----+--------+ +-----------+ +----------+ +-----------+
@@ -58,17 +58,30 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by
### Data Flow
-There are four independent paths for metric collection.
+There are five independent paths for metric collection.
-**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
+**Path 1: Pipeline-level metrics (`vllm_omni:*`)**
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
-- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
-- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
+- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm_omni:requests_success_total{finished_reason}`.
+- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter.
+- `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`.
-**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
+**Path 2: Image and diffusion metrics**
+
+Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice.
+
+- `stage_gen_time_s`, `stage_in_queue_s`, image workload, inference-step, and per-stage peak-memory values are emitted from `StageRequestStats` when a stage finishes. `image_pixels` and `image_count_total` are restricted to image outputs; `num_inference_steps` is restricted to diffusion stages.
+- Diffusion engine timings are accumulated per request by `OrchestratorAggregator.accumulate_diffusion_metrics`. Engine-side millisecond values are converted to seconds before `observe_modality_at_finalize(...)` dispatches the per-replica Histogram observations.
+- `request_queue_wait_s` is emitted once at final request completion from `pipeline_timings["queue_wait_ms"]`. A present value of `0` is a valid observation; a missing key produces no sample.
+- `stage_waiting_requests` is the sum of the latest scheduler snapshots for all live replicas in a stage. AR snapshots arrive through `SchedulerStats`; diffusion snapshots ride normal diffusion output metrics. A dead replica's snapshot is removed. Because this deliberately reuses existing message paths, the gauge can retain its last value while a replica produces no output.
+- `kv_wait_s` measures the AR scheduler interval from entering the KV-transfer-free wait state until extraction acknowledgement. The sample is carried to the Orchestrator through the existing engine output and labeled by connector type. Abort/error/removal paths discard incomplete timers rather than emitting partial waits.
+
+Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active.
+
+**Path 3: Audio modality metrics (`vllm_omni:audio_*`)**
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
@@ -76,7 +89,7 @@ There are four independent paths for metric collection.
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
-**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
+**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)**
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
@@ -87,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
-**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
+**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)**
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
@@ -97,9 +110,9 @@ The Orchestrator runs in a background thread. The API server (OmniBase) runs in
### Metric Registration and Lifecycle
-All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
+All `vllm_omni:*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
-The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
+The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm_omni:*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
## OmniPrometheusStatLogger Wrap
@@ -130,11 +143,11 @@ A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestr
## Gating: `--log-stats`
-All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
+All metrics — both the omni-specific `vllm_omni:*` families and the upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
Behavior with `--log-stats=off` (default):
-- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
+- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm_omni:*` stay materialized but never have data written to them.
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
@@ -162,10 +175,43 @@ second.
| Metric | Type | Labels | Description |
|--------|------|--------|-------------|
-| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
-| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
-| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
-| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+| `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
+| `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
+| `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
+| `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
+
+### Image and diffusion service-level metrics
+
+| Metric | Type | Labels | Description |
+|--------|------|--------|-------------|
+| `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing |
+| `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded |
+| `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage |
+| `vllm_omni:stage_in_queue_s` | Histogram | `model_name`, `stage` | Diffusion scheduler wait from enqueue to initial execution admission |
+| `vllm_omni:num_inference_steps` | Histogram | `model_name` | Denoising-step distribution for diffusion stages |
+| `vllm_omni:image_count_total` | Counter | `model_name` | Cumulative number of generated image outputs |
+| `vllm_omni:image_pixels` | Histogram | `model_name` | Pixel-count distribution for image outputs |
+| `vllm_omni:peak_memory_mb` | Gauge | `model_name`, `stage` | Positive per-stage peak-memory samples reported by engine outputs |
+| `vllm_omni:requests_failed_total` | Counter | `model_name`, `reason` | Request failures with bounded reason values (`client_abort`, `client_disconnect`, `stage_error`, `unknown`) |
+| `vllm_omni:kv_wait_s` | Histogram | `model_name`, `connector_type` | Completed AR KV-transfer wait intervals, grouped by connector backend |
+
+### Diffusion execution breakdown
+
+Labels: `{model_name, stage, replica}`.
+
+| Metric | Type | Description |
+|--------|------|-------------|
+| `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request |
+| `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count |
+| `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time |
+| `vllm_omni:diffusion_postprocess_s` | Histogram | Diffusion output postprocessing time |
+| `vllm_omni:vae_decode_s` | Histogram | VAE decode time when the engine reports it |
+| `vllm_omni:diffusion_forward_s` | Histogram | Forward-only denoise-loop time when pipeline profiling is enabled |
+| `vllm_omni:diffusion_kv_load_s` | Histogram | Diffusion-side KV receive/load time when upstream KV is used |
+| `vllm_omni:image_ttfp_s` | Histogram | Image-stage submit to first materialized image output |
+| `vllm_omni:denoise_step_latency_s` | Histogram | Mean diffusion forward time per denoise step for image output stages |
+
+The optional breakdown and memory families are sparse by design: when an engine does not produce the source measurement, no sample is emitted. This differs from queue timings, where an explicitly measured zero is meaningful and is retained.
### Audio (7)
@@ -173,13 +219,13 @@ Labels: `{model_name, stage, replica}` plus the listed extra label.
| Metric | Type | Extra label | Description |
|--------|------|-------------|-------------|
-| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
-| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
-| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
+| `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
+| `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
+| `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
+| `vllm_omni:audio_frames_total` | Counter | — | Cumulative audio frames generated |
+| `vllm_omni:audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
+| `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
+| `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
### Cross-stage transfer (4)
@@ -187,10 +233,10 @@ Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
| Metric | Type | Description |
|--------|------|-------------|
-| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
-| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
+| `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
+| `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
+| `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
+| `vllm_omni:transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
### LLM stage-level (wrapped `vllm:*`)
@@ -203,7 +249,7 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency,
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
- Sizes use the `_bytes` suffix.
-- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
+- All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`.
## Logging vs. Prometheus
diff --git a/docs/design/module/diffusion/offloader.md b/docs/design/module/diffusion/offloader.md
index 919b1f39a..7ac5d2cc1 100644
--- a/docs/design/module/diffusion/offloader.md
+++ b/docs/design/module/diffusion/offloader.md
@@ -18,7 +18,7 @@ validation_paths:
- tests/diffusion/offloader/**
upstream_refs:
- torch.nn.Module.to
-last_reviewed: 2026-07-16
+last_reviewed: 2026-08-25
---
# Diffusion offloader
@@ -48,6 +48,11 @@ and correctness required by the selected execution mode.
**Rule:** Retained host and device copies MUST respect configured limits and
provide deterministic teardown.
+Allocator-cache retention MUST remain local to an explicit component owner,
+declare both cache and physical-free-memory bounds, release conservatively
+when telemetry is unavailable, and force release on failure or memory
+pressure. It MUST NOT replace unconditional executor shutdown cleanup.
+
## Safe-change guide
Test enabled and disabled paths, repeated execution, asynchronous transfer,
diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md
index 5bb0ea0df..f3f07a43f 100644
--- a/docs/design/module/engine_orchestration.md
+++ b/docs/design/module/engine_orchestration.md
@@ -37,6 +37,7 @@ validation_paths:
- tests/engine/test_async_omni_engine_input.py
- tests/engine/test_async_omni_engine_outputs.py
- tests/engine/test_async_omni_engine_abort.py
+ - tests/e2e/offline_inference/test_qwen3_omni.py
- tests/engine/test_async_omni_engine_stage_init.py
- tests/engine/test_orchestrator.py
- tests/engine/test_orchestrator_error_handling.py
@@ -66,6 +67,12 @@ the boundary affected by the in-flight stage client/process refactor in
[#5441](https://github.com/vllm-project/vllm-omni/pull/5441). Names and
responsibilities proposed only by that PR are not current contracts.
+The orchestration loop also has an opt-in event-driven mode
+(`VLLM_OMNI_EVENT_DRIVEN_ORCH=1`, default off) proposed in
+[#5221](https://github.com/vllm-project/vllm-omni/pull/5221). It changes poll
+cadence only: the routing, ordering, and terminal-state contracts below hold
+identically on both loops.
+
## Ownership boundary
This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
@@ -107,6 +114,15 @@ Numbers become append-only after normative promotion.
Test routing, output ordering, queue correlation, cancellation, failure
propagation, shutdown ordering, and representative multi-stage execution.
+AR abort of a final-stage LLM request must deliver a terminal output with
+`finish_reason="abort"` and the generated prefix so collocated training can
+resume; diffusion abort remains whole-sample retry. Aborting a parallel-sampling
+child must drop `parent_requests` once no children remain. EngineCore control
+RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must
+propagate worker exceptions rather than returning `{"supported": False}`.
+When a request has multiple final output stages, only the last abort message
+is request-terminal. Output-processor abort state is committed only after the
+physical EngineCore abort succeeds.
Changes that cross into stage runtime or public error behavior require review
from that contract's owners.
diff --git a/docs/design/module/entrypoints.md b/docs/design/module/entrypoints.md
index 160f12514..fcdfe23ea 100644
--- a/docs/design/module/entrypoints.md
+++ b/docs/design/module/entrypoints.md
@@ -43,6 +43,7 @@ depends_on:
validation_paths:
- tests/entrypoints/test_omni_entrypoints.py
- tests/entrypoints/test_async_omni.py
+ - tests/entrypoints/test_async_omni_pause_sleep_routing.py
- tests/entrypoints/test_async_omni_duplex.py
- tests/entrypoints/test_serve.py
- tests/entrypoints/test_stream_finish_reason.py
@@ -119,7 +120,13 @@ promotion.
Test request validation, protocol conversion, model-adapter routing,
streaming, session identity, disconnect/cancellation, and error mapping for
-each affected offline or serving entrypoint.
+each affected offline or serving entrypoint. Sleep must wait for in-flight
+`generate()` admission before EngineCore offload; `wake_up` does not resume
+admission — callers must `resume_generation()`. Sleeping tags are tracked per
+stage so `wake_up(stage_ids=[0])` does not skip a later `wake_up(stage_ids=[1])`.
+Streaming input pumps take an admission slot immediately before each EngineCore
+ADD or update, not while waiting for the next client chunk. Frontend abort
+keeps `request_states` until `generate()` consumes the terminal output.
## Promotion gate
diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md
new file mode 100644
index 000000000..ac92711fc
--- /dev/null
+++ b/docs/design/module/host_weight_runtime.md
@@ -0,0 +1,341 @@
+---
+title: Host Weight Runtime
+kind: module
+status: draft
+owners:
+ - "@hsliuustc0106"
+primary_code_paths:
+ - vllm_omni/host_weight_runtime/**
+related_code_paths:
+ - vllm_omni/diffusion/model_loader/**
+ - vllm_omni/diffusion/offloader/**
+depends_on:
+ - model_integration.md
+ - quantization.md
+validation_paths:
+ - tests/host_weight_runtime/**
+upstream_refs:
+ - https://github.com/vllm-project/vllm-omni/issues/6414
+last_reviewed: 2026-08-25
+---
+
+# Host Weight Runtime
+
+The Host Weight Runtime is a loader-adjacent, transport-independent foundation
+for reusing immutable, runtime-ready model weights. It separates the semantic
+question—"are these exactly the weights this loader requested?"—from the
+physical question—"where is the validated host representation stored?"
+
+The runtime can save repeated checkpoint loading, transformation, and peak host
+memory when several workers need the same final representation. It also creates
+a second, versioned path for reconstructing a valid model, so a cache miss or
+recoverable store failure normally falls back to the canonical loader.
+
+This module does not enable host-weight caching for any model by itself. A model
+integration must still provide an exact identity, a producer, and a restorer.
+For feature behavior and consumer integration requirements, see the
+[Host Weight Runtime feature design](../feature/host_weight_runtime.md).
+
+## Component boundary
+
+```mermaid
+flowchart LR
+ LA["Model loader adapter"] -->|"exact identity + policy"| RT["HostWeightRuntime"]
+ RT -->|"lookup / get_or_build"| ST["HostWeightStore"]
+ PR["WeightProducer"] -->|"runtime-ready tensors"| ST
+ ST -->|"validated mmap views"| LE["HostWeightLease"]
+ LE --> RE["WeightRestorer"]
+ RE --> MO["Restored host model"]
+ MO --> TR["DLO or another GPU transport"]
+
+ CS["Canonical model source"] --> PR
+ LA -. "fallback" .-> CS
+```
+
+The contracts have deliberately narrow ownership:
+
+- `HostWeightRuntime` applies resolution and fallback policy and emits one
+ terminal resolution report. Explicit post-load publication emits a separate
+ report and closes its store-returned lease before returning; it does not
+ revise the completed resolution.
+ The runtime does not own a model or call the canonical loader.
+- `HostWeightStore` owns artifact lookup, coordinated construction,
+ publication, validation, quarantine, and lifecycle.
+- `HostWeightLease` owns process-local tensor views, mapped-file resources, and
+ the shared artifact lock that stabilizes their lifetime.
+- `WeightProducer` creates one declared final-layout representation. It writes
+ only through a store-scoped writer and cannot publish paths directly.
+- `WeightRestorer` validates a lease-to-model plan without mutation. Its
+ one-shot `commit() -> None` is the sole model-mutating phase. Concrete
+ restorers remain with the model or quantization integration.
+- GPU transport owns page registration, private staging, H2D scheduling, and
+ lease release. DLO must not parse manifests or manage store artifacts.
+
+## Exact semantic identity
+
+Storage lookup is exact; the store never coerces one representation into
+another. `WeightArtifactIdentity` includes:
+
+- immutable model revision and source-file fingerprint;
+- component and ownership boundary;
+- runtime representation, dtype, and representation metadata;
+- final tensor layout and semantic parallel coordinates;
+- static adaptation identity; and
+- producer version, implementation fingerprint, manifest schema, and restorer
+ schema.
+
+Loader integrations may persist content digests for canonical source shards in
+the selected storage domain. Reuse requires the same path, inode, size,
+timestamps, symlink target, and cache-record checksum. Corrupt records are
+rebuilt, and coordination or cache-I/O failure falls back to hashing the
+canonical source directly. This is an identity-computation optimization, not a
+second source of truth or an artifact substitution rule.
+
+The identity excludes transfer policy such as registered mmap versus private
+pinned staging. Those paths move the same host representation and must not
+produce duplicate artifacts.
+
+Parallel coordinates follow semantics rather than process topology:
+
+- TP size and rank are included when they change tensor contents or shape.
+- DP rank is excluded for replicated weights.
+- SP size and backend are included when they affect layout. SP rank is excluded
+ only when all SP ranks consume an identical representation.
+- PP, EP, or future parallel modes must be added by the model adapter whenever
+ they change component ownership or tensor layout.
+
+Dynamic LoRA overlays are not cached as base weights. A statically merged
+adapter is a distinct artifact identity with its own fingerprint.
+
+## Deterministic resolution
+
+The loader first resolves the canonical source and computes the exact requested
+identity. The runtime then follows this order:
+
+```text
+validated local artifact
+ -> exact remote artifact materialized locally (future)
+ -> registered local producer
+ -> canonical loader fallback, if policy permits
+```
+
+An already published local artifact wins because it avoids network transfer and
+is already in node-accessible storage. "Wins" means the highest-priority valid
+representation, not whichever backing responds first.
+
+Configuration modes are:
+
+- `disabled`: return the canonical-direct outcome without probing a root,
+ credentials, topology, or identity;
+- `preferred`: recoverable store failures may return canonical fallback; and
+- `required`: fail when no exact lease can be acquired.
+
+Failure-to-action policy is centralized in `HostWeightRuntime`. A clean miss or
+invalid cache artifact may fall back in preferred mode. A typed failure may do
+so only when it is marked retryable; unsupported capabilities, identity
+collisions, producer failures, and publication failures remain visible as
+startup failures.
+
+`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it
+is not a hard wall-clock deadline for in-process validation, production, or
+atomic publication. Once this V1 implementation becomes the producer, the
+synchronous producer and publication run to completion. A hung producer blocks
+its owning process and requires external process supervision; enforceable
+producer cancellation requires a future process-isolated producer contract. A
+waiter timing out never cancels another process's valid build.
+
+Remote providers are represented by protocols but are explicitly unsupported
+in the first implementation. Pre-load `resolve()` invokes only
+`PRE_LOAD_SAFE` producers. A loader that completed canonical materialization
+may explicitly call `publish_after_load()` with a `POST_LOAD_ONLY` producer
+when `allow_post_load_publish` is enabled.
+
+`allow_local_build` controls producer use during pre-load resolution;
+`allow_post_load_publish` independently controls the explicit post-load entry
+point. This permits lookup-only warm starts that populate a missing artifact
+only after the canonical loader has established the current model.
+
+Post-load publication is synchronous in V1, but its outcome is independent of
+the already completed resolution. A store failure is returned in a separate
+publication report and cannot turn a valid canonical model into a failed
+resolution. On success, the runtime closes the validated store lease before
+returning the report. Post-load publication only warms future startups; it does
+not restore, rebind, or otherwise mutate the canonically loaded model serving
+the current startup.
+
+## Restoration transaction boundary
+
+`WeightRestorer.plan_restore()` is validation-only: it must not modify the model
+or lease. The returned plan may be committed exactly once, `commit()` returns
+`None`, and commit is the sole phase allowed to mutate the model. Registration,
+staging, and GPU transfer remain transport responsibilities rather than restore
+side effects.
+
+If planning fails, preferred-mode canonical fallback may reuse the untouched
+model. If commit begins and then fails, the model is considered partially
+hydrated and must be discarded; canonical fallback must construct a fresh model
+instance instead of reusing it.
+
+## Filesystem domain and artifact lifecycle
+
+A storage domain describes physical locality and policy; it is not part of
+semantic identity. The first implementation supports node- or NUMA-scoped
+roots on verified local disk filesystems or tmpfs. The detected kernel
+filesystem type is preserved in `domain.json` and inspection output. Known
+remote filesystems and unknown filesystem types are rejected rather than being
+silently treated as local disk, because cross-node page cache and advisory-lock
+semantics do not satisfy this backend's contract. A NUMA root alone does not
+guarantee page placement; CPU pinning, local first-touch/prefault, and residency
+evidence belong to a later topology-aware integration.
+
+```text
+/
+├── domain.json
+├── domain-policy.json
+├── locks/
+│ ├── .build.lock
+│ └── .artifact.lock
+├── artifacts//
+│ ├── READY.json
+│ ├── manifest.json
+│ └── *.safetensors
+├── tmp/
+├── quarantine/
+└── deny/
+```
+
+Lock files live outside replaceable artifact directories. Every cooperative
+mutation follows one lock order: build lock, then exclusive artifact lock.
+Lookup acquires a shared artifact lock, validates `READY` and the manifest,
+opens payloads with non-following descriptor-based access, creates mmap-backed
+views, and transfers that shared lock to the lease.
+
+Lease acquisition rechecks the producer and restorer schemas against the exact
+identity. Malformed safetensors payloads are reported as typed invalid
+artifacts and denied rather than escaping the resolution policy as raw parser
+exceptions.
+
+Each `FileLock` records its creator PID. A forked child closes an inherited lock
+descriptor without issuing `LOCK_UN`, because `flock` state belongs to the
+shared open-file description and unlocking in the child would also release the
+parent's active lease lock.
+
+Publication follows this lifecycle:
+
+```text
+ABSENT -> invisible temporary build -> fully hashed and fsynced
+ -> atomic rename -> READY
+
+READY -> denied -> quarantined -> explicit removal
+```
+
+`READY.json` is written last in the temporary directory. Readers therefore see
+either a complete validated artifact or a miss, never an in-progress build.
+Failed process-local builds are left invisible and a later build owner removes
+their stale temporary directories. Explicit cleanup is nonblocking and reports
+an active builder or lease instead of unlinking live mappings.
+
+## Validation and capacity
+
+Every publication computes full file and tensor hashes. Local lookup supports:
+
+- `manifest_and_metadata`, which validates identity, immutable publication
+ metadata, file descriptors, sizes, tensor keys, shapes, dtypes, and strides;
+ and
+- `full_checksum`, which additionally reads and hashes every payload file.
+
+An ordered producer may compute file and tensor digests while payload bytes are
+written and overlap each closed payload's `fsync` with later producer work. An
+unordered producer retains parallel readback hashing. In both cases manifest
+and `READY` publication wait for every payload digest and `fsync`; the atomic
+publication and validation contract is unchanged.
+
+Filesystem verity is modeled but unsupported. Invalid artifacts receive an
+external deny marker before their shared lookup lock is released; a build owner
+then quarantines the immutable directory under the exclusive lock.
+Lookup checks the marker both before validation and after constructing the
+lease so a weaker overlapping lookup cannot return a lease after a stronger
+validation has denied the artifact.
+
+If publishing the external deny marker fails, lookup returns a retryable
+storage failure and marks the artifact directory with a persistent invalidation
+mode before releasing the shared lock. Cooperating readers check that mode both
+while opening and immediately before returning a lease. Lookup then makes a
+bounded attempt to take the exclusive artifact lock and quarantine the entry;
+the quarantine transition removes the internal mode marker. If a cooperating
+builder replaces the marked inode while the locks are exchanged, containment
+preserves that new publication instead of quarantining it. If neither the
+external marker nor the directory-mode marker can be persisted, cross-process
+exclusion cannot be guaranteed on an unavailable filesystem and the domain
+must remain failed rather than treating the artifact as an ordinary invalid
+cache entry.
+
+New publications carry the internal invalidation mode through the staging
+rename. Hardening clears it and is the logical authority transition for a
+process crash; the temporary- and artifact-parent syncs then make the rename
+durable, with an artifact-inode sync first making the hardened mode durable. A
+process that dies before hardening leaves a persistently denied
+entry, while one that dies after successful hardening leaves an authoritative
+ready entry. Any synchronous failure after rename is marked invalid and
+contained in quarantine before `PUBLICATION_FAILED` is returned, so that
+nonretryable result cannot later resolve as a ready hit. If the filesystem
+refuses both the inode marker and the quarantine move, the store instead returns
+a retryable storage failure with `outcome_uncertain` details; as with deny-marker
+double failure, cross-process exclusion cannot be guaranteed until storage
+recovers. Once the artifact-parent sync succeeds, a later artifact- or
+build-lock release error is logged and lookup resolves the committed entry
+instead of reporting a contradictory publication failure. The same authority
+rule applies when a competing publisher already installed byte-identical
+semantic content; staging-cleanup errors are retried without downgrading that
+validated ready entry to publication failure.
+
+Before widening a hardened artifact for a lifecycle move, the store durably
+sets the same invalidation mode. A process crash before rename therefore leaves
+the writable source denied. After rename, the store re-hardens the destination
+but retains the marker until both parent-directory syncs make the move durable;
+only then does it durably clear the marker. An interrupted move therefore leaves
+a marked transition that the next exact-key lifecycle operation re-hardens.
+Explicit cleanup uses the locked move and retries exact-key
+`.cleanup.*` tombstones left by an interrupted removal; a later build performs
+the same tombstone reconciliation before producing replacement content.
+
+An operational failure while inspecting a noncooperative competing publication
+is a retryable storage failure. The competing entry remains authoritative and
+is not mislabeled as corrupt or quarantined merely because its manifest could
+not be read transiently.
+
+`inspect_domain()` and `inspect_artifact()` return typed, read-only snapshots.
+Inspection validation does not create a deny marker. Reported lock activity is
+point-in-time evidence only; kernel locks are not a persistent owner registry.
+
+Capacity policy covers all store-owned bytes, including ready, temporary, and
+quarantined data. The local writer preflights `max_artifact_bytes`,
+`max_store_bytes`, and `min_free_bytes`, and preallocates payloads where the
+filesystem supports it. `ENOSPC` remains a normal typed store failure because
+concurrent preflight is inherently racy. Automatic eviction and strict
+concurrent reservations are not implemented.
+
+For tmpfs, artifact bytes are host-memory consumption and may consume swap;
+they must not be reported operationally as ordinary disk capacity. The store
+owns files and reports capacity. The kernel controls page-cache residency, and
+transport controls registration or page locking.
+
+## First implementation scope
+
+The foundation implements a CPU local filesystem store and
+`SINGLE_PROCESS` production for one exact artifact. It intentionally adds no:
+
+- public CLI or DLO configuration;
+- BF16, FP8, or model-specific producer/restorer;
+- change to DLO AllGather or no-AllGather behavior;
+- CUDA registration, pinned staging, or H2D scheduling;
+- remote provider, group-cohort producer, artifact bundle, eviction, or
+ cross-node coordination.
+
+Consumer PRs must prove ordinary-loader parity, warm-hit work avoidance,
+parallel-layout correctness, clean teardown, shared-backing evidence, and
+representation-specific startup and memory behavior. Generic foundation tests
+remain CPU-only and make no production performance claim.
+
+The complete rationale and staged rollout are tracked in
+[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414).
diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md
index 1d8d7ab73..2e848a49c 100644
--- a/docs/design/module/model_integration.md
+++ b/docs/design/module/model_integration.md
@@ -42,7 +42,7 @@ and execution behavior to stable runtime contracts.
### MODEL-INV-001: Registration is explicit
**Rule:** A model integration MUST declare how its model class, loader, input
-processor, and stage configuration are selected.
+processor, and pipeline configuration are selected.
### MODEL-INV-002: Model code does not route stages
diff --git a/docs/design/module/vllm_omni_config.md b/docs/design/module/vllm_omni_config.md
index 085e4cee6..1ee303b8c 100644
--- a/docs/design/module/vllm_omni_config.md
+++ b/docs/design/module/vllm_omni_config.md
@@ -10,16 +10,15 @@ owners:
primary_code_paths:
- vllm_omni/config/**
- vllm_omni/deploy/**
- - vllm_omni/model_executor/stage_configs/**
-related_code_paths:
- - vllm_omni/platforms/*/stage_configs/**
+ - vllm_omni/model_executor/models/*/pipeline.py
+related_code_paths: []
depends_on: []
validation_paths:
- tests/config/**
upstream_refs:
- vllm.config
-last_reviewed: 2026-08-07
-last_verified_commit: 3d7fc3b9ba3cac88d579d4dc35b78b0b641675fc
+last_reviewed: 2026-08-11
+last_verified_commit: e356708da4cd39992405686dacfc32b18fcccc7f
---
# vLLM-Omni configuration
@@ -39,11 +38,66 @@ The PR review proposal to capture environment-derived values in configuration
objects at initialization is an open design question, not a current
invariant.
+## Current structured-stage reuse behavior
+
+The current RFC #4021 implementation uses upstream vLLM config classes as the
+schema source for AR and generation stages:
+
+- `OmniStageLoadConfig` inherits `vllm.config.LoadConfig`.
+- `OmniStageCacheConfig` inherits `vllm.config.CacheConfig`.
+- `OmniStageSchedulerConfig` inherits `vllm.config.SchedulerConfig`.
+- `OmniStageParallelConfig` inherits `vllm.config.ParallelConfig`.
+
+This inheritance preserves the existing load, cache, scheduler, and parallel
+concern boundaries while reducing duplicate downstream declarations. It does
+not, by itself, make every inherited field an effective vLLM-Omni runtime
+option. Construction and runtime consumption are separate surfaces:
+
+| Surface | Current behavior |
+| --- | --- |
+| Structured schema | Upstream dataclass fields are inherited. Their defaults, default factories, and applicable Pydantic validation can participate when the structured object is constructed. |
+| Omni input ownership | Pipeline/deploy/CLI construction continues to accept only fields with an existing structured owner. Direct construction of an inherited sub-config can expose a wider upstream schema. |
+| Engine projection | For AR and generation stages, reusable fields are discovered by intersecting each upstream config dataclass with upstream `EngineArgs`. Only constructor-explicit values are emitted; inherited defaults remain deferred to terminal materialization. |
+| Terminal materialization | The engine-owning process constructs the final upstream `VllmConfig` and performs model-, platform-, rank-, port-, and backend-dependent initialization. |
+
+The generated projection maps include known upstream naming differences such
+as `cache_dtype` to `kv_cache_dtype`, `policy` to `scheduling_policy`, and
+`data_parallel_master_ip` to `data_parallel_address`. A field added to both an
+upstream concern config and `EngineArgs` therefore enters the AR/generation
+projection without requiring a second downstream allowlist. Explicit inherited
+fields that have no projection raise an error rather than being silently
+dropped.
+
+Ownership exclusions remain deliberate. Stage topology owns `scheduler_cls`,
+cache owns `disable_hybrid_kv_cache_manager`, and Omni runtime owns
+`distributed_executor_backend` and `worker_cls`; vLLM's private API-process
+fields are terminal internals. These exclusions prevent one input from being
+constructed or projected through two config concerns. Effective-engine-argument
+tests cover both the dynamic mapping and this exclusion boundary.
+
+`CompilationConfig` and `ProfilerConfig` are direct structured-stage fields:
+`stage.compilation_config` and `stage.profiler_config`. They accept mapping/YAML
+inputs and complete upstream value objects. Mapping inputs are validated and
+materialized as their concrete upstream types when the structured config is
+constructed, so downstream consumers see one resolved representation.
+Prebuilt upstream objects retain their type across the structured projection
+boundary.
+
+Quantization remains on its existing Omni-owned transport contract. This
+reuse step does not adopt upstream `QuantizationConfigArgs`, because the
+current engine-specific quantization split needs to be resolved separately.
+
+Diffusion keeps its existing engine-field ownership and projection surface.
+Shared Python inheritance must not be treated as evidence that an LLM-only
+parallel input is effective for a diffusion stage.
+
## Safe-change guide while deferred
Changes should trace every affected producer to its runtime consumer and test
-the structured, legacy, CLI, and deployment paths that are actually supported.
-Do not infer a stable contract from this placeholder page.
+the structured pipeline, CLI, and deployment paths that are actually supported.
+For inherited vLLM fields, tests should separately cover constructor behavior
+and effective `EngineArgs` projection. Do not infer a stable contract from
+this placeholder page.
## Promotion gate
diff --git a/docs/design/qwen3_omni_tts_performance_optimization.md b/docs/design/qwen3_omni_tts_performance_optimization.md
index bc1e29a7f..ff783a02a 100644
--- a/docs/design/qwen3_omni_tts_performance_optimization.md
+++ b/docs/design/qwen3_omni_tts_performance_optimization.md
@@ -111,7 +111,7 @@ This post walks through each optimization in the same order they are typically e
For both Qwen3-Omni and Qwen3-TTS, batching is a pipeline-level optimization:
-- Requests are grouped per stage using `runtime.max_batch_size`
+- Requests are grouped per stage using `stages[].max_num_seqs`
- Each stage executes batch inference with its own scheduler/worker
- Stage outputs are routed to downstream stages with per-request mapping preserved
@@ -171,7 +171,7 @@ At concurrency 10, batching alone brings Qwen3-TTS RTF from 2.19 (slower than re
In decode-heavy serving, repeatedly launching many small kernels from CPU can become a visible overhead. CUDA Graph reduces this overhead by capturing and replaying stable execution graphs.
-In stage configs, this is represented by `enforce_eager: false` for stages where graph capture is desired (Thinker/Talker), while Code2Wav keeps eager mode depending on stage behavior.
+In deploy configs, this is represented by `enforce_eager: false` for stages where graph capture is desired (Thinker/Talker), while Code2Wav keeps eager mode depending on stage behavior.
### CUDA Graph results on top of batching
@@ -320,7 +320,7 @@ In the async chunk pipeline, the standard `codec_chunk_frames` is 25 (each chunk
**Dynamic initial chunk sizing (default behavior):**
-Rather than using a fixed initial chunk size, vLLM-Omni dynamically selects it based on current server load. The initial chunk size is chosen from power-of-2 steps [2, 4, 8, 16] based on load factor (`active_requests / max_batch_size`):
+Rather than using a fixed initial chunk size, vLLM-Omni dynamically selects it based on current server load. The initial chunk size is chosen from power-of-2 steps [2, 4, 8, 16] based on load factor (`active_requests / max_num_seqs`):
| Server load | Initial chunk frames | Rationale |
| --- | --- | --- |
@@ -349,16 +349,15 @@ This overrides the dynamic calculation for that request.
**Config (server-side):**
```yaml
-runtime:
- connectors:
- connector_of_shared_memory:
- name: SharedMemoryConnector
- extra:
- codec_streaming: true
- codec_chunk_frames: 25 # standard chunk size (~2s of audio)
- codec_left_context_frames: 25
- # initial chunk is computed dynamically by default
- # set initial_codec_chunk_frames: 2 to force a fixed value
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ codec_streaming: true
+ codec_chunk_frames: 25 # standard chunk size (~2s of audio)
+ codec_left_context_frames: 25
+ # initial chunk is computed dynamically by default
+ # set initial_codec_chunk_frames: 2 to force a fixed value
```
The 64 ms TTFP result reported above for Qwen3-TTS at concurrency 1 uses the dynamic initial chunk, which picks `initial_codec_chunk_frames=2` at low load. At higher concurrency the dynamic sizing increases the initial chunk to maintain decode efficiency.
@@ -408,7 +407,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
Notes:
-- `runtime.max_batch_size` controls stage-level batching.
+- `stages[].max_num_seqs` controls stage-level batching.
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
@@ -426,33 +425,26 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
```yaml
async_chunk: true
-stage_args:
+stages:
- stage_id: 0 # thinker
- runtime:
- max_batch_size: 64
- engine_args:
- enforce_eager: false
- max_num_batched_tokens: 32768
- custom_process_next_stage_input_func: >-
- vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+ max_num_seqs: 64
+ enforce_eager: false
+ max_num_batched_tokens: 32768
- stage_id: 1 # talker
- runtime:
- max_batch_size: 64
- engine_args:
- enforce_eager: false
- max_num_batched_tokens: 32768
- custom_process_next_stage_input_func: >-
- vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+ max_num_seqs: 64
+ enforce_eager: false
+ max_num_batched_tokens: 32768
- stage_id: 2 # code2wav
- runtime:
- max_batch_size: 64
- engine_args:
- enforce_eager: true
- max_num_batched_tokens: 51200
+ max_num_seqs: 64
+ enforce_eager: true
+ max_num_batched_tokens: 51200
```
+The async handoff processors are fixed by the registered Qwen3-Omni
+`PipelineConfig`; deploy YAML only selects `async_chunk` and runtime sizing.
+
#### Reproduce Qwen3-Omni benchmarks
```bash
@@ -483,7 +475,7 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
The default config (`qwen3_tts.yaml`) enables the full optimization stack:
-- Batching with `max_batch_size: 10` on the Talker stage
+- Batching with `max_num_seqs: 10` on the Talker stage
- CUDA Graph on the Talker (`enforce_eager: false`)
- Async chunk with streaming transport
@@ -500,33 +492,29 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
```yaml
async_chunk: true
-stage_args:
+stages:
- stage_id: 0 # Talker (AR decoder)
- runtime:
- max_batch_size: 10
- engine_args:
- enforce_eager: false
- max_num_batched_tokens: 512
- custom_process_next_stage_input_func: >-
- vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+ max_num_seqs: 10
+ enforce_eager: false
+ max_num_batched_tokens: 512
- stage_id: 1 # Code2Wav (vocoder)
- runtime:
- max_batch_size: 1
- engine_args:
- enforce_eager: true
- max_num_batched_tokens: 8192
+ max_num_seqs: 1
+ enforce_eager: true
+ max_num_batched_tokens: 8192
-runtime:
- connectors:
- connector_of_shared_memory:
- name: SharedMemoryConnector
- extra:
- codec_streaming: true
- codec_chunk_frames: 25
- codec_left_context_frames: 25
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ codec_streaming: true
+ codec_chunk_frames: 25
+ codec_left_context_frames: 25
```
+The async Talker-to-Code2Wav processor is defined by the registered Qwen3-TTS
+pipeline rather than the deploy YAML.
+
#### Reproduce Qwen3-TTS benchmarks
```bash
diff --git a/docs/features/README.md b/docs/features/README.md
index 0dee558a4..7c3e89676 100644
--- a/docs/features/README.md
+++ b/docs/features/README.md
@@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow.
| Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
| Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |
+The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
+describes the experimental Qwen3-Omni runtime topology and its current
+integration boundaries. It is not yet a supported launch recipe.
+
Some runtime designs are deliberately not promoted as standalone User Guide
features yet:
-- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
+- Generic, OmniConnector-based
+ [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
+ yet a standalone User Guide feature. Its
[OmniConnector implementations](../design/index.md#communication) describe
- topology and transport contracts. Practical deployment configuration remains
- under [Pipeline and deploy configurations](../configuration/stage_configs.md).
+ transport contracts, while practical connector configuration remains under
+ [Pipeline and deploy configurations](../configuration/stage_configs.md).
- [Async Chunk](../design/feature/async_chunk.md) and
[Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
are model- and pipeline-dependent. Use the selected model's deploy
@@ -48,10 +54,10 @@ defines the shared configuration and backend extension points.
| Goal | User guide | Related design contract |
| --- | --- | --- |
| Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
-| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
+| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.md) |
| Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
| Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
-| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
+| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
| Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
| Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
| Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md
index 78d680f26..887dfa769 100644
--- a/docs/features/custom_pipeline.md
+++ b/docs/features/custom_pipeline.md
@@ -101,7 +101,7 @@ outputs = omni.generate(
)
# Access custom trajectory data
-output = outputs[0].request_output
+output = outputs[0]
print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
print(f"Trajectory latents shape: {output.latents.shape}")
```
diff --git a/docs/features/pd_disaggregation.md b/docs/features/pd_disaggregation.md
new file mode 100644
index 000000000..ae8497e81
--- /dev/null
+++ b/docs/features/pd_disaggregation.md
@@ -0,0 +1,140 @@
+# Prefill-Decode (PD) Disaggregation (experimental)
+
+!!! warning "Experimental design reference"
+ PD disaggregation is only partially integrated on the current `main` branch.
+ The runtime contains PD detection and routing scaffolding, but there is no
+ supported deploy overlay or end-to-end validated launch recipe yet. Treat this
+ page as a description of the current design, not as production guidance.
+
+Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
+stages: a prefill worker that processes the prompt and produces KV cache, and a
+decode worker that imports that cache and generates tokens. The Talker and
+Code2Wav stages remain downstream; PD does not disaggregate the entire
+Qwen3-Omni pipeline.
+
+## Current design
+
+### Runtime topology
+
+```mermaid
+flowchart LR
+ client["Client request"] --> orchestrator["Omni Orchestrator"]
+ orchestrator -->|"original prompt max_tokens = 1"| prefill["Thinker prefill KV producer"]
+ prefill -->|"KV blocks"| connector[("vLLM KV connector")]
+ prefill -->|"KV routing metadata"| orchestrator
+ orchestrator -->|"original prompt remote KV parameters"| decode["Thinker decode KV consumer"]
+ connector -->|"remote KV load"| decode
+ decode -->|"generated text"| client
+```
+
+The orchestrator owns the request lifecycle, while the vLLM KV connector owns
+the bulk KV transfer. The original prompt is submitted to both Thinker stages:
+the decode worker still needs the request tokens and multimodal features to
+construct its request, while the imported KV cache prevents it from repeating
+the completed prefill computation.
+
+For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
+
+| Stage | Role | Output or handoff |
+| --- | --- | --- |
+| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
+| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
+| 2 | Talker | Converts combined Thinker conditioning into codec codes |
+| 3 | Code2Wav | Converts codec codes into audio |
+
+Text remains a final output of the decode stage, and audio remains a final
+output of Code2Wav.
+
+### Request lifecycle
+
+1. At startup, PD detection looks for one stage marked `is_prefill_only` and
+ one downstream stage marked `is_decode_only`. The decode stage must list the
+ prefill stage as an input source.
+2. The orchestrator clones the Thinker sampling parameters for prefill, forces
+ `max_tokens=1`, clears stop conditions, and sets producer-side
+ `kv_transfer_params`. This makes prefill finish by length after exporting KV.
+3. When prefill finishes, the orchestrator captures its KV routing metadata and
+ Qwen multimodal output under the original request ID.
+4. The orchestrator submits the original prompt to decode with consumer-side
+ `kv_transfer_params`. Those parameters identify the transfer, prefill engine,
+ bootstrap endpoint, and, when required by the connector, remote request.
+5. The decode engine imports the KV cache and continues generation. The normal
+ Omni stage path then carries Thinker conditioning through Talker and
+ Code2Wav.
+
+### Control and data paths
+
+PD uses two different connector layers. They solve different handoffs and are
+not interchangeable.
+
+| Path | Payload | Owner |
+| --- | --- | --- |
+| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
+| PD data path | Thinker attention KV blocks | vLLM KV connector |
+| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
+
+### Detection and validation
+
+The current runtime accepts one logical prefill/decode pair per pipeline. Each
+stage may use tensor parallelism, but both sides of the pair must satisfy these
+invariants:
+
+- both stages define `kv_transfer_config`
+- prefill uses `kv_producer` or `kv_both`
+- decode uses `kv_consumer` or `kv_both`
+- both stages use the same KV connector
+- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
+- both stages use the same `tensor_parallel_size`
+
+These checks happen before request processing so an inconsistent PD topology
+fails during engine initialization.
+
+## What is implemented today
+
+| Area | Current state |
+| --- | --- |
+| Pair detection and topology validation | Implemented in the shared PD helper |
+| Prefill sampling preparation | Implemented for sync and async entrypoints |
+| Prefill-to-decode routing | Implemented in the current Orchestrator |
+| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
+| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
+| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
+| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
+| Performance | No non-regression or throughput/latency result is established for the current path |
+
+Because these pieces are still converging, the previous legacy `stage_args` YAML
+and `vllm serve` command have been removed. The serve CLI now accepts the
+deploy-based `stages:` schema, so retaining the legacy recipe would imply a
+launch path that current `main` rejects.
+
+## Configuration contract under development
+
+The intended deploy integration starts from the
+[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
+and expands its frozen three-stage pipeline into the four stages shown above.
+A complete integration must:
+
+- split the Thinker definition without duplicating the Talker or Code2Wav model
+ contracts
+- remap downstream stage IDs, input sources, and stage connectors
+- inject producer and consumer `kv_transfer_config` values into the two Thinker
+ stages
+- allocate separate device sets for prefill and decode
+- preserve both final outputs: decode text and Code2Wav audio
+- provide connector-specific routing metadata without hard-coding one connector
+ contract
+
+The required number of GPUs depends on tensor-parallel sizes and whether Talker
+and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
+for decode, and one for the two downstream audio stages.
+
+## Implementation map
+
+- PD detection, validation, and sampling-parameter preparation:
+ [`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
+- startup detection and connector metadata:
+ [`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
+- request-scoped prefill-to-decode routing:
+ [`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
+- Qwen3 downstream conditioning helpers:
+ [Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md
index b49cee794..448ca81f2 100644
--- a/docs/getting_started/installation/README.md
+++ b/docs/getting_started/installation/README.md
@@ -1,7 +1,7 @@
# Installation
!!! important
- vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.26.x requires vLLM 0.26.x; the stable v0.26.0 instructions pin vLLM 0.26.0.
+ vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0.
vLLM-Omni supports the following hardware platforms:
diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md
index 59d2aa482..196b6a638 100644
--- a/docs/getting_started/installation/gpu.md
+++ b/docs/getting_started/installation/gpu.md
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md
index a68bfdb58..a8f795861 100644
--- a/docs/getting_started/installation/gpu/cuda.inc.md
+++ b/docs/getting_started/installation/gpu/cuda.inc.md
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
```
-The 0.26.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.
+The 0.28.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
@@ -66,12 +66,12 @@ 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.26.0
+git checkout v0.28.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
-#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
- vllm/vllm-omni:v0.26.0 \
+ vllm/vllm-omni:v0.28.0 \
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -124,7 +124,7 @@ 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.26.0 \
+ --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \
-t vllm-omni-cuda .
```
diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md
index c44407299..e0f1614ae 100644
--- a/docs/getting_started/installation/gpu/rocm.inc.md
+++ b/docs/getting_started/installation/gpu/rocm.inc.md
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library
```bash
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
```
-The pre-built 0.26.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.28.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
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,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.26.0
+git checkout v0.28.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
- vllm/vllm-omni-rocm:v0.26.0 \
+ vllm/vllm-omni-rocm:v0.28.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
- vllm/vllm-omni-rocm:v0.26.0
+ vllm/vllm-omni-rocm:v0.28.0
```
# --8<-- [end:pre-built-images]
diff --git a/docs/getting_started/installation/gpu/xpu.inc.md b/docs/getting_started/installation/gpu/xpu.inc.md
index 8822e3a7c..b9a6e117c 100644
--- a/docs/getting_started/installation/gpu/xpu.inc.md
+++ b/docs/getting_started/installation/gpu/xpu.inc.md
@@ -23,6 +23,19 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
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:`
+base image, which Docker pulls automatically. To target a different vLLM release,
+pass `--build-arg VLLM_VERSION=`. If that tag has not been published yet, build
+the base from upstream's own Dockerfile first and point the build at it:
+
+```bash
+git clone --depth 1 --branch "$VLLM_VERSION" https://github.com/vllm-project/vllm /tmp/vllm
+DOCKER_BUILDKIT=1 docker build -t vllm-openai-xpu:local --target vllm-openai \
+ -f /tmp/vllm/docker/Dockerfile.xpu /tmp/vllm
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
+ --build-arg VLLM_BASE=vllm-openai-xpu:local .
+```
+
#### Launch the docker image
##### Launch with OpenAI API Server
diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md
index 0d8323760..036c2da17 100644
--- a/docs/getting_started/quickstart.md
+++ b/docs/getting_started/quickstart.md
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
-uv pip install vllm==0.26.0 --torch-backend=auto
+uv pip install vllm==0.28.0 --torch-backend=auto
# On ROCm
-uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
+uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
!!! 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.
- If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+ If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
## Offline Inference
@@ -47,7 +47,7 @@ if __name__ == "__main__":
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
- images = outputs[0].request_output.images
+ images = outputs[0].images
images[0].save("coffee.png")
```
@@ -66,7 +66,7 @@ from vllm_omni.entrypoints.omni import Omni
if __name__ == "__main__":
omni = Omni(
model="Tongyi-MAI/Z-Image-Turbo",
- # stage_configs_path="./stage-config.yaml", # See below
+ # deploy_config="./deploy-config.yaml", # Optional deploy override
)
prompts = [
"a cup of coffee on a table",
@@ -75,8 +75,7 @@ if __name__ == "__main__":
]
omni_outputs = omni.generate(prompts)
for i_prompt, prompt_output in enumerate(omni_outputs):
- this_request_output = prompt_output.request_output
- this_images = this_request_output.images
+ this_images = prompt_output.images
for i_image, image in enumerate(this_images):
image.save(f"p{i_prompt}-img{i_image}.jpg")
print("saved to", f"p{i_prompt}-img{i_image}.jpg")
diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py
index 8ddb63561..4b72f9a95 100644
--- a/docs/mkdocs/hooks/generate_examples.py
+++ b/docs/mkdocs/hooks/generate_examples.py
@@ -37,7 +37,6 @@ GENERAL_EXAMPLE_SLUGS = frozenset(
"text_to_speech",
"text_to_video",
"x_to_text",
- "x_to_video_audio",
}
)
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index b9686047a..4333c9934 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -21,43 +21,54 @@ th {
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen1) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
+| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen1) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
-| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | — |
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | — |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | ✅︎ | — |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `LingBotWorldCausalDMDPipeline` | LingBot-World 2.0 v1 (experimental) | `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` | Experimental | | | | — |
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `SanaWmPipeline` | SANA-WM | `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-WM.md) |
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | — |
+| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | ✅︎ | — |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
+| `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
+| `SanaVideoPipeline` | SANA-Video-2B T2V (native 480p/720p; Diffusers adapter validated at both resolutions) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
+| `SanaImageToVideoPipeline` | SANA-Video-2B I2V (native and Diffusers adapter validated at 480p/720p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/NVIDIA/SANA-Video-2B.md) |
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
-| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | | — |
+| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
-| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — |
+| `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -66,13 +77,11 @@ th {
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | — |
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — |
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — |
-| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — |
-| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — |
-| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
+| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — |
+| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | — |
+| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | ✅︎ | — |
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | — |
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | — |
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | — |
@@ -86,17 +95,18 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
-| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
-|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
The `Recipe` column links the corresponding deployment recipe. `Published` is
preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the
@@ -109,6 +119,10 @@ support metadata until their recipe is audited.
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
+See the [LingBot-World 2.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Robbyant/LingBot-World-2.0.md)
+for offline and experimental realtime usage, supported controls, validation,
+and current limitations.
+
H3 MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
Other AMD SKUs (e.g. MI325X) are not listed until their
diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md
index fef28142e..a38512342 100644
--- a/docs/serving/image_edit_api.md
+++ b/docs/serving/image_edit_api.md
@@ -123,7 +123,14 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res
}
],
"output_format": null,
- "size": null
+ "size": null,
+ "metrics": {
+ "stage_durations": {
+ "queue_wait_ms": 1.0,
+ "stage_0_gen_ms": 2.0
+ },
+ "peak_memory_mb": 1024.0
+ }
}
```
diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md
index 0deea7f8c..98fb7f380 100644
--- a/docs/serving/speech_api.md
+++ b/docs/serving/speech_api.md
@@ -131,7 +131,7 @@ Content-Type: application/json
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
-| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
+| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. |
| `ref_text` | string | null | Transcript of reference audio |
| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
@@ -141,7 +141,20 @@ The response shape depends on the streaming parameters:
**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the
complete clip as binary audio data with an appropriate `Content-Type` header (e.g.
-`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported.
+`audio/wav`). Because the raw-bytes body has no JSON carrier, successful
+non-streaming responses from non-diffusion speech servers report usage through
+response headers:
+
+| Header | Description |
+| --- | --- |
+| `x-vllm-omni-input-tokens` | Total input tokens (`text_tokens` + `audio_tokens`). |
+| `x-vllm-omni-output-tokens` | Generated codec/audio tokens. |
+| `x-vllm-omni-total-tokens` | `input_tokens` + `output_tokens`. |
+| `x-vllm-omni-input-text-tokens` | Tokens from the synthesized text (`input` plus `instructions`). |
+| `x-vllm-omni-input-audio-tokens` | Reference-audio codec frames, non-zero only for in-context voice cloning. |
+
+Diffusion-mode speech servers route through a separate response path and do not
+emit these headers.
**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
WAV) as they are decoded.
@@ -675,7 +688,7 @@ for result in response.json()["results"]:
|-----------|--------|---------|-------------|
| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
-All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
+All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed.
For best throughput, set both stages' `max_num_seqs` above 1 via `--stage-overrides`. On the current Qwen3-TTS CustomVoice benchmark, stage 1 performed best at `max_num_seqs: 10`:
@@ -798,6 +811,51 @@ If you encounter OOM errors:
Use `/v1/audio/voices` to list available voices for the loaded model.
+## Orchestration Loop (experimental)
+
+Multi-stage omni deployments route stage outputs through a single orchestrator
+loop. By default that loop polls every stage replica on a 1 ms cadence. An
+opt-in event-driven mode replaces the poll with one reader task per live stage
+replica awaiting its client directly, and switches the serving-side
+final-output drain to a condition-variable wakeup at the same time.
+
+**Configuration (environment variables):**
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
+
+Set it on the process that runs the orchestrator (stage 0 of an omni
+deployment) before starting the server:
+
+```bash
+export VLLM_OMNI_EVENT_DRIVEN_ORCH=1
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+ --omni \
+ --port 8091
+```
+
+The server logs the selected loop mode and its reader/poller counts once at
+startup, so you can confirm which loop is live.
+
+Routing, output ordering, and terminal-state behavior are identical on both
+loops; only the poll cadence changes. Leaving the variable unset keeps the
+legacy poll loop, which is the supported default.
+
+**Known limitations:**
+
+- The measured serving A/B (idle CPU 2.43% to 0.07%; TTFP p99 -32% at
+ concurrency 8) predates the rebuild on the per-replica fault-isolation work
+ in [#4285](https://github.com/vllm-project/vllm-omni/pull/4285). That work
+ changed dead-replica handling and reader/poller lifecycle rather than the
+ steady-state output path, and the parity suite covers it, but the serving
+ A/B has not been re-run on the current head.
+- The diffusion-poller branch is covered by unit tests only. Deployments whose
+ stages all run as standard engine cores never exercise it, including GLM-TTS,
+ which deploys its DiT without `stage_type: diffusion`.
+- Concurrency 1 and 32 measured at parity with the legacy loop. At 32 the
+ latency is admission-bound, which this mode does not address.
+
## Development
Enable debug logging:
diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md
index 88f74affc..688f5fe9c 100644
--- a/docs/serving/video_stream_api.md
+++ b/docs/serving/video_stream_api.md
@@ -10,7 +10,7 @@ Each server instance runs a single model specified at startup with `vllm serve <
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
- --deploy-config vllm_omni/deploy/qwen3_omni.yaml \
+ --deploy-config vllm_omni/deploy/qwen3_omni_moe.yaml \
--omni \
--port 8000 \
--trust-remote-code
diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md
index 6a94f9dab..7fc7d0ba8 100644
--- a/docs/serving/videos_api.md
+++ b/docs/serving/videos_api.md
@@ -161,11 +161,29 @@ are not implemented yet. Models may expose additional V2V controls through
`extra_params`. For example, Cosmos3 supports
`condition_frame_indexes_vision` and `condition_video_keep` to select which
decoded reference frames are used as clean conditioning. Cosmos3 transfer mode
-also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus
-transfer options such as `control_path`, `control_guidance`,
-`control_guidance_interval`, `num_video_frames_per_chunk`,
-`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
-Cosmos3 recipe for complete examples.
+also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints. Each hint
+may specify its own `control_path` and `control_weight`. Request-level transfer
+options include
+`control_guidance`, `control_guidance_interval`,
+`emphasize_control_in_prompt`, `num_video_frames_per_chunk`,
+`num_conditional_frames`, `show_control_condition`, and `show_input`. Transfer
+uses its transfer-specific system prompt and, by default, appends a
+control-adherence directive to the positive prompt. It adds duration/FPS and
+resolution metadata to both CFG branches but does not add a negative prompt
+automatically. The Cosmos3 recipe includes an optional reference negative
+prompt and shows how to pass it. Set `emphasize_control_in_prompt`,
+`use_duration_template`, or `use_resolution_template` to `false` to disable the
+corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or
+`none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete
+examples.
+
+HTTP redirects for `image_reference.image_url` follow vLLM's
+`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
+`1` (the default) to allow redirects or `0` to reject them. A redirect target
+can differ from the original host, and vLLM-Omni does not yet fully implement
+upstream vLLM's media URL allowlist protection. Deployments should therefore
+treat remote media URLs as untrusted and choose this setting as part of their
+URL access policy.
### Speech-to-Video
diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md
index b6fd13b12..e2a3b9bfc 100644
--- a/docs/user_guide/diffusion/attention_backends.md
+++ b/docs/user_guide/diffusion/attention_backends.md
@@ -1,494 +1,195 @@
# Diffusion Attention Backends
-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
+Use this page to select and configure a diffusion attention backend. Backend-
+specific installation, tuning, and compatibility details live in separate
+guides so that the selection contract stays easy to scan.
-## Overview
+Diffusion attention backend selection applies to DiT attention in image and
+video generation models. It does **not** change autoregressive LLM attention,
+which uses vLLM's attention selector.
-Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
+For the internal selector, registry, and platform contract, see
+[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
-This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
+## Choose a guide
-The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
+| Need | Guide |
+| --- | --- |
+| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
+| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
+| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
+| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
+| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
-## Backend Options
-
-| Value | Notes |
-|---|---|
-| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
-| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
-| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
-| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
-| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
-| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
-| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
-| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
-| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
-| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |
+## Backend options
+| Value | Family | Primary use | Detail |
+| --- | --- | --- | --- |
+| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
+| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
+| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
+| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
+| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
+| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
+| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
+| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
+| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
+| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
## Configuration
-Diffusion attention backends can be configured three ways, in priority order:
+Backend selection is resolved in this order:
-1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
-2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
-3. **Platform default** — used when nothing is configured.
+1. `--diffusion-attention-config` per-role configuration.
+2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
+ global default.
+3. The current platform's default.
-`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
+`--diffusion-attention-backend` is shorthand for
+`--diffusion-attention-config.default.backend`. Do not pass it together with an
+explicit `default.backend` in the structured configuration.
### Global default
-Set the default backend for every diffusion attention layer:
-
```bash
-# CLI flag
-vllm-omni serve --diffusion-attention-backend SAGE_ATTN
+vllm-omni serve --diffusion-attention-backend FLASH_ATTN
-# Environment variable (also recognized for backwards compatibility)
-export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
+# Backwards-compatible environment variable
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
```
### Per-role configuration
-Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
+Roles are declared by each diffusion model. Common categories are `self` and
+`cross`; a model may also use a more specific role such as
+`ltx2.audio_to_video`. Resolution order is:
-1. Exact `per_role[role]` match
-2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
-3. `default`
-4. Platform default
-
-Use vLLM-style dotted flags or one JSON blob:
+1. Exact `per_role[role]` match.
+2. Category `per_role[role_category]` match.
+3. `default`.
+4. Platform default.
```bash
# Dotted flags
vllm-omni serve \
- --diffusion-attention-config.default.backend FLASH_ATTN \
- --diffusion-attention-config.per_role.cross.backend TORCH_SDPA
+ --diffusion-attention-config.default.backend FLASH_ATTN \
+ --diffusion-attention-config.per_role.cross.backend TORCH_SDPA
-# JSON
+# Equivalent JSON
vllm-omni serve \
- --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
+ --diffusion-attention-config \
+ '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
```
-A backend that needs configuration exposes it as a typed field on the spec:
-`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
-`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
-
-```bash
---diffusion-attention-config.default.backend TRTLLM_ATTN \
---diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
-```
-
-### Programmatic API
-
-When constructing `OmniDiffusionConfig` directly:
+### Python API
```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
+from vllm_omni.diffusion.data import (
+ AttentionConfig,
+ AttentionSpec,
+ OmniDiffusionConfig,
+)
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
default=AttentionSpec(backend="FLASH_ATTN"),
- per_role={
- "cross": AttentionSpec(backend="TORCH_SDPA"),
- },
+ per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
),
...,
)
```
-A plain dict is also accepted and normalized to `AttentionConfig`.
+Backend-specific typed blocks are documented with their consumers:
-#### FlashInfer QK16/V8 quantized attention
+- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
+ and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
+- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
+- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
+- `fastvideo_vsa_topk`: [FastVideo VSA](attention_backends/fastvideo_vsa.md#choose-top-k).
-`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
-
-```python
-config = OmniDiffusionConfig(
- diffusion_attention_config=AttentionConfig(
- default=AttentionSpec(
- backend="FLASHINFER_ATTN",
- quant=AttnQuantSpec(
- dtype_qk="bfloat16",
- dtype_vo="fp8_e4m3",
- ),
- ),
- ),
- ...,
-)
-```
-
-`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
-
-## Platform Defaults
+## Platform defaults
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
-Auto-route preference, in order:
+The CUDA auto-route preference is:
-1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
-2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
-3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
-4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
-5. `TORCH_SDPA` — last resort
+1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
+ available, `head_dim=128`, and the model declares a compatible packed or
+ mask-free path.
+2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
+3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
+4. `FLASH_ATTN` when a compatible package is installed.
+5. `TORCH_SDPA`.
-`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.
+`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
+(sm_120/sm_121), for other head dimensions, or for paths that require an
+unsupported mask.
-`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.
+### Hopper, Ada, and Ampere
-The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
+The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
+to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
-**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
+Other platforms validate an explicit backend and choose their own default
+through the platform implementation. Check the startup log to confirm the
+resolved backend.
-### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
+## Choosing a backend manually
-Auto-route preference:
+Override the platform default when you need:
-1. `FLASH_ATTN` — when `flash-attn` is installed
-2. `TORCH_SDPA` — fallback
+- a correctness reference (`TORCH_SDPA`);
+- a backend-specific workaround;
+- training/rollout kernel alignment (Hub backends); or
+- an explicitly validated sparse or quantized speedup.
-`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
+The startup log prints the resolved backend and whether it came from explicit
+configuration or platform defaulting. If no resolution message appears, check
+earlier logs for diffusion-stage initialization failures.
-## TRTLLM_ATTN Backend and Skip-Softmax
+## Reference benchmark
-`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
-that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
-(algorithm: [design doc](../../design/feature/skip_softmax.md)).
+The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
+with the same prompt and seed across runs. Treat them as reference results, not
+portable guarantees.
-Enable it through the typed `skip_softmax` block on the attention spec:
-
-| Key | Valid values | Meaning |
-|---|---|---|
-| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
-| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
-| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
-
-```bash
-vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
- --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
- "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
-```
-
-Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
-
-AttentionConfig(
- default=AttentionSpec(
- backend="TRTLLM_ATTN",
- skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
- ),
-)
-```
-
-**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
-as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
-sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
-`skip_softmax` block.
-
-Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
-`TRTLLM_ATTN` raises rather than silently degrading.
-
-## TRTLLM_ATTN SAGE Quantization
-
-On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
-FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
-in low precision. Enable it through the typed `quant` block on the attention spec:
-
-| Key | Valid values | Meaning |
-|---|---|---|
-| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
-| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
-| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
-
-V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
-neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
-
-The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
-fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
-`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
-`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
-
-```bash
-vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
- --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
- "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
-```
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
-
-AttentionConfig(
- default=AttentionSpec(
- backend="TRTLLM_ATTN",
- quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
- ),
-)
-```
-
-**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
-availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
-(B300); `int8` QK kernels are compiled for **SM100 only**.
-
-## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
-
-`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
-128-token block of keys, ranks them per query block, and computes attention against only the
-top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
-spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
-
-Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
-first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
-nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
-exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
-delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
-
-Enable it through the typed `block_sparse` block on the attention spec, shared by every
-block-sparse backend:
-
-| Key | Valid values | Meaning |
-|---|---|---|
-| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
-| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
-| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
-
-```bash
-vllm-omni serve MiniMaxAI/MiniMax-H3 \
- --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
- "block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
-```
-
-Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
-
-AttentionConfig(
- default=AttentionSpec(
- backend="RAINFUSION_ATTN",
- block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
- ),
-)
-```
-
-**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
-cheapest way to recover structure, because the early high-noise steps decide global layout while the
-later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
-`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
-
-Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
-incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
-blocks — use Ulysses SP (`ring_degree=1`).
-
-### Which geometries run sparse
-
-MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
-necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
-does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
-promoted rows remain visible to every query and key.
-
-For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
-128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
-MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
-and kernel tiling aligned, at the cost of retaining more blocks.
-
-8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
-creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
-grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
-that implements protected video tails for rf_v2.
-
-## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
-
-Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
-
-| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
-|---|---|---|---|---|
+| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
+| --- | --- | ---: | ---: | ---: |
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
-Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
+Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
+were close enough that users should benchmark their exact model and shape.
-## Known Limitations
+## Compatibility anchors
-### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
+The following headings preserve links to sections that moved into dedicated
+guides.
-LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
+## TRTLLM_ATTN Backend and Skip-Softmax
-**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
+See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
- --model Lightricks/LTX-2 ...
-```
+## TRTLLM_ATTN SAGE Quantization
-### FlashAttention-4 on Blackwell
+See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
-Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
+## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
-```bash
-pip install 'vllm-omni[fa4]'
-```
-
-`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
-when it is unavailable. Version `4.0.0b18` is required; earlier beta
-wheels had known JIT failures on Blackwell.
-
-## Choosing a Backend Manually
-
-### When to override the default
-
-- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
-- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
-- **Workaround for known issues**: see Known Limitations above.
-
-### Verifying which backend is in use
-
-The startup log prints one of:
-
-```
-Using diffusion attention backend 'CUDNN_ATTN' # explicit override
-Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
-Defaulting to diffusion attention backend SDPA # nothing else available
-```
-
-If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
+See [RainFusion](attention_backends/rainfusion.md).
## SageAttention Installation
-vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
-
-Build from source:
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention
-
-export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
-pip install . --no-build-isolation
-```
-
-Quick check:
-
-```bash
-python -c "import sageattention; print(sageattention.__file__)"
-```
+See [SageAttention](attention_backends/sage.md#installation).
## SageAttention3 Installation
-vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
-
-Build from source:
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention/sageattention3_blackwell
-python setup.py install
-```
-
-Quick check:
-
-```bash
-python -c "import sageattn3; print(sageattn3.__file__)"
-```
-
-Notes:
-
-- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
-- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
+See [SageAttention](attention_backends/sage.md#sageattention3-installation).
## HuggingFace Kernels Hub Backends
-To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
-
-The following backend options are supported:
-- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
-- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
-
-### Installation
-
-To use these backends, you must install the `kernels` library:
-
-```bash
-pip install kernels==0.14.1
-```
-
-If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
-
-### Usage
-
-Select a Hub backend using the global CLI flag or environment variables:
-
-```bash
-# Environment variable
-export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
-
-# CLI flag
-vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB
-```
-
-## Usage Examples
-
-### Default (auto-route)
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
- --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
- --prompt "A dog running across a field of golden wheat." \
- --height 480 --width 832 --num-frames 33 \
- --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
- --output hv15.mp4
-```
-
-On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
-
-### Explicit backend selection
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
- --model Lightricks/LTX-2 \
- --prompt "A dog running across a field of golden wheat." \
- --height 480 --width 832 --num-frames 33 \
- --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
- --output ltx2.mp4
-```
-
-### SageAttention (lossy)
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
- --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
- --prompt "A dog running across a field of golden wheat." \
- --height 480 --width 832 --num-frames 33 \
- --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
- --tensor-parallel-size 2 \
- --output hv15_sage.mp4
-```
-
-Example: Wan2.2 TI2V 5B
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
- --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
- --prompt "A dog running across a field of golden wheat." \
- --height 704 --width 1280 --num-frames 49 \
- --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
- --tensor-parallel-size 2 \
- --output outputs/wan22_sage.mp4
-```
-
-### Enable SageAttention3
-
-Example:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
- --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
- --prompt "A dog running across a field of golden wheat." \
- --height 480 --width 832 --num-frames 33 \
- --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
- --tensor-parallel-size 2 \
- --output outputs/hv15_sage3.mp4
-```
+See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md
new file mode 100644
index 000000000..030550dac
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md
@@ -0,0 +1,110 @@
+# Dense Attention Backends
+
+Use these backends when you want dense attention without backend-specific
+sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
+the fastest compatible kernel for your model, shape, and hardware.
+
+For selection precedence and per-role configuration, see the
+[attention backend overview](../attention_backends.md).
+
+## `TORCH_SDPA`
+
+`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
+dispatcher choose the implementation. It is always available and is the most
+conservative reference when validating another backend.
+
+```bash
+vllm-omni serve --diffusion-attention-backend TORCH_SDPA
+```
+
+## `FLASH_ATTN`
+
+`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
+prefers the CuTe-based FlashAttention 4 path and falls back to
+FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
+Ampere it is the preferred automatic route when a compatible package is
+installed.
+
+```bash
+vllm-omni serve --diffusion-attention-backend FLASH_ATTN
+```
+
+### FlashAttention 4 on Blackwell
+
+Install the optional CUDA 13 extra:
+
+```bash
+pip install 'vllm-omni[fa4]'
+```
+
+Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
+Blackwell. If the CuTe path is unavailable, the backend falls back to the
+compatible FlashAttention 3 or 2 path.
+
+## `CUDNN_ATTN`
+
+`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
+for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
+9.5 or newer is available and the higher-priority TRTLLM route is not
+compatible.
+
+```bash
+vllm-omni serve --diffusion-attention-backend CUDNN_ATTN
+```
+
+### LTX-2.0 limitation
+
+LTX-2 audio attention has a symbolic head dimension during `torch.compile`
+tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
+aborts compilation. This is tracked in
+[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
+
+Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
+ python examples/offline_inference/text_to_video/text_to_video.py \
+ --model Lightricks/LTX-2 ...
+```
+
+## `FLASHINFER_ATTN`
+
+`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
+option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
+is installed but cuDNN is too old for `CUDNN_ATTN`.
+
+```bash
+vllm-omni serve --diffusion-attention-backend FLASHINFER_ATTN
+```
+
+### FlashInfer quantized attention
+
+The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
+in FP16 or BF16 and use FP8 E4M3 for V:
+
+```python
+from vllm_omni.diffusion.data import (
+ AttentionConfig,
+ AttentionSpec,
+ AttnQuantSpec,
+ OmniDiffusionConfig,
+)
+
+config = OmniDiffusionConfig(
+ diffusion_attention_config=AttentionConfig(
+ default=AttentionSpec(
+ backend="FLASHINFER_ATTN",
+ quant=AttnQuantSpec(
+ dtype_qk="bfloat16",
+ dtype_vo="fp8_e4m3",
+ ),
+ ),
+ ),
+ ...,
+)
+```
+
+`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
+configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
+schema is also consumed by TRTLLM, but each backend validates its own allowed
+fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md
new file mode 100644
index 000000000..2439633d4
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md
@@ -0,0 +1,125 @@
+# FastVideo VSA
+
+FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
+self-attention by partitioning the post-patch latent grid into spatiotemporal
+blocks. For every query block, VSA scores the key/value blocks and computes
+attention only against the selected top-k blocks.
+
+The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
+
+VSA is a CUDA-only, explicitly selected backend. It requires the
+`fastvideo-kernel` package and currently supports non-causal self-attention
+with equal query and key/value sequence lengths. Unsupported shapes, masks,
+dtypes, sequence-parallel execution, or kernel failures fall back to
+`TORCH_SDPA` and emit a warning with the reason.
+
+## Enable the backend
+
+For online serving, select the backend with the existing attention backend
+flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
+for every query block:
+
+```bash
+vllm-omni serve \
+ --diffusion-attention-backend FASTVIDEO_VSA \
+ --fastvideo-vsa-topk 64
+```
+
+The backwards-compatible environment variable selects the backend with the
+default `topk=64`:
+
+```bash
+export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
+vllm-omni serve
+```
+
+To tune top-k, pass the CLI backend and top-k flags together as shown above,
+or use the equivalent structured configuration:
+
+```bash
+vllm-omni serve \
+ --diffusion-attention-config \
+ '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
+```
+
+Do not combine `--diffusion-attention-backend` with an explicit
+`diffusion_attention_config.default.backend`. The top-k value must be positive
+and is valid only when the default backend is `FASTVIDEO_VSA`.
+
+For a deploy YAML stage, use either the shorthand fields:
+
+```yaml
+stages:
+ - stage_id: 0
+ diffusion_attention_backend: FASTVIDEO_VSA
+ fastvideo_vsa_topk: 64
+```
+
+or the structured configuration:
+
+```yaml
+stages:
+ - stage_id: 0
+ diffusion_attention_config:
+ default:
+ backend: FASTVIDEO_VSA
+ fastvideo_vsa_topk: 64
+```
+
+## Choose top-k
+
+At runtime the backend logs the sequence shape and derived block count:
+
+```text
+FASTVIDEO_VSA routing: seq_len=27280, dit_seq_shape=(31, 22, 40),
+block_size=(4, 8, 8), num_blocks=120, topk=64, keep_ratio=53.3%,
+checkpoint_mode=native, route=VSA
+```
+
+Use `num_blocks` as the upper bound when tuning top-k. Top-k is the number of
+key/value blocks retained for each query block, not the number of tokens or
+the total number of blocks processed by the layer.
+
+- A smaller top-k increases sparsity and may reduce attention computation, but
+ it can remove relevant blocks and reduce visual quality or temporal
+ consistency. Routing and padding overhead can also make a smaller value
+ slower for some shapes.
+- A larger top-k retains more context and generally approaches dense-attention
+ quality, but reduces the potential speedup and increases memory traffic.
+- `topk > num_blocks` is invalid for the runtime shape and falls back to SDPA.
+- For a native checkpoint, `topk == num_blocks` routes to SDPA because scoring
+ every block provides no sparsity benefit.
+- For a FastVideo DMD checkpoint, `topk == num_blocks` stays on the VSA
+ all-block path so the checkpoint keeps its trained compensation semantics.
+
+There is no universally optimal value. Resolution, frame count, GPU, kernel
+version, and checkpoint all affect both quality and latency. Start from the
+logged `num_blocks`, test several keep ratios on the target workload, and
+compare output quality against the same checkpoint running dense attention.
+
+## Checkpoint behavior
+
+The backend does not expose a user-selectable gate mode. Wan checkpoints that
+contain learned `to_gate_compress` weights use that projection automatically.
+When those weights are absent, the unused projection is removed and VSA uses
+the sparse branch without learned compensation.
+
+FastVideo DMD checkpoints use their fixed distilled timestep schedule. Native
+Wan checkpoints keep their normal scheduler and inference-step configuration;
+selecting VSA does not turn a native checkpoint into a distilled model.
+
+## Verify routing and fallback
+
+Check the startup and first-forward logs instead of assuming that selecting
+the backend guarantees sparse execution:
+
+- `route=VSA` means top-k block selection is active.
+- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
+ blocks through the VSA kernel.
+- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
+ executed; the warning includes the reason.
+
+The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
+standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
+and no active sequence-parallel context. NPU and XPU paths do not execute the
+FastVideo VSA CUDA kernel.
diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md
new file mode 100644
index 000000000..946786b97
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md
@@ -0,0 +1,32 @@
+# Hugging Face Hub Attention Backends
+
+Hub backends load the same published kernels commonly used with Hugging Face
+Diffusers. They are useful when training and serving must use matching kernel
+implementations to reduce numerical drift and sampling divergence.
+
+| Backend | Kernel | Platform |
+| --- | --- | --- |
+| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
+| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
+
+## Installation
+
+```bash
+pip install kernels==0.14.1
+```
+
+If `kernels` is unavailable, vLLM-Omni warns and falls back to the
+corresponding local FlashAttention implementation. On CUDA devices older than
+Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
+
+## Usage
+
+```bash
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
+
+# Equivalent CLI selection
+vllm-omni serve --diffusion-attention-backend FLASH_ATTN_HUB
+```
+
+For per-role kernel alignment, configure the Hub backend on only the relevant
+role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
diff --git a/docs/user_guide/diffusion/attention_backends/rainfusion.md b/docs/user_guide/diffusion/attention_backends/rainfusion.md
new file mode 100644
index 000000000..00f4e37a1
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/rainfusion.md
@@ -0,0 +1,80 @@
+# RainFusion Attention
+
+`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
+attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
+block, and attends to the highest-scoring blocks after arranging video tokens
+in `(t, h, w)` order.
+
+Only the video segment is sparse. Prefix rows such as text, visual conditions,
+and audio, plus first-frame blocks, remain dense. Unsupported calls—including
+warmup steps, skipped layers, missing video geometry, and video segments below
+32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
+globally.
+
+## Configuration
+
+| Key | Valid values | Meaning |
+| --- | --- | --- |
+| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
+| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
+| `end_step` | integer, `>= 0` | Number of final denoise steps kept dense (tail fallback) |
+| `precision` | `"bf16"`, `"fp8"`, `"mix"` | Kernel precision mode; default `"bf16"`. Requires MindIE-SD with `sparse_attention(precision=...)` |
+| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+ --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
+ "block_sparse":{"sparsity":0.8,"start_step":0}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
+
+config = AttentionConfig(
+ default=AttentionSpec(
+ backend="RAINFUSION_ATTN",
+ block_sparse=BlockSparseSpec(
+ sparsity=0.8,
+ start_step=0,
+ skip_layers="0-1",
+ ),
+ ),
+)
+```
+
+Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
+early-step window is usually the cheapest way to recover global structure;
+use layer exclusions only after a same-seed dense comparison identifies
+sensitive blocks.
+
+## Requirements and compatibility
+
+RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
+platform raises. It is incompatible with ring sequence parallelism because
+the kernel needs the complete key sequence for ranking; use Ulysses sequence
+parallelism with `ring_degree=1`.
+
+The `precision` knob requires a MindIE-SD release whose `sparse_attention`
+accepts `precision=`; older releases accept it through `**kwargs` but
+silently ignore it and run the BF16 path. `RAINFUSION_ATTN` raises at runtime
+when a non-`bf16` precision is requested without that support, so pin a
+compatible MindIE-SD release when using `precision="mix"` or `"fp8"`.
+
+## Geometry handling
+
+RainFusion handles arbitrary video grids by rearranging video spatially and,
+when necessary, promoting a real-video suffix to the always-kept prefix. It
+does not pad because `rf_v2` does not consume a padding attention mask.
+
+For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
+The implementation promotes 2,976 rows and leaves 59,520 sparse rows
+(`465 x 128`) so the mask and kernel tiling remain aligned.
+
+Spatial grids aligned to 8x8 are still preferable. A latent height or width
+not divisible by 8 creates an always-kept suffix and reduces realized
+sparsity. At the image level, that means width and height divisible by 256.
+Protected video tails require a compatible MindIE-SD release.
+
+For common configuration and selector behavior, see the
+[attention backend overview](../attention_backends.md) and the
+[backend selection design](../../../design/feature/attention_backend_selection.md).
diff --git a/docs/user_guide/diffusion/attention_backends/sage.md b/docs/user_guide/diffusion/attention_backends/sage.md
new file mode 100644
index 000000000..1e1fa5b60
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/sage.md
@@ -0,0 +1,63 @@
+# SageAttention
+
+SageAttention backends provide lossy low-precision attention for diffusion
+models. Validate output quality against `TORCH_SDPA` at the same seed before
+using either backend in production.
+
+## `SAGE_ATTN`
+
+`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
+accumulation.
+
+### Installation
+
+Install SageAttention into the same environment as vLLM-Omni:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention
+export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
+pip install . --no-build-isolation
+```
+
+Verify the installation:
+
+```bash
+python -c "import sageattention; print(sageattention.__file__)"
+```
+
+Select it globally:
+
+```bash
+vllm-omni serve --diffusion-attention-backend SAGE_ATTN
+```
+
+## `SAGE_ATTN_3`
+
+`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
+
+### SageAttention3 installation
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention/sageattention3_blackwell
+python setup.py install
+```
+
+Verify the installation:
+
+```bash
+python -c "import sageattn3; print(sageattn3.__file__)"
+```
+
+```bash
+vllm-omni serve --diffusion-attention-backend SAGE_ATTN_3
+```
+
+`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
+GPU. Its kernel assumes the query-head count equals the key/value-head count.
+GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
+correctness.
+
+For common configuration and platform routing, see the
+[attention backend overview](../attention_backends.md).
diff --git a/docs/user_guide/diffusion/attention_backends/trtllm.md b/docs/user_guide/diffusion/attention_backends/trtllm.md
new file mode 100644
index 000000000..dd36ba61e
--- /dev/null
+++ b/docs/user_guide/diffusion/attention_backends/trtllm.md
@@ -0,0 +1,103 @@
+# TRTLLM Attention
+
+`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
+BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
+
+For common selection and per-role configuration, see the
+[attention backend overview](../attention_backends.md).
+
+## Requirements
+
+Selecting `TRTLLM_ATTN` requires:
+
+- datacenter Blackwell, sm_100 or sm_103;
+- `head_dim=128`;
+- FlashInfer; and
+- a model-declared packed or mask-free compatible path.
+
+It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
+explicit selection raises instead of silently changing the backend.
+
+## Dense mode
+
+Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
+
+```bash
+vllm-omni serve --diffusion-attention-backend TRTLLM_ATTN
+```
+
+## Skip-Softmax
+
+Skip-Softmax trades some fidelity for speed by skipping low-contribution
+attention work. See the [feature design](../../../design/feature/skip_softmax.md)
+for the algorithm and implementation contract.
+
+| Key | Valid values | Meaning |
+| --- | --- | --- |
+| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
+| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
+| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+ --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
+ "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
+
+config = AttentionConfig(
+ default=AttentionSpec(
+ backend="TRTLLM_ATTN",
+ skip_softmax=SkipSoftmaxSpec(
+ target_sparsity=0.65,
+ disabled_until_timestep=0.86,
+ ),
+ ),
+)
+```
+
+Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
+sparsity only after comparing quality against dense output at the same seed.
+The benefit grows with sequence length because only attention is accelerated.
+
+## SAGE quantization
+
+TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
+matrix multiplications use low precision.
+
+| Key | Valid values | Meaning |
+| --- | --- | --- |
+| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
+| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
+| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
+
+V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
+Only block sizes with compiled kernels are accepted.
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+ --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
+ "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
+
+config = AttentionConfig(
+ default=AttentionSpec(
+ backend="TRTLLM_ATTN",
+ quant=AttnQuantSpec(
+ dtype_qk="fp8_e4m3",
+ q_block_size=1,
+ k_block_size=16,
+ ),
+ ),
+)
+```
+
+This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
+available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
+The shared `AttnQuantSpec` is backend typed: values intended for
+`FLASHINFER_ATTN` are rejected here and vice versa.
diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md
index 02ab9f074..93ac5e48e 100644
--- a/docs/user_guide/diffusion/cpu_offload.md
+++ b/docs/user_guide/diffusion/cpu_offload.md
@@ -1,389 +1,74 @@
# CPU Offloading for Diffusion Models
-## Overview
+vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
+the coarsest strategy that meets your memory target: finer-grained streaming
+saves more device memory but adds more transfer and setup overhead.
-vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
+For the shared factory, discovery, and lifecycle contract, see the
+[CPU Offloading design](../../design/feature/offloader/README.md).
-1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
-2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
+## Choose a strategy
-Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
+| Strategy | Device residency | Parallel scope | Guide |
+| --- | --- | --- | --- |
+| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
+| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
+| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
+All strategies use pinned host memory for faster transfers where applicable.
+Configuration priority is:
+
+1. Distributed layerwise offloading.
+2. Layerwise offloading.
+3. Model-level offloading.
+
+Treat the flags as mutually exclusive. If more than one is enabled, the
+higher-priority strategy is selected.
+
+## Quick selection
+
+- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
+ encoders and DiTs is enough to fit and phase-boundary transfers are
+ acceptable.
+- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
+ DiTs where block transfers can overlap computation.
+- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md)
+ when a multi-rank deployment also needs bounded device residency and,
+ optionally, sharded host weights.
+
+## Supported models
+
+| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
+| --- | --- | --- | :---: | :---: | :---: | --- |
+| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
+| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
+| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
+| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
+
+Model-level support requires discoverable DiT and encoder components.
+Layerwise support requires transformer block topology. Distributed support
+reuses that topology but still requires validation for the model, checkpoint,
+and parallel configuration.
+
+## Compatibility anchors
+
+These headings preserve links to sections that moved into dedicated guides.
## Model-level (Sequential) Offloading
-### How It Works
-
-Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
-
-- **When encoders run**: DiT transformer is offloaded to CPU
-- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
-- **VAE**: Stays resident on GPU
-
-Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
-
-### Usage
-
-**Python API:**
-```python
-from vllm_omni import Omni
-
-m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
-```
-
-**CLI:**
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
-```
-
-### To Support a Model
-
-Implement the `SupportsComponentDiscovery` protocol to declare which
-submodules serve as pipeline components (used by offloading, HSDP
-sharding, and other framework features):
-
-```python
-from typing import ClassVar
-from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
-
-class MyPipeline(nn.Module, SupportsComponentDiscovery):
- _dit_modules: ClassVar[list[str]] = ["transformer"]
- _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
- _vae_modules: ClassVar[list[str]] = ["vae"]
- _resident_modules: ClassVar[list[str]] = [] # optional
-
- def __init__(self):
- super().__init__()
- self.transformer = ... # DiT — stays on GPU during denoising
- self.text_encoder = ... # Encoder — offloaded to CPU during denoising
- self.vision_model = ... # Encoder — offloaded to CPU during denoising
- self.vae = ... # VAE — always on GPU
-```
-
-- `_dit_modules`: attribute names of denoising submodules (kept on GPU
- during the diffusion loop).
-- `_encoder_modules`: attribute names of encoder/vision submodules
- (offloaded to CPU during the diffusion loop).
-- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
- part of the mutual exclusion hooks).
-- `_resident_modules`: attribute names of small submodules that must
- stay on GPU during layerwise offloading (e.g. embedders, connectors).
- Optional — defaults to `[]`.
-
-All attribute names support dotted paths for nested submodules
-(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
-
-Both DiT and encoder lists are needed because the offload hooks use
-mutual exclusion: when one group runs, the other moves to CPU.
-
-### Limitations
-- Cold start latency increases
-- Adds overhead from CPU-GPU transfers between encoder and denoising phases
-- Support single GPU only for now
-
-
-### Component offloading for split models (e.g. Cosmos3)
-
-Some models split their transformer into mutually-exclusive *components* that run
-in different phases of a single forward pass rather than as separate pipeline
-components -- e.g. Cosmos3's understanding (reasoner) component runs once per
-generation while the generation (generator) component runs every denoising step.
-Such models have no separate text encoder to swap against, so the transformer
-owns a small model-local offload path and wraps each phase with
-`with self._offload_context(name):`
-
-```python
-class Cosmos3VFMTransformer(nn.Module):
- def forward(self, ...):
- with self._offload_context("reasoner"):
- ... # understanding pass, runs once
- with self._offload_context("generator"):
- ... # denoising pass, runs every step
-```
-
-Model-level offloading then keeps exactly one component GPU-resident at a time
-(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
-pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
-transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
-works for these models too -- each component declares its own block container via
-`_layerwise_offload_blocks_attrs`.
-
+See [Model-Level Offloading](offloader/cpu_offload.md).
## Layerwise (Blockwise) Offloading
-### How It Works
-
-Layerwise offloading keeps only one transformer block on GPU at a time.
-
-As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
-
-**Execution Flow:**
-
-| Block | Pre-forward Hook | Forward | Post-forward Hook |
-|-------|------------------|---------|-------------------|
-| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
-| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
-| ... | ... | ... | ... |
-| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
-
-Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
-
-Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
-
-### Usage
-
-**Python API:**
-```python
-from vllm_omni import Omni
-
-# Text-to-video
-m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
-
-# Or image-to-video
-m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
-```
-
-**CLI:**
-```bash
-# Text-to-video
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
-
-# Or image-to-video
-vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
-```
-
-### To Support a Model
-
-Models must define the blocks attribute name for layerwise offloading:
-
-```python
-class WanTransformer3DModel(nn.Module):
- _layerwise_offload_blocks_attrs = ["blocks"] # Attribute names containing transformer blocks
-
- def __init__(self):
- self.blocks = nn.ModuleList([...]) # Transformer blocks
-```
-
-For models with multiple block types:
-
-```python
-class Flux2Transformer2DModel(nn.Module):
- _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
-```
-
-### Limitations
-- Cold start latency increases because offloaded components must be moved to CPU
- during setup; layerwise offload may add extra weight consolidation and pinning
- work.
-- Performance depends on compute cost and H2D bandwidth as well
-- Support single GPU only for now
-
+See [Layerwise Offloading](offloader/layerwise_offload.md).
## Distributed Layerwise Offloading
-### How It Works
-
-Distributed layerwise offloading extends single-GPU layerwise offloading to
-multi-device deployments. In the default mode, each rank stores only
-**1/dp_size** of the model weights in host memory; full layer weights are
-reconstructed at runtime via **AllGather** on a dedicated communication stream,
-overlapped with computation via a fixed double-buffer scheme. The
-`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
-weights independently and does not shard weights across ranks.
-
-For the implementation design and compatibility matrix, see
-[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
-
-**Key features:**
-- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
- reconstructed per-layer via `all_gather_into_tensor`
-- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
- time, regardless of model size
-- **DP multi-concurrency**: N concurrent requests processed in parallel
- (AllGather only gathers weight shards, which are request-independent)
-- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
- page cache, eliminating O(dp_size × model_size) RSS during model creation
-- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
- (CANN/HCCL) via vLLM-Omni's platform abstraction
-
-**Execution Flow (per device):**
-
-```
-Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
-H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
-AllGather: [AG N+1] [AG N+2]
-Slot Usage: Slot0: Layer N Slot1: Layer N+1
-```
-
-### Usage
-
-**CLI:**
-```bash
-# 4× GPU/NPU with AllGather (recommended)
-vllm serve /path/to/model --omni \
- --enable-distributed-layerwise-offload \
- --data-parallel-size 4
-
-# Without DLO AllGather (each rank streams standard-loader rank-local weights)
-vllm serve /path/to/model --omni \
- --enable-distributed-layerwise-offload \
- --data-parallel-size 4 \
- --dlo-no-use-allgather
-
-# With SP instead of DP (long sequences)
-vllm serve /path/to/model --omni \
- --enable-distributed-layerwise-offload \
- --usp 4
-```
-
-**Python API:**
-```python
-from vllm_omni import Omni
-
-m = Omni(
- model="/path/to/model",
- enable_distributed_layerwise_offload=True,
- dlo_use_allgather=True, # default
-)
-```
-
-### CLI Flags
-
-| Flag | Description | Default |
-|------|-------------|---------|
-| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
-| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
-| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
-
-### How model weights are loaded (mmap path)
-
-When DLO + AllGather is active, the offloader:
-
-1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
- from the normally-created transformer
-2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
- initialization weights
-3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
- which return views into the OS page cache (shared across all ranks, 0 RSS)
-4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
- (e.g. Cosmos3's `time_embedder` → FP32)
-5. **Restores non-persistent buffers** from saved copies
-
-This approach requires **zero model-specific code changes** — no pipeline or
-transformer modifications are needed.
-
-### OffloadPlan (declarative topology metadata)
-
-Models can optionally declare an `OffloadPlan` class variable to provide
-topology metadata (block attribute names, submodules to offload) without
-any offload-specific logic:
-
-```python
-from vllm_omni.diffusion.offloader import OffloadPlan
-
-class MyPipeline(nn.Module):
- _dit_modules = ["transformer"]
- _offload_plan = OffloadPlan(
- block_attrs={"transformer": ("blocks",)},
- offload_submodules={"context_encoder": "layers"},
- )
-```
-
-When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
-and heuristic attribute search. This is backward-compatible — existing models
-work without any changes.
-
-### DP Multi-concurrency
-
-When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
-up to `dp_size` requests per denoise step. Each DP rank processes a different
-request while AllGather synchronizes only weight shards (request-independent).
-
-**Requirements for concurrent requests:**
-- `num_inference_steps` must be specified explicitly (None is not allowed)
-- All concurrent requests must have the same `num_inference_steps` value
-- These constraints exist because AllGather is a collective that requires
- every rank to participate at each denoise step
-
-### Limitations
-- Online quantization (FP8) is incompatible with mmap loading — falls back to
- regular loading only when the model is outside the DLO+AllGather mmap path;
- online quantization with DLO+AllGather is rejected. Use
- `--dlo-no-use-allgather` or disable online quantization.
-- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
- weight-loader callbacks are bypassed. The no-AllGather path retains the
- standard TP loader and is experimental.
-- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
- with no-AllGather is accepted at configuration level, but has limited
- end-to-end validation.
-- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
-
-**Module Discovery**
-
-The offloader discovers pipeline components in two ways:
-
-1. **Protocol-based** (preferred): If the pipeline implements
- `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
- `_vae_modules`, and `_resident_modules` class variables are used
- directly. All attribute names support dotted paths (e.g.
- `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
-
-2. **Fallback attribute scan**: Otherwise, the offloader scans for
- well-known attribute names:
- - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
- - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
- - **VAE**: `vae`, `audio_vae`
-
-**Hook System**
-
-Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
-
-**Backend Architecture**
-
-```
-OffloadBackend (base class)
-├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
-│ (delegates to a pipeline's enable_omni_model_cpu_offload
-│ for split models like Cosmos3)
-├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
-│ (single-GPU, full weights on host)
-└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
- (multi-GPU, sharded AllGather or rank-local mode)
-```
-
-Factory function `get_offload_backend()` selects the appropriate backend based on
-configuration.
-
-For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
-`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
-reasoner/generator components inside the model forward pass.
-
-
-## Supported Models
-
-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
-| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
-| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
-| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
-| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
-| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
-| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
-| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
-| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
-
-**Notes:**
-- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
-- Distributed Layerwise Offloading reuses the layerwise block topology, but the
- model and parallelism combination still needs validation. AllGather mode
- additionally requires a compatible checkpoint/mmap layout; no-AllGather mode
- retains the standard loader's rank-local tensors. See the
- [DLO design](../../design/feature/distributed_layerwise_offload.md) and
- [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
- for usage examples.
+See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md).
diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md
index 9611a8210..f9fbc1be0 100644
--- a/docs/user_guide/diffusion/execution_modes.md
+++ b/docs/user_guide/diffusion/execution_modes.md
@@ -101,8 +101,16 @@ set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
`--max-num-seqs >1`. See the
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
for its validated configuration. Helios supports single-request step execution only: use
-`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
-documentation and source for the latest support status.
+`--step-execution --max-num-seqs 1` for Helios. MiniMax H3 supports step-wise
+continuous batching by packing co-batched requests into one sequence that keeps
+a separate attention document per request; that layout needs a backend which
+honors the packed `cu_seqlens` metadata, so run it with
+`--diffusion-attention-backend FLASH_ATTN` (other backends stay correct but
+fall back to one transformer forward per request). Batching H3 does not improve
+its throughput — see the measured numbers in the
+[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
+— so keep `--max-num-seqs 1` unless you need step-level scheduling. Consult the
+selected pipeline's documentation and source for the latest support status.
## Streaming Output
@@ -207,18 +215,16 @@ outputs = omni.generate(
)
```
-## Stage Configuration
+## Deploy Configuration
Use the equivalent engine arguments in a deployment YAML:
```yaml
-stage_args:
+stages:
- stage_id: 0
- stage_type: diffusion
- engine_args:
- step_execution: false
- max_num_seqs: 4
- request_batch_max_wait_ms: 20
+ step_execution: false
+ max_num_seqs: 4
+ request_batch_max_wait_ms: 20
```
For step execution, set `step_execution: true` and remove
diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md
index 44705129f..3c8582af0 100644
--- a/docs/user_guide/diffusion/lora.md
+++ b/docs/user_guide/diffusion/lora.md
@@ -12,15 +12,17 @@ LoRA adapters are lightweight, model-specific fine-tuning weights that can be ap
## LoRA Adapter Format
### PEFT (Parameter-Efficient Fine-Tuning) format (default)
+
A typical PEFT-format LoRA adapter directory structure:
-```
+```text
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
The `adapter_config.json` file contains metadata about the LoRA adapter, including:
+
- `r`: LoRA rank
- `lora_alpha`: LoRA alpha scaling factor
- `target_modules`: List of module names to apply LoRA to
@@ -67,7 +69,7 @@ For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or
For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
| Pipeline | Supported distilled LoRA repo | Notes |
-|----------|--------------------------------|-------|
+| ---------- | -------------------------------- | ------- |
| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
@@ -240,6 +242,37 @@ Notes:
- Output quality and speed depend on the replacement checkpoints and sampling params you choose.
- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path .safetensors .safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
+## MiniMax-H3 adapter-declared schedules
+
+MiniMax-H3 supports two few-step mechanisms that must not be conflated:
+
+- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
+ `model_index.json`. Requests must pass `num_inference_steps` as the interval
+ count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
+- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and
+ enforce `flow_shift=6`, `audio_flow_shift=3`.
+- **Runtime native LoRA**: FlashGen-style artifacts declare
+ `key_format=minimax-h3-native` and embed `base_schedule` in safetensors
+ metadata. When active, the adapter schedule overrides the base checkpoint and
+ requests must use the interval-count contract (`num_inference_steps=4`).
+ Request-mode generation may omit the field and take the count from the adapter
+ schedule; step execution requires it explicitly, because the step scheduler
+ derives the total step count from the request before the adapter schedule is
+ known.
+
+Native artifacts also declare `qkv_layout=grouped`. The H3 loader reorders fused
+`qkv_proj` LoRA rows with the same `_reorder_grouped_qkv_to_qkv` helper used for
+base-weight loading, then binds the packed Q/K/V slices through the legacy PEFT
+manager without modifying `DiffusionLoRAManager`.
+
+Both runtime adapters run with distributed layerwise offload in request-mode
+generation, where the manager keeps the LoRA A/B buffers resident on the compute
+device while DLO streams the base blocks. MiniMax-H3 step execution still
+rejects DLO, so `--step-execution` cannot be combined with
+`--enable-distributed-layerwise-offload`. Model-level CPU offload and standard
+layerwise offload are rejected in every mode, because the dynamic LoRA tensors
+are neither parameters nor registered buffers and therefore do not participate
+in those weight lifecycles.
## See Also
diff --git a/docs/user_guide/diffusion/offloader/cpu_offload.md b/docs/user_guide/diffusion/offloader/cpu_offload.md
new file mode 100644
index 000000000..558985e66
--- /dev/null
+++ b/docs/user_guide/diffusion/offloader/cpu_offload.md
@@ -0,0 +1,81 @@
+# Model-Level Offloading
+
+Model-level, or sequential, offloading keeps only the pipeline component group
+currently executing on the accelerator. It is the simplest offload strategy
+and is selected with `--enable-cpu-offload`.
+
+## How it works
+
+Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
+
+- before an encoder runs, the DiT moves to CPU;
+- before a DiT runs, encoders and other DiTs move to CPU; and
+- VAE modules remain on the accelerator.
+
+Pinned host memory reduces transfer overhead. Transfers occur at phase
+boundaries, so cold-start and encoder-to-denoiser transitions become slower.
+
+## Usage
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+ model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+ enable_cpu_offload=True,
+)
+```
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+ --omni --enable-cpu-offload
+```
+
+## Model integration
+
+Pipelines should implement `SupportsComponentDiscovery`:
+
+```python
+from typing import ClassVar
+
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+
+
+class MyPipeline(nn.Module, SupportsComponentDiscovery):
+ _dit_modules: ClassVar[list[str]] = ["transformer"]
+ _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
+ _vae_modules: ClassVar[list[str]] = ["vae"]
+ _resident_modules: ClassVar[list[str]] = []
+```
+
+All entries may be dotted paths. DiT and encoder lists are both required for
+mutual exclusion. VAE modules are pinned but not swapped; resident modules are
+small modules that must stay on the accelerator for layerwise paths.
+
+## Split-model components
+
+Some models, such as Cosmos3, split one transformer into mutually exclusive
+components that run in different phases. The pipeline exposes
+`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
+contexts:
+
+```python
+class Cosmos3VFMTransformer(nn.Module):
+ def forward(self, ...):
+ with self._offload_context("reasoner"):
+ ...
+ with self._offload_context("generator"):
+ ...
+```
+
+This preserves the same invariant—exactly one component is device resident—
+while reusing sequential `.to()` movers.
+
+## Limitations
+
+- Single device only.
+- Higher cold-start latency.
+- Transfers between encoder and denoising phases add latency.
+
+See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
+for lifecycle and extension invariants.
diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
new file mode 100644
index 000000000..45ff8df2d
--- /dev/null
+++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
@@ -0,0 +1,285 @@
+# Distributed Layerwise Offloading
+
+Distributed layerwise offloading (DLO) extends block streaming to multi-device
+deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
+the host weights and reconstructs each layer at runtime. Without AllGather,
+each rank streams a complete block independently. Compatible TP1 deployments
+can share checkpoint-backed host pages among processes on the same node;
+otherwise DLO streams the ordinary loader's rank-local tensors.
+
+See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
+for the implementation contract and compatibility matrix.
+
+## Execution model
+
+DLO overlaps three operations with a fixed two-block device buffer:
+
+```text
+Compute stream: [Layer N] [Layer N+1] [Layer N+2]
+H2D stream: [H2D shard N+1] [H2D shard N+2]
+AllGather: [AG N+1] [AG N+2]
+Slots: slot 0: Layer N slot 1: Layer N+1
+```
+
+AllGather communicates only request-independent weight shards, so data-
+parallel ranks may process different requests concurrently.
+
+## Usage
+
+```bash
+# Four ranks with sharded host weights and AllGather
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --data-parallel-size 4
+
+# Standard-loader rank-local weights, without DLO AllGather
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --data-parallel-size 4 \
+ --dlo-no-use-allgather
+
+# Sequence parallel deployment
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --usp 4
+```
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+ model="/path/to/model",
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=True,
+)
+```
+
+## Flags
+
+| Flag | Meaning | Default |
+| --- | --- | --- |
+| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
+| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
+| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
+| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` |
+| `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset |
+| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` |
+
+## Host-weight loading
+
+The diffusion loader chooses host storage before DLO is enabled. It first
+attempts to build a complete, validated direct-checkpoint mmap plan. If names,
+coverage, shape, dtype, topology, or loader-callback compatibility cannot be
+proven, it runs the ordinary model loader instead. DLO consumes that result and
+does not make a second checkpoint-compatibility decision.
+
+The shared-mmap optimization in this phase is supported only with TP1. TP
+greater than one falls back before model mutation to ordinary TP-aware loading.
+DLO may still consume those TP-local tensors, but this is a compatibility path:
+it does not share checkpoint-backed runtime weights across DP replicas and
+provides no shared-mmap host-memory guarantee.
+
+The mmap plan skips only dedicated DiT weight sources. Other component sources,
+such as a text encoder loaded through the shared diffusion loader, continue to
+use their ordinary component loader. A checkpoint source that mixes DiT and
+non-DiT weights falls back completely rather than leaving an unplanned
+component uninitialized.
+
+With direct checkpoint mmap, the loader:
+
+1. saves non-persistent buffers such as RoPE frequencies;
+2. moves the normally created transformer to the meta device;
+3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
+4. applies any loader-owned bounded layout adapters while packing blocks;
+5. restores saved non-persistent buffers; and
+6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
+ hooks.
+
+For AllGather with a group larger than one, each process copies only its
+persistent shard and then releases the source mapping. For no-AllGather, each
+process keeps the mapping open and packs complete blocks through two bounded
+pinned staging slots. Processes mapping the same files on one node share the
+immutable pages; no-AllGather still performs a complete-block H2D copy in each
+process.
+
+When the effective DLO group size is one, `dlo_use_allgather=True` does not
+perform a collective and uses the same rank-local transfer behavior.
+
+### Final-layout Host Weight Runtime
+
+HWR is an opt-in startup optimization for models that declare the final-layout
+BF16 restore contract. Use it only with no-AllGather DLO:
+
+The validated BF16 model contracts currently cover MiniMax H3 and
+`black-forest-labs/FLUX.2-klein-4B`. FLUX.2-klein-9B shares the same model
+class but has not been validated against this contract. FLUX.2-dev, online
+FP8, HSDP, LoRA/adapted weights, and non-default load formats remain outside
+the validated scope.
+
+```bash
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --dlo-no-use-allgather \
+ --host-weight-runtime-mode preferred \
+ --host-weight-runtime-root /var/cache/vllm-omni/hwr
+```
+
+#### Choosing a mode
+
+| Mode | Exact local artifact hit | Miss or recoverable artifact/store problem | Intended use |
+| --- | --- | --- | --- |
+| `disabled` | HWR is not consulted | Use the existing checkpoint-mmap or ordinary-loader path | Default compatibility path |
+| `preferred` | Restore the final-layout artifact | Load canonically, serve with those tensors, and attempt to publish an artifact for the next startup | Normal deployment and store population |
+| `required` | Restore the final-layout artifact | Fail startup without canonical DiT fallback or publication | Enforce a pre-populated store in controlled rollouts or CI |
+
+Both enabled modes still fail on non-retryable configuration, identity, or
+compatibility errors. `preferred` is a fallback policy for a cache miss or a
+recoverable cache problem; it does not hide an invalid deployment.
+
+#### Populating a store for `required` mode
+
+`required` is deliberately consume-only. PR2 does not include a separate
+prewarm command, so populate each node-local storage domain with one matching
+`preferred` producer cohort:
+
+1. Choose a persistent, writable root visible to every diffusion worker in the
+ storage domain. Do not use a process-private temporary directory.
+2. Start the deployment with `preferred`, using the exact model revision,
+ dtype, TP size, SP configuration, and other weight-layout settings intended
+ for serving. No inference request is needed; publication happens during
+ model startup. Wait for the engine to become healthy, then shut it down
+ cleanly.
+3. Restart with the same arguments and root, changing only the mode to
+ `required`. A successful startup proves that every worker acquired a valid
+ artifact; a miss, corrupt artifact, or incompatible identity fails startup.
+4. Repeat the population start on every node or storage domain because this
+ store is node-local.
+
+For example:
+
+```bash
+# First startup: canonically load and populate the exact artifacts.
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --dlo-no-use-allgather \
+ --host-weight-runtime-mode preferred \
+ --host-weight-runtime-root /var/cache/vllm-omni/hwr
+
+# After a healthy startup and clean shutdown, enforce cache hits.
+vllm serve /path/to/model --omni \
+ --enable-distributed-layerwise-offload \
+ --dlo-no-use-allgather \
+ --host-weight-runtime-mode required \
+ --host-weight-runtime-root /var/cache/vllm-omni/hwr
+```
+
+Include the same TP and SP layout flags in both commands. DP rank and DP size
+are excluded from artifact identity, so the population and serving DP sizes may
+differ and equivalent DP replicas share artifacts. TP rank is included, so a
+TP-N deployment normally needs N rank-specific artifacts in each storage
+domain; launching the matching TP cohort creates that set. If the model
+revision or layout changes, run `preferred` again for the new identity before
+returning to `required`.
+
+On a cold start, the canonical loader remains authoritative and publishes a
+validated final-layout artifact for later workers. A warm start restores the
+DiT final tensors without ordinary DiT materialization. DLO then attempts to
+register the immutable mapping for direct H2D; if registration is unavailable,
+it streams through the same two bounded host staging slots. The artifact
+identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards,
+and distinct SP layouts do not alias one another. Publication failure remains
+separate from a `preferred` serving startup; the next `required` startup
+provides the explicit artifact-availability check.
+
+For local canonical checkpoints, the first eligible worker may hash source
+shards to establish immutable identity. HWR caches those digests in the same
+node-local storage domain and validates file metadata before reuse, so later
+workers normally avoid repeating that read. Cold BF16 publication also hashes
+ordered payloads as they are written and overlaps payload durability work with
+later shards; this changes startup work only, not artifact contents or runtime
+H2D behavior.
+
+#### Registered direct H2D
+
+Registration is attempted automatically only for an eligible warm HWR hit when
+`pin_cpu_memory` is enabled. A successful path is:
+
+```text
+shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel
+```
+
+It removes the recurrent mmap-to-private-staging CPU copy and does not allocate
+the two host staging slots. It does not reduce the H2D payload or make GPU
+kernels access host memory directly.
+
+CUDA requires read-only host-registration capability for the immutable HWR
+mapping. Unsupported capability, a positive registration limit smaller than
+the complete page-aligned mapping, or a safely rolled-back registration error
+falls back to bounded staging. Programmatic configurations can set
+`pin_cpu_memory=False` to disable registration explicitly. A successful
+registration locks the complete mapped range in host memory for that worker's
+lifetime, so use
+`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is
+required.
+
+Shutdown drains H2D work, releases hook/source references, unregisters every
+range, and only then closes the HWR lease. AllGather and direct checkpoint mmap
+retain their existing transfer paths.
+
+When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does
+not resolve HWR sources or construct its store.
+
+## Declarative topology
+
+Models may declare an `OffloadPlan` instead of embedding offload logic:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+
+class MyPipeline(nn.Module):
+ _dit_modules = ["transformer"]
+ _offload_plan = OffloadPlan(
+ block_attrs={"transformer": ("blocks",)},
+ offload_submodules={"context_encoder": "layers"},
+ )
+```
+
+When no plan exists, discovery falls back to
+`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
+
+## Data-parallel concurrency
+
+With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
+up to `dp_size` requests per denoising step. Every concurrent request must set
+the same explicit `num_inference_steps`; `None` is rejected because every rank
+must enter each collective.
+
+## Limitations
+
+- Direct checkpoint mmap currently requires TP1. TP greater than one is
+ outside the Phase A shared-mmap support scope and falls back before model
+ mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
+ and eligible no-AllGather configurations may use HWR final-layout artifacts,
+ but direct checkpoint mmap provides no shared-mmap host-memory guarantee for
+ TP greater than one.
+- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
+ AllGather has limited end-to-end validation.
+- Per-tensor online FP8 linears use the ordinary loader and can run with either
+ DLO transfer path. With AllGather, every rank temporarily materializes the
+ complete FP8 model in host memory before DLO retains only its shard. Other
+ online quantization methods require no-AllGather until their runtime layouts
+ are validated.
+- Resident leading layers require `--dlo-no-use-allgather` and a model
+ `OffloadPlan` that declares eligible `resident_dit_paths`.
+- DP concurrency requires an explicit, identical inference-step count.
+
+Sharing quantized or otherwise unvalidated transformed layouts through a
+normalized HWR producer is a follow-up design in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
+the current BF16 final-layout path.
+
+See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
+for an end-to-end example.
diff --git a/docs/user_guide/diffusion/offloader/layerwise_offload.md b/docs/user_guide/diffusion/offloader/layerwise_offload.md
new file mode 100644
index 000000000..c88b67fcc
--- /dev/null
+++ b/docs/user_guide/diffusion/offloader/layerwise_offload.md
@@ -0,0 +1,65 @@
+# Layerwise Offloading
+
+Layerwise, or blockwise, offloading keeps one transformer block on the
+accelerator and prefetches the next block while the current block computes.
+It is best suited to compute-heavy video DiTs whose block execution can hide
+host-to-device transfers.
+
+## Execution flow
+
+Each block has a pre-forward and post-forward hook. Parameters are consolidated
+in pinned host tensors and rematerialized for execution on a dedicated copy
+stream.
+
+| Block | Pre-forward hook | Forward | Post-forward hook |
+| --- | --- | --- | --- |
+| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
+| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
+| ... | ... | ... | ... |
+| last block | Prefetch block 0 | Compute last block | Free last block |
+
+Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
+remain device resident.
+
+## Usage
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+ model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+ enable_layerwise_offload=True,
+)
+```
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+ --omni --enable-layerwise-offload
+```
+
+## Model integration
+
+Transformer classes declare containers of executable blocks:
+
+```python
+class WanTransformer3DModel(nn.Module):
+ _layerwise_offload_blocks_attrs = ["blocks"]
+
+
+class Flux2Transformer2DModel(nn.Module):
+ _layerwise_offload_blocks_attrs = [
+ "transformer_blocks",
+ "single_transformer_blocks",
+ ]
+```
+
+See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
+the discovery and hook invariants. `OffloadPlan` is a separate declarative
+topology path for distributed layerwise offload.
+
+## Limitations
+
+- Single device only.
+- Setup consolidates and pins block parameters, increasing cold-start time.
+- Performance depends on block compute time and host-to-device bandwidth;
+ lightweight blocks may not hide transfers.
diff --git a/docs/user_guide/diffusion/parallelism/hsdp.md b/docs/user_guide/diffusion/parallelism/hsdp.md
index 96a357c86..9491ed9f5 100644
--- a/docs/user_guide/diffusion/parallelism/hsdp.md
+++ b/docs/user_guide/diffusion/parallelism/hsdp.md
@@ -64,6 +64,41 @@ omni = Omni(
)
```
+### Combined with Online FP8 Quantization
+
+HSDP can be combined with online FP8 quantization for diffusion models that
+support both features. HSDP shards the quantized weights across GPUs to reduce
+the per-GPU weight memory footprint.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+ model="",
+ quantization="fp8",
+ parallel_config=DiffusionParallelConfig(
+ use_hsdp=True,
+ hsdp_shard_size=2,
+ ),
+)
+```
+
+For online serving:
+
+```bash
+vllm serve \
+ --omni \
+ --quantization fp8 \
+ --use-hsdp \
+ --hsdp-shard-size 2
+```
+
+The model must support both HSDP and online FP8 quantization. HSDP remains
+incompatible with Tensor Parallelism. See the
+[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and
+hardware.
+
---
## Example Script
@@ -117,7 +152,8 @@ In `DiffusionParallelConfig`:
**Constraints:**
- `hsdp_replicate_size × hsdp_shard_size == world_size`
-- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+- HSDP cannot be used with Tensor Parallelism, ordinary Data Parallelism,
+ Pipeline Parallelism, or Expert Parallelism.
---
@@ -147,3 +183,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co
1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
+4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint
diff --git a/docs/user_guide/diffusion/regional_compilation.md b/docs/user_guide/diffusion/regional_compilation.md
index 469a62b2c..448f19518 100644
--- a/docs/user_guide/diffusion/regional_compilation.md
+++ b/docs/user_guide/diffusion/regional_compilation.md
@@ -24,7 +24,7 @@ stages:
For an experimental whole-transformer compile scope, set
`--diffusion-compile-granularity full` or use
-`diffusion_compile_granularity: full` in the stage configuration. Full scope may
+`diffusion_compile_granularity: full` in the deploy configuration. Full scope may
still contain graph breaks; it does not force one graph. It is rejected when
HSDP, sequence parallelism, CPU offload, or layerwise offload is enabled. Use
regional scope with those features.
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 7c664e88b..c0b29cf6e 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -106,7 +106,10 @@ The following tables show which models support each feature:
> Notes:
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
+> 1. CPU Offload has three strategies: model-level, layerwise, and distributed
+> layerwise. The tables below show **layerwise support** only. Split models
+> like Cosmos3 swap their reasoner/generator components for model-level
+> offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.
### ImageGen
@@ -121,11 +124,11 @@ The following tables show which models support each feature:
| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **HiDream-O1-Image** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅* |
| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -144,7 +147,7 @@ The following tables show which models support each feature:
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
> 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
@@ -158,20 +161,29 @@ The following tables show which models support each feature:
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
+| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
-| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **LongCat-Video-Avatar-1.5** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MiniMax-H3** | ✅ (FL2VA) | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
+| **SANA-WM** | ❌ | ❌ | ❌5 | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+
+> Notes:
+> 5. SANA-WM cannot support sequence parallelism: its bidirectional gated delta
+> recurrence carries state across frames, so a rank cannot denoise a slice of
+> the token sequence in isolation. Doing so would need a distributed scan or
+> an all-gather before every GDN block. The remaining ❌ columns are simply
+> unvalidated on this model, not known-broken.
> **Step execution note:** Helios supports single-request step execution only;
> use `max_num_seqs=1`.
**Frame Interpolation Support**
-- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
-- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines; SANA-Video-2B native text-to-video and image-to-video pipelines
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM; SANA-Video-2B Diffusers-adapter text-to-video and image-to-video pipelines
### AudioGen
@@ -214,7 +226,7 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. The CPU Offloading (Layerwise) row describes local layerwise offload.
Multi-device Distributed Layerwise Offload has a separate topology and
- compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
+ compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md).
5. The compatibility matrix uses FP8 as the representative quantization method.
6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
diff --git a/docs/user_guide/examples/offline_inference/glm_image.md b/docs/user_guide/examples/offline_inference/glm_image.md
index 09d927ff7..267819cc5 100644
--- a/docs/user_guide/examples/offline_inference/glm_image.md
+++ b/docs/user_guide/examples/offline_inference/glm_image.md
@@ -43,7 +43,7 @@ if __name__ == "__main__":
),
],
)
- outputs[0].request_output.images[0].save("output.png")
+ outputs[0].images[0].save("output.png")
```
## Image-to-Image (Image Editing)
@@ -72,7 +72,7 @@ if __name__ == "__main__":
),
],
)
- outputs[0].request_output.images[0].save("output.png")
+ outputs[0].images[0].save("output.png")
```
## Generation Parameters
diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md
index db7a31ce5..ff57063b8 100644
--- a/docs/user_guide/examples/offline_inference/image_to_video.md
+++ b/docs/user_guide/examples/offline_inference/image_to_video.md
@@ -4,8 +4,8 @@ Source **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md
index f631cd962..24cab3a3c 100644
--- a/docs/user_guide/examples/offline_inference/text_to_image.md
+++ b/docs/user_guide/examples/offline_inference/text_to_image.md
@@ -55,7 +55,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
- images = outputs[0].request_output.images
+ images = outputs[0].images
images[0].save("coffee.png")
```
@@ -179,7 +179,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
- output.request_output.images[0].save(f"{i}.jpg")
+ output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -207,7 +207,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
- output.request_output.images[0].save(f"{i}.jpg")
+ output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md
index ad7777014..a6544c1e5 100644
--- a/docs/user_guide/examples/offline_inference/text_to_speech.md
+++ b/docs/user_guide/examples/offline_inference/text_to_speech.md
@@ -69,7 +69,7 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
### Notes
- Output: 48 kHz mono WAV.
-- Stage config: `vllm_omni/model_executor/stage_configs/voxcpm2.yaml` (default).
+- Deploy config: `vllm_omni/deploy/voxcpm2.yaml` (default).
---
@@ -152,7 +152,7 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
--text "Hello, this is a streaming test." \
--streaming
```
-Streaming requires `async_chunk: true` in the stage config.
+Streaming requires `async_chunk: true` in the deploy config.
### Notes
- Output: 44.1 kHz mono WAV.
@@ -274,7 +274,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
--streaming \
--output-dir /tmp/out_stream
```
-Streaming requires `async_chunk: true` in the stage config.
+Streaming requires `async_chunk: true` in the deploy config.
### Batched decoding
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`:
diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md
index 1a2082f82..ec61bc749 100644
--- a/docs/user_guide/examples/offline_inference/text_to_video.md
+++ b/docs/user_guide/examples/offline_inference/text_to_video.md
@@ -16,6 +16,8 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
+| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE |
+| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | ~23.6 GiB |
## Local CLI Usage
@@ -122,6 +124,31 @@ python text_to_video.py \
--output quick_test.mp4
```
+### SANA-Video-2B
+
+```bash
+python text_to_video.py \
+ --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ --model-class-name SanaVideoPipeline \
+ --prompt "A cinematic tracking shot of a sailboat crossing the ocean at sunset." \
+ --height 480 \
+ --width 832 \
+ --num-frames 81 \
+ --num-inference-steps 50 \
+ --guidance-scale 6.0 \
+ --extra-body '{"motion_score": 30}' \
+ --fps 16 \
+ --output sana_video_480p.mp4
+```
+
+For the 720p checkpoint, switch the model to
+`Efficient-Large-Model/SANA-Video_2B_720p_diffusers`, use
+`--height 704 --width 1280`, and write to a different output path. SANA-Video
+accepts `motion_score`, `clean_caption`, and `use_resolution_binning` through
+`--extra-body`. For image-to-video, use the shared
+`examples/offline_inference/image_to_video/image_to_video.py` example with
+`--model-class-name SanaImageToVideoPipeline` and `--image `.
+
## Key Arguments
### Common
diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md
deleted file mode 100644
index d5cd0ea46..000000000
--- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md
+++ /dev/null
@@ -1,115 +0,0 @@
-# X-To-Video-Audio
-
-Source .
-
-
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
-
-## Local CLI Usage
-### Download the Model locally
-Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
-
-```bash
-python download_dreamid_omni.py --output-dir ./dreamid_omni
-```
-After download, the model directory will look like this:
-
-```
-dreamid_omni/
-├── DreamID-Omni/
-│ ├── dreamid_omni.safetensors
-├── MMAudio/
-│ ├── ext_weights/
-│ │ ├── best_netG.pt
-│ │ ├── v1-16.pth
-├── Wan2.2-TI2V-5B/
-│ ├── google/*
-│ ├── models_t5_umt5-xxl-enc-bf16.pth
-│ ├── Wan2.2_VAE.pth
-│
-├── model_index.json # create by download_dreamid_omni.py
-```
-
-### Run the Inference
-```python
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt "Two people walking together and singing happily" \
- --image-path ./example0.png ./example1.png \
- --audio-path ./example0.wav ./example1.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --output out_dreamid_omni_twoip.mp4
-```
-In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
-The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt "Two people walking together and singing happily" \
- --image-path ./example0.png ./example1.png \
- --audio-path ./example0.wav ./example1.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --use-hsdp \
- --hsdp-shard-size 2 \
- --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \
- --image-path 9.png \
- --audio-path 9.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --output out_dreamid_omni_oneip.mp4
-```
-
-
-Key arguments:
-- `--prompt`: text description (string).
-- `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
-- `--image-path`: path to the input image list.
-- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
-- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--num-inference-steps`: number of denoising steps (defaults 45).
-- `--video-negative-prompt`: negative prompt for video generation.
-- `--audio-negative-prompt`: negative prompt for audio generation.
-- `--enable-cpu-offload`: enable CPU offload (defaults False).
-
-## Example materials
-
-??? abstract "download_dreamid_omni.py"
- ``````py
- --8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
- ``````
-??? abstract "x_to_video_audio.py"
- ``````py
- --8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
- ``````
diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md
index b4c05bdbe..f92463f84 100644
--- a/docs/user_guide/examples/online_serving/qwen3_omni.md
+++ b/docs/user_guide/examples/online_serving/qwen3_omni.md
@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file
```
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+ --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
### Launch individual stages (stage-based CLI)
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
@@ -271,7 +283,7 @@ The script supports the following arguments:
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
-If you have custom stage configs file:
+If you have a custom deploy config file:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
```
diff --git a/docs/user_guide/examples/online_serving/text_to_audio.md b/docs/user_guide/examples/online_serving/text_to_audio.md
deleted file mode 100644
index 1b6dbc463..000000000
--- a/docs/user_guide/examples/online_serving/text_to_audio.md
+++ /dev/null
@@ -1,193 +0,0 @@
-# Text-To-Audio
-
-Source .
-
-This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
-
-## Supported Models
-
-| Model | Description |
-|-------|-------------|
-| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
-
-## Start Server
-
-### Basic Start
-
-```bash
-vllm serve stabilityai/stable-audio-open-1.0 \
- --host 0.0.0.0 \
- --port 8091 \
- --gpu-memory-utilization 0.9 \
- --trust-remote-code \
- --enforce-eager \
- --omni
-```
-
-## API Calls
-
-### Method 1: Using curl
-
-```bash
-# Run all curl examples
-bash curl_examples.sh
-
-# Or execute directly
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "The sound of a cat purring",
- "audio_length": 10.0
- }' --output cat.wav
-```
-
-### Method 2: Using Python Client
-
-```bash
-cd examples/online_serving/stable_audio
-
-# Simple generation
-python stable_audio_client.py \
- --text "The sound of a cat purring"
-
-# With custom duration
-python stable_audio_client.py \
- --text "A dog barking" \
- --audio_length 5.0
-
-# With all parameters
-python stable_audio_client.py \
- --text "Thunder and rain" \
- --audio_length 15.0 \
- --negative_prompt "Low quality" \
- --guidance_scale 7.0 \
- --num_inference_steps 100 \
- --seed 42 \
- --output thunder.wav
-```
-
-The Python client supports the following command-line arguments:
-
-- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
-- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
-- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
-- `--audio_start`: Audio start time in seconds (default: `0.0`)
-- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
-- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
-- `--num_inference_steps`: Number of inference steps (default: `100`)
-- `--seed`: Random seed for reproducibility (default: `None`)
-- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
-- `--output`: Output file path (default: `stable_audio_output.wav`)
-
-### Method 3: Using Python httpx
-
-```python
-import httpx
-
-response = httpx.post(
- "http://localhost:8091/v1/audio/generate",
- json={
- "input": "The sound of ocean waves crashing on a beach",
- "audio_length": 10.0,
- "negative_prompt": "Low quality, distorted",
- "guidance_scale": 7.0,
- "num_inference_steps": 100,
- },
- timeout=300.0,
-)
-
-with open("ocean.wav", "wb") as f:
- f.write(response.content)
-```
-
-## Request Format
-
-### Simple Generation
-
-```json
-{
- "input": "The sound of ocean waves"
-}
-```
-
-### Generation with Parameters
-
-```json
-{
- "input": "A piano playing a gentle melody",
- "audio_length": 10.0,
- "negative_prompt": "Low quality, distorted, noisy",
- "guidance_scale": 8.0,
- "num_inference_steps": 150,
- "seed": 42,
- "response_format": "wav"
-}
-```
-
-## API Reference
-
-### Endpoint
-
-```
-POST /v1/audio/generate
-Content-Type: application/json
-```
-
-### Generation Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text prompt describing the audio to generate |
-| `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
-| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
-| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
-| `audio_start` | float | 0.0 | Audio start time in seconds |
-| `negative_prompt` | string | null | Text describing what to avoid in generation |
-| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
-| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
-| `seed` | int | null | Random seed for reproducible generation |
-
-### Response Format
-
-Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
-
-## Tuning Tips
-
-1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
-2. **Quality vs Speed**:
- - 50 steps: Fast, decent quality (quick previews)
- - 100 steps: Good balance (general purpose)
- - 150+ steps: High quality, slower (final / critical audio)
-3. **Guidance Scale**:
- - Lower (3 - 5): More creative / varied output
- - Default (7): Good balance
- - Higher (10+): Strict adherence to the prompt
-4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
-5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
-
-## File Description
-
-| File | Description |
-|------|-------------|
-| `curl_examples.sh` | Curl examples covering common use cases |
-| `stable_audio_client.py` | Python client with full CLI argument support |
-
-## Troubleshooting
-
-1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
-2. **Connection refused**: Make sure the server is running on the correct port.
-3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
-4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
-5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
-
-## Example materials
-
-??? abstract "stable_audio_client.py"
- ``````py
- --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
- ``````
-??? abstract "curl_examples.sh"
- ``````sh
- --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
- ``````
diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md
index c6a613723..5d4ccbf3b 100644
--- a/docs/user_guide/examples/online_serving/text_to_speech.md
+++ b/docs/user_guide/examples/online_serving/text_to_speech.md
@@ -419,7 +419,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
"response_format": "pcm"
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
-Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
+Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
### Streaming WebSocket
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
diff --git a/docs/user_guide/fault_injection_reliability_matrix.md b/docs/user_guide/fault_injection_reliability_matrix.md
index 28271affa..4a2fa0a16 100644
--- a/docs/user_guide/fault_injection_reliability_matrix.md
+++ b/docs/user_guide/fault_injection_reliability_matrix.md
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
## Source of Conclusions
-The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
+The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, and **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`.
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
diff --git a/docs/user_guide/feature_compatibility.md b/docs/user_guide/feature_compatibility.md
index ea05fe81e..eeb087510 100644
--- a/docs/user_guide/feature_compatibility.md
+++ b/docs/user_guide/feature_compatibility.md
@@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
-- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
+- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level.
- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md
index ddca3ea19..010ef4ba7 100644
--- a/docs/user_guide/quantization/fp8.md
+++ b/docs/user_guide/quantization/fp8.md
@@ -80,20 +80,22 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
## Model Type Support
-### Diffusion Model (Qwen-Image, Wan2.2)
+### Diffusion Models
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
+| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
+| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
-| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan deploy config for multi-stage runs | None | |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
-| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
+| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -136,6 +138,39 @@ outputs = omni.generate(
)
```
+### Global and per-component scope
+
+A plain, global quantization configuration is passed unchanged to every
+quantization-aware component constructed by a pipeline. This includes an
+eligible encoder when that encoder is implemented with vLLM quantizable
+layers; it does not rewrite arbitrary `torch.nn` modules. A structured
+component map is the only way to narrow that scope. Pipeline integrations that
+do not yet expose an encoder through the quantization factory remain DiT-only.
+
+For a pipeline that exposes both a transformer and a quantization-aware text
+encoder, the scope is:
+
+| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
+|---------------|-------------|--------------|-------------------------------------------------|
+| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
+| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
+| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
+
+Use `quantization_config` for component-selective Python configuration. For
+example, quantize only the text decoder:
+
+```python
+omni = Omni(
+ model="",
+ quantization_config={"text_encoder": {"method": "fp8"}},
+)
+```
+
+Component keys are runtime prefixes exposed by the pipeline integration; common
+keys include `transformer` and `text_encoder`. Entries may be combined, and
+`ignored_layers` can keep named eligible layers in checkpoint precision. Check
+the model recipe for supported components and their runtime prefixes.
+
CLI:
```bash
@@ -144,6 +179,24 @@ python text_to_image.py --model --quantization fp8 --ignored-layers
vllm serve --omni --quantization fp8
```
+### Distributed layerwise offload
+
+Per-tensor online FP8 linears can be combined with DLO's default AllGather path:
+
+```bash
+vllm serve \
+ --omni \
+ --quantization fp8 \
+ --enable-distributed-layerwise-offload
+```
+
+Direct checkpoint mmap is not used for this combination. Each rank runs the
+ordinary loader to create the final FP8 weights and scales before DLO shards
+them, so startup temporarily materializes the complete FP8 model in host memory
+on every rank. Runtime DLO host residency is sharded normally. Use
+`--dlo-no-use-allgather` when the synchronized AllGather request-wave contract
+does not fit the workload.
+
## Parameters
| Parameter | Type | Default | Description |
diff --git a/docs/user_guide/quantization/msmodelslim.md b/docs/user_guide/quantization/msmodelslim.md
index c03628b9d..ea42f6250 100644
--- a/docs/user_guide/quantization/msmodelslim.md
+++ b/docs/user_guide/quantization/msmodelslim.md
@@ -92,4 +92,4 @@ For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
2. The `ascend` quantization method expects weights produced by the Ascend
tooling; it is not a load-time CUDA quantizer.
3. Keep the quantized checkpoint aligned with the same model architecture and
- stage config used for BF16 inference.
+ pipeline and deploy config used for BF16 inference.
diff --git a/docs/user_guide/quantization/quantized_kvcache.md b/docs/user_guide/quantization/quantized_kvcache.md
index e560032d3..99a90e08a 100644
--- a/docs/user_guide/quantization/quantized_kvcache.md
+++ b/docs/user_guide/quantization/quantized_kvcache.md
@@ -83,23 +83,23 @@ Online serving:
vllm serve --omni --diffusion-kv-cache-dtype fp8
```
-Stage config:
+Deploy config:
```yaml
-stage_args:
+stages:
- stage_id: 0
- stage_type: diffusion
- engine_args:
- model_stage: dit
- diffusion_kv_cache_dtype: "fp8"
- diffusion_kv_cache_skip_steps: "0,1"
- diffusion_kv_cache_skip_layers: "0-2"
+ diffusion_kv_cache_dtype: "fp8"
+ diffusion_kv_cache_skip_steps: "0,1"
+ diffusion_kv_cache_skip_layers: "0-2"
```
-Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and
-`kv_cache_skip_layers` are still accepted when constructing
-`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*`
-names for new configs.
+The `model_stage` and diffusion execution type belong to the registered
+`PipelineConfig`; the deploy YAML only carries runtime overrides.
+
+The legacy keyword aliases `kv_cache_dtype`, `kv_cache_skip_steps`, and
+`kv_cache_skip_layers` remain accepted when constructing
+`OmniDiffusionConfig` directly. They are not deploy YAML fields; prefer the
+`diffusion_*` names for new code.
## Parameters
diff --git a/docs/user_guide/quantization/svdquant.md b/docs/user_guide/quantization/svdquant.md
new file mode 100644
index 000000000..4fbd8040d
--- /dev/null
+++ b/docs/user_guide/quantization/svdquant.md
@@ -0,0 +1,50 @@
+# SVDQuant W4A4
+
+## Overview
+
+[SVDQuant](https://arxiv.org/abs/2411.05007) combines four-bit weights
+and activations with a small low-rank branch that corrects part of the
+quantization error. vLLM-Omni consumes an offline-quantized checkpoint;
+it does not calibrate the model while loading.
+
+## Checkpoint contract
+
+Place the following entry in the diffusion transformer's `config.json`:
+
+```json
+{
+ "quantization_config": {
+ "quant_method": "svdquant",
+ "rank": 32,
+ "precision": "nvfp4",
+ "act_unsigned": false,
+ "modules_to_not_convert": []
+ }
+}
+```
+
+For a quantized linear with input size `K`, output size `N`, and correction
+rank `R`, the checkpoint stores:
+
+| Suffix | Shape | dtype |
+| --- | --- | --- |
+| `qweight` | `(N, K / 2)` | `int8` (two packed FP4 values per byte) |
+| `wscales` | `(K / 16, N)` | `float8_e4m3fn` |
+| `proj_down` | `(K, R)` | `bfloat16` |
+| `proj_up` | `(N, R)` | `bfloat16` |
+| `smooth_factor` | `(K,)` | `bfloat16` |
+| `wcscales` | `(N,)` | `bfloat16` |
+| `wtscale` | `(1,)` | `bfloat16` |
+
+`K` must be divisible by 16 on every tensor-parallel rank. Set `wcscales` to
+ones when no per-output correction is needed. Modules listed in
+`modules_to_not_convert` keep their checkpoint precision.
+
+## Runtime support
+
+The compatibility path accepts BF16 inputs and executes an NVFP4 GEMM followed
+by the BF16 rank correction. It supports vLLM's FlashInfer, CUTLASS, and FBGEMM
+NVFP4 tensor layouts; incompatible forced backends fail during model loading.
+SM103 is the currently validated and enabled hardware target. Native fusion of
+the NVFP4 GEMM and rank correction is separate from this checkpoint-loading
+contract.
diff --git a/examples/__init__.py b/examples/__init__.py
new file mode 100644
index 000000000..78a85387e
--- /dev/null
+++ b/examples/__init__.py
@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""vLLM-Omni example scripts, importable as a package.
+
+Keep this ``__init__.py``: with the upstream vLLM installed editable (which also
+places an ``examples`` package on ``sys.path``), a regular package at this repo
+root wins over the upstream ``examples`` namespace during ``import examples``.
+"""
diff --git a/examples/offline_inference/audex/README.md b/examples/offline_inference/audex/README.md
index 793b69eb9..4b570b06a 100644
--- a/examples/offline_inference/audex/README.md
+++ b/examples/offline_inference/audex/README.md
@@ -17,7 +17,7 @@ Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
pair contract) lives in `common.py`; each task script keeps only its own
prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
-Open, AudioX) use the generic
+Open) use the generic
`examples/offline_inference/text_to_audio/text_to_audio.py` instead.
## text_to_speech.py — text → speech
diff --git a/examples/offline_inference/audex/text_to_audio.py b/examples/offline_inference/audex/text_to_audio.py
index c357819c3..e81a53bb3 100644
--- a/examples/offline_inference/audex/text_to_audio.py
+++ b/examples/offline_inference/audex/text_to_audio.py
@@ -7,7 +7,7 @@ thinker generates interleaved 4-codebook RVQ tokens under an
RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
mono WAV.
-(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
+(For DIFFUSION text-to-audio models — Stable Audio Open — use
``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
TTA is a two-stage autoregressive pipeline with a different interface.)
diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py
index bac385d8e..29ccb3539 100644
--- a/examples/offline_inference/aura_omni/end2end.py
+++ b/examples/offline_inference/aura_omni/end2end.py
@@ -24,7 +24,6 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
DEFAULT_QWEN3_TTS_REF_TEXT,
@@ -228,7 +227,6 @@ def parse_args():
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
- nullify_stage_engine_defaults(parser)
return parser.parse_args()
diff --git a/examples/offline_inference/covo_audio/README.md b/examples/offline_inference/covo_audio/README.md
index 80a822fa2..083666ebd 100644
--- a/examples/offline_inference/covo_audio/README.md
+++ b/examples/offline_inference/covo_audio/README.md
@@ -2,7 +2,8 @@
## Setup
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+Use `--deploy-config` for deployment overrides such as stage memory allocation.
+See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/).
> **Note**
> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
@@ -41,7 +42,7 @@ python end2end.py -m /path/to/Covo-Audio-Chat --output-dir ./my_output
| `--sampling-rate` | | `16000` | Sampling rate for audio loading (Hz) |
| `--output-dir` | | `./output_audio` | Output directory for generated files |
| `--num-prompts` | | `1` | Number of prompts to generate |
-| `--stage-configs-path` | | (auto) | Path to stage configs YAML file |
+| `--deploy-config` | | (auto) | Optional path to a deploy YAML override |
| `--log-stats` | | `false` | Enable detailed statistics logging |
| `--stage-init-timeout` | | `300` | Stage initialization timeout (seconds) |
| `--batch-timeout` | | `5` | Batching timeout (seconds) |
diff --git a/examples/offline_inference/covo_audio/end2end.py b/examples/offline_inference/covo_audio/end2end.py
index 5b7ac5fef..b5050f62c 100644
--- a/examples/offline_inference/covo_audio/end2end.py
+++ b/examples/offline_inference/covo_audio/end2end.py
@@ -59,7 +59,7 @@ def main(args):
omni = Omni(
model=args.model_name,
- stage_configs_path=args.stage_configs_path,
+ deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
batch_timeout=args.batch_timeout,
@@ -167,10 +167,10 @@ def parse_args():
help="Sampling rate for audio loading (default: 16000).",
)
parser.add_argument(
- "--stage-configs-path",
+ "--deploy-config",
type=str,
default=None,
- help="Path to stage configs YAML file.",
+ help="Path to a deploy config YAML file.",
)
parser.add_argument(
"--log-stats",
diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py
index 5f4311050..851d18e18 100644
--- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py
+++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py
@@ -53,7 +53,6 @@ from typing import Any
import torch
from PIL import Image
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -139,12 +138,6 @@ async def main():
# ---- Torch setup ----
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
- parallel_config = DiffusionParallelConfig(
- ulysses_degree=args.ulysses_degree,
- ring_degree=args.ring_degree,
- cfg_parallel_size=args.cfg_parallel_size,
- tensor_parallel_size=args.tensor_parallel_size,
- )
# ---- Cache Config ----
if args.cache_backend == "cache_dit":
@@ -171,7 +164,10 @@ async def main():
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
- parallel_config=parallel_config,
+ ulysses_degree=args.ulysses_degree,
+ ring_degree=args.ring_degree,
+ cfg_parallel_size=args.cfg_parallel_size,
+ tensor_parallel_size=args.tensor_parallel_size,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
diffusion_load_format="dummy",
diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py
new file mode 100644
index 000000000..b3a4ae0d1
--- /dev/null
+++ b/examples/offline_inference/diffusion/lingbot_world_v2.py
@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Generate a LingBot-World v2 video from an image and camera trajectory.
+
+The official checkpoint is licensed separately under CC BY-NC-SA and is
+restricted to non-commercial use. This vLLM-Omni integration remains
+Apache-2.0 licensed.
+
+Example:
+ python examples/offline_inference/diffusion/lingbot_world_v2.py \
+ --prompt "The camera moves slowly forward through the scene." \
+ --image /path/to/first_frame.png \
+ --action-dir /path/to/actions/forward \
+ --output lingbot_world_v2.mp4
+
+``--action-dir`` must contain ``poses.npy`` and ``intrinsics.npy``. The
+example configures its parent as the trusted action root and sends the
+contained directory through ``extra_args["action_path"]``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_NUM_INFERENCE_STEPS = 4
+_MAX_SEQUENCE_LENGTH = 512
+_MAX_PIXEL_AREA = 480 * 832
+_MAX_RAW_FRAMES = 117
+_MAX_SOURCE_CAMERA_FRAMES = 4096
+_TEMPORAL_COMPRESSION = 4
+_LATENT_FRAMES_PER_BLOCK = 3
+_SPATIAL_ALIGNMENT = 16
+
+
+@dataclass(frozen=True)
+class LingBotPaths:
+ image: Path
+ action_dir: Path
+ action_root: Path
+ action_relative: Path
+ camera_frames: int
+ output: Path
+
+
+def _validate_camera_arrays(action_dir: Path) -> int:
+ """Validate the bounded camera contract before starting the engine."""
+
+ try:
+ poses = np.load(action_dir / "poses.npy", mmap_mode="r", allow_pickle=False)
+ intrinsics = np.load(action_dir / "intrinsics.npy", mmap_mode="r", allow_pickle=False)
+ except (OSError, TypeError, ValueError):
+ raise ValueError("Camera files must be valid numeric poses.npy and intrinsics.npy arrays.") from None
+ if poses.ndim != 3 or poses.shape[1:] != (4, 4):
+ raise ValueError("poses.npy must have shape [frames, 4, 4].")
+ if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,):
+ raise ValueError("intrinsics.npy must have shape [frames, 4].")
+ if poses.shape[0] != intrinsics.shape[0]:
+ raise ValueError("poses.npy and intrinsics.npy must contain the same frame count.")
+ if not 1 <= poses.shape[0] <= _MAX_SOURCE_CAMERA_FRAMES:
+ raise ValueError("Camera arrays must contain between 1 and 4096 source frames.")
+ if poses.dtype.kind not in "fiu" or intrinsics.dtype.kind not in "fiu":
+ raise ValueError("Camera arrays must contain real numeric values.")
+ if not np.isfinite(poses).all() or not np.isfinite(intrinsics).all():
+ raise ValueError("Camera arrays must contain only finite values.")
+ return int(poses.shape[0])
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Generate a LingBot-World v2 video from a first-frame image and camera trajectory."
+ )
+ parser.add_argument("--model", default=_MODEL, help="Official Hugging Face model ID or local checkpoint path.")
+ parser.add_argument("--prompt", required=True, help="Text description of the scene and desired motion.")
+ parser.add_argument("--image", required=True, help="Path to the first-frame image.")
+ parser.add_argument(
+ "--action-dir",
+ required=True,
+ help="Camera action directory containing poses.npy and intrinsics.npy.",
+ )
+ parser.add_argument("--height", type=int, default=480, help="Output height; must be divisible by 16.")
+ parser.add_argument("--width", type=int, default=832, help="Output width; must be divisible by 16.")
+ parser.add_argument(
+ "--num-frames",
+ type=int,
+ default=9,
+ help="Raw output frame count. Nine frames are one three-latent-frame causal block.",
+ )
+ parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+ parser.add_argument(
+ "--tensor-parallel-size",
+ type=int,
+ default=1,
+ help="Number of GPUs used for tensor parallelism inside the DiT.",
+ )
+ parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.")
+ parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.")
+ parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.")
+ parser.add_argument(
+ "--enforce-eager",
+ action="store_true",
+ help="Disable compilation and force eager execution.",
+ )
+ return parser.parse_args(argv)
+
+
+def resolve_cli_paths(args: argparse.Namespace) -> LingBotPaths:
+ """Canonicalize local inputs and construct a safe trusted-root mapping."""
+
+ image = Path(args.image).expanduser().resolve()
+ if not image.is_file():
+ raise ValueError("--image must point to an existing file.")
+
+ action_dir = Path(args.action_dir).expanduser().resolve()
+ if not action_dir.is_dir():
+ raise ValueError("--action-dir must point to an existing directory.")
+ if not all((action_dir / name).is_file() for name in ("poses.npy", "intrinsics.npy")):
+ raise ValueError("--action-dir must contain poses.npy and intrinsics.npy.")
+ camera_frames = _validate_camera_arrays(action_dir)
+
+ action_root = action_dir.parent.resolve()
+ try:
+ action_relative = action_dir.relative_to(action_root)
+ except ValueError:
+ raise ValueError("--action-dir must be contained by its canonical trusted action root.") from None
+ if not action_relative.parts:
+ raise ValueError("--action-dir must name a directory contained by a trusted action root.")
+
+ output = Path(args.output).expanduser().resolve()
+ return LingBotPaths(
+ image=image,
+ action_dir=action_dir,
+ action_root=action_root,
+ action_relative=action_relative,
+ camera_frames=camera_frames,
+ output=output,
+ )
+
+
+def _positive_finite(value: float, flag: str) -> float:
+ result = float(value)
+ if not math.isfinite(result) or result <= 0:
+ raise ValueError(f"{flag} must be a positive finite number.")
+ return result
+
+
+def build_omni_kwargs(
+ args: argparse.Namespace,
+ paths: LingBotPaths,
+) -> dict[str, Any]:
+ """Build Omni configuration without overriding checkpoint class discovery."""
+
+ if args.tensor_parallel_size <= 0:
+ raise ValueError("--tensor-parallel-size must be a positive integer.")
+ flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+ model_path = Path(args.model).expanduser()
+ model = str(model_path.resolve()) if model_path.exists() else args.model
+ return {
+ "model": model,
+ "flow_shift": flow_shift,
+ "tensor_parallel_size": args.tensor_parallel_size,
+ "enforce_eager": args.enforce_eager,
+ "model_config": {"lingbot_action_root": str(paths.action_root)},
+ }
+
+
+def build_request(
+ args: argparse.Namespace,
+ paths: LingBotPaths,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+ """Build the prompt and fixed v1 causal-DMD sampling contract."""
+
+ prompt_text = args.prompt.strip()
+ if not prompt_text:
+ raise ValueError("--prompt must contain non-whitespace text.")
+ if args.height <= 0 or args.width <= 0:
+ raise ValueError("--height and --width must be positive integers.")
+ if args.height % _SPATIAL_ALIGNMENT or args.width % _SPATIAL_ALIGNMENT:
+ raise ValueError("--height and --width must each be divisible by 16.")
+ if args.height * args.width > _MAX_PIXEL_AREA:
+ raise ValueError("--height * --width must not exceed 480 * 832 pixels.")
+ if args.num_frames <= 0:
+ raise ValueError("--num-frames must be a positive integer.")
+ if args.num_frames > _MAX_RAW_FRAMES:
+ raise ValueError("--num-frames must not exceed 117 raw frames.")
+ if args.num_frames > paths.camera_frames:
+ raise ValueError("Camera arrays must contain at least --num-frames frames.")
+ if (args.num_frames - 1) % _TEMPORAL_COMPRESSION:
+ raise ValueError("--num-frames must satisfy (num_frames - 1) divisible by 4.")
+ latent_frames = (args.num_frames - 1) // _TEMPORAL_COMPRESSION + 1
+ if latent_frames % _LATENT_FRAMES_PER_BLOCK:
+ raise ValueError("--num-frames must map to whole three-frame latent blocks.")
+ if args.fps <= 0:
+ raise ValueError("--fps must be a positive integer.")
+ flow_shift = _positive_finite(args.flow_shift, "--flow-shift")
+
+ prompt = {
+ "prompt": prompt_text,
+ "multi_modal_data": {"image": str(paths.image)},
+ }
+ sampling_kwargs = {
+ "height": args.height,
+ "width": args.width,
+ "num_frames": args.num_frames,
+ "num_inference_steps": _NUM_INFERENCE_STEPS,
+ "max_sequence_length": _MAX_SEQUENCE_LENGTH,
+ "seed": args.seed,
+ "fps": args.fps,
+ "extra_args": {
+ "action_path": paths.action_relative.as_posix(),
+ "flow_shift": flow_shift,
+ },
+ }
+ return prompt, sampling_kwargs
+
+
+def extract_video_array(outputs: Any) -> np.ndarray:
+ """Unwrap one standard Omni diffusion video into ``[frames, H, W, C]``."""
+
+ if not isinstance(outputs, list) or len(outputs) != 1:
+ raise ValueError("Expected exactly one Omni output for the single LingBot request.")
+ output = outputs[0]
+ request_output = getattr(output, "request_output", None)
+ if request_output is not None and getattr(request_output, "images", None):
+ output = request_output
+
+ images = getattr(output, "images", None)
+ if not images:
+ multimodal_output = getattr(output, "multimodal_output", None) or {}
+ images = multimodal_output.get("video")
+ if images is None:
+ raise ValueError("No video frames found in Omni output.")
+ if isinstance(images, list) and len(images) == 1:
+ images = images[0]
+
+ video: np.ndarray = np.asarray(images)
+ if video.ndim == 5:
+ if video.shape[0] != 1:
+ raise ValueError(f"Expected one generated video, got batch shape {video.shape}.")
+ video = video[0]
+ if video.ndim != 4 or video.shape[-1] not in (3, 4):
+ raise ValueError(f"Expected video shape [frames, height, width, channels], got {video.shape}.")
+ return video
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+ args = parse_args(argv)
+ paths = resolve_cli_paths(args)
+ prompt, sampling_kwargs = build_request(args, paths)
+ paths.output.parent.mkdir(parents=True, exist_ok=True)
+
+ # Keep vLLM imports below pure CLI/path validation so ``--help`` and the
+ # helper tests work on machines without the native vLLM runtime.
+ from diffusers.utils import export_to_video
+
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ omni_kwargs = build_omni_kwargs(args, paths)
+ omni = Omni(**omni_kwargs)
+ try:
+ outputs = omni.generate(
+ prompt,
+ OmniDiffusionSamplingParams(**sampling_kwargs),
+ use_tqdm=False,
+ )
+ finally:
+ omni.close()
+
+ frames = extract_video_array(outputs)
+ export_to_video(frames, str(paths.output), fps=args.fps)
+ print(f"Saved {frames.shape[0]} frames to {paths.output}")
+ return paths.output
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py
new file mode 100644
index 000000000..86461f9f8
--- /dev/null
+++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py
@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Run LingBot-World 2.0 as an in-process realtime AR-Diffusion session.
+
+This example exercises the internal session API. It is not a public HTTP or
+WebSocket protocol. Each JSONL event advances the world by one three-latent
+frame AR block and writes the latent plus its identity metadata. The current
+scoped implementation supports at most ten ticks per generation epoch because
+its image condition is bounded to 117 pixel frames; reset or create a session
+to start a new world.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import math
+import time
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
+_CAMERA_ACTION_SCHEMA = "lingbot.camera_actions.v1"
+_FRAMES_PER_BLOCK = 3
+# ((117 pixel frames - 1) / VAE temporal factor 4 + 1) / 3 latent frames.
+_MAX_REALTIME_TICKS = 10
+
+
+def _camera_event_data(frames: list[list[str]]) -> dict[str, Any]:
+ """Build the event-side script consumed by LingBotCameraControlReducer."""
+ return {"mode": "script", "frames": frames}
+
+
+def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description="Run an in-process realtime LingBot-World 2.0 session.")
+ parser.add_argument("--model", default=_MODEL, help="Hugging Face model ID or local checkpoint path.")
+ parser.add_argument("--image", required=True, help="Initial RGB image.")
+ parser.add_argument("--prompt", required=True, help="Initial scene prompt.")
+ parser.add_argument(
+ "--events",
+ required=True,
+ help="JSONL file with one prompt/action event per AR block; at most 10 events.",
+ )
+ parser.add_argument("--output-dir", required=True, help="Directory for chunk latents and metadata.")
+ parser.add_argument("--session-id", default="lingbot-world", help="Persistent world session identifier.")
+ parser.add_argument("--height", type=int, default=480)
+ parser.add_argument("--width", type=int, default=832)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--gpu-memory-fraction", type=float, default=0.1)
+ parser.add_argument("--tensor-parallel-size", type=int, default=1)
+ parser.add_argument("--enforce-eager", action="store_true")
+ return parser.parse_args(argv)
+
+
+def _load_events(path: Path) -> list[dict[str, Any]]:
+ events: list[dict[str, Any]] = []
+ for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1):
+ line = raw_line.strip()
+ if not line:
+ continue
+ try:
+ value = json.loads(line)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"events line {line_number} is not valid JSON: {exc.msg}.") from None
+ if not isinstance(value, dict):
+ raise ValueError(f"events line {line_number} must be a JSON object.")
+ event_id = value.get("event_id")
+ if isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0:
+ raise ValueError(f"events line {line_number} requires a non-negative integer event_id.")
+ prompt = value.get("prompt")
+ if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()):
+ raise ValueError(f"events line {line_number} prompt must be a non-empty string.")
+ frames = value.get("frames")
+ if frames is not None:
+ if not isinstance(frames, list) or len(frames) != _FRAMES_PER_BLOCK:
+ raise ValueError(f"events line {line_number} frames must contain exactly three lists.")
+ for frame in frames:
+ if not isinstance(frame, list) or any(not isinstance(action, str) for action in frame):
+ raise ValueError(f"events line {line_number} frames must contain only action strings.")
+ if prompt is None and frames is None:
+ raise ValueError(f"events line {line_number} must update prompt and/or frames.")
+ events.append({"event_id": event_id, "prompt": prompt, "frames": frames})
+ if not events:
+ raise ValueError("events file must contain at least one event.")
+ if len(events) > _MAX_REALTIME_TICKS:
+ raise ValueError(
+ "events file must contain at most "
+ f"{_MAX_REALTIME_TICKS} events because the current LingBot realtime "
+ "image-condition horizon is 117 pixel frames."
+ )
+ event_ids = [event["event_id"] for event in events]
+ if event_ids != sorted(set(event_ids)):
+ raise ValueError("event_id values must be unique and strictly increasing.")
+ return events
+
+
+def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path]:
+ image = Path(args.image).expanduser().resolve()
+ events = Path(args.events).expanduser().resolve()
+ output_dir = Path(args.output_dir).expanduser().resolve()
+ if not image.is_file():
+ raise ValueError("--image must point to an existing file.")
+ if not events.is_file():
+ raise ValueError("--events must point to an existing JSONL file.")
+ if not args.prompt.strip():
+ raise ValueError("--prompt must contain non-whitespace text.")
+ if args.height <= 0 or args.width <= 0 or args.height % 16 or args.width % 16:
+ raise ValueError("--height and --width must be positive multiples of 16.")
+ if args.tensor_parallel_size <= 0:
+ raise ValueError("--tensor-parallel-size must be positive.")
+ if not math.isfinite(args.gpu_memory_fraction) or not 0 < args.gpu_memory_fraction <= 1:
+ raise ValueError("--gpu-memory-fraction must be in (0, 1].")
+ return image, events, output_dir
+
+
+async def run(argv: Sequence[str] | None = None) -> Path:
+ args = parse_args(argv)
+ image, events_path, output_dir = _validate_args(args)
+ events = _load_events(events_path)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Imports stay below pure input validation so --help and helper tests do not
+ # require a CUDA-enabled vLLM installation.
+ import torch
+
+ from vllm_omni.diffusion.models.lingbot_world.actions import LingBotCameraControlReducer
+ from vllm_omni.entrypoints.async_omni import AsyncOmni
+ from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+ from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSessionEvent,
+ ARDiffusionSessionManager,
+ ARDiffusionWorkerLifecycle,
+ )
+ from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionControlInput
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ engine = AsyncOmni(
+ model=args.model,
+ engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine",
+ enforce_eager=args.enforce_eager,
+ tensor_parallel_size=args.tensor_parallel_size,
+ max_num_seqs=1,
+ model_config={
+ "ar_diffusion_height": args.height,
+ "ar_diffusion_width": args.width,
+ "ar_diffusion_kv_config": {
+ "gpu_memory_fraction": args.gpu_memory_fraction,
+ "warmup_cudagraph": True,
+ },
+ },
+ )
+ sampling = OmniDiffusionSamplingParams(
+ height=args.height,
+ width=args.width,
+ num_frames=9,
+ num_inference_steps=4,
+ max_sequence_length=512,
+ seed=args.seed,
+ output_type="latent",
+ extra_args={"flow_shift": 5.0},
+ )
+ consumer = ARDiffusionOmniTickConsumer(
+ engine,
+ prompt_provider=lambda tick: {
+ "prompt": tick.prompt,
+ "multi_modal_data": {"image": str(image)},
+ },
+ sampling_params_list=[sampling],
+ diffusion_stage_id=0,
+ )
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=ARDiffusionWorkerLifecycle(engine, stage_ids=[0], timeout=180.0),
+ max_pending_events=32,
+ control_reducer_factory=LingBotCameraControlReducer,
+ )
+ session = await manager.create_session(args.session_id)
+ measurements: list[dict[str, Any]] = []
+ try:
+ for chunk_index, event in enumerate(events):
+ controls = ()
+ if event["frames"] is not None:
+ controls = (
+ ARDiffusionControlInput(
+ track="camera",
+ schema=_CAMERA_ACTION_SCHEMA,
+ data=_camera_event_data(event["frames"]),
+ ),
+ )
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=event["event_id"],
+ prompt=event["prompt"]
+ if event["prompt"] is not None
+ else (args.prompt if chunk_index == 0 else None),
+ controls=controls,
+ )
+ )
+ started = time.perf_counter()
+ output = await session.next_chunk()
+ elapsed = time.perf_counter() - started
+ if len(output.images) != 1 or not isinstance(output.images[0], torch.Tensor):
+ raise RuntimeError("Expected one latent tensor from each realtime LingBot chunk.")
+ latent = output.images[0].detach().float().cpu()
+ metadata = consumer.chunk_metadata(output).to_dict()
+ torch.save(latent, output_dir / f"chunk_{chunk_index:03d}.pt")
+ (output_dir / f"chunk_{chunk_index:03d}.json").write_text(
+ json.dumps(metadata, indent=2, sort_keys=True) + "\n"
+ )
+ measurements.append(
+ {
+ "chunk_index": chunk_index,
+ "latency_seconds": elapsed,
+ "shape": list(latent.shape),
+ "finite": bool(torch.isfinite(latent).all()),
+ "metadata": metadata,
+ }
+ )
+ print(json.dumps(measurements[-1], sort_keys=True), flush=True)
+ finally:
+ try:
+ await manager.close_session(args.session_id)
+ finally:
+ engine.shutdown()
+
+ summary_path = output_dir / "summary.json"
+ summary_path.write_text(json.dumps({"chunks": measurements}, indent=2, sort_keys=True) + "\n")
+ return summary_path
+
+
+def main(argv: Sequence[str] | None = None) -> Path:
+ return asyncio.run(run(argv))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md
index 19fa8c69e..6cf834142 100644
--- a/examples/offline_inference/dynin_omni/README.md
+++ b/examples/offline_inference/dynin_omni/README.md
@@ -107,4 +107,4 @@ python /examples/offline_inference/dynin_omni/end2end.py \
- If you want to force local config resolution, pass `--dynin-config-path `.
- If you see the warning
`max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
- reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).
+ reduce `max_num_batched_tokens` in the deploy config (for example, `4096` in the CI config).
diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py
index 0cab46fee..74d33fe01 100644
--- a/examples/offline_inference/dynin_omni/end2end.py
+++ b/examples/offline_inference/dynin_omni/end2end.py
@@ -876,7 +876,7 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES)
parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.")
parser.add_argument(
- "--stage-config-path",
+ "--deploy-config",
type=str,
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
help="Path to deploy config yaml.",
@@ -1394,8 +1394,8 @@ def main() -> None:
from vllm_omni.entrypoints.omni import Omni
- stage_config_path = str(Path(args.stage_config_path).expanduser())
- omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype)
+ deploy_config_path = str(Path(args.deploy_config).expanduser())
+ omni = Omni(model=model_source, deploy_config=deploy_config_path, dtype=args.dtype)
sampling_params_list = [
SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False)
for _ in range(omni.num_stages)
diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md
index 8b90e6b7f..b4d5af2c7 100644
--- a/examples/offline_inference/hunyuan_image3/README.md
+++ b/examples/offline_inference/hunyuan_image3/README.md
@@ -7,17 +7,17 @@ YAMLs under `vllm_omni/deploy/`.
| File | Topology | Default use |
| :--- | :--- | :--- |
-| `vllm_omni/deploy/hunyuan_image3.yaml` | AR + DiT | Default for `text2img` and `img2img`. |
+| `vllm_omni/deploy/hunyuan_image_3_moe.yaml` | AR + DiT | Default for `text2img` and `img2img`. |
| `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. |
| `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. |
-The example chooses a deploy config automatically when `--deploy-config` and
-`--stage-configs-path` are both omitted:
+The example chooses a deploy config automatically when `--deploy-config` is
+omitted:
| `--modality` | `mode` passed to Omni | Default deploy |
| :--- | :--- | :--- |
-| `text2img` | `text-to-image` | `hunyuan_image3.yaml` |
-| `img2img` | `image-editing` | `hunyuan_image3.yaml` |
+| `text2img` | `text-to-image` | `hunyuan_image_3_moe.yaml` |
+| `img2img` | `image-editing` | `hunyuan_image_3_moe.yaml` |
| `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` |
| `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` |
@@ -35,8 +35,8 @@ path; the image endpoints infer the image task from the endpoint and payload.
| Online scenario | Server deploy | Request |
| :--- | :--- | :--- |
-| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. |
-| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/edits`. |
+| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. |
+| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/edits`. |
| Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. |
| DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. |
@@ -96,7 +96,7 @@ Override the default full AR + DiT deploy explicitly:
python examples/offline_inference/hunyuan_image3/end2end.py \
--model tencent/HunyuanImage-3.0-Instruct \
--modality text2img \
- --deploy-config vllm_omni/deploy/hunyuan_image3.yaml \
+ --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \
--prompts "A cute cat"
```
@@ -117,8 +117,7 @@ python end2end.py --modality text2img \
| Argument | Description |
| :--- | :--- |
-| `--deploy-config` | Preferred config path for unified deploy YAMLs. |
-| `--stage-configs-path` | Legacy stage config path, kept only for compatibility. Prefer `--deploy-config`. |
+| `--deploy-config` | Optional path for unified deploy YAML overrides. |
| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. |
| `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. |
| `--steps` | Number of diffusion inference steps for image generation. |
@@ -132,7 +131,7 @@ python end2end.py --modality text2img \
- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy.
- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`.
-- The old HunyuanImage3 YAMLs under `model_executor/stage_configs/` and `platforms/*/stage_configs/` have been folded into the deploy YAMLs.
+- Platform-specific overrides are included in the unified deploy YAMLs.
## Prompt Format
diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py
index 9a2350156..fbe4289be 100644
--- a/examples/offline_inference/hunyuan_image3/end2end.py
+++ b/examples/offline_inference/hunyuan_image3/end2end.py
@@ -73,7 +73,6 @@ def parse_args():
)
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
- parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
parser.add_argument(
"--stream",
action="store_true",
@@ -146,13 +145,7 @@ def main():
else:
bot_task = args.bot_task
- if args.deploy_config is not None and args.stage_configs_path is not None:
- raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
-
- deploy_config = args.deploy_config
- stage_configs_path = args.stage_configs_path
- if deploy_config is None and stage_configs_path is None:
- deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
+ deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
omni_kwargs = {
"model": args.model,
@@ -168,10 +161,7 @@ def main():
if additional_config is not None:
omni_kwargs["additional_config"] = additional_config
- if deploy_config is not None:
- omni_kwargs["deploy_config"] = deploy_config
- else:
- omni_kwargs["stage_configs_path"] = stage_configs_path
+ omni_kwargs["deploy_config"] = deploy_config
omni = Omni(**omni_kwargs)
@@ -272,10 +262,7 @@ def main():
print(f" Modality: {args.modality}")
print(f" Prompt task: {task}")
print(f" Bot task: {bot_task}")
- if deploy_config is not None:
- print(f" Deploy config: {deploy_config}")
- else:
- print(f" Stage config: {stage_configs_path}")
+ print(f" Deploy config: {deploy_config}")
print(f" Num stages: {omni.num_stages}")
if args.modality in ("text2img", "img2img"):
print(f" Inference steps: {args.steps}")
diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index ae84353f7..136d99857 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/examples/offline_inference/image_to_image/image_edit.py
@@ -96,7 +96,6 @@ from typing import Any
import torch
from PIL import Image
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
@@ -452,14 +451,6 @@ def main():
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
- parallel_config = DiffusionParallelConfig(
- ulysses_degree=args.ulysses_degree,
- ring_degree=args.ring_degree,
- cfg_parallel_size=args.cfg_parallel_size,
- tensor_parallel_size=args.tensor_parallel_size,
- enable_expert_parallel=args.enable_expert_parallel,
- )
-
# Configure cache based on backend type
cache_config = None
if args.cache_backend == "cache_dit":
@@ -490,7 +481,11 @@ def main():
vae_use_tiling=args.vae_use_tiling,
cache_backend=args.cache_backend,
cache_config=cache_config,
- parallel_config=parallel_config,
+ ulysses_degree=args.ulysses_degree,
+ ring_degree=args.ring_degree,
+ cfg_parallel_size=args.cfg_parallel_size,
+ tensor_parallel_size=args.tensor_parallel_size,
+ enable_expert_parallel=args.enable_expert_parallel,
enable_cpu_offload=args.enable_cpu_offload,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md
index c8f7fe617..02f233642 100644
--- a/examples/offline_inference/image_to_video/README.md
+++ b/examples/offline_inference/image_to_video/README.md
@@ -1,7 +1,7 @@
# Image-To-Video
This shared example generates videos from images with VACE, Wan2.2, LTX-2,
-HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
+HunyuanVideo-1.5, SANA-Video, Cosmos3, and other compatible pipelines.
- `image_to_video.py`: command-line script for single video generation with advanced options.
@@ -27,6 +27,8 @@ This folder provides a unified CLI script for image-to-video generation using vL
| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
+| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480 x 832 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; Wan VAE |
+| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704 x 1280 | 81 | 50 | 6.0 | Native `SanaImageToVideoPipeline`; LTX-2 Video VAE |
!!! info
Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
@@ -96,7 +98,7 @@ if __name__ == "__main__":
),
)
from diffusers.utils import export_to_video
- frames = outputs[0].request_output.images
+ frames = outputs[0].images
export_to_video(frames, "quick_test_i2v.mp4", fps=16)
```
@@ -227,6 +229,48 @@ python image_to_video.py \
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, T2V, defaults, and advanced options.
+### SANA-Video-2B
+
+The checkpoint metadata selects the upstream T2V pipeline, so pass
+`SanaImageToVideoPipeline` explicitly for I2V:
+
+```bash
+python image_to_video.py \
+ --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ --model-class-name SanaImageToVideoPipeline \
+ --image cherry_blossom.jpg \
+ --prompt "Cherry blossoms sway in the breeze as petals drift past the camera." \
+ --negative-prompt "blurry, low quality, temporal artifacts" \
+ --height 480 \
+ --width 832 \
+ --num-frames 81 \
+ --num-inference-steps 50 \
+ --guidance-scale 6.0 \
+ --extra-body '{"motion_score": 30}' \
+ --fps 16 \
+ --seed 42 \
+ --output sana_video_i2v_480p.mp4
+```
+
+Use the 720p checkpoint with `--height 704 --width 1280`. In this example,
+81 frames at 16 FPS means the standard approximately five-second SANA-Video
+checkpoint request. It does not enable the separate LongSANA/LongLive
+minute-scale autoregressive workflow.
+
+The native `SanaImageToVideoPipeline` supports both checkpoint variants. The
+480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper and the 720p
+path uses `DistributedAutoencoderKLLTX2Video`; both wrap their corresponding
+Diffusers VAE implementations. The native denoising loop intentionally keeps
+Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`.
+
+The Diffusers adapter has completed end-to-end I2V serving validation at both
+480p and 720p. For online adapter serving, select the matching checkpoint and
+use `WIDTH=1280 HEIGHT=704` for the 720p request. Adapter T2V is also validated
+at both resolutions; see the recipe for the exact backend matrix.
+
+See the [SANA-Video recipe](../../../recipes/NVIDIA/SANA-Video-2B.md) for
+online serving, native/adapter validation boundaries, and hardware notes.
+
## Advanced Features
### CFG Parallel
diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py
index 222502e8a..65e7389f0 100644
--- a/examples/offline_inference/image_to_video/image_to_video.py
+++ b/examples/offline_inference/image_to_video/image_to_video.py
@@ -1,15 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
-HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE.
+HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE.
Supports:
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
- LTX2 image-to-video pipeline
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
+- SANA-Video 2B: first-frame latent conditioning at 480p or 720p
- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning
Usage:
@@ -38,6 +39,13 @@ Usage:
--image input.jpg --prompt "A cat playing with yarn" \
--flow-shift 5.0 --guidance-scale 6.0
+ # SANA-Video 2B I2V (480p)
+ python image_to_video.py --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ --model-class-name SanaImageToVideoPipeline \
+ --image input.jpg --prompt "A cat turns toward the camera." \
+ --height 480 --width 832 --num-frames 81 --num-inference-steps 50 \
+ --guidance-scale 6.0
+
# Cosmos3 I2V (image conditioning)
python image_to_video.py --model nvidia/Cosmos3-Nano \
--image input.jpg --prompt "The scene comes to life with smooth, natural motion." \
@@ -56,7 +64,6 @@ import numpy as np
import PIL.Image
import torch
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -68,6 +75,7 @@ from vllm_omni.model_extras import (
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
+ should_preserve_reference_image_size,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -103,17 +111,27 @@ def build_image_to_video_prompt(
return result
+def _validate_video_output_type(output_type: str) -> None:
+ if output_type not in {"image", "video"}:
+ raise ValueError(
+ f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
+ )
+
+
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=(
"Generate a video from one or more images "
- "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)."
+ "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
)
)
parser.add_argument(
"--model",
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
- help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).",
+ help=(
+ "Diffusers I2V model ID or local path "
+ "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, SANA-Video, Cosmos3, or Wan2.1 VACE)."
+ ),
)
parser.add_argument(
"--model-class-name",
@@ -377,6 +395,7 @@ def main():
is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
+ is_sana = model_class_name == "SanaImageToVideoPipeline" or "sana-video" in model_name
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
@@ -434,6 +453,17 @@ def main():
512 * 768,
32,
)
+ elif is_sana:
+ is_720p = "720p" in model_name
+ d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+ 16,
+ 6.0,
+ 81,
+ 50,
+ 5.0,
+ (704 * 1280) if is_720p else (480 * 832),
+ 32 if is_720p else 16,
+ )
else: # Wan2.2 / HunyuanVideo-1.5
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
@@ -454,7 +484,13 @@ def main():
media_inputs: dict[str, Any] = {}
if image is not None:
- media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+ preserve_image_size = should_preserve_reference_image_size(
+ model_class_name,
+ model=args.model,
+ )
+ media_inputs["image"] = (
+ image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+ )
if last_image is not None:
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
if mask_image is not None:
@@ -496,17 +532,6 @@ def main():
}
profiler_enabled = args.profiler_config is not None
- parallel_config = DiffusionParallelConfig(
- ulysses_degree=args.ulysses_degree,
- ring_degree=args.ring_degree,
- cfg_parallel_size=args.cfg_parallel_size,
- tensor_parallel_size=args.tensor_parallel_size,
- vae_patch_parallel_size=args.vae_patch_parallel_size,
- use_hsdp=args.use_hsdp,
- hsdp_shard_size=args.hsdp_shard_size,
- hsdp_replicate_size=args.hsdp_replicate_size,
- pipeline_parallel_size=args.pipeline_parallel_size,
- )
omni_kwargs = dict(
model=args.model,
enable_layerwise_offload=args.enable_layerwise_offload,
@@ -517,7 +542,15 @@ def main():
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
enable_cpu_offload=args.enable_cpu_offload,
- parallel_config=parallel_config,
+ ulysses_degree=args.ulysses_degree,
+ ring_degree=args.ring_degree,
+ cfg_parallel_size=args.cfg_parallel_size,
+ tensor_parallel_size=args.tensor_parallel_size,
+ vae_patch_parallel_size=args.vae_patch_parallel_size,
+ use_hsdp=args.use_hsdp,
+ hsdp_shard_size=args.hsdp_shard_size,
+ hsdp_replicate_size=args.hsdp_replicate_size,
+ pipeline_parallel_size=args.pipeline_parallel_size,
enforce_eager=args.enforce_eager,
model_class_name=model_class_name,
cache_backend=args.cache_backend,
@@ -639,10 +672,7 @@ def main():
frames = frames[0] if frames else None
if isinstance(frames, OmniRequestOutput):
- if frames.final_output_type != "image":
- raise ValueError(
- f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
- )
+ _validate_video_output_type(frames.final_output_type)
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py
index d26910161..2b70bd749 100644
--- a/examples/offline_inference/lance/end2end.py
+++ b/examples/offline_inference/lance/end2end.py
@@ -82,9 +82,6 @@ def parse_args():
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
# and ``create_default_diffusion`` materializes the stage config.
- from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
- nullify_stage_engine_defaults(p)
return p.parse_args()
diff --git a/examples/offline_inference/longcat_video/end2end.py b/examples/offline_inference/longcat_video/end2end.py
index 8071dfc25..6e0c5593e 100644
--- a/examples/offline_inference/longcat_video/end2end.py
+++ b/examples/offline_inference/longcat_video/end2end.py
@@ -15,7 +15,6 @@ import imageio
import numpy as np
import torch
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.models.longcat_video.pipeline_longcat_video_avatar import (
prepare_longcat_video_avatar_model_for_omni,
)
@@ -236,14 +235,6 @@ def main():
dtype="bfloat16",
enforce_eager=args.enforce_eager,
additional_config=additional_config,
- parallel_config=DiffusionParallelConfig(
- ulysses_degree=1,
- ring_degree=1,
- cfg_parallel_size=1,
- tensor_parallel_size=1,
- vae_patch_parallel_size=1,
- pipeline_parallel_size=1,
- ),
)
audio_input: str | dict[str, str] | None = args.audio
diff --git a/examples/offline_inference/mimo_audio/README.md b/examples/offline_inference/mimo_audio/README.md
index 60f433a4f..126896259 100644
--- a/examples/offline_inference/mimo_audio/README.md
+++ b/examples/offline_inference/mimo_audio/README.md
@@ -19,7 +19,7 @@ MiMo-Audio provides multiple task variants for audio understanding and generatio
## Setup
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware.
### Environment Variables
@@ -189,7 +189,7 @@ Note: This task uses hardcoded message lists in the script.
### Other
-- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
+- If the model or deploy config fails to load, check the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
## Notes
diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md
index ee7e93fd0..29fd59a8d 100644
--- a/examples/offline_inference/ming_flash_omni/README.md
+++ b/examples/offline_inference/ming_flash_omni/README.md
@@ -14,7 +14,7 @@ For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Te
## Setup
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware.
When no `--deploy-config` is passed, the model registry auto-loads the full thinker+talker `vllm_omni/deploy/ming_flash_omni.yaml` (See [Omni-Speech](#omni-speech-thinker--talker)).
diff --git a/examples/offline_inference/minicpmo/README.md b/examples/offline_inference/minicpmo/README.md
index b80fc0aa6..b827a4e98 100644
--- a/examples/offline_inference/minicpmo/README.md
+++ b/examples/offline_inference/minicpmo/README.md
@@ -8,7 +8,7 @@ Two-stage pipeline: **thinker** (multimodal understanding) → **talker + Token2
- Install talker deps: `pip install stepaudio2-minicpmo`
- `--trust-remote-code` is always passed by `end2end.py` (required for MiniCPMO)
-- See [stage configuration docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory tuning
+- See [pipeline and deploy configuration docs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory tuning
## Run examples
diff --git a/examples/offline_inference/minicpmo/end2end.py b/examples/offline_inference/minicpmo/end2end.py
index e9c9019c0..8e536735d 100644
--- a/examples/offline_inference/minicpmo/end2end.py
+++ b/examples/offline_inference/minicpmo/end2end.py
@@ -335,7 +335,7 @@ def main(args):
print(f"use_tts (tts_bos): {use_tts}")
for stage_outputs in omni_generator:
- output = stage_outputs.request_output
+ output = stage_outputs
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -452,12 +452,6 @@ def parse_args():
default=None,
help="Path to a .txt file with one prompt per line (preferred).",
)
- parser.add_argument(
- "--stage-configs-path",
- type=str,
- default=None,
- help="Path to a stage configs file (deprecated; prefer --deploy-config).",
- )
parser.add_argument(
"--deploy-config",
type=str,
diff --git a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_lifecycle.py
similarity index 72%
rename from examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py
rename to examples/offline_inference/minimax_h3/dlo_lifecycle.py
index f8117b0f9..5630777cb 100644
--- a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py
+++ b/examples/offline_inference/minimax_h3/dlo_lifecycle.py
@@ -1,21 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
+"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
+
+The DLO modes accept any positive data-parallel size. With AllGather enabled,
+the script submits one complete DP wave and verifies recovery after an invalid
+wave. It also checks that every mode shuts down all worker processes cleanly.
Examples:
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
+ CUDA_VISIBLE_DEVICES=0,1,2,3 \
+ python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+ --model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
+
+ # Rank-local mmap storage, with no DLO collective or DP wave scheduling.
+ VLLM_WORKER_MULTIPROC_METHOD=spawn \
+ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
- python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
- --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
+ python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+ --model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
- python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
- --model /path/to/MiniMax-H3/FL2VA --mode request
+ python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
+ --model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
"""
from __future__ import annotations
@@ -40,14 +51,33 @@ DEFAULT_PROMPTS = (
)
+def positive_int(value: str) -> int:
+ parsed = int(value)
+ if parsed < 1:
+ raise argparse.ArgumentTypeError("must be a positive integer")
+ return parsed
+
+
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
- choices=("dlo-dp2", "request"),
+ choices=("dlo", "dlo-no-allgather", "request"),
required=True,
- help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
+ help="DLO (with or without AllGather) or ordinary non-DLO request mode",
+ )
+ parser.add_argument(
+ "--dp-size",
+ type=positive_int,
+ default=2,
+ help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
+ )
+ parser.add_argument(
+ "--tp-size",
+ type=positive_int,
+ default=2,
+ help="Tensor-parallel size used only by request mode (default: 2)",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
@@ -60,10 +90,11 @@ def parse_args() -> argparse.Namespace:
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+ is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
common: dict[str, Any] = {
"model": args.model,
"trust_remote_code": True,
- "num_gpus": 2,
+ "num_gpus": args.dp_size if is_dlo else args.tp_size,
"usp": 1,
"ring": 1,
"vae_parallel_mode": "tile",
@@ -74,22 +105,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
- if args.mode == "dlo-dp2":
+ if is_dlo:
common.update(
tensor_parallel_size=1,
- data_parallel_size=2,
+ data_parallel_size=args.dp_size,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
- dlo_use_allgather=True,
+ dlo_use_allgather=args.mode == "dlo",
dlo_resident_layers=0,
)
else:
common.update(
- tensor_parallel_size=2,
+ tensor_parallel_size=args.tp_size,
data_parallel_size=1,
- text_encoder_tp_size=2,
- vae_patch_parallel_size=2,
+ text_encoder_tp_size=args.tp_size,
+ vae_patch_parallel_size=args.tp_size,
enable_distributed_layerwise_offload=False,
)
return common
@@ -172,7 +203,8 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
"engine_kwargs": engine_kwargs(args),
}
try:
- if args.mode == "dlo-dp2":
+ uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
+ if uses_dp_wave:
started = time.perf_counter()
asymmetric = await asyncio.gather(
generate_one(
@@ -182,12 +214,15 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
prompt="",
seed=1001,
),
- generate_one(
- engine,
- args,
- request_id="invalid-peer",
- prompt=DEFAULT_PROMPTS[0],
- seed=1002,
+ *(
+ generate_one(
+ engine,
+ args,
+ request_id=f"invalid-peer-{index}",
+ prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
+ seed=1001 + index,
+ )
+ for index in range(1, args.dp_size)
),
return_exceptions=True,
)
@@ -195,18 +230,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
summary["asymmetric_errors"] = [
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
]
- if len(summary["asymmetric_errors"]) != 2:
- raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
+ if len(summary["asymmetric_errors"]) != args.dp_size:
+ raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
started = time.perf_counter()
- request_count = 2 if args.mode == "dlo-dp2" else 1
+ request_count = args.dp_size if uses_dp_wave else 1
outputs = await asyncio.gather(
*(
generate_one(
engine,
args,
request_id=f"recovery-{index}",
- prompt=DEFAULT_PROMPTS[index],
+ prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
seed=2000 + index,
)
for index in range(request_count)
diff --git a/examples/offline_inference/nemotron_voicechat/end2end.py b/examples/offline_inference/nemotron_voicechat/end2end.py
index b020e5bad..ac5fac376 100644
--- a/examples/offline_inference/nemotron_voicechat/end2end.py
+++ b/examples/offline_inference/nemotron_voicechat/end2end.py
@@ -67,7 +67,7 @@ def main() -> None:
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
parser.add_argument(
- "--stage-configs-path",
+ "--deploy-config",
default=None,
help=(
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
@@ -134,8 +134,8 @@ def main() -> None:
]
omni_kwargs: dict = {}
- if args.stage_configs_path:
- omni_kwargs["stage_configs_path"] = args.stage_configs_path
+ if args.deploy_config:
+ omni_kwargs["deploy_config"] = args.deploy_config
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
outputs = omni.generate(inputs, sampling_params_list)
diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md
index 059c10226..490711d34 100644
--- a/examples/offline_inference/qwen2_5_omni/README.md
+++ b/examples/offline_inference/qwen2_5_omni/README.md
@@ -1,7 +1,7 @@
# Qwen2.5-Omni: Offline inference
## Setup
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation for your hardware.
## Run examples
diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md
index 2f4ee11cc..d216042b4 100644
--- a/examples/offline_inference/qwen3_omni/README.md
+++ b/examples/offline_inference/qwen3_omni/README.md
@@ -1,7 +1,8 @@
# Qwen3-Omni: Offline inference
## Setup
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+Use `--deploy-config` for deployment overrides such as stage memory allocation.
+See the [pipeline and deploy configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/).
## Run examples
@@ -70,8 +71,8 @@ For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
start **before** the upstream stage (Thinker) finishes -- use the async_chunk
example. This requires:
-1. A deploy config YAML with ``async_chunk: true`` (e.g.
- ``qwen3_omni_moe.yaml``).
+1. A deploy YAML with ``async_chunk: true`` (for example, an overlay based on
+ ``vllm_omni/deploy/qwen3_omni_moe.yaml``).
2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
config).
@@ -97,11 +98,11 @@ bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
python end2end_async_chunk.py --query-type text --modalities text
```
-#### Custom stage config
+#### Custom deploy config
```bash
python end2end_async_chunk.py \
--query-type use_audio \
- --deploy-config /path/to/your_deploy_config.yaml
+ --deploy-config /path/to/your_async_chunk.yaml
```
> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py
index 9434b3058..faa200bc9 100644
--- a/examples/offline_inference/qwen3_omni/end2end.py
+++ b/examples/offline_inference/qwen3_omni/end2end.py
@@ -484,10 +484,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
- "--stage-configs-path",
+ "--deploy-config",
type=str,
default=None,
- help="Path to a stage configs file.",
+ help="Path to a deploy config YAML.",
)
parser.add_argument(
"--video-path",
diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py
index c224586c4..c4570abae 100644
--- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py
+++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py
@@ -386,8 +386,8 @@ async def run_all(args):
# argparse defaults do not silently override deploy YAML values.
async_omni = AsyncOmni.from_cli_args(args, model=args.model)
- # Use default sampling params from stage config (they are pre-configured
- # in the YAML for each stage).
+ # Use default sampling params from the resolved pipeline and deploy
+ # config.
#
# NOTE: Since we do not set the sampling params directly, .generate in
# will automatically set the output kind to delta, since this is what
diff --git a/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py
index 964cd4874..62dd29720 100644
--- a/examples/offline_inference/speech_to_video/speech_to_video.py
+++ b/examples/offline_inference/speech_to_video/speech_to_video.py
@@ -23,7 +23,6 @@ from typing import Any
import torch
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -250,17 +249,6 @@ def main():
"scm_steps_policy": "dynamic",
}
- parallel_config = DiffusionParallelConfig(
- tensor_parallel_size=args.tensor_parallel_size,
- cfg_parallel_size=args.cfg_parallel_size,
- vae_patch_parallel_size=args.vae_patch_parallel_size,
- ring_degree=args.ring_degree,
- ulysses_degree=args.ulysses_degree,
- use_hsdp=args.use_hsdp,
- hsdp_shard_size=args.hsdp_shard_size,
- hsdp_replicate_size=args.hsdp_replicate_size,
- )
-
profiler_enabled = args.profiler_config is not None
omni_kwargs = dict(
@@ -271,7 +259,14 @@ def main():
vae_use_tiling=args.vae_use_tiling,
enable_cpu_offload=args.enable_cpu_offload,
enable_layerwise_offload=args.enable_layerwise_offload,
- parallel_config=parallel_config,
+ tensor_parallel_size=args.tensor_parallel_size,
+ cfg_parallel_size=args.cfg_parallel_size,
+ vae_patch_parallel_size=args.vae_patch_parallel_size,
+ ring_degree=args.ring_degree,
+ ulysses_degree=args.ulysses_degree,
+ use_hsdp=args.use_hsdp,
+ hsdp_shard_size=args.hsdp_shard_size,
+ hsdp_replicate_size=args.hsdp_replicate_size,
enforce_eager=args.enforce_eager,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
cache_backend=args.cache_backend,
diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md
index abd21f371..7457863f5 100644
--- a/examples/offline_inference/step_audio2/README.md
+++ b/examples/offline_inference/step_audio2/README.md
@@ -29,7 +29,7 @@ Step-Audio2 is a two-stage audio model:
**Notes:**
- Single GPU mode requires high VRAM due to both stages sharing memory
- Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs
-- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config
+- VRAM usage can be adjusted via `gpu_memory_utilization` in the deploy config
## Performance Benchmark
@@ -197,11 +197,6 @@ This mode:
### Advanced Options
```bash
-# Use custom stage configuration
-python end2end.py --query-type audio_to_text \
- --model stepfun-ai/Step-Audio-2-mini \
- --stage-configs-path /path/to/custom_config.yaml
-
# Use custom deploy configuration
python end2end.py --query-type audio_to_text \
--model stepfun-ai/Step-Audio-2-mini \
diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py
index f175db750..418a7f129 100644
--- a/examples/offline_inference/step_audio2/end2end.py
+++ b/examples/offline_inference/step_audio2/end2end.py
@@ -232,15 +232,8 @@ def main(args):
raise ValueError(f"Unknown query type: {args.query_type}")
# Initialize vLLM-Omni with Step-Audio2
- # Resolve deploy config path. Keep explicit stage_configs_path for legacy
- # custom configs, but use bundled deploy configs by default.
- deploy_config_path = None
- stage_config_path = None
- if args.stage_configs_path and args.deploy_config:
- raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive")
- if args.stage_configs_path:
- stage_config_path = args.stage_configs_path
- elif args.deploy_config:
+ # Resolve an explicit deploy config or use the bundled task default.
+ if args.deploy_config:
deploy_config_path = args.deploy_config
else:
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy"
@@ -262,10 +255,7 @@ def main(args):
"ray_address": args.ray_address,
"trust_remote_code": True,
}
- if stage_config_path is not None:
- omni_kwargs["stage_configs_path"] = stage_config_path
- else:
- omni_kwargs["deploy_config"] = deploy_config_path
+ omni_kwargs["deploy_config"] = deploy_config_path
omni_llm = Omni(**omni_kwargs)
@@ -418,12 +408,6 @@ def parse_args():
" - HuggingFace ID: stepfun-ai/Step-Audio-2-mini"
),
)
- parser.add_argument(
- "--stage-configs-path",
- type=str,
- default=None,
- help="Path to a legacy stage config YAML file (default: use bundled deploy config)",
- )
parser.add_argument(
"--deploy-config",
type=str,
diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md
index 89ffc8bd8..f096d96c8 100644
--- a/examples/offline_inference/text_to_audio/README.md
+++ b/examples/offline_inference/text_to_audio/README.md
@@ -1,11 +1,10 @@
# Text-To-Audio
-A unified script for text/video-to-audio generation. Supported models:
+A unified script for text-to-audio generation. Supported models:
| Model | Tasks | Notes |
|-------|-------|-------|
| `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` |
-| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` |
The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
@@ -50,42 +49,12 @@ python text_to_audio.py \
--output stable_audio_output.wav
```
-### AudioX
-
-AudioX supports six tasks. Sampler and reference knobs (declared in
-`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body`
-JSON flag, routed into sampling `extra_args`.
-
-Text tasks (`t2a` / `t2m`):
-
-```bash
-python text_to_audio.py \
- --model zhangj1an/AudioX --task t2a \
- --prompt "Fireworks burst twice, followed by a clock ticking." \
- --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \
- --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
- --output t2a.wav
-```
-
-Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`:
-
-```bash
-python text_to_audio.py \
- --model zhangj1an/AudioX --task tv2a \
- --prompt "drum beating sound and human talking" \
- --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
- --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
- --output tv2a.wav
-```
-
Key arguments:
- `--prompt`: text description (string).
-- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`.
-- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`).
-- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX).
-- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX).
-- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`.
+- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s`).
+- `--audio-length`: audio duration in seconds (audio length for Stable Audio).
+- `--extra-body`: JSON dict of model-specific knobs, merged into sampling `extra_args`.
- `--negative-prompt`: negative prompt for classifier-free guidance.
- `--seed`: integer seed for deterministic generation.
- `--guidance-scale`: classifier-free guidance scale.
diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py
index eff9c9d1a..094a1c66e 100644
--- a/examples/offline_inference/text_to_audio/text_to_audio.py
+++ b/examples/offline_inference/text_to_audio/text_to_audio.py
@@ -2,15 +2,13 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
-Example script for text/video-to-audio generation with diffusion audio models.
+Example script for text-to-audio generation with diffusion audio models.
This script supports:
* Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio.
- * AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m.
-Model-specific generation knobs are declared in `vllm_omni/model_extras` and
-routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so
-the same script serves both models without per-model branching in the engine.
+Model-specific generation knobs can be passed via `--extra-body` and are
+merged into `sampling_params.extra_args`.
Usage:
# Stable Audio Open
@@ -18,15 +16,6 @@ Usage:
python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache
-
- # AudioX (text-to-audio / text-to-music)
- python text_to_audio.py --model zhangj1an/AudioX --task t2a \
- --prompt "Fireworks burst twice, then a clock ticking." \
- --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0
- # AudioX video-conditioned (v2*/tv2*) — pass a video file/URL
- python text_to_audio.py --model zhangj1an/AudioX --task tv2a \
- --prompt "drum beating sound and human talking" --video clip.mp4 \
- --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}'
"""
import argparse
@@ -38,17 +27,11 @@ from typing import Any
import numpy as np
import torch
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
-AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
-AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
-
def parse_extra_body(value: str) -> dict[str, Any]:
"""Parse a JSON object string for --extra-body."""
@@ -66,7 +49,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--model",
default="stabilityai/stable-audio-open-1.0",
- help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.",
+ help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0.",
)
parser.add_argument(
"--prompt",
@@ -97,19 +80,13 @@ def parse_args() -> argparse.Namespace:
"--audio-start",
type=float,
default=0.0,
- help=(
- "Audio start offset in seconds. Applies to both models "
- "(audio_start_in_s for Stable Audio, seconds_start for AudioX)."
- ),
+ help="Audio start offset in seconds (maps to audio_start_in_s).",
)
parser.add_argument(
"--audio-length",
type=float,
default=10.0,
- help=(
- "Audio duration in seconds. Maps to audio length for Stable Audio "
- "(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX."
- ),
+ help=("Audio duration in seconds. Maps to audio length for Stable Audio (max ~47s for stable-audio-open-1.0)."),
)
parser.add_argument(
"--num-inference-steps",
@@ -123,41 +100,20 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of audio waveforms to generate for the given prompt.",
)
- parser.add_argument(
- "--task",
- type=str,
- default=None,
- choices=list(AUDIOX_TASKS),
- help=(
- "[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. "
- "v2*/tv2* require --video. Ignored for Stable Audio Open."
- ),
- )
- parser.add_argument(
- "--video",
- type=str,
- default=None,
- help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.",
- )
parser.add_argument(
"--extra-body",
type=parse_extra_body,
default=None,
help=(
- "Model-specific generation params as a JSON object, e.g. "
- '\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered '
- "against the model's declared extra_body_params (see vllm_omni/model_extras), "
- "so unknown keys are silently dropped. Values here override the equivalent flags."
+ "Model-specific generation params as a JSON object, merged into "
+ "sampling extra_args. Values here override the equivalent flags."
),
)
parser.add_argument(
"--output",
type=str,
default=None,
- help=(
- "Path to save the generated audio (WAV format). Defaults to "
- "audiox_output.wav for AudioX and stable_audio_output.wav otherwise."
- ),
+ help=("Path to save the generated audio (WAV format). Defaults to stable_audio_output.wav."),
)
parser.add_argument(
"--sample-rate",
@@ -276,8 +232,6 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
def main():
args = parse_args()
- model_name = str(args.model).lower() if args.model else ""
- is_audiox = "audiox" in model_name
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
cache_config = None
if args.cache_backend == "tea_cache":
@@ -286,11 +240,9 @@ def main():
}
print(f"\n{'=' * 60}")
- print("Text/Video-to-Audio Generation")
+ print("Text-to-Audio Generation")
print(f"{'=' * 60}")
print(f" Model: {args.model}")
- if args.task is not None:
- print(f" Task: {args.task}")
print(f" Prompt: {args.prompt}")
print(f" Negative prompt: {args.negative_prompt}")
print(f" Duration: {args.audio_length}s")
@@ -306,28 +258,18 @@ def main():
print(f" Seed: {args.seed}")
print(f"{'=' * 60}\n")
- parallel_config = DiffusionParallelConfig(
+ omni_kwargs: dict[str, Any] = dict(
+ model=args.model,
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
- )
-
- # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio
- # is resolved from its config.
- omni_kwargs: dict[str, Any] = dict(
- model=args.model,
- parallel_config=parallel_config,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
enable_cpu_offload=args.enable_cpu_offload,
enable_layerwise_offload=args.enable_layerwise_offload,
)
- if is_audiox:
- omni_kwargs["model_class_name"] = "AudioXPipeline"
omni = Omni(**omni_kwargs)
- model_class_name = get_model_class_name(omni)
- declared_extra_body_params = get_extra_body_params(model_class_name)
diffusion_params = OmniDiffusionSamplingParams(
generator=generator,
@@ -337,38 +279,19 @@ def main():
seed=args.seed,
)
- # Negative prompt is opt-in for both models (off by default). Including a
+ # Negative prompt is opt-in (off by default). Including a
# non-empty negative prompt enables the separate CFG conditioning branch.
prompt: dict[str, Any] = {"prompt": args.prompt}
if args.negative_prompt:
prompt["negative_prompt"] = args.negative_prompt
- if model_class_name == "AudioXPipeline":
- if args.task is None:
- raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).")
- if args.task in AUDIOX_VIDEO_TASKS and not args.video:
- raise ValueError(f"AudioX task {args.task!r} requires --video.")
- # AudioX core knobs are surfaced as shared flags; sampler/reference knobs
- # (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and
- # are filtered against the declared extra_body contract.
- user_extra: dict[str, Any] = {
- "audiox_task": args.task,
- "seconds_start": args.audio_start,
- "seconds_total": args.audio_length,
- "video_path": args.video,
- }
- if args.extra_body:
- user_extra.update(args.extra_body)
- apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
- else:
- # Stable Audio Open path (backward-compatible).
- audio_end_in_s = args.audio_start + args.audio_length
- diffusion_params.extra_args = {
- "audio_start_in_s": args.audio_start,
- "audio_end_in_s": audio_end_in_s,
- }
- if args.extra_body:
- diffusion_params.extra_args.update(args.extra_body)
+ audio_end_in_s = args.audio_start + args.audio_length
+ diffusion_params.extra_args = {
+ "audio_start_in_s": args.audio_start,
+ "audio_end_in_s": audio_end_in_s,
+ }
+ if args.extra_body:
+ diffusion_params.extra_args.update(args.extra_body)
generation_start = time.perf_counter()
@@ -379,7 +302,7 @@ def main():
print(f"Total generation time: {generation_time:.2f} seconds")
- output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav")
+ output = args.output or "stable_audio_output.wav"
output_path = Path(output)
output_path.parent.mkdir(parents=True, exist_ok=True)
suffix = output_path.suffix or ".wav"
diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md
index b71c998b3..cfe05255d 100644
--- a/examples/offline_inference/text_to_image/README.md
+++ b/examples/offline_inference/text_to_image/README.md
@@ -74,7 +74,7 @@ if __name__ == "__main__":
omni = Omni(model="Qwen/Qwen-Image")
prompt = "a cup of coffee on the table"
outputs = omni.generate(prompt)
- images = outputs[0].request_output.images
+ images = outputs[0].images
images[0].save("coffee.png")
```
@@ -265,7 +265,7 @@ if __name__ == "__main__":
]
outputs = omni.generate(prompts)
for i, output in enumerate(outputs):
- output.request_output.images[0].save(f"{i}.jpg")
+ output.images[0].save(f"{i}.jpg")
```
!!! info
@@ -301,7 +301,7 @@ if __name__ == "__main__":
}
])
for i, output in enumerate(outputs):
- output.request_output.images[0].save(f"{i}.jpg")
+ output.images[0].save(f"{i}.jpg")
```
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index 360c726fd..2ee09e178 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/examples/offline_inference/text_to_image/text_to_image.py
@@ -89,15 +89,9 @@ def parse_args() -> argparse.Namespace:
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
"black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
"black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
- "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
+ "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
)
- parser.add_argument(
- "--stage-configs-path",
- type=str,
- default=None,
- help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
- )
parser.add_argument(
"--deploy-config",
type=str,
@@ -475,8 +469,6 @@ def main():
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
if args.enforce_eager is not None:
omni_kwargs["enforce_eager"] = args.enforce_eager
- if args.stage_configs_path:
- omni_kwargs["stage_configs_path"] = args.stage_configs_path
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if use_nextstep:
@@ -515,8 +507,8 @@ def main():
print(f" Image size: {args.width}x{args.height}")
if args.lora_path:
print(f" LoRA: scale={args.lora_scale}")
- if args.stage_configs_path:
- print(f" stage-configs-path: {args.stage_configs_path}")
+ if args.deploy_config:
+ print(f" deploy-config: {args.deploy_config}")
print(f"{'=' * 60}\n")
# Build LoRA request when --lora-path is set
diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md
index d246f9e45..92e08fab5 100644
--- a/examples/offline_inference/text_to_speech/README.md
+++ b/examples/offline_inference/text_to_speech/README.md
@@ -16,6 +16,7 @@ list of supported architectures across all modalities, see
|---|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
+| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -23,6 +24,7 @@ list of supported architectures across all modalities, see
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
+| dots.tts | `rednote-hilab/dots.tts-soar` | single (native AR) | — (not wired yet) | — | — | 48 kHz |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
| IndexTTS-2.5 | native `checkpoints/` bundle | 2 (AR talker + EnhancedCodec + S2Mel DiT + BigVGAN) | ✓ (required) | — | multilingual (`--lang`) + emotion control | 22.05 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -154,7 +156,7 @@ python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
--text "Hello, this is a streaming test." \
--streaming
```
-Streaming requires `async_chunk: true` in the stage config.
+Streaming requires `async_chunk: true` in the deploy config.
### Notes
- Output: 44.1 kHz mono WAV.
@@ -369,7 +371,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
--streaming \
--output-dir /tmp/out_stream
```
-Streaming requires `async_chunk: true` in the stage config.
+Streaming requires `async_chunk: true` in the deploy config.
### Word Timestamps
Word-level timestamps are currently a serving-path feature: launch
@@ -378,13 +380,13 @@ Word-level timestamps are currently a serving-path feature: launch
example is not provided in this release.
### Batched decoding
-The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
+The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, create an overlay based on `vllm_omni/deploy/qwen3_tts.yaml` and pass it with `--deploy-config`:
```bash
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
--query-type CustomVoice \
--txt-prompts examples/offline_inference/text_to_speech/qwen3_tts/benchmark_prompts.txt \
--batch-size 4 \
- --stage-configs-path /path/to/qwen3_tts_batched.yaml
+ --deploy-config /path/to/qwen3_tts_batched.yaml
```
`--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…).
@@ -394,6 +396,57 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
---
+## Gepard-1.0
+
+Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
+
+### Prerequisites
+The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
+
+```bash
+uv pip install "nemo_toolkit[tts]==2.7.3" \
+ --overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
+```
+
+The overrides are what make that command work. `nemo_toolkit[tts]` declares
+`transformers~=4.57.0`, so a plain install downgrades the version
+`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
+which `import vllm` fails. Both packages are named here rather than left to resolve on
+their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
+that would make the second repair an accident of the first. `--constraint` cannot be
+used instead, since it narrows a version range rather than overriding a dependency a
+package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
+against transformers 5.x, so nothing is lost by holding it there.
+
+(The same clash is why this is documented here rather than declared in
+`pyproject.toml`: an extra resolves together with the base dependencies, so
+`vllm-omni[gepard]` would be unsatisfiable rather than opt-in.)
+
+On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
+consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
+sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
+dies inside `profile_run` long after every import has succeeded, which reads as a model
+failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
+here.
+
+### Quick start (zero-shot, default voice)
+```bash
+python examples/offline_inference/text_to_speech/gepard/end2end.py \
+ --text "Hello, this is a Gepard demo."
+```
+
+### Voice cloning
+Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
+
+### Notes
+- Output: 22.05 kHz mono WAV.
+- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
+- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
+- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
+- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
+- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
+- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
+
## VoxCPM2
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
@@ -430,6 +483,31 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
---
+## dots.tts
+
+Single-stage native AR TTS at 48 kHz (rednote-hilab). Pipeline: `Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR loopback → AudioVAE (streaming decode)`. Same "vLLM-native base LM + side-path computation" pattern as VoxCPM2, with a plain Qwen2 backbone instead of MiniCPM4 and no FSQ / residual LM stage.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
+ --model rednote-hilab/dots.tts-soar \
+ --text "Hello, this is a test of dots TTS running on vLLM Omni."
+```
+
+### Voice cloning
+Not wired in this release — generation is zero-shot only. The CAM++ x-vector speaker encoder weights load, but `end2end.py` has no `--ref-audio`/`--ref-text` flags yet.
+
+### Streaming
+The AudioVAE decoder has an internal streaming path (`init_stream_state` / `stream_step` / `stream_flush`) used to avoid boundary artifacts between 160 ms patches, but it is not yet exposed through an online serving endpoint or example.
+
+### Notes
+- Output: 48 kHz mono WAV.
+- Deploy config: `vllm_omni/deploy/dots_tts.yaml` (auto-loaded by HF `model_type`).
+- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default. `dots.tts-base` shares the same architecture but is unvalidated in this repo. `dots.tts-mf` (MeanFlow, 2-4 step) is not supported yet.
+- Known limitation: no CUDA graph capture and no batched side-path computation yet, so concurrent requests do not currently scale (each request's DiT Euler steps run serially). See `recipes/rednote-hilab/dots.tts.md` for details and the roadmap.
+
+---
+
## IndexTTS-2 and IndexTTS-2.5
Both versions use a 2-stage TTS pipeline and produce 22.05 kHz mono audio. Stage 0 is a GPT AR talker. IndexTTS-2 Stage 1 uses RepCodec semantic embeddings, GPT latent, S2Mel CFM/DiT, and BigVGAN. IndexTTS-2.5 replaces RepCodec with EnhancedCodec and uses the official code-only path (`use_gpt_latent=false`). Every request requires reference audio for zero-shot voice cloning.
@@ -538,100 +616,3 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
- `--num-prompts N` replicates the prompt for performance measurement.
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
- Run `--help` for the full argument surface.
-
----
-
-## SoulX-Singer
-
-Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
-
-### Prerequisites
-
-Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
-
-```bash
-# 1. DiT weights
-export BASE=path/to/SoulX-Singer
-export PREPROCESS=path/to/SoulX-Singer-Preprocess
-export SVC_DIR=path/to/SoulX-Singer-svc
-
-huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
-
-# 2. Preprocess weights (required)
-huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
-export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
-
-# 3. SVS / SVC view directories
-mkdir -p "$SVC_DIR"
-cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
-mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
-
-cat > "$BASE/config.json" <<'EOF'
-{
- "model_type": "soulxsinger",
- "architectures": ["SoulXSingerPipeline"],
- "max_num_seqs": 1
-}
-EOF
-
-cat > "$SVC_DIR/config.json" <<'EOF'
-{
- "model_type": "soulxsinger",
- "architectures": ["SoulXSingerSVCPipeline"],
- "max_num_seqs": 1
-}
-EOF
-```
-
-`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
-
-**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
-
-```bash
-pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
-```
-
-Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
-
-```bash
-# install optional dependencies:
-pip install -e ".[soulx-svs]"
-```
-
-English SVS adds NeMo ASR and NLTK data; pass `--language English`:
-
-```bash
-pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
-python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
-```
-
-**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
-
-### Quick start
-
-```bash
-# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
-python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
- --model "$BASE" \
- --preprocess-weights-dir "$PREPROCESS" \
- --control score \
- --num-inference-steps 32 \
- -o output.wav
-
-python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
- --model "$SVC_DIR" \
- --preprocess-weights-dir "$PREPROCESS" \
- --svc \
- --num-inference-steps 32 \
- -o output_svc.wav
-```
-
-`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
-
-### Notes
-
-- Output: 24 kHz mono WAV; batch only.
-- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
-- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
-
----
diff --git a/examples/offline_inference/text_to_speech/dots_tts/end2end.py b/examples/offline_inference/text_to_speech/dots_tts/end2end.py
new file mode 100644
index 000000000..8705f03f3
--- /dev/null
+++ b/examples/offline_inference/text_to_speech/dots_tts/end2end.py
@@ -0,0 +1,125 @@
+"""Offline dots.tts inference example (single-stage native AR pipeline).
+
+Uses the single-stage native AR config (vllm_omni/deploy/dots_tts.yaml).
+Reference checkpoint: rednote-hilab/dots.tts-soar.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+SAMPLE_RATE = 48_000
+
+
+def parse_args():
+ parser = TrackingArgumentParser(description="Offline dots.tts native AR inference")
+ parser.add_argument(
+ "--model",
+ type=str,
+ default="rednote-hilab/dots.tts-soar",
+ help="dots.tts model path or HuggingFace repo ID.",
+ )
+ parser.add_argument(
+ "--text",
+ type=str,
+ default="Hello, this is a test of dots TTS running on vLLM Omni.",
+ help="Text to synthesize.",
+ )
+ parser.add_argument(
+ "--output-dir",
+ type=str,
+ default="output_audio",
+ help="Directory for output WAV files.",
+ )
+ parser.add_argument(
+ "--deploy-config",
+ type=str,
+ default=None,
+ help="Override the deploy config path. If unset, auto-loads "
+ "vllm_omni/deploy/dots_tts.yaml based on the HF model_type.",
+ )
+ return parser.parse_args()
+
+
+def extract_audio(multimodal_output: dict) -> torch.Tensor:
+ """Extract the final complete audio tensor from multimodal output.
+
+ The output processor concatenates per-step delta tensors under
+ ``model_outputs``. Falls back to ``audio`` for backwards compat.
+ """
+ audio = multimodal_output.get("model_outputs")
+ if audio is None:
+ audio = multimodal_output.get("audio")
+ if audio is None:
+ raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}")
+
+ if isinstance(audio, list):
+ valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+ if not valid:
+ raise ValueError("Audio list is empty or all elements are None.")
+ return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+
+ return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def main():
+ args = parse_args()
+
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ engine = Omni(
+ model=args.model,
+ deploy_config=args.deploy_config,
+ )
+
+ from transformers import AutoTokenizer
+
+ from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
+ build_dots_tts_prompt,
+ )
+
+ tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+
+ prompt = build_dots_tts_prompt(
+ tokenizer=tokenizer,
+ text=args.text,
+ )
+
+ print(f"Model : {args.model}")
+ print(f"Text : {args.text}")
+ print(f"Output dir : {output_dir}")
+ print(f"Prompt len : {len(prompt['prompt_token_ids'])} tokens")
+
+ t_start = time.perf_counter()
+ outputs = engine.generate([prompt])
+ elapsed = time.perf_counter() - t_start
+
+ request_output = outputs[0]
+ mm = request_output.outputs[0].multimodal_output
+ audio = extract_audio(mm)
+
+ duration = audio.numel() / SAMPLE_RATE
+ rtf = elapsed / duration if duration > 0 else float("inf")
+
+ output_path = output_dir / "output.wav"
+ sf.write(str(output_path), audio.numpy(), SAMPLE_RATE, format="WAV")
+
+ print(f"Saved : {output_path}")
+ print(f"Duration : {duration:.2f}s")
+ print(f"Inference : {elapsed:.2f}s")
+ print(f"RTF : {rtf:.3f}")
+
+
+if __name__ == "__main__":
+ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+ main()
diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py
index 18b4d79f8..e2e5738fd 100644
--- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py
+++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py
@@ -38,17 +38,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
logger = logging.getLogger(__name__)
DEFAULT_MODEL = "fishaudio/s2-pro"
-DEFAULT_STAGE_CONFIG = os.path.join(
- os.path.dirname(__file__),
- "..",
- "..",
- "..",
- "..",
- "vllm_omni",
- "model_executor",
- "stage_configs",
- "fish_speech_s2_pro.yaml",
-)
def build_prompt(
@@ -154,8 +143,6 @@ def _build_omni_kwargs(args, model_name: str) -> dict:
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
- else:
- omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG
return omni_kwargs
@@ -278,12 +265,6 @@ def parse_args():
default=None,
help="Deploy config YAML path (under vllm_omni/deploy/).",
)
- parser.add_argument(
- "--stage-configs-path",
- type=str,
- default=None,
- help="Path to stage configs YAML (legacy, prefer --deploy-config).",
- )
parser.add_argument(
"--output-dir",
default="output_audio",
diff --git a/examples/offline_inference/text_to_speech/gepard/end2end.py b/examples/offline_inference/text_to_speech/gepard/end2end.py
new file mode 100644
index 000000000..2f6b61d89
--- /dev/null
+++ b/examples/offline_inference/text_to_speech/gepard/end2end.py
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
+
+Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
+frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
+waveform. Zero-shot: the default learned voice, no reference audio.
+
+Generation length and seed are stage settings rather than flags here. One
+output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
+frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
+
+Usage:
+ python end2end.py --text "Hello, this is Gepard speaking."
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+
+import vllm_omni
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+from vllm_omni import Omni # noqa: E402
+
+MODEL = "nineninesix/gepard-1.0"
+SAMPLE_RATE = 22050
+# Absolute: the loader hands the path straight to OmegaConf, which resolves a
+# bare name against the current working directory, not the package.
+DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
+
+
+def build_request(text: str) -> dict:
+ """Build an Omni request payload for Gepard (zero-shot, no ref audio).
+
+ The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
+ a bare text prompt would have no speaker slots and no SOS.
+ """
+ from transformers import AutoTokenizer
+
+ from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+ from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+
+ # From the checkpoint's own sidecar, which is what the worker builds its
+ # config from too. A layout that disagrees with the trained one does not
+ # fail, it just degrades the speech.
+ cfg = GepardConfig.from_checkpoint(MODEL)
+ tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
+ prompt_token_ids = build_gepard_prompt_ids(
+ tokenizer(text, add_special_tokens=False)["input_ids"],
+ config=cfg,
+ )
+ return {
+ "prompt_token_ids": prompt_token_ids,
+ "additional_information": {"text": [text]},
+ }
+
+
+def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
+ sf.write(path, waveform.float().numpy(), sample_rate)
+ seconds = waveform.numel() / sample_rate
+ print(
+ f" Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
+ f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
+ )
+
+
+def main(args) -> None:
+ omni = Omni(
+ model=MODEL,
+ deploy_config=args.deploy_config,
+ stage_init_timeout=args.stage_init_timeout,
+ init_timeout=args.init_timeout,
+ )
+
+ # No explicit SamplingParams: one supplied here would replace the stage
+ # defaults rather than merge over them, dropping the pipeline's stop token.
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ print(f"Synthesizing: {args.text!r}")
+ inputs = build_request(text=args.text)
+
+ for i, stage_outputs in enumerate(omni.generate(inputs)):
+ for j, out in enumerate(stage_outputs.outputs):
+ mm = out.multimodal_output
+ if mm is None:
+ print(f" [req {i}] No audio output.")
+ continue
+ # The consolidation path renames "model_outputs" to "audio".
+ audio = mm.get("audio")
+ if audio is None:
+ audio = mm.get("model_outputs")
+ if isinstance(audio, list):
+ audio = audio[0] if len(audio) else None
+ if audio is None:
+ print(f" [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
+ continue
+ sr_tensor = mm.get("sr")
+ sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
+ save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
+
+ print("Done.")
+
+
+def parse_args():
+ parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
+ parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
+ parser.add_argument(
+ "--output-dir",
+ default=os.path.join(
+ os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
+ "gepard_output",
+ ),
+ help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
+ )
+ parser.add_argument(
+ "--deploy-config",
+ default=DEFAULT_DEPLOY_CONFIG,
+ help=(
+ "Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
+ "stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
+ "Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
+ ),
+ )
+ # The upstream defaults are too tight for a cold start.
+ parser.add_argument("--stage-init-timeout", type=int, default=900)
+ parser.add_argument("--init-timeout", type=int, default=1800)
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ main(parse_args())
diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py
index 6624c2b35..3f66ed8c3 100644
--- a/examples/offline_inference/text_to_speech/glm_tts/end2end.py
+++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py
@@ -145,7 +145,7 @@ def main(args):
omni = Omni(
model=args.model,
- stage_configs_path=deploy_config_path,
+ deploy_config=deploy_config_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
diff --git a/examples/offline_inference/text_to_speech/ming_tts/README.md b/examples/offline_inference/text_to_speech/ming_tts/README.md
index 3fdfaf4d5..87f2df1f5 100644
--- a/examples/offline_inference/text_to_speech/ming_tts/README.md
+++ b/examples/offline_inference/text_to_speech/ming_tts/README.md
@@ -131,7 +131,7 @@ python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
--enforce-eager
```
-`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
+`--streaming` uses `AsyncOmni` and the async chunk deploy config. It currently
supports one prompt per process invocation; use blocking mode for
`--num-prompts > 1`.
diff --git a/examples/offline_inference/text_to_speech/ming_tts/end2end.py b/examples/offline_inference/text_to_speech/ming_tts/end2end.py
index 1a6991e40..89d5fda73 100644
--- a/examples/offline_inference/text_to_speech/ming_tts/end2end.py
+++ b/examples/offline_inference/text_to_speech/ming_tts/end2end.py
@@ -13,6 +13,7 @@ import yaml
from transformers import AutoTokenizer
from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
KEY_CFG,
KEY_MAX_DECODE_STEPS,
@@ -21,7 +22,6 @@ from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
SAMPLE_RATE,
)
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
-from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
try:
from .runner import (
@@ -142,7 +142,8 @@ def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
return _load_speaker_embedding(args.speaker_embedding)
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
return None
- embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
+ # Runs before the engine is built, so nothing has pre-fetched campplus.onnx.
+ embs = SpeakerEmbeddingExtractor(args.model, allow_download=True).extract_many(paths)
if not embs:
raise RuntimeError("Speaker extraction produced no embeddings")
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
index 7c0ce125e..75f562af6 100644
--- a/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
+++ b/examples/offline_inference/text_to_speech/qwen3_tts/end2end.py
@@ -490,10 +490,10 @@ def parse_args():
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument(
- "--stage-configs-path",
+ "--deploy-config",
type=str,
default=None,
- help="Path to a stage configs file.",
+ help="Path to a deploy config YAML.",
)
parser.add_argument(
"--audio-path",
diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py
index 62f5e2caa..7b282c560 100644
--- a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py
+++ b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py
@@ -32,11 +32,10 @@ from end2end import _estimate_prompt_len
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
-def _default_stage_config() -> str:
+def _default_deploy_config() -> str:
repo_root = Path(__file__).resolve().parents[4]
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
@@ -109,7 +108,7 @@ def main(args: Any) -> None:
"output_dir",
):
omni_kwargs.pop(key, None)
- omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
+ omni_kwargs["deploy_config"] = args.deploy_config or _default_deploy_config()
omni_kwargs["log_stats"] = args.log_stats
omni = Omni(**omni_kwargs)
@@ -156,7 +155,7 @@ def parse_args() -> Any:
default=None,
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
)
- parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
+ parser.add_argument("--deploy-config", default=None, help="Qwen3-TTS deploy YAML")
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
@@ -164,7 +163,6 @@ def parse_args() -> Any:
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
- nullify_stage_engine_defaults(parser)
return parser.parse_args()
diff --git a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py
deleted file mode 100644
index 14dd5e145..000000000
--- a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py
+++ /dev/null
@@ -1,162 +0,0 @@
-"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
-
-Usage:
- python benchmark.py --model /path/to/SoulX-Singer --svs \\
- --prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
- --target-audio ../SoulX-Singer/example/audio/music.mp3 \\
- --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
- -o benchmark.wav
-
- python benchmark.py ... --enable-diffusion-pipeline-profiler
-"""
-
-import argparse
-import statistics
-import time
-from pathlib import Path
-
-import soundfile as sf
-from end2end import (
- SVC_DEPLOY_CONFIG,
- SVS_DEPLOY_CONFIG,
- add_inference_args,
- build_sampling,
- extract_audio,
- resolve_preprocess_weights_dir,
-)
-
-from vllm_omni.entrypoints.omni import Omni
-
-
-def _audio_duration_sec(outputs) -> float:
- audio_np, sr = extract_audio(outputs)
- return audio_np.size / sr
-
-
-def _stage_durations_from_output(outputs) -> dict[str, float]:
- omni_out = outputs[0]
- durations = getattr(omni_out, "stage_durations", None) or {}
- if durations:
- return dict(durations)
- ro = omni_out
- if ro is not None:
- return dict(getattr(ro, "stage_durations", None) or {})
- return {}
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
- add_inference_args(parser)
- parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
- parser.add_argument("--runs", type=int, default=3, help="Measured runs")
- parser.add_argument(
- "--enable-diffusion-pipeline-profiler",
- action="store_true",
- help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
- )
- parser.add_argument(
- "--enforce-eager",
- action="store_true",
- help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
- )
- parser.add_argument(
- "--dtype",
- type=str,
- default=None,
- choices=["bfloat16", "bf16", "float16", "fp16", "half"],
- help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
- )
- parser.add_argument(
- "--output",
- "-o",
- type=str,
- default=None,
- help="Save WAV from the last measured run (after warmup)",
- )
- parser.add_argument(
- "--save-each-run",
- action="store_true",
- help="With --output, also write measured runs as _run1.wav, _run2.wav, ...",
- )
- return parser.parse_args()
-
-
-def main() -> None:
- args = parse_args()
- preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
- deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
- mode = "SVS" if args.svs else "SVC"
-
- omni_kwargs: dict = {
- "model": args.model,
- "deploy_config": deploy_config,
- "async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
- }
- if args.enable_diffusion_pipeline_profiler:
- omni_kwargs["enable_diffusion_pipeline_profiler"] = True
- if args.enforce_eager:
- omni_kwargs["enforce_eager"] = True
- if args.dtype is not None:
- omni_kwargs["dtype"] = args.dtype
-
- compile_mode = "eager" if args.enforce_eager else "torch.compile"
- dtype_label = args.dtype or "from deploy YAML"
- print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
-
- omni = Omni(**omni_kwargs)
- kind = "svs" if args.svs else "svc"
- sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
- prompt = {"prompt_token_ids": [0]}
-
- latencies_ms: list[float] = []
- rtfs: list[float] = []
- last_stages: dict[str, float] = {}
- last_measured_outputs = None
-
- total_runs = args.warmup + args.runs
- for run_idx in range(total_runs):
- is_warmup = run_idx < args.warmup
- label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
- t0 = time.perf_counter()
- outputs = list(omni.generate([prompt], sampling))
- elapsed_ms = (time.perf_counter() - t0) * 1000.0
- audio_sec = _audio_duration_sec(outputs)
- rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
- print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
- if not is_warmup:
- latencies_ms.append(elapsed_ms)
- rtfs.append(rtf)
- last_stages = _stage_durations_from_output(outputs)
- last_measured_outputs = outputs
- if args.output and args.save_each_run:
- measured_idx = run_idx - args.warmup + 1
- out_path = Path(args.output)
- run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
- audio_np, sr = extract_audio(outputs)
- sf.write(str(run_path), audio_np, sr)
- print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
-
- omni.close()
-
- if args.output and last_measured_outputs is not None:
- audio_np, sr = extract_audio(last_measured_outputs)
- out_path = Path(args.output)
- out_path.parent.mkdir(parents=True, exist_ok=True)
- sf.write(str(out_path), audio_np, sr)
- print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
-
- if latencies_ms:
- print("\n=== Summary (measured runs) ===")
- print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
- print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
- if last_stages:
- print("\nStage durations (last measured run):")
- for name, value in sorted(last_stages.items()):
- if name.endswith("_ms"):
- print(f" {name}: {value:.1f} ms")
- else:
- print(f" {name}: {value * 1000.0:.1f} ms")
-
-
-if __name__ == "__main__":
- main()
diff --git a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py
deleted file mode 100644
index a1bb112af..000000000
--- a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py
+++ /dev/null
@@ -1,285 +0,0 @@
-"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
-
-Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
-
-Usage:
- # SVS (default)
- python end2end.py --model /path/to/SoulX-Singer \\
- --prompt-audio path/to/prompt/audio.mp3 \\
- --target-audio path/to/target/audio.mp3 \\
- --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
- -o output.wav
-
- # SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
- python end2end.py --model /path/to/SoulX-Singer-svc \\
- --prompt-audio path/to/prompt/audio.mp3 \\
- --target-audio path/to/target/audio.mp3 \\
- --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
- -o output_svc.wav
-"""
-
-import argparse
-import os
-from pathlib import Path
-
-import numpy as np
-import soundfile as sf
-
-from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-REPO_ROOT = Path(__file__).resolve().parents[4]
-SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
-SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
-DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
-DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
-_SAMPLE_RATE = 24000
-
-
-def _require_paths(paths: dict[str, str | None]) -> None:
- missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
- if missing:
- raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
-
-
-def resolve_preprocess_weights_dir(explicit: str | None) -> str:
- if explicit:
- path = Path(explicit).expanduser().resolve()
- if not (path / "rmvpe" / "rmvpe.pt").is_file():
- raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
- return str(path)
- env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
- if env:
- path = Path(env).expanduser().resolve()
- if (path / "rmvpe" / "rmvpe.pt").is_file():
- return str(path)
- raise FileNotFoundError(
- "SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
- "SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
- "(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
- )
-
-
-def build_sampling(
- args: argparse.Namespace,
- *,
- preprocess_weights_dir: str,
- kind: str,
-) -> OmniDiffusionSamplingParams:
- extra: dict = {
- "prompt_audio": os.path.abspath(args.prompt_audio),
- "target_audio": os.path.abspath(args.target_audio),
- "preprocess_weights_dir": preprocess_weights_dir,
- "vocal_sep": args.vocal_sep,
- "auto_shift": args.auto_shift,
- "pitch_shift": args.pitch_shift,
- }
- if kind == "svs":
- extra["language"] = args.language
- extra["control"] = args.control
- if args.prompt_metadata_path:
- extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
- if args.target_metadata_path:
- extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
- if args.prompt_metadata_path or args.target_metadata_path:
- extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
-
- extra = validate_soulx_extra_args(kind, extra)
-
- return OmniDiffusionSamplingParams(
- num_inference_steps=args.num_inference_steps,
- guidance_scale=args.guidance_scale,
- seed=args.seed,
- extra_args=extra,
- )
-
-
-def extract_audio(outputs) -> tuple[np.ndarray, int]:
- out = outputs[0]
- mm = getattr(out, "multimodal_output", None)
- if mm is None:
- ro = out
- if ro is not None:
- mm = getattr(ro, "multimodal_output", None)
- if mm is None and ro.outputs:
- mm = getattr(ro.outputs[0], "multimodal_output", None)
- if not mm or "audio" not in mm:
- raise RuntimeError("No audio in multimodal_output")
-
- audio_val = mm["audio"]
- if isinstance(audio_val, list):
- chunks = [
- np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
- for chunk in audio_val
- if chunk is not None
- ]
- audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
- elif hasattr(audio_val, "cpu"):
- audio_np = audio_val.detach().cpu().numpy().squeeze()
- else:
- audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
-
- sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
- if isinstance(sr_val, list) and sr_val:
- sr_val = sr_val[-1]
- if hasattr(sr_val, "item"):
- sr_val = sr_val.item()
- return audio_np.astype(np.float32, copy=False), int(sr_val)
-
-
-def add_inference_args(parser: argparse.ArgumentParser) -> None:
- mode = parser.add_mutually_exclusive_group()
- mode.add_argument(
- "--svc",
- action="store_true",
- help="Assert SVC mode (must match config.json architectures)",
- )
- parser.add_argument(
- "--model",
- type=str,
- required=True,
- help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
- )
- parser.add_argument(
- "--deploy-config",
- type=str,
- default=None,
- help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
- )
- parser.add_argument(
- "--prompt-audio",
- type=str,
- default=str(DEFAULT_PROMPT_AUDIO),
- help="Prompt / reference audio (mp3 or wav)",
- )
- parser.add_argument(
- "--target-audio",
- type=str,
- default=str(DEFAULT_TARGET_AUDIO),
- help="Target content audio (mp3 or wav)",
- )
- parser.add_argument(
- "--prompt-metadata-path",
- type=str,
- default=None,
- help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
- )
- parser.add_argument(
- "--target-metadata-path",
- type=str,
- default=None,
- help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
- )
- parser.add_argument(
- "--prompt-wav-path",
- type=str,
- default=None,
- help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
- )
- parser.add_argument(
- "--preprocess-weights-dir",
- type=str,
- default=None,
- help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
- )
- parser.add_argument(
- "--language",
- type=str,
- default="Mandarin",
- help="SVS lyrics language (Mandarin, Cantonese, English)",
- )
- parser.add_argument(
- "--control",
- type=str,
- default="melody",
- choices=["score", "melody"],
- help="SVS control mode",
- )
- parser.add_argument(
- "--vocal-sep",
- action=argparse.BooleanOptionalAction,
- default=False,
- help="Run vocal separation during preprocess",
- )
- parser.add_argument(
- "--auto-shift",
- action=argparse.BooleanOptionalAction,
- default=True,
- help="Auto pitch shift to align target pitch with prompt",
- )
- parser.add_argument("--pitch-shift", type=int, default=0)
- parser.add_argument("--num-inference-steps", type=int, default=32)
- parser.add_argument("--guidance-scale", type=float, default=3.0)
- parser.add_argument(
- "--seed",
- type=int,
- default=42,
- help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
- )
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
- add_inference_args(parser)
- parser.add_argument("--output", "-o", type=str, default="output.wav")
- args = parser.parse_args()
- return args
-
-
-def main() -> None:
- args = parse_args()
- kind = resolve_soulx_kind(args.model)
-
- if not (args.svc ^ (kind == "svs")):
- raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
-
- required = {
- "prompt_audio": args.prompt_audio,
- "target_audio": args.target_audio,
- }
- if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
- raise ValueError(
- "SVS precomputed metadata requires both --prompt-metadata-path and "
- "--target-metadata-path (or neither, for online preprocess)."
- )
- if args.prompt_metadata_path:
- required["prompt_metadata_path"] = args.prompt_metadata_path
- required["target_metadata_path"] = args.target_metadata_path
- if args.prompt_wav_path:
- required["prompt_wav_path"] = args.prompt_wav_path
- _require_paths(required)
- preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
- deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
- mode = kind.upper()
-
- print(f"Loading SoulX-Singer {mode} from {args.model}")
- print(f" deploy: {deploy_config}")
- print(f" prompt: {args.prompt_audio}")
- print(f" target: {args.target_audio}")
-
- omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
- sampling = build_sampling(
- args,
- preprocess_weights_dir=preprocess_dir,
- kind=kind,
- )
- prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
-
- print("Running generate (inline preprocess + DiT)...")
- outputs = list(omni.generate([prompt], sampling))
- if not outputs:
- raise RuntimeError("No output from omni.generate")
-
- err = getattr(outputs[0], "error", None)
- if err is not None:
- raise RuntimeError(f"Pipeline error: {err}")
-
- audio_np, sr = extract_audio(outputs)
- sf.write(args.output, audio_np, sr)
- print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
- omni.close()
-
-
-if __name__ == "__main__":
- main()
diff --git a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py
index c9225c212..e67e9c350 100644
--- a/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py
+++ b/examples/offline_inference/text_to_speech/voxtral_tts/end2end.py
@@ -304,7 +304,7 @@ def parse_args() -> Namespace:
"--cfg-alpha",
type=float,
default=None,
- help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).",
+ help="CFG alpha for flow-matching guidance (default: use the deploy config value, typically 1.2).",
)
parser.add_argument(
"--quantization",
diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md
index 339c7dac9..0cb75a8d4 100644
--- a/examples/offline_inference/text_to_video/text_to_video.md
+++ b/examples/offline_inference/text_to_video/text_to_video.md
@@ -14,6 +14,8 @@ A unified script for text-to-video generation. Supports multiple models with mod
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
+| `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | 480x832 | 81 | 50 | 6.0 | BF16 DiT + FP32 Wan VAE wrapper |
+| `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` | 704x1280 | 81 | 50 | 6.0 | BF16 DiT + LTX-2 Video VAE wrapper |
## Local CLI Usage
@@ -151,6 +153,40 @@ python text_to_video.py \
--output quick_test.mp4
```
+### SANA-Video-2B
+
+Select the native T2V pipeline explicitly:
+
+```bash
+python text_to_video.py \
+ --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ --model-class-name SanaVideoPipeline \
+ --prompt "A cinematic tracking shot of a sailboat crossing the ocean at sunset." \
+ --height 480 \
+ --width 832 \
+ --num-frames 81 \
+ --num-inference-steps 50 \
+ --guidance-scale 6.0 \
+ --extra-body '{"motion_score": 30}' \
+ --fps 16 \
+ --output sana_video_480p.mp4
+```
+
+For the 720p checkpoint, switch the model to
+`Efficient-Large-Model/SANA-Video_2B_720p_diffusers`, use
+`--height 704 --width 1280`, and write to a different output path. Both native
+checkpoint variants are supported. The Diffusers adapter T2V compatibility
+path is also validated at both resolutions through online serving.
+
+The native 480p path uses vLLM-Omni's `DistributedAutoencoderKLWan` wrapper;
+the native 720p path uses `DistributedAutoencoderKLLTX2Video`. Both wrappers
+retain their underlying Diffusers autoencoders. The pipeline intentionally
+uses Diffusers' checkpoint-compatible `DPMSolverMultistepScheduler`.
+
+For image-to-video, use the shared
+[`image_to_video.py`](../image_to_video/README.md#sana-video-2b) example with
+`--model-class-name SanaImageToVideoPipeline` and `--image `.
+
### Cosmos3
```bash
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 129ca860a..02383c6e4 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/examples/offline_inference/text_to_video/text_to_video.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import argparse
import json
@@ -10,7 +10,6 @@ from typing import Any
import numpy as np
import torch
-from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -115,12 +114,32 @@ _MODEL_PRESETS = {
"fps": 24,
"output": "ltx23_output.mp4",
},
+ "sana_480p": {
+ "height": 480,
+ "width": 832,
+ "num_frames": 81,
+ "num_inference_steps": 50,
+ "guidance_scale": 6.0,
+ "fps": 16,
+ "output": "sana_video_480p.mp4",
+ },
+ "sana_720p": {
+ "height": 704,
+ "width": 1280,
+ "num_frames": 81,
+ "num_inference_steps": 50,
+ "guidance_scale": 6.0,
+ "fps": 16,
+ "output": "sana_video_720p.mp4",
+ },
}
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
model_lower = model.lower()
class_lower = (model_class_name or "").lower()
+ if "sana-video" in model_lower or "sana_video" in model_lower or "sanavideo" in class_lower:
+ return _MODEL_PRESETS["sana_720p" if "720p" in model_lower else "sana_480p"]
if "lingbot" in model_lower or "lingbotvideo" in class_lower:
return _MODEL_PRESETS["lingbot"]
if "ltx" in class_lower or "ltx" in model_lower:
@@ -155,6 +174,13 @@ def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict
return result
+def _validate_video_output_type(output_type: str) -> None:
+ if output_type not in {"image", "video"}:
+ raise ValueError(
+ f"Unexpected output type '{output_type}', expected 'video' or legacy 'image' for video generation."
+ )
+
+
def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
"""Normalize decoded frames according to the pipeline's output contract."""
if not tensor.is_floating_point():
@@ -462,17 +488,6 @@ def main():
"scm_steps_policy": "dynamic",
}
- # Configure parallel settings
- parallel_config = DiffusionParallelConfig(
- ulysses_degree=args.ulysses_degree,
- ring_degree=args.ring_degree,
- cfg_parallel_size=args.cfg_parallel_size,
- tensor_parallel_size=args.tensor_parallel_size,
- vae_patch_parallel_size=args.vae_patch_parallel_size,
- pipeline_parallel_size=args.pipeline_parallel_size,
- enable_expert_parallel=args.enable_expert_parallel,
- )
-
profiler_enabled = args.profiler_config is not None
omni_kwargs = dict(
@@ -481,7 +496,13 @@ def main():
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
enable_cpu_offload=args.enable_cpu_offload,
- parallel_config=parallel_config,
+ ulysses_degree=args.ulysses_degree,
+ ring_degree=args.ring_degree,
+ cfg_parallel_size=args.cfg_parallel_size,
+ tensor_parallel_size=args.tensor_parallel_size,
+ vae_patch_parallel_size=args.vae_patch_parallel_size,
+ pipeline_parallel_size=args.pipeline_parallel_size,
+ enable_expert_parallel=args.enable_expert_parallel,
enforce_eager=args.enforce_eager,
model_class_name=model_class_name,
cache_backend=args.cache_backend,
@@ -570,6 +591,8 @@ def main():
num_frames=args.num_frames,
extra_args=extra_args,
)
+ if args.frame_rate is not None:
+ sampling_kwargs["frame_rate"] = args.frame_rate
if args.guidance_scale_high is not None:
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
@@ -604,10 +627,7 @@ def main():
frames = frames[0] if frames else None
if isinstance(frames, OmniRequestOutput):
- if frames.final_output_type != "image":
- raise ValueError(
- f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
- )
+ _validate_video_output_type(frames.final_output_type)
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
diff --git a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py b/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py
deleted file mode 100644
index 2f66d5f77..000000000
--- a/examples/offline_inference/x_to_video_audio/download_dreamid_omni.py
+++ /dev/null
@@ -1,108 +0,0 @@
-import argparse
-import fcntl
-import os
-import site
-import subprocess
-import tempfile
-import time
-from pathlib import Path
-
-from huggingface_hub import snapshot_download
-
-DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
-DEPENDENCY_BRANCH = "omni"
-CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
-LOCK_FILE = CACHE_DIR / ".install.lock"
-DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
-
-
-def download_dependency():
- CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
- with open(LOCK_FILE, "w") as f:
- fcntl.flock(f, fcntl.LOCK_EX)
- if not DEPENDENCY_DIR.exists():
- print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
- subprocess.run(
- ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
- check=True,
- )
- print("Download finished.")
- fcntl.flock(f, fcntl.LOCK_UN)
-
- # write .pth to site-packages
- site_packages = Path(site.getsitepackages()[0])
- pth_file = site_packages / "vllm_omni_dependency.pth"
- pth_file.write_text(str(DEPENDENCY_DIR))
- print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
-
-
-def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
- """Download files from HF repo and log time + destination."""
- if os.path.exists(local_dir):
- print(f"Directory {local_dir} already exists. Skipping download.")
- return
- print(f"Starting download from {repo_id} into {local_dir}")
- start_time = time.time()
-
- snapshot_download(
- repo_id=repo_id,
- local_dir=local_dir,
- local_dir_use_symlinks=False,
- allow_patterns=allow_patterns,
- )
-
- elapsed = time.time() - start_time
- print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
-
-
-def main(output_dir: str):
- # Wan2.2
- wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
- timed_download(
- repo_id="Wan-AI/Wan2.2-TI2V-5B",
- local_dir=wan_dir,
- allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
- )
-
- # MMAudio
- mm_audio_dir = os.path.join(output_dir, "MMAudio")
- timed_download(
- repo_id="hkchengrex/MMAudio",
- local_dir=mm_audio_dir,
- allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
- )
-
- dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
-
- timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
-
- # Now we construct the config file
- import json
-
- data = {
- "_class_name": "DreamIDOmniPipeline",
- }
-
- with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
- json.dump(data, f, indent=2)
-
- print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
-
- transformer_dir = os.path.join(output_dir, "transformer")
- os.makedirs(transformer_dir, exist_ok=True)
- with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
- json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
- print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
-
- # now we download the dependency code
- download_dependency()
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser(description="Download models from Hugging Face")
- parser.add_argument(
- "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
- )
- args = parser.parse_args()
- main(args.output_dir)
diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md
deleted file mode 100644
index 8305c785f..000000000
--- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md
+++ /dev/null
@@ -1,163 +0,0 @@
-# X-To-Video-Audio
-
-The `DreamID-Omni` pipeline generates short videos from text, image and video.
-
-## Local CLI Usage
-### Download the Model locally
-Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
-
-```bash
-python download_dreamid_omni.py --output-dir ./dreamid_omni
-```
-After download, the model directory will look like this:
-
-```
-dreamid_omni/
-├── DreamID-Omni/
-│ ├── dreamid_omni.safetensors
-├── MMAudio/
-│ ├── ext_weights/
-│ │ ├── best_netG.pt
-│ │ ├── v1-16.pth
-├── Wan2.2-TI2V-5B/
-│ ├── google/*
-│ ├── models_t5_umt5-xxl-enc-bf16.pth
-│ ├── Wan2.2_VAE.pth
-│
-├── model_index.json
-└── transformer/
- └── config.json # create by download_dreamid_omni.py
-```
-
-### Run the Inference
-```python
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt "Two people walking together and singing happily" \
- --image-path ./example0.png ./example1.png \
- --audio-path ./example0.wav ./example1.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --output out_dreamid_omni_twoip.mp4
-```
-In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
-The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
-For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
-
-```python
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt "Two people walking together and singing happily" \
- --image-path ./example0.png ./example1.png \
- --audio-path ./example0.wav ./example1.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --use-hsdp \
- --hsdp-shard-size 2 \
- --output out_dreamid_omni_twoip.mp4
-```
-
-
-You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
-
-For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing.
-
-```python
-# Example usage for oneip, ref media from the official repo DreamID-Omni
-python x_to_video_audio.py \
- --model /path/to/dreamid_omni \
- --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \
- --image-path 9.png \
- --audio-path 9.wav \
- --video-negative-prompt "jitter, bad hands, blur, distortion" \
- --audio-negative-prompt "robotic, muffled, echo, distorted" \
- --cfg-parallel-size 2 \
- --num-inference-steps 45 \
- --height 704 \
- --width 1280 \
- --output out_dreamid_omni_oneip.mp4
-```
-
-
-## MagiHuman (`--model-type magi-human`)
-
-MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
-T5Gemma text encoder. A detailed text prompt is the only required input; an optional
-image and/or audio file may be supplied for conditioning. Natively supports Tensor
-Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
-the MoE weights and the text encoder.
-
-> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
-> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
-
-Text-only generation:
-
-```bash
-python x_to_video_audio.py \
- --model-type magi-human \
- --model /path/to/daVinci-MagiHuman \
- --prompt "A young woman with long, wavy golden blonde hair... " \
- --tensor-parallel-size 4 \
- --height 256 --width 448 \
- --num-inference-steps 8 \
- --seed 52 \
- --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
- --output output_magihuman.mp4
-```
-
-With optional image and audio conditioning:
-
-```bash
-python x_to_video_audio.py \
- --model-type magi-human \
- --model /path/to/daVinci-MagiHuman \
- --prompt "A young woman..." \
- --tensor-parallel-size 4 \
- --height 256 --width 448 \
- --num-inference-steps 8 \
- --seed 52 \
- --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
- --output output_magihuman.mp4
-```
-
-MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
-in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
-
-- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
- because the audio length then determines the number of frames).
-- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
- the audio drives the frame count and is encoded as a latent condition into the DiT
- (`is_a2v` mode). Omit for pure text-to-video+audio generation.
-- `image_path`: path to an image file for visual conditioning. Applied at both the
- base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
-- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
- when these are omitted.
-- `sr_num_inference_steps`: denoising steps for the SR stage.
-
-## DreamID-Omni arguments
-
-Key arguments:
-- `--prompt`: text description (string).
-- `--model`: path to the model local directory.
-- `--height/--width`: output resolution (defaults 704 * 1024).
-- `--image-path`: path to the input image list.
-- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
-- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
-- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
-- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
-- `--hsdp-replicate-size`: number of HSDP replica groups.
-- `--num-inference-steps`: number of denoising steps (defaults 45).
-- `--video-negative-prompt`: negative prompt for video generation.
-- `--audio-negative-prompt`: negative prompt for audio generation.
-- `--enable-cpu-offload`: enable CPU offload (defaults False).
-- `--cache-backend`: enable `cache_dit` for acceleration.
-- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
- (VAEs, the T5 text encoder, norms and modulation stay bf16.)
diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py
deleted file mode 100644
index 9512ac65c..000000000
--- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py
+++ /dev/null
@@ -1,284 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import argparse
-import json
-import re
-import time
-from typing import Any
-
-import numpy as np
-from PIL import Image
-from vllm.multimodal.media.audio import load_audio
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
-AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
- parser.add_argument("--model", required=True, help="Model ckpt root directory.")
- parser.add_argument(
- "--model-type",
- default="dreamid-omni",
- choices=["dreamid-omni", "magi-human"],
- help="Model type.",
- )
- parser.add_argument("--prompt", default=None, help="Text prompt.")
-
- parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
- parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
- parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
- parser.add_argument(
- "--extra-body",
- type=str,
- default=None,
- help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
- 'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
- )
-
- parser.add_argument("--height", type=int, default=704, help="Video height.")
- parser.add_argument("--width", type=int, default=1280, help="Video width.")
- parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
- parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
- parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
- parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
- parser.add_argument(
- "--cfg-parallel-size",
- type=int,
- default=1,
- choices=[1, 2, 3, 4],
- help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
- )
- parser.add_argument(
- "--tensor-parallel-size",
- type=int,
- default=1,
- help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
- )
- parser.add_argument(
- "--video-negative-prompt",
- default="jitter, bad hands, blur, distortion",
- help="Negative prompt for video.",
- )
- parser.add_argument(
- "--audio-negative-prompt",
- default="robotic, muffled, echo, distorted",
- help="Negative prompt for audio.",
- )
- parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
- parser.add_argument(
- "--quantization",
- type=str,
- default=None,
- choices=["fp8", "int8"],
- help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
- )
- parser.add_argument(
- "--enable-cpu-offload",
- action="store_true",
- default=False,
- help="Enable CPU offloading for diffusion models.",
- )
- parser.add_argument(
- "--enable-layerwise-offload",
- action="store_true",
- help="Enable layerwise (blockwise) offloading on DiT modules.",
- )
- parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit"], help="Cache backend.")
- parser.add_argument(
- "--use-hsdp",
- action="store_true",
- help="Enable HSDP for DreamID-Omni fused transformer blocks.",
- )
- parser.add_argument(
- "--hsdp-shard-size",
- type=int,
- default=1,
- help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
- )
- parser.add_argument(
- "--hsdp-replicate-size",
- type=int,
- default=1,
- help="Number of HSDP replica groups. Default 1 means pure sharding.",
- )
- return parser.parse_args()
-
-
-def load_image_and_audio(image_paths, audio_paths):
- image = []
- audio = []
-
- for path in image_paths:
- with Image.open(path) as img:
- img = img.convert("RGB")
- image.append(img)
-
- for path in audio_paths:
- audio_array, sr = load_audio(path, sr=16000)
- audio_array = audio_array[int(sr * 1) : int(sr * 3)]
- audio.append(audio_array)
- return image, audio
-
-
-def _extract_peak_memory_mb(result: Any) -> float:
- """Pull worker-reported peak VRAM (MiB) generation result.
-
- Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
- """
- if isinstance(result, list):
- result = result[0] if result else None
- if result is None:
- return 0.0
- val = getattr(result, "peak_memory_mb", 0.0)
- if not val:
- inner = result
- if isinstance(inner, list):
- inner = inner[0] if inner else None
- val = getattr(inner, "peak_memory_mb", 0.0)
- try:
- return float(val or 0.0)
- except (TypeError, ValueError):
- return 0.0
-
-
-def main() -> None:
- args = parse_args()
- if args.prompt is None and args.prompt_file is None:
- raise ValueError("Either --prompt or --prompt-file must be provided.")
-
- text_prompt = args.prompt
- if args.prompt_file:
- with open(args.prompt_file) as f:
- text_prompt = json.load(f)
- text_prompt = re.sub(
- r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
- ).strip()
- text_prompt = re.sub(
- r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
- ).strip()
- text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
- text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
-
- if args.model_type == "magi-human":
- prompt = text_prompt
- extra_args = json.loads(args.extra_body) if args.extra_body else {}
- sampling_params = OmniDiffusionSamplingParams(
- height=args.height,
- width=args.width,
- num_inference_steps=args.num_inference_steps,
- seed=args.seed,
- extra_args=extra_args,
- )
- else:
- image, audio = load_image_and_audio(args.image_path, args.audio_path)
-
- prompt = {
- "prompt": text_prompt,
- "video_negative_prompt": args.video_negative_prompt,
- "audio_negative_prompt": args.audio_negative_prompt,
- "multi_modal_data": {"image": image, "audio": audio},
- }
-
- sampling_params = OmniDiffusionSamplingParams(
- height=args.height,
- width=args.width,
- num_inference_steps=args.num_inference_steps,
- seed=args.seed,
- extra_args={
- "solver_name": args.solver_name,
- "shift": args.shift,
- },
- )
-
- parallel_config = DiffusionParallelConfig(
- cfg_parallel_size=args.cfg_parallel_size,
- tensor_parallel_size=args.tensor_parallel_size,
- use_hsdp=args.use_hsdp,
- hsdp_shard_size=args.hsdp_shard_size,
- hsdp_replicate_size=args.hsdp_replicate_size,
- )
-
- cache_config = None
- if args.cache_backend == "cache_dit":
- cache_config = {
- "Fn_compute_blocks": 1,
- "Bn_compute_blocks": 0,
- "max_warmup_steps": 4,
- "max_cached_steps": 20,
- "residual_diff_threshold": 0.24,
- "max_continuous_cached_steps": 3,
- "enable_taylorseer": False,
- "taylorseer_order": 1,
- "scm_steps_mask_policy": None,
- "scm_steps_policy": "dynamic",
- }
-
- omni_kwargs = dict(
- model=args.model,
- parallel_config=parallel_config,
- model_type=args.model_type,
- enable_cpu_offload=args.enable_cpu_offload,
- enable_layerwise_offload=args.enable_layerwise_offload,
- cache_backend=args.cache_backend,
- cache_config=cache_config,
- )
- if args.quantization is not None:
- omni_kwargs["quantization"] = args.quantization
- omni = Omni(**omni_kwargs)
- start = time.perf_counter()
- outputs = omni.generate(prompt, sampling_params)
- elapsed = time.perf_counter() - start
-
- if not outputs:
- raise RuntimeError("No output returned from the model.")
- result = outputs[0]
- if not result.images:
- raise RuntimeError("No video frames found in model output.")
- generated_video = result.images[0]
- mm = result.multimodal_output or {}
- generated_audio = mm.get("audio")
-
- if args.model_type == "magi-human":
- # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
- fps = float(mm.get("fps", 25))
- sample_rate = int(mm.get("audio_sample_rate", 24000))
- frames = generated_video
- else:
- fps = float(mm.get("fps", 24))
- sample_rate = int(mm.get("audio_sample_rate", 16000))
- # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
- # mux_video_audio_bytes expects (F, H, W, C) uint8.
- if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
- raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
- frames = generated_video.transpose(1, 2, 3, 0)
- frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
-
- audio_np = None
- if generated_audio is not None:
- audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32)
-
- output_path = args.output
- video_bytes = mux_video_audio_bytes(
- frames,
- audio_np,
- fps=float(fps),
- audio_sample_rate=sample_rate,
- )
- with open(output_path, "wb") as f:
- f.write(video_bytes)
- print(f"Saved generated video to {output_path}")
- print(f"Total time: {elapsed:.2f}s")
- peak_mb = _extract_peak_memory_mb(outputs)
- if peak_mb:
- print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-
-
-if __name__ == "__main__":
- main()
diff --git a/examples/online_serving/audiox/README.md b/examples/online_serving/audiox/README.md
deleted file mode 100644
index bf068b3eb..000000000
--- a/examples/online_serving/audiox/README.md
+++ /dev/null
@@ -1,64 +0,0 @@
-# AudioX online serving
-
-Launches the `AudioXPipeline` behind vLLM-Omni's OpenAI-compatible chat endpoint and provides a
-minimal Python client that covers all six tasks (`t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`).
-
-## Start the server
-
-```bash
-cd examples/online_serving/audiox
-bash run_server.sh # defaults: MODEL=zhangj1an/AudioX, PORT=8099
-```
-
-Environment overrides: `MODEL`, `PORT`, `DIFFUSION_ATTENTION_BACKEND`.
-
-## Call from Python
-
-```bash
-# text-to-audio
-python openai_chat_client.py --task t2a \
- --prompt "Fireworks burst twice, followed by a period of silence before a clock begins ticking." \
- --output t2a.wav
-
-# text-to-music
-python openai_chat_client.py --task t2m \
- --prompt "Uplifting ukulele tune for a travel vlog" \
- --output t2m.wav
-
-# video-to-audio (no text)
-python openai_chat_client.py --task v2a --video path/to/clip.mp4 --output v2a.wav
-
-# text+video-to-audio
-python openai_chat_client.py --task tv2a \
- --prompt "drum beating sound and human talking" \
- --video path/to/clip.mp4 \
- --output tv2a.wav
-```
-
-The client sends:
-
-- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`,
- `seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server
- routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's
- `sampling_params.extra_args`.
-- For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files)
-
-## curl
-
-```bash
-curl -sS -X POST http://localhost:8099/v1/chat/completions \
- -H 'Content-Type: application/json' \
- -d '{
- "model": "zhangj1an/AudioX",
- "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
- "extra_body": {
- "num_inference_steps": 250,
- "guidance_scale": 7.0,
- "seed": 42,
- "audiox_task": "t2m",
- "seconds_total": 10.0,
- "sigma_min": 0.3,
- "sigma_max": 500.0
- }
- }' > t2m.json
-```
diff --git a/examples/online_serving/audiox/openai_chat_client.py b/examples/online_serving/audiox/openai_chat_client.py
deleted file mode 100755
index fd9aa8cc5..000000000
--- a/examples/online_serving/audiox/openai_chat_client.py
+++ /dev/null
@@ -1,119 +0,0 @@
-#!/usr/bin/env python3
-"""AudioX OpenAI-compatible chat client.
-
-AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the
-chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content
-item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed)
-are sent as flattened keys under ``extra_body``; the server routes the declared knobs into
-``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``).
-
-Usage:
- python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav
- python openai_chat_client.py --task tv2a --prompt "drum beating" --video clip.mp4 -o tv2a.wav
-"""
-
-from __future__ import annotations
-
-import argparse
-import base64
-import io
-import mimetypes
-import sys
-from pathlib import Path
-
-import requests
-import soundfile
-import torch
-
-VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
-TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"})
-
-
-def _to_data_url(path: str) -> str:
- mime, _ = mimetypes.guess_type(path)
- mime = mime or "video/mp4"
- with open(path, "rb") as f:
- data = base64.b64encode(f.read()).decode("ascii")
- return f"data:{mime};base64,{data}"
-
-
-def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
- audio = audio.to(torch.float32)
- audio = audio / audio.abs().max().clamp(min=1e-8)
- path.parent.mkdir(parents=True, exist_ok=True)
- # soundfile expects channels-last (T, C); project convention is (C, T).
- soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
-
-
-def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
- for choice in body.get("choices", []):
- audio_obj = choice.get("message", {}).get("audio")
- if not (isinstance(audio_obj, dict) and audio_obj.get("data")):
- continue
- data, sr = soundfile.read(io.BytesIO(base64.b64decode(audio_obj["data"])), dtype="float32", always_2d=True)
- return torch.from_numpy(data).transpose(0, 1), sr
- brief = {k: v for k, v in body.items() if k != "choices"}
- raise RuntimeError(f"no audio in response message.audio: {brief}")
-
-
-def main() -> int:
- p = argparse.ArgumentParser(description="AudioX OpenAI chat client")
- p.add_argument("--task", required=True, choices=["t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m"])
- p.add_argument("--prompt", "-p", default="", help="Text prompt (required for t2*/tv2*).")
- p.add_argument("--video", help="Video path or URL (required for v2*/tv2*).")
- p.add_argument("--output", "-o", default="audiox_out.wav")
- p.add_argument("--server", "-s", default="http://localhost:8099")
- p.add_argument("--model", default="zhangj1an/AudioX")
- p.add_argument("--steps", type=int, default=250)
- p.add_argument("--guidance-scale", type=float, default=7.0)
- p.add_argument("--seed", type=int, default=42)
- p.add_argument("--seconds-total", type=float, default=10.0)
- p.add_argument("--seconds-start", type=float, default=0.0)
- p.add_argument("--sigma-min", type=float, default=0.03)
- p.add_argument("--sigma-max", type=float, default=1000.0)
- args = p.parse_args()
-
- if args.task in VIDEO_TASKS and not args.video:
- print(f"ERROR: task {args.task!r} requires --video", file=sys.stderr)
- return 2
- if args.task in TEXT_TASKS and not args.prompt.strip() and args.task not in {"v2a", "v2m"}:
- print(f"ERROR: task {args.task!r} requires --prompt", file=sys.stderr)
- return 2
-
- content: list[dict] = [{"type": "text", "text": args.prompt}]
- if args.task in VIDEO_TASKS:
- vurl = args.video if args.video.startswith(("http://", "https://")) else _to_data_url(args.video)
- content.append({"type": "video_url", "video_url": {"url": vurl}})
-
- payload = {
- "model": args.model,
- "messages": [{"role": "user", "content": content}],
- "extra_body": {
- "num_inference_steps": args.steps,
- "guidance_scale": args.guidance_scale,
- "seed": args.seed,
- "audiox_task": args.task,
- "seconds_start": args.seconds_start,
- "seconds_total": args.seconds_total,
- "sigma_min": args.sigma_min,
- "sigma_max": args.sigma_max,
- },
- }
-
- print(f"POST {args.server}/v1/chat/completions task={args.task} steps={args.steps}")
- r = requests.post(
- f"{args.server}/v1/chat/completions",
- headers={"Content-Type": "application/json"},
- json=payload,
- timeout=600,
- )
- r.raise_for_status()
- audio, sr = _decode_audio_from_response(r.json())
- _save_wav(audio, Path(args.output), sr)
- dur = audio.shape[-1] / sr
- print(f"saved {args.output} sr={sr}Hz duration={dur:.2f}s channels={audio.shape[0]}")
- return 0
-
-
-if __name__ == "__main__":
- raise SystemExit(main())
diff --git a/examples/online_serving/audiox/run_server.sh b/examples/online_serving/audiox/run_server.sh
deleted file mode 100755
index d1a259bfd..000000000
--- a/examples/online_serving/audiox/run_server.sh
+++ /dev/null
@@ -1,16 +0,0 @@
-#!/bin/bash
-# AudioX online serving startup script.
-
-MODEL="${MODEL:-zhangj1an/AudioX}"
-PORT="${PORT:-8099}"
-DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-FLASH_ATTN}"
-
-echo "Starting AudioX server..."
-echo "Model: $MODEL"
-echo "Port: $PORT"
-echo "Diffusion attention backend: $DIFFUSION_ATTENTION_BACKEND"
-
-DIFFUSION_ATTENTION_BACKEND="$DIFFUSION_ATTENTION_BACKEND" \
- vllm serve "$MODEL" --omni \
- --model-class-name AudioXPipeline \
- --port "$PORT"
diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md
index e6d9b070f..7fac613df 100644
--- a/examples/online_serving/image_to_video/README.md
+++ b/examples/online_serving/image_to_video/README.md
@@ -1,5 +1,19 @@
# Image-To-Video
+This example demonstrates how to deploy image-to-video models for online
+video generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Model ID |
+|-------|----------|
+| Wan2.2 I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
+| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` |
+| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` |
+
+For SANA-Video 2B image-to-video online serving, see the
+[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md).
+
This example demonstrates online image-to-video generation with vLLM-Omni.
The startup script defaults to Wan2.2, but `MODEL` can select any supported
image-to-video model. The existing curl helper remains a concrete Wan2.2
diff --git a/examples/online_serving/image_to_video/run_curl_sana_video.sh b/examples/online_serving/image_to_video/run_curl_sana_video.sh
new file mode 100755
index 000000000..35806335f
--- /dev/null
+++ b/examples/online_serving/image_to_video/run_curl_sana_video.sh
@@ -0,0 +1,67 @@
+#!/bin/bash
+# Submit and download a standard SANA-Video-2B image-to-video request.
+
+set -euo pipefail
+
+INPUT_IMAGE="${INPUT_IMAGE:-}"
+BASE_URL="${BASE_URL:-http://localhost:8099}"
+OUTPUT_PATH="${OUTPUT_PATH:-sana_video_i2v.mp4}"
+WIDTH="${WIDTH:-832}"
+HEIGHT="${HEIGHT:-480}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
+
+if [ -z "$INPUT_IMAGE" ] || [ ! -f "$INPUT_IMAGE" ]; then
+ echo "Set INPUT_IMAGE to an existing image file."
+ exit 1
+fi
+
+create_response="$(
+ curl -sS -X POST "${BASE_URL}/v1/videos" \
+ -H "Accept: application/json" \
+ -F "prompt=A cat turns toward the camera with smooth, natural motion." \
+ -F "negative_prompt=blurry, low quality, temporal artifacts" \
+ -F "input_reference=@${INPUT_IMAGE}" \
+ -F "width=${WIDTH}" \
+ -F "height=${HEIGHT}" \
+ -F "num_frames=81" \
+ -F "fps=16" \
+ -F "num_inference_steps=50" \
+ -F "guidance_scale=6.0" \
+ -F "seed=42" \
+ -F 'extra_params={"motion_score":30}'
+)"
+
+video_id="$(echo "${create_response}" | jq -r '.id')"
+if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
+ echo "Failed to create SANA-Video job:"
+ echo "${create_response}" | jq .
+ exit 1
+fi
+
+while true; do
+ status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
+ status="$(echo "${status_response}" | jq -r '.status')"
+
+ case "${status}" in
+ queued|in_progress)
+ echo "Video job ${video_id} status: ${status}"
+ sleep "${POLL_INTERVAL}"
+ ;;
+ completed)
+ break
+ ;;
+ failed)
+ echo "SANA-Video generation failed:"
+ echo "${status_response}" | jq .
+ exit 1
+ ;;
+ *)
+ echo "Unexpected status response:"
+ echo "${status_response}" | jq .
+ exit 1
+ ;;
+ esac
+done
+
+curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
+echo "Saved SANA-Video to ${OUTPUT_PATH}"
diff --git a/examples/online_serving/image_to_video/run_server_sana_video.sh b/examples/online_serving/image_to_video/run_server_sana_video.sh
new file mode 100755
index 000000000..257cb2b00
--- /dev/null
+++ b/examples/online_serving/image_to_video/run_server_sana_video.sh
@@ -0,0 +1,13 @@
+#!/bin/bash
+# SANA-Video-2B native image-to-video serving.
+
+set -euo pipefail
+
+MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}"
+PORT="${PORT:-8099}"
+
+vllm serve "$MODEL" \
+ --omni \
+ --model-class-name SanaImageToVideoPipeline \
+ --dtype bfloat16 \
+ --port "$PORT"
diff --git a/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh
new file mode 100755
index 000000000..bf0573529
--- /dev/null
+++ b/examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+# SANA-Video-2B image-to-video serving through the Diffusers adapter.
+# Both 480p and 720p checkpoints are validated. For 720p:
+# MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers bash run_server_sana_video_diffusers.sh
+# INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 bash run_curl_sana_video.sh
+
+set -euo pipefail
+
+MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}"
+PORT="${PORT:-8099}"
+
+vllm serve "$MODEL" \
+ --omni \
+ --model-class-name SanaImageToVideoPipeline \
+ --diffusion-load-format diffusers \
+ --diffusion-attention-backend TORCH_SDPA \
+ --dtype bfloat16 \
+ --port "$PORT"
diff --git a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py
index 3030fa3e3..f2aa159cf 100644
--- a/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py
+++ b/examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py
@@ -328,7 +328,7 @@ def run_multimodal_generation(args) -> None:
extra_body = {
"sampling_params_list": sampling_params_list
- # Optional, it has a default setting in stage_configs of the corresponding model.
+ # Optional; defaults come from the resolved pipeline and deploy config.
}
# Build messages list (filter None so concurrent tasks get valid structure)
diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md
index 6ed7d8c32..96c021a44 100644
--- a/examples/online_serving/minicpmo/README.md
+++ b/examples/online_serving/minicpmo/README.md
@@ -32,12 +32,14 @@ GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
profile admits at most four concurrent sequences per stage.
| deploy config | GPUs | Notes |
-|---|---|---|
+| --- | --- | --- |
| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
-| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
+
+Every profile sets `session_mode: duplex`, so native full-duplex is served
+from the same process as `/v1/chat/completions`.
The split pipeline preserves native-duplex epoch/turn identity, segment text,
turn completion, reference voice, and terminal-audio metadata through
@@ -55,8 +57,8 @@ vllm serve openbmb/MiniCPM-o-4_5 \
```
For local ModelScope checkpoints, replace `openbmb/MiniCPM-o-4_5` with the
-checkpoint path. To start the experimental native duplex backend, use
-`vllm_omni/deploy/minicpmo_4_5_duplex.yaml`.
+checkpoint path. For native full-duplex, connect `/v1/realtime?duplex=1` on
+this same server.
### Per-stage overrides
@@ -145,7 +147,7 @@ assistant template, so it is not an apples-to-apples accuracy run.
## Run the Realtime duplex CLI demo
-After the duplex backend is running, stream one WAV through the Realtime
+After the server is running, stream one WAV through the Realtime
WebSocket endpoint:
```bash
@@ -157,6 +159,33 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
--output-dir /tmp/minicpmo_realtime_duplex_demo
```
+Video input uses the same session. PyAV demuxes JPEG frames at `--video-fps`
+(default 1.0) and vLLM `load_audio` extracts a 16 kHz mono WAV unless
+`--input-wav` overrides the soundtrack. Frames keep their capture size
+(`--frame-max-side 0`); the server `process_image` normalizes at 448.
+
+`--stack-frames N` raises the visual refresh rate the way official duplex does:
+each 1 s unit also samples the `N-1` sub-frames captured inside it and tiles
+them into a single composite image sent next to that unit's base frame. The
+audio cadence never changes — a unit is always one second — and the wire always
+carries 2 images per unit however large `N` is, because the sub-frames share one
+composite. Official uses 5 for high refresh rate mode:
+
+```bash
+python examples/online_serving/minicpmo/realtime_duplex_demo.py \
+ --url ws://localhost:8099/v1/realtime?duplex=1 \
+ --model openbmb/MiniCPM-o-4_5 \
+ --input-video /path/to/clip.mp4 \
+ --stack-frames 5 \
+ --ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
+ --output-dir /tmp/minicpmo_realtime_duplex_video_demo
+```
+
+Detail inside a composite is capped by `scale_resolution=448` at
+`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for
+stacked frames, which the duplex adapter does not implement yet. Reading small
+text or digits out of a wide scene is limited by that, not by frame timing.
+
## Open the experimental browser client
The browser UI serves the page and proxies the same-origin Realtime WebSocket to
@@ -190,6 +219,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
two-response WAV, its `--input-sha256`, and an
`--expect-second-response-substring` value.
+## Run Omni-DuplexEval
+
+Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
+endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
+served from the same vLLM-Omni environment. This validates the vLLM-Omni
+duplex client and local judge integration; it does not reproduce the paper's
+original MiniCPM-o inference implementation.
+
+Start the duplex generation server:
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
+ --deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
+ --served-model-name openbmb/MiniCPM-o-4_5 \
+ --host 0.0.0.0 --port 8099
+```
+
+Start a separate multimodal judge. The allowed path must contain any local
+videos passed with `--judge-video-mode video_url`:
+
+```bash
+vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
+ --served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
+ --allowed-local-media-path /data/omni-duplex-eval \
+ --host 0.0.0.0 --port 8000
+```
+
+Generate, evaluate, and summarize:
+
+```bash
+vllm bench omni-duplex-eval --omni generate \
+ --url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
+ --model openbmb/MiniCPM-o-4_5 \
+ --ref-audio /data/ref.wav \
+ --dataset Hothan/Omni-DuplexEval \
+ --concurrency 2 \
+ --response-root /data/omni-duplex-eval/responses
+
+vllm bench omni-duplex-eval --omni evaluate \
+ --dataset Hothan/Omni-DuplexEval \
+ --response-root /data/omni-duplex-eval/responses \
+ --score-root /data/omni-duplex-eval/scores \
+ --judge-base-url http://127.0.0.1:8000 \
+ --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+ --judge-video-mode video_url \
+ --eval-workers 4
+
+vllm bench omni-duplex-eval --omni summarize \
+ --score-root /data/omni-duplex-eval/scores
+```
+
+The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
+splits, not dataset configurations. Pass one with `--split`, or omit it to
+run all nine splits. `--limit 1` is useful for a smoke test. Realtime
+generation records `clock=media`; artifacts generated with
+`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
+them unless `--allow-invalid-clock` is explicit.
+
## Related examples
- [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)
diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py
index b6c861747..b12c20a56 100644
--- a/examples/online_serving/minicpmo/realtime_duplex_demo.py
+++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Minimal single-input MiniCPM-o 4.5 Realtime duplex demo.
Run this after starting the duplex server. Strict lifecycle, overlap, and
@@ -9,7 +12,9 @@ from __future__ import annotations
import argparse
import asyncio
import base64
+import io
import json
+import math
import sys
import time
from pathlib import Path
@@ -28,6 +33,17 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
wait_for,
write_pcm16_wav,
)
+from vllm_omni.experimental.fullduplex.client import chunk_period_ms as _chunk_period_ms # noqa: E402
+from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
+ has_residual_model_unit as _has_residual_model_unit,
+)
+from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
+ reference_audio_data_url as _ref_audio_data_url,
+)
+from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
+ concat_frames_b64,
+ unit_subframe_offsets,
+)
class _StreamingOutputWriter:
@@ -133,31 +149,29 @@ def _latest_model_decision(
return decision
-def _chunk_period_ms(events: list[dict[str, object]]) -> int:
- for event in reversed(events):
- session = event.get("session")
- if not isinstance(session, dict):
- continue
- capabilities = session.get("capabilities")
- if not isinstance(capabilities, dict):
- continue
- chunk_period_ms = capabilities.get("chunk_period_ms")
- if isinstance(chunk_period_ms, int) and chunk_period_ms > 0:
- return chunk_period_ms
- return 1000
-
-
-def _has_residual_model_unit(pcm16: bytes, *, chunk_period_ms: int) -> bool:
- unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_period_ms // 1000
- return bool(unit_bytes > 0 and len(pcm16) % unit_bytes)
-
-
def _response_in_progress(events: list[dict[str, object]]) -> bool:
return sum(event.get("type") == "response.created" for event in events) > sum(
event.get("type") == "response.done" for event in events
)
+async def _commit_input_after_playback_checkpoint(client: RealtimeDuplexClient) -> float:
+ """Commit input only after preserving already-played assistant history.
+
+ Native duplex can finish several responses while the user is still
+ streaming. Under the ``ack_only`` playback policy, acknowledging those
+ responses after the final input commit is intentionally rejected: it could
+ append an old assistant message after the newly committed user message.
+ Checkpoint playback first so the server commits each assistant item at its
+ original history position. A second acknowledgement after the commit can
+ then safely advance response-local playback that drained in the meantime.
+ """
+ await client.acknowledge_playback()
+ commit_sent_at_s = time.monotonic()
+ await client.commit()
+ return commit_sent_at_s
+
+
def _event_count_after(
events: list[dict[str, object]],
event_type: str,
@@ -168,21 +182,258 @@ def _event_count_after(
return sum(event.get("type") == event_type for event in events[index + 1 :])
-def _ref_audio_data_url(path: str | None) -> str | None:
- if path is None:
- return None
- ref_path = Path(path).expanduser()
- return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
+def _probe_video_frames_and_fps(video_path: Path) -> tuple[int, float]:
+ """Frame count and average FPS via PyAV (same approach as Daily-Omni)."""
+ import av
+
+ with av.open(str(video_path)) as container:
+ stream = container.streams.video[0]
+ avg_fps = float(stream.average_rate) if stream.average_rate else 0.0
+ num_frames = stream.frames or 0
+ if num_frames <= 0:
+ duration: float | None = None
+ if stream.duration is not None and stream.time_base is not None:
+ duration = float(stream.duration * stream.time_base)
+ elif container.duration is not None:
+ duration = container.duration / av.time_base
+ if duration and avg_fps:
+ num_frames = int(math.floor(duration * avg_fps))
+ if num_frames <= 0:
+ stream.thread_type = "AUTO"
+ num_frames = sum(1 for _ in container.decode(stream))
+ if num_frames <= 0:
+ raise ValueError(f"No decodable video frames in {video_path}")
+ return num_frames, avg_fps or 1.0
+
+
+def _upright(image, rotation: int):
+ """Apply a container display-matrix rotation, in degrees, to ``image``.
+
+ PyAV hands back the stored raster and leaves the display matrix for the
+ caller, unlike the ffmpeg CLI which autorotates. Phone clips routinely
+ carry +/-90, so skipping this feeds the model a sideways frame: arrows and
+ digits stop being readable and the model mirrors the motion it reports.
+ """
+ from PIL import Image
+
+ quarter_turns = {
+ 90: Image.Transpose.ROTATE_90,
+ 180: Image.Transpose.ROTATE_180,
+ 270: Image.Transpose.ROTATE_270,
+ }
+ turn = quarter_turns.get(int(rotation) % 360)
+ return image.transpose(turn) if turn is not None else image
+
+
+def _decode_video_frames_rgb(video_path: Path, frame_idx: list[int]) -> list:
+ """Decode presentation-order ``frame_idx`` into upright RGB ``PIL.Image``s."""
+ import av
+
+ wanted = sorted({int(i) for i in frame_idx})
+ if not wanted:
+ return []
+
+ decoded: dict[int, object] = {}
+ with av.open(str(video_path)) as container:
+ stream = container.streams.video[0]
+ stream.thread_type = "AUTO"
+ cursor = 0
+ for position, frame in enumerate(container.decode(stream)):
+ if wanted[cursor] > position:
+ continue
+ image = _upright(frame.to_image(), getattr(frame, "rotation", 0) or 0)
+ while cursor < len(wanted) and wanted[cursor] <= position:
+ decoded[wanted[cursor]] = image
+ cursor += 1
+ if cursor >= len(wanted):
+ break
+
+ if not decoded:
+ raise ValueError(f"Decoded no frames from {video_path} for indices {wanted[:8]}")
+ last_image = decoded[max(decoded)]
+ return [decoded.get(int(i), last_image) for i in frame_idx]
+
+
+def _resize_max_side(image, max_side: int):
+ width, height = image.size
+ longest = max(width, height)
+ if longest <= max_side or max_side <= 0:
+ return image
+ scale = max_side / float(longest)
+ new_size = (max(1, int(round(width * scale))), max(1, int(round(height * scale))))
+ return image.resize(new_size)
+
+
+def _jpeg_bytes(image, *, quality: int = 95) -> bytes:
+ buf = io.BytesIO()
+ image.convert("RGB").save(buf, format="JPEG", quality=quality)
+ return buf.getvalue()
+
+
+def _to_pcm16_bytes(audio_np) -> bytes:
+ import numpy as np
+
+ arr = np.asarray(audio_np).reshape(-1)
+ if arr.dtype != np.int16:
+ arr = (np.clip(arr, -1.0, 1.0) * 32767.0).astype(np.int16)
+ return arr.tobytes()
+
+
+def _pcm16_wav_duration_s(path: Path) -> float:
+ pcm16 = read_pcm16_wav(path)
+ return len(pcm16) / float(PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+
+
+def _extract_video_frames(
+ video_path: Path,
+ work_dir: Path,
+ *,
+ fps: float,
+ max_side: int,
+ duration_s: float,
+ stack_frames: int = 1,
+) -> tuple[list[str], list[str | None]]:
+ """Demux JPEG frames only. Frame count follows ``duration_s``, not soundtrack.
+
+ Returns the base frames (one per model unit) and the parallel stacked track.
+ With ``stack_frames=N`` each unit also samples its ``N-1`` interior
+ sub-frames and tiles them into one composite, which is how official duplex
+ raises the visual refresh rate while keeping one second of audio per unit.
+ A unit whose interior falls past the end of the clip stacks nothing.
+ """
+ num_frames, avg_fps = _probe_video_frames_and_fps(video_path)
+ sample_fps = max(float(fps), 1e-6)
+ unit_s = 1.0 / sample_fps
+ timeline_s = max(float(duration_s), unit_s)
+ num_samples = max(1, int(math.ceil(timeline_s * sample_fps - 1e-9)))
+ offsets = unit_subframe_offsets(stack_frames, unit_s=unit_s)
+
+ def frame_at(seconds: float) -> int:
+ # Official get_video_frame_audio_segments uses truncating int(), not
+ # round(), so 29.97 fps lands on the same presentation index.
+ return min(int(seconds * avg_fps), num_frames - 1)
+
+ base_idx = [frame_at(unit * unit_s) for unit in range(num_samples)]
+ sub_times = [
+ [unit * unit_s + offset for offset in offsets if unit * unit_s + offset < timeline_s]
+ for unit in range(num_samples)
+ ]
+ sub_idx = [frame_at(seconds) for times in sub_times for seconds in times]
+ images = _decode_video_frames_rgb(video_path, base_idx + sub_idx)
+ sub_images = iter(images[len(base_idx) :])
+
+ frame_dir = work_dir / "frames"
+ frame_dir.mkdir(parents=True, exist_ok=True)
+ for stale in frame_dir.glob("*.jpg"):
+ stale.unlink()
+ frames_b64: list[str] = []
+ stacked_b64: list[str | None] = []
+ for index, (image, times) in enumerate(zip(images[: len(base_idx)], sub_times), start=1):
+ jpeg = _jpeg_bytes(_resize_max_side(image, max_side))
+ (frame_dir / f"frame_{index:05d}.jpg").write_bytes(jpeg)
+ frames_b64.append(base64.b64encode(jpeg).decode("ascii"))
+ interior = [_resize_max_side(next(sub_images), max_side) for _ in times]
+ if not interior:
+ stacked_b64.append(None)
+ continue
+ composite = concat_frames_b64(interior)
+ (frame_dir / f"stack_{index:05d}.jpg").write_bytes(base64.b64decode(composite))
+ stacked_b64.append(composite)
+ return frames_b64, stacked_b64
+
+
+def _extract_video_soundtrack(video_path: Path, work_dir: Path) -> Path:
+ """Write a 16 kHz mono WAV from the video soundtrack."""
+ try:
+ from vllm.multimodal.media.audio import load_audio
+ except ImportError as exc: # pragma: no cover
+ raise SystemExit("--input-video needs vLLM audio helpers (load_audio)") from exc
+
+ audio_np, sr = load_audio(str(video_path), sr=PCM16_SAMPLE_RATE, mono=True)
+ work_dir.mkdir(parents=True, exist_ok=True)
+ audio_path = work_dir / "input_16k.wav"
+ write_pcm16_wav(audio_path, _to_pcm16_bytes(audio_np), sample_rate_hz=int(sr))
+ return audio_path
+
+
+def _extract_video_input(
+ video_path: Path,
+ work_dir: Path,
+ *,
+ fps: float,
+ max_side: int,
+ stack_frames: int,
+) -> tuple[Path, list[str], list[str | None]]:
+ """Demux a video into duplex inputs: 16 kHz mono WAV plus JPEG frames.
+
+ Use this only when the video soundtrack is the audio source. When the
+ caller already has ``--input-wav``, extract frames against that WAV
+ duration instead so a silent clip cannot fail ``load_audio``.
+ """
+ audio_path = _extract_video_soundtrack(video_path, work_dir)
+ frames_b64, stacked_b64 = _extract_video_frames(
+ video_path,
+ work_dir,
+ fps=fps,
+ max_side=max_side,
+ duration_s=_pcm16_wav_duration_s(audio_path),
+ stack_frames=stack_frames,
+ )
+ return audio_path, frames_b64, stacked_b64
+
+
+def _resolve_duplex_av_inputs(
+ *,
+ input_wav: str | None,
+ input_video: str | None,
+ work_dir: Path,
+ fps: float,
+ max_side: int,
+ stack_frames: int = 1,
+) -> tuple[str, list[str], list[str | None]]:
+ """Return ``(wav_path, frames_b64, stacked_b64)``. External WAV wins over video audio."""
+ if input_video:
+ video_path = Path(input_video)
+ if input_wav:
+ frames, stacked = _extract_video_frames(
+ video_path,
+ work_dir,
+ fps=fps,
+ max_side=max_side,
+ duration_s=_pcm16_wav_duration_s(Path(input_wav)),
+ stack_frames=stack_frames,
+ )
+ return input_wav, frames, stacked
+ extracted_wav, frames, stacked = _extract_video_input(
+ video_path,
+ work_dir,
+ fps=fps,
+ max_side=max_side,
+ stack_frames=stack_frames,
+ )
+ return str(extracted_wav), frames, stacked
+ if not input_wav:
+ raise SystemExit("provide --input-wav, --input-video, or both")
+ return input_wav, [], []
async def run_demo(args: argparse.Namespace) -> dict[str, object]:
- input_pcm16 = read_pcm16_wav(Path(args.input_wav))
- if not input_pcm16:
- raise ValueError("input WAV has no audio")
-
output_dir = Path(args.output_dir)
stream_writer = _StreamingOutputWriter(output_dir)
+ input_wav, video_frames, stacked_video_frames = _resolve_duplex_av_inputs(
+ input_wav=args.input_wav,
+ input_video=args.input_video,
+ work_dir=output_dir / "video_input",
+ fps=args.video_fps,
+ max_side=args.frame_max_side,
+ stack_frames=args.stack_frames,
+ )
+
+ input_pcm16 = read_pcm16_wav(Path(input_wav))
+ if not input_pcm16:
+ raise ValueError("input WAV has no audio")
+
url = build_realtime_url(
args.url,
args.model,
@@ -200,10 +451,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
timeout_s=args.timeout_s,
)
stream_event_cursor = len(client.events.events)
- await client.stream_pcm16(
+ frames_sent = await client.stream_pcm16(
input_pcm16,
chunk_ms=args.chunk_ms,
realtime=not args.no_realtime_pacing,
+ video_frames=video_frames,
+ stacked_video_frames=stacked_video_frames,
)
commit_event_cursor = len(client.events.events)
stream_decision = _latest_model_decision(client.events.events, stream_event_cursor)
@@ -212,8 +465,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
chunk_period_ms=_chunk_period_ms(client.events.events),
)
wait_for_post_commit_decision = False
- commit_sent_at_s = time.monotonic()
- await client.commit()
+ commit_sent_at_s = await _commit_input_after_playback_checkpoint(client)
wait_error: str | None = None
committed_index: int | None = None
post_commit_decision: str | None = None
@@ -326,6 +578,10 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
committed_index,
),
},
+ "input_video": args.input_video,
+ "video_frames_extracted": len(video_frames),
+ "video_frames_sent": frames_sent,
+ "video_stacked_frames_extracted": sum(frame is not None for frame in stacked_video_frames),
"audio_bytes": len(audio),
"audio_chunk_count": len(stream_writer.audio_chunk_paths),
"audio_chunk_files": [str(path) for path in stream_writer.audio_chunk_paths],
@@ -372,7 +628,48 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
parser.add_argument("--session-id")
- parser.add_argument("--input-wav", required=True)
+ parser.add_argument(
+ "--input-wav",
+ help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
+ )
+ parser.add_argument(
+ "--input-video",
+ help=(
+ "Video to stream as omni duplex input. PyAV demuxes frames and "
+ "vLLM load_audio extracts a 16 kHz mono WAV (unless --input-wav overrides the audio)."
+ ),
+ )
+ parser.add_argument(
+ "--video-fps",
+ type=float,
+ default=1.0,
+ help="Frame extraction rate. The wire contract carries ~1 frame per second of audio.",
+ )
+ parser.add_argument(
+ "--frame-max-side",
+ type=int,
+ default=0,
+ help=(
+ "Optional client-side downscale of the longest side, in pixels. "
+ "0 (default) keeps the capture resolution so the server "
+ "process_image can normalize at scale_resolution=448 from a "
+ "higher-quality source. The web demo never resizes on the client."
+ ),
+ )
+ parser.add_argument(
+ "--stack-frames",
+ type=int,
+ default=1,
+ help=(
+ "Frames sampled per 1 s audio unit. 1 (default) sends the unit's "
+ "frame alone. N > 1 also samples the N-1 sub-frames inside the unit "
+ "and tiles them into one composite sent next to the base frame, so "
+ "motion (arrows, doors, digits changing) is visible at N fps while "
+ "the audio stays 1 s per unit. Official uses 5 for high refresh "
+ "rate mode; the wire carries 2 images per unit at any N because the "
+ "sub-frames share one composite."
+ ),
+ )
parser.add_argument(
"--ref-audio",
required=True,
@@ -388,7 +685,7 @@ def parse_args() -> argparse.Namespace:
"--temperature",
type=float,
default=None,
- help="Stage0 sampling temperature; response-required soft-interrupt uses 0.0.",
+ help="Stage 0 sampling temperature; omit to preserve the model default.",
)
parser.add_argument("--no-realtime-pacing", action="store_true")
parser.add_argument("--require-audio", action="store_true")
diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py
index 63a3f9ec4..fd74ebe71 100644
--- a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py
+++ b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py
@@ -415,7 +415,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
audio_path=audio_path,
)
extra_body = {
- # Optional, it has default settings in stage configs. you can override them here.
+ # Optional; defaults come from the resolved pipeline and deploy config.
}
if args.query_type == "use_audio_in_video":
extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
diff --git a/examples/online_serving/personaplex/README.md b/examples/online_serving/personaplex/README.md
index cf105dd6c..3b737d275 100644
--- a/examples/online_serving/personaplex/README.md
+++ b/examples/online_serving/personaplex/README.md
@@ -38,7 +38,7 @@ full-duplex control plane:
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
/path/to/personaplex-7b-v1 \
--omni \
- --stage-configs-path vllm_omni/deploy/personaplex.yaml
+ --deploy-config vllm_omni/deploy/personaplex.yaml
```
Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz
diff --git a/examples/online_serving/pi0/README.md b/examples/online_serving/pi0/README.md
new file mode 100644
index 000000000..1f424a88b
--- /dev/null
+++ b/examples/online_serving/pi0/README.md
@@ -0,0 +1,104 @@
+# π0 (Pi-Zero) VLA — OpenPI realtime serving
+
+[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
+model from Physical Intelligence: multi-camera images + a language instruction +
+robot proprioceptive state → a continuous action chunk via flow-matching denoising
+(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
+websocket protocol at `/v1/realtime/robot/openpi`.
+
+## Install extras
+
+The core `pip install -e .` does not include the OpenPI client used here:
+
+- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
+ repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
+download it (`MODEL=lerobot/pi0_base`) or point at a local copy
+(`MODEL=/path/to/pi0_base`).
+
+## Run the server
+
+```bash
+vllm serve lerobot/pi0_base --omni --port 8000 \
+ --served-model-name pi0 \
+ --deploy-config vllm_omni/deploy/pi0.yaml \
+ --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
+metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
+
+### Commonly adjusted `model_config` keys
+
+These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
+
+| Key | Default | Meaning |
+|---|---|---|
+| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
+| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
+| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
+| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
+| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
+| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
+| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
+| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
+| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
+
+The `policy_server_config` block below them is the OpenPI handshake metadata
+advertised to the client; keep its `action_horizon` / `action_dim` /
+`image_resolution` in sync with the `model_config` values above. The e2e test
+`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
+live server and asserts the advertised metadata matches these `pi0.yaml` values.
+
+## Run the client
+
+```bash
+python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
+ --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
+cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
+proprioceptive state to drive a robot.
+
+### Observation format
+
+The client sends a flat dict per inference:
+
+```python
+{
+ "observation.images.base_0_rgb": np.uint8[H, W, 3],
+ "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
+ "observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
+ "state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
+ "prompt": "pick up the red block",
+ "session_id": "", # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
+(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
+version-stable across transformers releases (the version-stability checks in
+`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
+lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
+
+## Limitations
+
+- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
+ fully supported (state passes through, actions are returned in the model's space).
+ Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
+ stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
+ not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
+ added before its actions are in real-world units.
diff --git a/examples/online_serving/pi0/openpi_client.py b/examples/online_serving/pi0/openpi_client.py
new file mode 100644
index 000000000..f2154b33d
--- /dev/null
+++ b/examples/online_serving/pi0/openpi_client.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
+
+Connects to a running π0 server (see README.md), sends a single robot
+observation (3 cameras + state + language instruction), and prints the returned
+continuous action chunk ``[action_horizon, action_dim]``.
+
+The server speaks the OpenPI websocket protocol:
+ connect -> server sends msgpack(policy_server_config metadata)
+ infer -> client sends msgpack(obs), server replies msgpack(actions ndarray)
+
+Usage:
+ python openpi_client.py --host 127.0.0.1 --port 8000 \
+ --prompt "pick up the red block"
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+import uuid
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+ # Avoid shadowing the installed ``openpi_client`` package with this dir.
+ example_dir = str(Path(__file__).resolve().parent)
+ removed = sys.path and sys.path[0] == example_dir
+ if removed:
+ sys.path.pop(0)
+ try:
+ from openpi_client.websocket_client_policy import WebsocketClientPolicy
+ finally:
+ if removed:
+ sys.path.insert(0, example_dir)
+except ImportError as exc: # pragma: no cover - runtime dependency guard
+ raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
+
+logger = logging.getLogger("pi0_openpi_client")
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+
+# π0 / pi0_base camera identities (must match the server's image_feature_keys).
+CAMERA_KEYS = (
+ "observation.images.base_0_rgb",
+ "observation.images.left_wrist_0_rgb",
+ "observation.images.right_wrist_0_rgb",
+)
+STATE_DIM = 32
+
+
+def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
+ """Build a single π0 observation with blank cameras + zero state.
+
+ Replace the zeros with real camera frames (HWC uint8) and proprioceptive
+ state to drive an actual robot.
+ """
+ obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
+ obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
+ obs["prompt"] = prompt
+ obs["session_id"] = session_id
+ return obs
+
+
+def _ws_uri(host: str, port: int, path: str) -> str:
+ """Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
+ ``host`` (it has no separate path arg)."""
+ if host.startswith(("ws://", "wss://")):
+ return host
+ return f"ws://{host}:{port}{path}"
+
+
+def _as_action_array(response: Any) -> np.ndarray:
+ if isinstance(response, dict) and response.get("type") == "error":
+ raise RuntimeError(f"Inference failed: {response.get('message', response)}")
+ return np.asarray(response, dtype=np.float32)
+
+
+def main() -> int:
+ parser = argparse.ArgumentParser(description="π0 OpenPI client")
+ parser.add_argument("--host", default=DEFAULT_HOST)
+ parser.add_argument("--port", type=int, default=DEFAULT_PORT)
+ parser.add_argument("--path", default=DEFAULT_PATH)
+ parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+ parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
+ args = parser.parse_args()
+ logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+ policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
+ metadata = dict(policy.get_server_metadata())
+ logger.info("Server metadata: %s", metadata)
+
+ session_id = str(uuid.uuid4())
+ for step in range(args.num_steps):
+ obs = _make_dummy_obs(args.prompt, session_id)
+ actions = _as_action_array(policy.infer(obs))
+ logger.info(
+ "[step %d] actions shape=%s mean=%.4f std=%.4f",
+ step,
+ actions.shape,
+ float(actions.mean()),
+ float(actions.std()),
+ )
+ if not np.isfinite(actions).all():
+ raise RuntimeError("Server returned non-finite actions.")
+ logger.info("Done — π0 OpenPI round-trip OK.")
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh
index ed09bac14..d0ed84b85 100644
--- a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh
+++ b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh
@@ -51,7 +51,7 @@ code2wav_sampling_params='{
"detokenize": true,
"repetition_penalty": 1.1
}'
-# Above is optional, it has a default setting in stage_configs of the corresponding model.
+# The block above is optional; defaults come from the resolved pipeline and deploy config.
# Define URLs for assets
MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md
index 74a18af32..926a58682 100644
--- a/examples/online_serving/qwen3_omni/README.md
+++ b/examples/online_serving/qwen3_omni/README.md
@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_deploy_config.yaml
```
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+ --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
use `--stage-overrides` on top of the default config:
@@ -513,7 +525,7 @@ The script supports the following arguments:
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
-If you have custom stage configs file:
+If you have a custom deploy config file:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
```
diff --git a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh
index 399c65481..0a059a03e 100644
--- a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh
+++ b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh
@@ -50,7 +50,7 @@ code2wav_sampling_params='{
"detokenize": true,
"repetition_penalty": 1.1
}'
-# Above is optional, it has a default setting in stage_configs of the corresponding model.
+# The block above is optional; defaults come from the resolved pipeline and deploy config.
# Define URLs for assets
MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
diff --git a/examples/online_serving/replica_data_parallel/bench_replica_dp.py b/examples/online_serving/replica_data_parallel/bench_replica_dp.py
index 977f475ac..d17899b78 100755
--- a/examples/online_serving/replica_data_parallel/bench_replica_dp.py
+++ b/examples/online_serving/replica_data_parallel/bench_replica_dp.py
@@ -15,7 +15,7 @@ so an output is a deterministic function of its own input and nothing else.
Measuring replica scaling
-------------------------
This client is replica-agnostic. Vary the *server's* replica count
-(`runtime.num_replicas` in the stage config; see run_server.sh) and keep the
+(`stages[].num_replicas` in the deploy config; see run_server.sh) and keep the
client workload fixed, then compare throughput:
N=1 -> baseline
diff --git a/examples/online_serving/stable_audio/README.md b/examples/online_serving/stable_audio/README.md
deleted file mode 100644
index 954ac4c85..000000000
--- a/examples/online_serving/stable_audio/README.md
+++ /dev/null
@@ -1,234 +0,0 @@
-# Stable Audio Online Serving
-
-Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
-
-## Features
-
-- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
-- **Flexible control**: Adjust audio length, guidance scale, inference steps
-- **Quality control**: Use negative prompts to avoid unwanted characteristics
-- **Reproducible**: Set random seed for deterministic generation
-
-## Quick Start
-
-### 1. Start the Server
-
-```bash
-vllm serve stabilityai/stable-audio-open-1.0 \
- --host 0.0.0.0 \
- --port 8091 \
- --gpu-memory-utilization 0.9 \
- --trust-remote-code \
- --enforce-eager \
- --omni
-```
-
-### 2. Generate Audio
-
-#### Using curl
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "The sound of a cat purring",
- "audio_length": 10.0
- }' --output cat.wav
-```
-
-#### Using Python Client
-
-```bash
-python stable_audio_client.py \
- --text "The sound of a cat purring" \
- --audio_length 10.0 \
- --output cat.wav
-```
-
-#### Using Bash Script
-
-```bash
-bash curl_examples.sh
-```
-
-## API Reference
-
-### Endpoint
-
-```
-POST /v1/audio/generate
-```
-
-### Request Body
-
-```json
-{
- "input": "Text description of the audio",
- "audio_length": 10.0,
- "audio_start": 0.0,
- "negative_prompt": "Low quality",
- "guidance_scale": 7.0,
- "num_inference_steps": 100,
- "seed": 42,
- "response_format": "wav"
-}
-```
-
-### Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text prompt describing the audio to generate |
-| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
-| `audio_start` | float | 0.0 | Audio start time in seconds |
-| `negative_prompt` | string | null | Text describing what to avoid in generation |
-| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
-| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
-| `seed` | int | null | Random seed for reproducibility |
-| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
-
-### Response
-
-Returns audio data in the requested format (default: WAV).
-
-## Usage Examples
-
-### Basic Generation
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "The sound of ocean waves"
- }' --output ocean.wav
-```
-
-### Custom Duration
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "A dog barking",
- "audio_length": 5.0
- }' --output dog_5s.wav
-```
-
-### High Quality with Negative Prompt
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "A piano playing a gentle melody",
- "audio_length": 10.0,
- "negative_prompt": "Low quality, distorted, noisy",
- "guidance_scale": 8.0,
- "num_inference_steps": 150
- }' --output piano_hq.wav
-```
-
-### Reproducible Generation
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "Thunder and rain sounds",
- "audio_length": 15.0,
- "seed": 42
- }' --output thunder.wav
-```
-
-### Quick Generation (Fewer Steps)
-
-For faster generation with slightly lower quality:
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "Birds chirping in a forest",
- "audio_length": 8.0,
- "num_inference_steps": 50
- }' --output birds_quick.wav
-```
-
-## Python Client Examples
-
-### Simple Generation
-
-```bash
-python stable_audio_client.py \
- --text "The sound of a cat purring"
-```
-
-### Custom Parameters
-
-```bash
-python stable_audio_client.py \
- --text "Thunder and rain" \
- --audio_length 15.0 \
- --negative_prompt "Low quality" \
- --guidance_scale 7.0 \
- --num_inference_steps 100 \
- --seed 42 \
- --output thunder.wav
-```
-
-### Different Output Format
-
-```bash
-python stable_audio_client.py \
- --text "Guitar playing" \
- --response_format mp3 \
- --output guitar.mp3
-```
-
-## Tips
-
-1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
-2. **Quality vs Speed**:
- - 50 steps: Fast, decent quality
- - 100 steps: Good balance (default)
- - 150+ steps: High quality, slower
-3. **Guidance Scale**:
- - Lower (3-5): More creative/varied
- - Default (7): Good balance
- - Higher (10+): More literal to prompt
-4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
-5. **Seeds**: Use same seed for reproducible results
-
-## Performance
-
-| Inference Steps | Quality | Speed | Use Case |
-|----------------|---------|-------|----------|
-| 50 | Good | Fast | Quick previews |
-| 100 (default) | Very Good | Medium | Production |
-| 150+ | Excellent | Slow | Final/critical audio |
-
-## Troubleshooting
-
-### Server not responding
-- Check if server is running: `curl http://localhost:8091/health`
-- Check server logs for errors
-
-### Audio quality issues
-- Increase `num_inference_steps` (e.g., 150)
-- Add negative prompts: `"Low quality, distorted, noisy"`
-- Increase `guidance_scale` for more prompt adherence
-
-### Generation timeout
-- Reduce `num_inference_steps`
-- Reduce `audio_length`
-- Check GPU memory with `nvidia-smi`
-
-### Wrong audio length
-- Ensure `audio_length` is within model limits (~47s max)
-- Adjust `audio_start` if trimming is needed
-
-## See Also
-
-- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
-- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
-- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)
diff --git a/examples/online_serving/stable_audio/curl_examples.sh b/examples/online_serving/stable_audio/curl_examples.sh
deleted file mode 100755
index 16a12ef1c..000000000
--- a/examples/online_serving/stable_audio/curl_examples.sh
+++ /dev/null
@@ -1,54 +0,0 @@
-#!/bin/bash
-# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
-
-# Example 1: Simple request with default parameters
-echo "Example 1: Simple request with default parameters"
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "The sound audience clapping and cheering in a stadium"
- }' --output stadium.wav
-
-# Example 2: Request with custom audio_length
-echo "Example 2: Custom audio length (5 seconds)"
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "The sound of a dog barking",
- "audio_length": 5.0
- }' --output dog_5s.wav
-
-# Example 3: Request with negative prompt for quality control
-echo "Example 3: With negative prompt"
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "A piano playing a gentle melody",
- "audio_length": 10.0,
- "negative_prompt": "Low quality, distorted, noisy"
- }' --output piano.wav
-
-# Example 4: Full control with all parameters
-echo "Example 4: Full control (custom length, guidance, steps, seed)"
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "Thunder and rain sounds",
- "audio_length": 15.0,
- "negative_prompt": "Low quality",
- "guidance_scale": 7.0,
- "num_inference_steps": 100,
- "seed": 42
- }' --output thunder_rain.wav
-
-# Example 5: Quick generation with fewer steps (faster but lower quality)
-echo "Example 5: Quick generation (fewer steps)"
-curl -X POST http://localhost:8091/v1/audio/generate \
- -H "Content-Type: application/json" \
- -d '{
- "input": "Ocean waves crashing on a beach",
- "audio_length": 8.0,
- "num_inference_steps": 50
- }' --output ocean.wav
-
-echo "All examples completed!"
diff --git a/examples/online_serving/stable_audio/stable_audio_client.py b/examples/online_serving/stable_audio/stable_audio_client.py
deleted file mode 100755
index 7f28a5245..000000000
--- a/examples/online_serving/stable_audio/stable_audio_client.py
+++ /dev/null
@@ -1,170 +0,0 @@
-#!/usr/bin/env python3
-"""
-OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
-
-This script demonstrates how to use the OpenAI-compatible speech API
-to generate audio from text using Stable Audio models.
-
-Examples:
- # Simple generation
- python stable_audio_client.py --text "The sound of a cat purring"
-
- # With custom duration
- python stable_audio_client.py --text "A dog barking" --audio_length 5.0
-
- # With all parameters
- python stable_audio_client.py --text "Thunder and rain" \
- --audio_length 15.0 \
- --negative_prompt "Low quality" \
- --guidance_scale 7.0 \
- --num_inference_steps 100 \
- --seed 42 \
- --output thunder.wav
-"""
-
-import argparse
-import sys
-
-import requests
-
-
-def parse_args():
- parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
- parser.add_argument(
- "--api_url",
- default="http://localhost:8091/v1/audio/generate",
- help="API endpoint URL",
- )
- parser.add_argument(
- "--text",
- default="The sound of a cat purring",
- help="Text prompt for audio generation",
- )
- parser.add_argument(
- "--audio_length",
- type=float,
- default=10.0,
- help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
- )
- parser.add_argument(
- "--audio_start",
- type=float,
- default=0.0,
- help="Audio start time in seconds",
- )
- parser.add_argument(
- "--negative_prompt",
- default="Low quality",
- help="Negative prompt for classifier-free guidance",
- )
- parser.add_argument(
- "--guidance_scale",
- type=float,
- default=7.0,
- help="Guidance scale for diffusion (higher = more adherence to prompt)",
- )
- parser.add_argument(
- "--num_inference_steps",
- type=int,
- default=100,
- help="Number of inference steps (higher = better quality, slower)",
- )
- parser.add_argument(
- "--seed",
- type=int,
- default=None,
- help="Random seed for reproducibility",
- )
- parser.add_argument(
- "--output",
- default="stable_audio_output.wav",
- help="Output file path",
- )
- parser.add_argument(
- "--response_format",
- default="wav",
- choices=["wav", "mp3", "flac", "pcm"],
- help="Audio output format",
- )
- return parser.parse_args()
-
-
-def generate_audio(args):
- """Generate audio using the API."""
-
- # Build request payload
- payload = {
- "input": args.text,
- "audio_length": args.audio_length,
- "audio_start": args.audio_start,
- "response_format": args.response_format,
- }
-
- # Add optional parameters
- if args.negative_prompt:
- payload["negative_prompt"] = args.negative_prompt
- if args.guidance_scale:
- payload["guidance_scale"] = args.guidance_scale
- if args.num_inference_steps:
- payload["num_inference_steps"] = args.num_inference_steps
- if args.seed is not None:
- payload["seed"] = args.seed
-
- print(f"\n{'=' * 60}")
- print("Stable Audio - Text-to-Audio Generation")
- print(f"{'=' * 60}")
- print(f"API URL: {args.api_url}")
- print(f"Prompt: {args.text}")
- print(f"Audio length: {args.audio_length}s")
- print(f"Negative prompt: {args.negative_prompt}")
- print(f"Guidance scale: {args.guidance_scale}")
- print(f"Inference steps: {args.num_inference_steps}")
- if args.seed is not None:
- print(f"Seed: {args.seed}")
- print(f"Output: {args.output}")
- print(f"{'=' * 60}\n")
-
- try:
- # Make the API request
- print("Generating audio...")
- response = requests.post(
- args.api_url,
- json=payload,
- headers={"Content-Type": "application/json"},
- timeout=300, # 5 minute timeout for long generations
- )
-
- # Check for errors
- if response.status_code != 200:
- print(f"Error: API returned status code {response.status_code}")
- print(f"Response: {response.text}")
- return False
-
- # Save the audio
- with open(args.output, "wb") as f:
- f.write(response.content)
-
- print(f"✓ Audio saved to {args.output}")
- print(f" File size: {len(response.content) / 1024:.1f} KB")
- return True
-
- except requests.exceptions.Timeout:
- print("Error: Request timed out. Try reducing inference steps or audio length.")
- return False
- except requests.exceptions.ConnectionError:
- print(f"Error: Could not connect to {args.api_url}")
- print("Make sure the server is running.")
- return False
- except Exception as e:
- print(f"Error: {e}")
- return False
-
-
-def main():
- args = parse_args()
- success = generate_audio(args)
- sys.exit(0 if success else 1)
-
-
-if __name__ == "__main__":
- main()
diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md
index 25e2a6785..ba2797dae 100644
--- a/examples/online_serving/step_audio2/README.md
+++ b/examples/online_serving/step_audio2/README.md
@@ -189,5 +189,5 @@ Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker
- Check server logs for errors
### Out of memory
-- Reduce `gpu_memory_utilization` in stage configs
+- Reduce `gpu_memory_utilization` in the deploy config
- Use a smaller batch size
diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md
index 1ea191a37..69a50540b 100644
--- a/examples/online_serving/text_to_speech/README.md
+++ b/examples/online_serving/text_to_speech/README.md
@@ -25,7 +25,6 @@ For the full list of supported architectures across all modalities, see
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
-| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) |
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead.
@@ -570,7 +569,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
"response_format": "pcm"
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
-Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
+Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
### Streaming WebSocket
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
@@ -715,65 +714,3 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
- Voice cloning is gated upstream and may require a recent `mistral_common`.
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
-
----
-
-## SoulX-Singer
-
-Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`).
-
-### Prerequisites
-
-Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory.
-
-### Launch
-
-```bash
-# SVS (default)
-export MODEL=/path/to/SoulX-Singer
-export PREPROCESS=/path/to/SoulX-Singer-Preprocess
-bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
-
-# SVC
-export MODE=svc
-export MODEL=/path/to/SoulX-Singer-svc
-bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh
-```
-
-Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly:
-```bash
-export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS
-vllm serve $MODEL --omni \
- --deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \
- --port 8192 --trust-remote-code --enforce-eager
-```
-
-### Sending requests
-
-Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`.
-
-```bash
-# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
-python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
- --prompt-audio /path/on/server/zh_prompt.mp3 \
- --target-audio /path/on/server/music.mp3 \
- --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \
- -o output.wav
-```
-
-Use precomputed metadata to skip online preprocess with following command:
-```bash
-python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \
- --prompt-metadata-path /path/on/server/zh_prompt.json \
- --target-metadata-path /path/on/server/music.json \
- --audio-path /path/on/server/zh_prompt.mp3 \
- -o output.wav
-```
-
-`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`.
-
-### Notes
-
-- Output: 24 kHz mono WAV; batch only.
-- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
-- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
diff --git a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py
deleted file mode 100644
index 91035e625..000000000
--- a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py
+++ /dev/null
@@ -1,186 +0,0 @@
-#!/usr/bin/env python3
-"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC).
-
-Sends prompt audio via ``input_audio`` and target accompaniment via
-``extra_args['target_audio']`` (server-local path). For integrated preprocess,
-also pass ``preprocess_weights_dir`` in ``extra_args``.
-
-Usage:
- python openai_chat_client.py \\
- --prompt-audio /path/on/server/zh_prompt.mp3 \\
- --target-audio /path/on/server/music.mp3 \\
- --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\
- -o output.wav
-"""
-
-from __future__ import annotations
-
-import argparse
-import base64
-import io
-import sys
-from pathlib import Path
-
-import requests
-import soundfile
-import torch
-
-
-def _audio_to_data_url(path: Path) -> str:
- with path.open("rb") as handle:
- data = base64.b64encode(handle.read()).decode("ascii")
- return f"data:audio/mpeg;base64,{data}"
-
-
-def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
- audio = audio.to(torch.float32)
- peak = audio.abs().max().clamp(min=1e-8)
- audio = audio / peak
- path.parent.mkdir(parents=True, exist_ok=True)
- soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
-
-
-def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]:
- for choice in body.get("choices", []):
- audio_obj = choice.get("message", {}).get("audio")
- if isinstance(audio_obj, dict) and audio_obj.get("data"):
- data, sr = soundfile.read(
- io.BytesIO(base64.b64decode(audio_obj["data"])),
- dtype="float32",
- always_2d=True,
- )
- return torch.from_numpy(data).transpose(0, 1), sr
- brief = {k: v for k, v in body.items() if k != "choices"}
- raise RuntimeError(f"no audio in response message.audio: {brief}")
-
-
-def main() -> int:
- repo_root = Path(__file__).resolve().parents[4]
- default_assets = repo_root / "tests" / "assets" / "soulxsinger"
-
- parser = argparse.ArgumentParser(description="SoulX-Singer online chat client")
- parser.add_argument("--port", type=int, default=8192)
- parser.add_argument("--model", default="Soul-AILab/SoulX-Singer")
- parser.add_argument(
- "--prompt-audio",
- default=str(default_assets / "zh_prompt.mp3"),
- help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)",
- )
- parser.add_argument(
- "--target-audio",
- default=str(default_assets / "music.mp3"),
- help="Target accompaniment path on the server (extra_args['target_audio'])",
- )
- parser.add_argument(
- "--prompt-metadata-path",
- default=None,
- help="SVS precomputed prompt metadata.json",
- )
- parser.add_argument(
- "--target-metadata-path",
- default=None,
- help="SVS precomputed target metadata.json",
- )
- parser.add_argument(
- "--audio-path",
- default=None,
- help="SVS prompt vocal wav for precomputed metadata",
- )
- parser.add_argument("--preprocess-weights-dir", default=None)
- parser.add_argument("--output", "-o", default="soulxsinger_out.wav")
- parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs")
- parser.add_argument("--language", default="Mandarin")
- parser.add_argument("--num-inference-steps", type=int, default=32)
- parser.add_argument("--guidance-scale", type=float, default=3.0)
- parser.add_argument(
- "--seed",
- type=int,
- default=42,
- help="Optional CFM seed. Omit for non-deterministic sampling.",
- )
- parser.add_argument(
- "--auto-shift",
- action=argparse.BooleanOptionalAction,
- default=True,
- help="Auto pitch shift (default: on, original upstream infer.sh)",
- )
- parser.add_argument(
- "--control",
- default="melody",
- choices=["melody", "score"],
- help="SVS control mode",
- )
- parser.add_argument("--vocal-sep", action="store_true")
- args = parser.parse_args()
-
- meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path)
- if any(meta_paths) and not all(meta_paths):
- print(
- "ERROR: precomputed metadata requires --prompt-metadata-path, "
- "--target-metadata-path, and --audio-path together.",
- file=sys.stderr,
- )
- return 2
-
- extra_args: dict = {
- "vocal_sep": args.vocal_sep,
- "auto_shift": args.auto_shift,
- "pitch_shift": 0,
- }
- if all(meta_paths):
- extra_args.update(
- {
- "prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()),
- "target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()),
- "audio_path": str(Path(args.audio_path).expanduser().resolve()),
- }
- )
- content = [{"type": "text", "text": "soulx-singer"}]
- else:
- prompt_path = Path(args.prompt_audio).expanduser().resolve()
- if not prompt_path.is_file():
- print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr)
- return 2
- extra_args["prompt_audio"] = str(prompt_path)
- extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve())
- if args.preprocess_weights_dir:
- extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve())
- content = [
- {"type": "text", "text": "soulx-singer"},
- {
- "type": "input_audio",
- "input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"},
- },
- ]
- if not args.svc:
- extra_args["language"] = args.language
- extra_args["control"] = args.control
-
- payload = {
- "model": args.model,
- "modalities": ["audio"],
- "messages": [{"role": "user", "content": content}],
- "num_inference_steps": args.num_inference_steps,
- "guidance_scale": args.guidance_scale,
- "extra_args": extra_args,
- }
- if args.seed is not None:
- payload["seed"] = args.seed
-
- print(f"POST http://localhost:{args.port}/v1/chat/completions")
- response = requests.post(
- f"http://localhost:{args.port}/v1/chat/completions",
- headers={"Content-Type": "application/json"},
- json=payload,
- timeout=1800,
- )
- response.raise_for_status()
- audio, sample_rate = _decode_audio_from_response(response.json())
- _save_wav(audio, Path(args.output), sample_rate)
- duration = audio.shape[-1] / sample_rate
- print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s")
- return 0
-
-
-if __name__ == "__main__":
- raise SystemExit(main())
diff --git a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh
deleted file mode 100755
index 3e20de39b..000000000
--- a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh
+++ /dev/null
@@ -1,39 +0,0 @@
-#!/bin/bash
-# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline).
-#
-# Usage:
-# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \
-#
-# Audio paths in client extra_args must be readable on the server host.
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
-
-MODEL="${MODEL:-Soul-AILab/SoulX-Singer}"
-MODE="${MODE:-svs}"
-PORT="${PORT:-8192}"
-GPUS="${GPUS:-0}"
-
-if [[ "$MODE" == "svc" ]]; then
- DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}"
-else
- DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}"
-fi
-
-echo "Starting SoulX-Singer server"
-echo " MODEL=$MODEL"
-echo " MODE=$MODE"
-echo " PORT=$PORT"
-echo " DEPLOY_CONFIG=$DEPLOY_CONFIG"
-echo " CUDA_VISIBLE_DEVICES=$GPUS"
-
-CUDA_VISIBLE_DEVICES="$GPUS" \
-vllm serve "$MODEL" \
- --omni \
- --deploy-config "$DEPLOY_CONFIG" \
- --host 0.0.0.0 \
- --port "$PORT" \
- --trust-remote-code \
- --enforce-eager
diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md
index 51512fe0c..cf303ccb1 100644
--- a/examples/online_serving/text_to_video/README.md
+++ b/examples/online_serving/text_to_video/README.md
@@ -11,6 +11,11 @@ This example demonstrates how to deploy text-to-video models for online video ge
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
| LTX-2 | `Lightricks/LTX-2` |
| Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` |
+| SANA-Video 2B (480p) | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` |
+| SANA-Video 2B (720p) | `Efficient-Large-Model/SANA-Video_2B_720p_diffusers` |
+
+For SANA-Video 2B text-to-video online serving, see the
+[SANA-Video 2B recipe](../../../recipes/NVIDIA/SANA-Video-2B.md).
## Wan2.2 T2V
diff --git a/examples/online_serving/text_to_video/run_curl_sana_video.sh b/examples/online_serving/text_to_video/run_curl_sana_video.sh
new file mode 100755
index 000000000..490a743e5
--- /dev/null
+++ b/examples/online_serving/text_to_video/run_curl_sana_video.sh
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+BASE_URL="${BASE_URL:-http://localhost:8091}"
+OUTPUT_PATH="${OUTPUT_PATH:-sana_video.mp4}"
+
+curl -sS -X POST "$BASE_URL/v1/videos/sync" \
+ -F "prompt=A cinematic tracking shot of a sailboat crossing the ocean at sunset" \
+ -F "height=480" \
+ -F "width=832" \
+ -F "num_frames=81" \
+ -F "fps=16" \
+ -F "num_inference_steps=50" \
+ -F "guidance_scale=6.0" \
+ -F "seed=42" \
+ -F 'extra_params={"motion_score":30}' \
+ --output "$OUTPUT_PATH"
diff --git a/examples/online_serving/text_to_video/run_server_sana_video.sh b/examples/online_serving/text_to_video/run_server_sana_video.sh
new file mode 100755
index 000000000..446a75239
--- /dev/null
+++ b/examples/online_serving/text_to_video/run_server_sana_video.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+# SANA-Video-2B text-to-video serving.
+
+MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}"
+PORT="${PORT:-8091}"
+
+vllm serve "$MODEL" \
+ --omni \
+ --model-class-name SanaVideoPipeline \
+ --dtype bfloat16 \
+ --port "$PORT"
diff --git a/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh
new file mode 100755
index 000000000..adfb125c0
--- /dev/null
+++ b/examples/online_serving/text_to_video/run_server_sana_video_diffusers.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+# SANA-Video through the black-box Diffusers backend adapter.
+
+MODEL="${MODEL:-Efficient-Large-Model/SANA-Video_2B_480p_diffusers}"
+PORT="${PORT:-8091}"
+
+vllm serve "$MODEL" \
+ --omni \
+ --diffusion-load-format diffusers \
+ --diffusion-attention-backend TORCH_SDPA \
+ --dtype bfloat16 \
+ --port "$PORT"
diff --git a/mkdocs.yml b/mkdocs.yml
index 1a3b5109f..8a0e322e3 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -79,8 +79,13 @@ plugins:
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
user_guide/quantization/index.md: user_guide/quantization/overview.md
+ configuration/pd_disaggregation.md: features/pd_disaggregation.md
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
+ design/feature/cpu_offloading.md: design/feature/offloader/README.md
+ design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
+ design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
+ design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
- glightbox
- git-revision-date-localized:
# exclude files
diff --git a/pyproject.toml b/pyproject.toml
index 0e91bfae6..92dcb1a05 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,8 +40,11 @@ dev = [
"pytest-asyncio==1.4.0",
"pytest-cov==7.1.0",
"pytest-mock==3.15.1",
+ "pytest-shard==0.1.2",
+ "pytest-xdist==3.8.0",
"websockets==16.1",
"pytest-subtests==0.15.0",
+ "filelock>=3.16.1,<4",
"datasets==3.6.0",
"mypy==1.11.1",
"pre-commit==4.0.1",
@@ -68,9 +71,11 @@ dev = [
"voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
"torchdiffeq==0.2.5", # for covo_audio Test
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
- "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
"s3tokenizer==0.3.0", # for tts test
+ # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
+ # whose top-level stepaudio2 package conflicts with this distribution.
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
+ "silero-vad==6.1", # Realtime server-VAD duplex E2E
]
demo = [
@@ -80,6 +85,10 @@ demo = [
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
]
+server-vad = [
+ "silero-vad==6.1",
+]
+
forced-aligner = [
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
]
@@ -132,13 +141,6 @@ docs = [
"pydantic",
]
-soulx-svs = [
- "funasr>=1.0.0",
- "g2pM>=0.1.2.5",
- "g2p-en>=2.1.0",
- "ToJyutping>=3.2.0",
-]
-
[project.urls]
Homepage = "https://github.com/vllm-project/vllm-omni"
Repository = "https://github.com/vllm-project/vllm-omni"
@@ -189,6 +191,7 @@ select = [
"N", # pep8-naming
"UP", # pyupgrade
"TID251", # flake8-tidy-imports.banned-api
+ "TC004", # flake8-type-checking: runtime-used imports
]
ignore = [
"E203", # whitespace before ':' (conflicts with black)
@@ -224,19 +227,14 @@ ignore = [
"torch.cuda.synchronize".msg = "Use torch.accelerator.synchronize"
[tool.mypy]
-python_version = "3.12, 3.13"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = true
-disallow_incomplete_defs = true
+python_version = "3.12"
+# Align with vLLM: silent follow for library code; missing third-party stubs
+# are ignored everywhere (including tests). tests/ still get
+# --follow-imports skip from tools/pre_commit/mypy.py.
+plugins = ["pydantic.mypy"]
+ignore_missing_imports = true
check_untyped_defs = true
-disallow_untyped_decorators = true
-no_implicit_optional = true
-warn_redundant_casts = true
-warn_unused_ignores = true
-warn_no_return = true
-warn_unreachable = true
-strict_equality = true
+follow_imports = "silent"
[tool.pytest.ini_options]
asyncio_mode = "auto"
@@ -246,14 +244,15 @@ python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"--strict-markers",
- "--strict-config"
+ "--strict-config",
+ "--dist=loadgroup"
]
markers = [
- # ci/cd required
- "core_model: L1&L2 tests (run in each PR)",
- "advanced_model: L3 level tests (run on each merge)",
- "full_model: L4 level tests (run nightly)",
- "local_model: tests (run locally)",
+ # ci/cd required — tag with [ci-level] for tests.helpers.mark.get_level_markers().
+ "core_model: [ci-level] L1&L2 tests (run in each PR)",
+ "advanced_model: [ci-level] L3 level tests (run on each merge)",
+ "full_model: [ci-level] L4 level tests (run nightly)",
+ "local_model: [ci-level] tests (run locally)",
# function module markers
"diffusion: Diffusion model tests",
"omni: Omni model tests",
@@ -263,44 +262,64 @@ markers = [
"sp: Sequence parallelism tests",
"example: Doc example code tests",
# platform markers
- "cpu: Tests that run on CPU",
- "gpu: Tests that run on GPU (auto-added)",
- "cuda: Tests that run on CUDA (auto-added)",
- "rocm: Tests that run on AMD/ROCm (auto-added)",
- "xpu: Tests that run on XPU (auto-added)",
- "npu: Tests that run on NPU/Ascend (auto-added)",
- "musa: Tests that run on MUSA/Moore Threads (auto-added)",
+ # Tag with [hardware-platform] for tests.helpers.mark.get_supported_platforms().
+ # Add [gpu] on accelerator platforms that should also get pytest.mark.gpu.
+ # cpu/gpu are pytest marks; hardware_test(res=...) rejects them as keys.
+ "cpu: [hardware-platform] Tests that run on CPU",
+ "gpu: [hardware-platform] Tests that run on GPU (auto-added)",
+ "cuda: [hardware-platform] [gpu] Tests that run on CUDA (auto-added)",
+ "rocm: [hardware-platform] [gpu] Tests that run on AMD/ROCm (auto-added)",
+ "xpu: [hardware-platform] [gpu] Tests that run on XPU (auto-added)",
+ "npu: [hardware-platform] Tests that run on NPU/Ascend (auto-added)",
+ "musa: [hardware-platform] [gpu] Tests that run on MUSA/Moore Threads (auto-added)",
# specified computation resources marks (auto-added)
- # Tag with [hardware-resource] for tests.helpers.mark.get_hardware_mark_list().
- "H100: [hardware-resource] Tests that require H100 GPU",
- "H200: [hardware-resource] Tests that require H200 GPU",
- "L4: [hardware-resource] Tests that require L4 GPU",
- "B200: [hardware-resource] Tests that require B200 GPU",
- "MI325: [hardware-resource] Tests that require MI325 GPU (AMD/ROCm)",
- "B60: [hardware-resource] Tests that require Intel Arc Pro B60 XPU",
- "S5000: [hardware-resource] Tests that require S5000 GPU (Moore Threads/MUSA)",
- "A2: [hardware-resource] Tests that require A2 NPU",
- "A3: [hardware-resource] Tests that require A3 NPU",
- "distributed_cuda: Tests that require multi cards on CUDA platform",
- "distributed_rocm: Tests that require multi cards on ROCm platform",
- "distributed_xpu: Tests that require multi cards on XPU platform",
- "distributed_npu: Tests that require multi cards on NPU platform",
- "distributed_musa: Tests that require multi cards on MUSA platform",
- "skipif_cuda: Skip if the num of CUDA cards is less than the required",
- "skipif_rocm: Skip if the num of ROCm cards is less than the required",
- "skipif_xpu: Skip if the num of XPU cards is less than the required",
- "skipif_npu: Skip if the num of NPU cards is less than the required",
- "skipif_musa: Skip if the num of MUSA cards is less than the required",
+ # [hardware-resource] + [cuda]/[rocm]/[xpu]/[npu]/[musa] for hardware_test(res=...).
+ # Do not put [gpu] on SKU lines: that would make gpu a res key.
+ "H100: [hardware-resource] [cuda] Tests that require H100 GPU",
+ "H800: [hardware-resource] [cuda] Tests that require H800 GPU",
+ "H200: [hardware-resource] [cuda] Tests that require H200 GPU",
+ "L4: [hardware-resource] [cuda] Tests that require L4 GPU",
+ "B200: [hardware-resource] [cuda] Tests that require B200 GPU",
+ "MI325: [hardware-resource] [rocm] Tests that require MI325 GPU (AMD/ROCm)",
+ "B60: [hardware-resource] [xpu] Tests that require Intel Arc Pro B60 XPU",
+ "S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)",
+ "A2: [hardware-resource] [npu] Tests that require A2 NPU",
+ "A3: [hardware-resource] [npu] Tests that require A3 NPU",
+ # Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts().
+ "cards_1: [hardware-cards] Tests that require 1 accelerator card",
+ "cards_2: [hardware-cards] Tests that require 2 accelerator cards",
+ "cards_3: [hardware-cards] Tests that require 3 accelerator cards",
+ "cards_4: [hardware-cards] Tests that require 4 accelerator cards",
+ "cards_5: [hardware-cards] Tests that require 5 accelerator cards",
+ "cards_6: [hardware-cards] Tests that require 6 accelerator cards",
+ "cards_7: [hardware-cards] Tests that require 7 accelerator cards",
+ "cards_8: [hardware-cards] Tests that require 8 accelerator cards",
# more detailed markers
"sp: Sequence parallelism tests (multi-GPU)",
- "slow: Slow tests (may skip in quick CI)",
+ "slow: [ci-level] Slow tests (may skip in quick CI)",
"benchmark: Benchmark tests",
+ "xdist: xdist-parallelizable tests",
]
filterwarnings = [
"ignore:.*does not have '__test__' attribute.*:UserWarning",
"ignore:.*does not have '__bases__' attribute.*:UserWarning",
]
+[tool.coverage.run]
+core = "sysmon"
+source = ["vllm_omni"]
+# Online-serving tests launch the server via subprocess.Popen and stop it with
+# SIGTERM. pytest-cov no longer instruments subprocesses implicitly, and a
+# SIGTERM-ed process discards its data, so without patch/sigterm a coverage run
+# over an online test measures only the pytest parent.
+patch = ["subprocess"]
+concurrency = ["multiprocessing", "thread"]
+sigterm = true
+
+[tool.coverage.report]
+show_missing = true
+skip_covered = true
+
[tool.typos.default]
extend-ignore-identifiers-re = [
".*_thw",
@@ -327,6 +346,8 @@ extend-ignore-identifiers-re = [
]
extend-ignore-re = [
"\\.\\.\\.n,d->\\.\\.\\.nd",
+ # setuptools-scm / git-describe local versions, e.g. 0.26.1.dev14+gae6577ea
+ "\\+g[0-9a-f]{7,}",
]
[tool.typos.default.extend-words]
ue = "ue"
@@ -351,3 +372,5 @@ Leary = "Leary"
fo = "fo"
ba = "ba"
SER = "SER"
+# NumPy ndarray.flags.writeable (canonical flag name, not "writable").
+writeable = "writeable"
diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md
index 4c2d75aca..3b4b6c594 100644
--- a/recipes/Bagel/BAGEL-7B-MoT.md
+++ b/recipes/Bagel/BAGEL-7B-MoT.md
@@ -57,14 +57,14 @@ move the diffusion stage to a second GPU in a custom deploy config.
#### Offline Commands
Run text-to-image with the shared offline example from the repository root.
-Pick the topology with `--stage-configs-path` and forward BAGEL-specific
+Pick the topology with `--deploy-config` and forward BAGEL-specific
generation parameters as a JSON object through `--extra-body`:
```bash
# Two-stage (Thinker + DiT), shares one GPU by default
python examples/offline_inference/text_to_image/text_to_image.py \
--model ByteDance-Seed/BAGEL-7B-MoT \
- --stage-configs-path vllm_omni/deploy/bagel.yaml \
+ --deploy-config vllm_omni/deploy/bagel.yaml \
--prompt "A beautiful sunset over mountains" \
--height 1024 \
--width 1024 \
@@ -79,7 +79,7 @@ python examples/offline_inference/text_to_image/text_to_image.py \
# Single-stage (DiT only, with internal LLM/ViT/VAE)
python examples/offline_inference/text_to_image/text_to_image.py \
--model ByteDance-Seed/BAGEL-7B-MoT \
- --stage-configs-path vllm_omni/deploy/bagel_single_stage.yaml \
+ --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
--prompt "A beautiful sunset over mountains" \
--height 1024 \
--width 1024 \
diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md
index b9bff789f..b1edf0b4e 100644
--- a/recipes/Boogu/Boogu-Image.md
+++ b/recipes/Boogu/Boogu-Image.md
@@ -1,11 +1,14 @@
# Boogu-Image
-> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
+> Text-to-image and image-editing online serving
+(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
## Summary
- Vendor: Boogu
-- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
+- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
+ `Boogu/Boogu-Image-0.1-Edit` (image editing), and
+ `Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
- Task: Text-to-image generation and text-guided image editing (TI2I)
- Mode: Online serving with the OpenAI-compatible API
- Maintainer: Community
@@ -13,19 +16,23 @@
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
-`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
+`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
+`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
required).
Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
-family. This recipe covers the Base text-to-image checkpoint, which pairs a
-Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
-Euler scheduler with time-shift. It handles photorealistic generation and
-Chinese/English text rendering.
+family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
+with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
+time-shift. It handles photorealistic generation and Chinese/English text
+rendering. The Edit and Edit-Turbo checkpoints use the same native
+`BooguImagePipeline` for image editing.
## References
- Upstream model card:
+- Edit model card:
+- Edit-Turbo 1K hotfix:
- Project page:
- GitHub:
- Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
@@ -93,7 +100,7 @@ curl -s http://localhost:8091/v1/chat/completions \
- **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
to trim peak VRAM if needed.
- **Key flags:**
- - `--omni` — enables vLLM-Omni diffusion serving.
+ - `--omni` — enables vLLM-Omni diffusion serving.
- **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
`text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
active whenever `guidance_scale > 1.0`.
@@ -111,13 +118,13 @@ the image-editing (TI2I) path activates automatically when a request carries a
reference image. The Base text-to-image path is unaffected (no reference image
is sent).
-#### Command
+### Command
```bash
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
```
-#### Verification
+### Verification
Edit an image with `/v1/images/edits` (the model-card example — change a photo
to a colored-pencil drawing). Diffusion parameters are plain multipart form
@@ -157,14 +164,14 @@ curl -s http://localhost:8091/v1/chat/completions \
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
```
-#### Notes
+### Notes
- **Single reference image:** only one input image is supported for now (the
upstream "Only support 1 reference image for now" limit).
- **Guidance semantics:**
- - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
+ - `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
`4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
- - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
+ - `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
(3 model predictions per step), steering more strongly toward the reference
image.
@@ -178,3 +185,69 @@ curl -s http://localhost:8091/v1/chat/completions \
supports Boogu-Image-Edit directly
(`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
`--guidance-scale-2`).
+
+## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
+
+Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
+declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
+The upstream 1K hotfix is recommended for stable results and is pinned below
+to avoid silently loading the older checkpoint from the repository's default
+revision.
+
+### Command
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
+ --omni \
+ --revision hotfix-1k-20260708 \
+ --port 8091
+```
+
+### Verification
+
+Edit an image with the distilled four-step settings:
+
+```bash
+curl -s http://localhost:8091/v1/images/edits \
+ -F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \
+ -F image="@input.png" \
+ -F prompt="Change the style to a colored pencil drawing." \
+ -F num_inference_steps=4 \
+ -F guidance_scale=1.0 \
+ -F seed=42 \
+ | jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png
+```
+
+Or via chat completions:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+ -H "Content-Type: application/json" \
+ -d '{
+ "messages": [
+ {"role": "user", "content": [
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,"}},
+ {"type": "text", "text": "Change the style to a colored pencil drawing."}
+ ]}
+ ],
+ "extra_body": {
+ "num_inference_steps": 4,
+ "guidance_scale": 1.0,
+ "seed": 42
+ }
+ }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png
+```
+
+### Notes
+
+- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
+ hotfix over the 1.5K variant for more stable results.
+- **Recommended settings:** `num_inference_steps=4` and
+ `guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
+ behavior is already baked into the distilled checkpoint. A scale of `1.0`
+ avoids applying additional classifier-free guidance; unlike the regular Edit
+ checkpoint, it should not be increased to `5.0`.
+- **Reference images:** the same single-reference and `align_res` behavior as
+ the regular Edit checkpoint applies.
+- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
+ CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
diff --git a/recipes/BosonAI/Higgs-Audio-V3-TTS.md b/recipes/BosonAI/Higgs-Audio-V3-TTS.md
index b924ac472..0391fe03c 100644
--- a/recipes/BosonAI/Higgs-Audio-V3-TTS.md
+++ b/recipes/BosonAI/Higgs-Audio-V3-TTS.md
@@ -45,7 +45,7 @@ MusicGen-style delay pattern).
**Online serving:**
```bash
-vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
+vllm serve bosonai/higgs-audio-v3-tts-4b \
--host 0.0.0.0 --port 8095 \
--trust-remote-code --omni
```
diff --git a/recipes/GLM/GLM-Image.md b/recipes/GLM/GLM-Image.md
index bd11558c9..df6b512cd 100644
--- a/recipes/GLM/GLM-Image.md
+++ b/recipes/GLM/GLM-Image.md
@@ -52,7 +52,7 @@ Start the server from the repository root:
vllm serve zai-org/GLM-Image --omni --port 8091
```
-To use the bundled stage config explicitly (same default as above):
+To use the bundled deploy config explicitly (same default as above):
```bash
vllm serve zai-org/GLM-Image \
@@ -106,7 +106,7 @@ Overall summary from the run’s metrics. Rough wall-time split: **Stage 0 (AR)*
#### Notes
- Memory usage: Roughly **~38 GiB + KV** on Stage 0 (AR) and **~20 GiB** on Stage 1 (DiT+VAE) per the user guide; two 80 GB cards match the default split.
-- Key flags: `--omni` is required; `--stage-configs-path` is optional unless you use a custom YAML (for example single-GPU).
+- Key flags: `--omni` is required; `--deploy-config` is optional unless you use a custom YAML (for example single-GPU).
- Keep **Transformers ≥ 5.5.1** (this recipe used **5.5.4**) so `glm_image` configs resolve; otherwise Stage 0 can fail at `ModelConfig` validation.
- Known limitations: This starter recipe follows the dual-GPU online path documented under `examples/online_serving/glm_image`. The first request may be slower due to warmup.
- Generation time: about **61 s** wall time end-to-end for the sample above (50 inference steps, 1024×1024).
diff --git a/recipes/Helios/Helios.md b/recipes/Helios/Helios.md
index fdc0660bb..f1c11ed92 100644
--- a/recipes/Helios/Helios.md
+++ b/recipes/Helios/Helios.md
@@ -79,7 +79,7 @@ Total generation time: seconds ( ms)
Saved generated video to helios_t2v_base.mp4
```
-#### Important flags and stage configs
+#### Important flags and deploy config
- `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by
this recipe.
@@ -88,7 +88,7 @@ Saved generated video to helios_t2v_base.mp4
- `--cache-backend cache_dit` enables the cache-dit acceleration path.
- `--enable-cache-dit-summary` prints cache-dit summary information after
diffusion forward passes.
-- No separate stage config is required for this offline recipe; the shared
+- No separate deploy config is required for this offline recipe; the shared
`text_to_video.py` example configures the pipeline through its arguments.
- Helios-specific knobs (declared in `vllm_omni/model_extras/helios.py`) are
passed via the generic `--extra-body` JSON flag:
diff --git a/recipes/HiDream-ai/HiDream-O1-Image.md b/recipes/HiDream-ai/HiDream-O1-Image.md
new file mode 100644
index 000000000..1864324b2
--- /dev/null
+++ b/recipes/HiDream-ai/HiDream-O1-Image.md
@@ -0,0 +1,113 @@
+# HiDream-O1-Image
+
+> HiDream-O1 text-to-image generation through the shared offline example
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-O1-Image`
+- Task: Text-to-image
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to generate native-resolution images with HiDream-O1 on one
+or two CUDA GPUs. The two-GPU command shards the Qwen3-VL attention and MLP
+weights, as well as the timestep and pixel projections.
+
+## References
+
+- Upstream model:
+ [`HiDream-ai/HiDream-O1-Image`](https://huggingface.co/HiDream-ai/HiDream-O1-Image)
+- Related offline example:
+ [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- Related issue:
+ [`vllm-project/vllm-omni#3479`](https://github.com/vllm-project/vllm-omni/issues/3479)
+
+## Hardware Support
+
+## GPU
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+ --model HiDream-ai/HiDream-O1-Image \
+ --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+ --height 2048 \
+ --width 2048 \
+ --num-inference-steps 50 \
+ --guidance-scale 5.0 \
+ --seed 42 \
+ --output hidream_o1_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- The pipeline supports BF16 text-to-image generation.
+- Requested sizes are matched to the closest resolution supported by the checkpoint.
+
+### 2x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+ --model HiDream-ai/HiDream-O1-Image \
+ --tensor-parallel-size 2 \
+ --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+ --height 2048 \
+ --width 2048 \
+ --num-inference-steps 50 \
+ --guidance-scale 5.0 \
+ --seed 42 \
+ --output hidream_o1_tp2_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_tp2_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- Tensor parallelism shards the attention, MLP, timestep, and pixel projection weights across both GPUs.
+- Image editing and image-conditioned generation are not supported.
diff --git a/recipes/LTX/LTX-2.5.md b/recipes/LTX/LTX-2.5.md
new file mode 100644
index 000000000..f0be6948b
--- /dev/null
+++ b/recipes/LTX/LTX-2.5.md
@@ -0,0 +1,217 @@
+# LTX-2.5
+
+> Text-to-video and first-frame image-to-video generation with synchronized audio
+
+## Overview
+
+vLLM-Omni supports the gated
+[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+checkpoint for T2V and first-frame I2V. Every output includes synchronized
+48 kHz stereo audio.
+
+The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
+repository is not directly loadable with `--model`. It supplies the canonical
+Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
+the Full/SFT two-stage pipeline, so accept both model licenses and authenticate
+before first use.
+
+The vLLM-Omni integration code is provided under the project's
+[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate
+artifacts covered by the
+[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md),
+not the vLLM-Omni license. Users are responsible for reviewing and complying
+with the model license before downloading or using the weights.
+
+## Choose a pipeline
+
+| Pipeline | Mode | Default size | Schedule |
+| --- | --- | ---: | ---: |
+| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
+| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
+| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
+| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
+
+Both two-stage pipelines generate at half resolution, apply the official x2
+latent upsampler, and run a three-step refinement stage. Select the class with
+`--model-class-name`; no `--task-type` flag is required. Supplying one initial
+image selects I2V, while omitting it selects T2V.
+
+## Video decoder
+
+LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in
+to the legacy convolutional VAE, set the startup-only stage-0 model extra
+`ltx2_use_conv_vae: true`. The choice applies to all requests and all four
+pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub
+during startup. `decode_timestep` and `decode_noise_scale` condition only the
+legacy ConvVAE and have no effect when DiffVAE is selected.
+
+For offline Python usage:
+
+```python
+omni = Omni(
+ model="Lightricks/LTX-2.5-Diffusers",
+ model_class_name="LTX2Pipeline",
+ stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
+)
+```
+
+For online serving:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+ --omni \
+ --model-class-name LTX2Pipeline \
+ --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'
+```
+
+Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving
+serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial
+dimension. DiffVAE also supports distributed VAE decode; see the
+[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md).
+DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder.
+
+## Prerequisites
+
+```bash
+hf auth login
+export MODEL=Lightricks/LTX-2.5-Diffusers
+```
+
+Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
+`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
+`libx264` encoder.
+
+## Hardware
+
+| GPU | Status | Recommended scope |
+| --- | --- | --- |
+| NVIDIA B300 | Verified | All four canonical pipelines |
+| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
+
+The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
+80 GB GPUs do not have enough safety margin for the canonical two-stage
+configuration; use a larger GPU or reduce the output configuration.
+
+### H100 80 GB FP8
+
+Use the distilled one-stage pipeline with FP8 and cuDNN attention:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+ --omni \
+ --model-class-name LTX2DistilledOneStagePipeline \
+ --quantization fp8 \
+ --diffusion-attention-backend CUDNN_ATTN \
+ --host 0.0.0.0 \
+ --port 8000
+```
+
+Use the T2V or I2V request below after the server is ready.
+
+## Offline inference
+
+Choose values from the pipeline table. For example, the distilled two-stage
+path uses:
+
+```bash
+export PIPELINE=LTX2DistilledTwoStagePipeline
+export WIDTH=1920
+export HEIGHT=1088
+export STEPS=8
+```
+
+### Text-to-video
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+ --model "${MODEL}" \
+ --model-class-name "${PIPELINE}" \
+ --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
+ --width "${WIDTH}" \
+ --height "${HEIGHT}" \
+ --num-frames 121 \
+ --num-inference-steps "${STEPS}" \
+ --frame-rate 24 \
+ --fps 24 \
+ --seed 42 \
+ --output ltx25-t2v.mp4
+```
+
+### First-frame image-to-video
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+ --model "${MODEL}" \
+ --model-class-name "${PIPELINE}" \
+ --image /absolute/path/to/first-frame.png \
+ --prompt "The red fox walks forward while the camera tracks alongside." \
+ --width "${WIDTH}" \
+ --height "${HEIGHT}" \
+ --num-frames 121 \
+ --num-inference-steps "${STEPS}" \
+ --frame-rate 24 \
+ --fps 24 \
+ --seed 42 \
+ --output ltx25-i2v.mp4
+```
+
+LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
+
+## Online serving
+
+Start one server for the selected pipeline:
+
+```bash
+vllm serve "${MODEL}" \
+ --omni \
+ --model-class-name "${PIPELINE}" \
+ --host 0.0.0.0 \
+ --port 8000 \
+ --stage-init-timeout 900
+```
+
+T2V request:
+
+```bash
+curl -sS --fail-with-body \
+ -X POST http://127.0.0.1:8000/v1/videos/sync \
+ -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
+ -o ltx25-online-t2v.mp4
+```
+
+For I2V, add exactly one first frame:
+
+```bash
+export FIRST_FRAME=/absolute/path/to/first-frame.png
+
+curl -sS --fail-with-body \
+ -X POST http://127.0.0.1:8000/v1/videos/sync \
+ -F "input_reference=@${FIRST_FRAME};type=image/png" \
+ -F 'prompt=The red fox walks forward while the camera tracks alongside.' \
+ -o ltx25-online-i2v.mp4
+```
+
+Restart the server after changing `PIPELINE` because each class selects a
+different weight layout and execution topology.
+
+## Constraints
+
+- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
+ are optional and use the selected pipeline defaults when omitted. `seed` is
+ optional and random when omitted.
+- `num_frames` must be `8k+1` when overridden.
+- One-stage width and height must be divisible by 32; two-stage final
+ dimensions must be divisible by 64.
+- Full/SFT pipelines accept negative prompts. Distilled pipelines are
+ positive-only and reject negative prompts.
+- I2V accepts exactly one initial image per prompt.
+- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
+
+## References
+
+- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
+- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
+- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)
diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md
index 3142f9fd1..ec4e8fb3b 100644
--- a/recipes/LTX/LTX-2.md
+++ b/recipes/LTX/LTX-2.md
@@ -4,11 +4,22 @@
## Pipelines
-| Model | Checkpoint | Task | `--model-class-name` | Batching |
-|---|---|---|---|---|
-| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
-| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
-| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
+| `--model-class-name` | Task | Required checkpoint repositories |
+|---|---|---|
+| `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
+| `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
+| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
+| `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` `Lightricks/LTX-2.3` |
+| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` `Lightricks/LTX-2.3` |
+
+Repositories in the table are download units. A full pipeline repository
+contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
+and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
+The runtime first searches the model root for sidecars under their official
+filenames, then downloads them from the matching Lightricks Hub repository.
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
@@ -16,10 +27,13 @@ image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
-`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
-or provide one initial image for I2V. As in the official LTX pipelines,
-distilled mode is selected explicitly rather than inferred from the checkpoint
-name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
+`LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
+then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
+merged distilled Transformer without upsampling, while
+`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
+T2V and I2V; select their class explicitly. The deprecated
+`LTX2DistilledPipeline` name remains an alias for
+`LTX2DistilledTwoStagePipeline`.
## API Migration
@@ -43,8 +57,8 @@ The consolidation also removes these registry names without aliases:
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
-| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
-| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
+| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
+| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
Passing any second positional argument now raises `TypeError`. These changes
affect direct Python callers and explicit `--model-class-name` overrides;
@@ -68,24 +82,28 @@ The model recipe also supplies the default negative prompt. Top-level
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
-## Distilled Two-Stage Defaults
+## Two-Stage Defaults
-| Parameter | Value |
-|---|---:|
-| Final width × height | 1536 × 1024 |
-| Stage 1 width × height | 768 × 512 |
-| Frames / frame rate | 121 / 24 |
-| Stage 1 / Stage 2 denoise steps | 8 / 3 |
-| Guidance | Fixed positive-only |
+| Parameter | Ordinary | Full-distilled |
+|---|---:|---:|
+| Final width × height | 1536 × 1024 | 1536 × 1024 |
+| Stage 1 width × height | 768 × 512 | 768 × 512 |
+| Frames / frame rate | 121 / 24 | 121 / 24 |
+| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
+| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
-The dimensions passed to the API are the final output dimensions and must be
-divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
-The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
-`8` when supplied.
+API dimensions are final dimensions and must be divisible by 64. All LTX
+requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
+LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
+`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
+sigmas and input latents.
+
+Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
+and automatically switches to dynamic LoRA when quantization is enabled.
## Serving
-Start either checkpoint:
+Start either one-stage checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -95,11 +113,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
-Start the distilled checkpoint with its two-stage entry:
+Start a two-stage checkpoint with its explicit class:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
- --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+ --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
+# LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
+vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
+ --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
+# Ordinary LTX-2.3 two-stage
+vllm serve diffusers/LTX-2.3-Diffusers --omni \
+ --model-class-name LTX2TwoStagePipeline \
+ --enable-layerwise-offload \
+ --stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
@@ -135,8 +161,9 @@ together with `input_reference`.
## Guidance
-LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
-cross-modality guidance, and rescaling:
+One-stage and ordinary Stage 1 support independent video/audio CFG,
+spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
+Distilled stages and ordinary Stage 2 are fixed positive-only.
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
@@ -172,9 +199,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
applied after all pass predictions are gathered.
-The default one-stage recipe has four Transformer passes per denoise step:
-`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
-configurations are therefore:
+The default one-stage recipe and ordinary Stage 1 have four Transformer passes
+per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
+(cross-modality). The useful balanced configurations are therefore:
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|---:|---:|---:|---|
@@ -195,9 +222,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
```
The total device count is the product of `cfg_parallel_size` and the other
-configured parallel dimensions. Distilled pipelines use positive-only
-guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
-them.
+configured parallel dimensions. Positive-only distilled phases do not benefit
+from `cfg_parallel_size > 1`.
### Python API
@@ -242,22 +268,22 @@ noted below.
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
-| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
+| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
-| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
+| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
-| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
+| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
-| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
-| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
+| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
+| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
-| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
+| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
@@ -273,13 +299,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
-| Override | One-stage | Distilled two-stage |
-|---|---|---|
-| Guidance | Supported | Fixed positive-only |
-| Negative prompt/embeddings | Supported | Rejected |
-| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
-| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
-| Video/audio latents | Supported | Rejected |
+| Override | One-stage | Ordinary two-stage | Distilled two-stage |
+|---|---|---|---|
+| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
+| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
+| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
+| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
+| Video/audio latents | Supported | Rejected | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -309,6 +335,10 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
+- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
+ Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
+ divisible request shape.
+- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
@@ -323,14 +353,15 @@ bundled offline CLI do not currently expose `sigmas`.
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
-- Current two-stage support is limited to the distilled checkpoint; official
- LTX two-stage/HQ execution remains out of scope.
+- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
+ remains out of scope.
## References
-
-
-
+-
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md
index 7056a6f8f..6dc5f12e9 100644
--- a/recipes/MammothModa2/MammothModa2.md
+++ b/recipes/MammothModa2/MammothModa2.md
@@ -81,7 +81,7 @@ MammothModa2 generation parameters as a JSON object through `--extra-body`:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model ./MammothModa2-Preview \
- --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
+ --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
--height 1024 \
--width 1024 \
@@ -132,6 +132,44 @@ ls -lh mammoth_t2i.png
python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
```
+### 1x AMD MI300X, MammothModa2 Preview
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Offline Commands
+
+The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+ --model bytedance-research/MammothModa2-Preview \
+ --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+ --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+ --height 1024 \
+ --width 1024 \
+ --seed 42 \
+ --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+ --enable-diffusion-pipeline-profiler \
+ --log-stats \
+ --output mammoth_t2i.png
+```
+
+#### Verification
+
+The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
+
+The output was a valid 1024 by 1024 RGB PNG.
+
## MammothModa2-Dev unified inference
MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
diff --git a/recipes/MiniMaxAI/MiniMax-H3-4090.md b/recipes/MiniMaxAI/MiniMax-H3-4090.md
index 1e22f5d2e..782c1483b 100644
--- a/recipes/MiniMaxAI/MiniMax-H3-4090.md
+++ b/recipes/MiniMaxAI/MiniMax-H3-4090.md
@@ -110,9 +110,10 @@ outputs of each pair agreed to within 115 bytes.
Four-GPU T2VA repeats: 273.8 s and 269.5 s, peaking at 15,560 MiB and
15,612 MiB. Four-GPU Ref2VA is a single successful end-to-end run (545.2 s,
16,448 MiB). Follow-up Ref2VA requests on this build completed diffusion but
-then hit the engine's hardcoded 30 s async output wait
-(`_ASYNC_OUTPUT_TIMEOUT` in `diffusion_engine.py`) during post-compute D2H,
-so they are not reported as repeats.
+then hit the engine's async output wait during post-compute D2H, so they are
+not reported as repeats. That wait was a hardcoded 30 s when these numbers were
+taken; it is now `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` (default 600 s), so a rerun on
+current main should not lose these repeats to the bound.
Relative to the two-GPU baseline on the same shape, four-GPU TP2×USP2 was about
1.6× faster for both tasks, while per-GPU peak HBM stayed in the mid-15 GiB
diff --git a/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
new file mode 100644
index 000000000..debc89d89
--- /dev/null
+++ b/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md
@@ -0,0 +1,82 @@
+# MiniMax H3 disaggregated text encoder
+
+This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
+its hidden states and token-role metadata to an encoder-free diffusion stage.
+The standard MiniMax H3 recipes remain single-stage and continue to load the
+text encoder inside the diffusion pipeline.
+
+## Start the server
+
+Choose the topology explicitly and load its deployment defaults:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+ --omni \
+ --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml
+```
+
+The default deployment assigns stage 0 to GPUs 0-1 with tensor parallel size
+2 and `max_num_seqs: 1`. Stage 1 uses GPUs 2-5 with diffusion tensor parallel
+size 1, Ulysses degree 4, and VAE patch parallel size 4. Adjust the
+`devices`, `tensor_parallel_size`, and stage 1 `parallel_config` values in a
+deployment override for the available hardware. Diffusion quantization,
+layerwise offload, distributed layerwise offload, VAE parallelism, and USP
+settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
+
+For example, this five-GPU topology assigns one GPU to the encoder and four
+to the diffusion stage. `--stage-overrides` keeps placement and parallelism
+scoped to the owning stage rather than broadcasting an override to both:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3,4 \
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+ --omni \
+ --deploy-config vllm_omni/deploy/minimax_h3_disaggregated.yaml \
+ --stage-overrides '{"0":{"devices":"0","tensor_parallel_size":1},"1":{"devices":"1,2,3,4","tensor_parallel_size":1,"ulysses_degree":4,"vae_patch_parallel_size":4}}'
+```
+
+For memory-constrained deployments, start from the CPU-offload or distributed
+layerwise-offload profiles in [MiniMax-H3.md](MiniMax-H3.md). Apply memory and
+quantization options only to Stage 1 with `--stage-overrides`; retain the
+encoder's BF16 configuration and the video/audio VAEs' FP32 precision. Select
+one offload strategy per deployment:
+
+```bash
+# Stage 1 model-level CPU offload.
+--stage-overrides '{"1":{"enable_cpu_offload":true}}'
+
+# Stage 1 distributed layerwise offload. Tune resident layers for available RAM.
+--stage-overrides '{"1":{"enable_distributed_layerwise_offload":true,"dlo_use_allgather":false,"dlo_resident_layers":20}}'
+
+# Stage 1 online FP8 quantization of the DiT only.
+--stage-overrides '{"1":{"diffusion_quantization_config":"{\"transformer\":{\"method\":\"fp8\"}}"}}'
+```
+
+The Stage 1 VAE patch-parallel options remain independent of offload and
+quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
+hardware-qualified profiles before combining these options.
+
+Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
+text-encoder weights. This H3 topology explicitly keeps its single-replica
+diffusion stage inline, avoiding serialization of decoded video through a
+subprocess. It expects both stages to run in one deployment; cross-node payload
+transport is outside this configuration.
+
+The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
+`fl2va`, and `ref2va`) are unchanged from the single-stage recipe.
+
+## Turbo LoRA
+
+MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
+`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
+requests that omit sampling controls do not inherit the 50-step base schedule:
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+ --omni \
+ --lora-path /path/to/MiniMax-H3-Turbo \
+ --deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
+```
+
+The base deployment intentionally retains 50 inference steps for non-LoRA
+quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md
new file mode 100644
index 000000000..36eb78a5d
--- /dev/null
+++ b/recipes/MiniMaxAI/MiniMax-H3-NPU-950PR.md
@@ -0,0 +1,195 @@
+# MiniMax-H3 on a single NPU 950PR
+
+This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
+NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
+1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
+[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
+
+## Capacity requirements
+
+| Resource | Requirement |
+| --- | ---: |
+| NPU | 1x NPU 950PR |
+| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
+| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
+| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
+| Container shared memory | 8 GiB, or a writable local filesystem for spill |
+
+The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
+sized for exactly one resident task partition at 1024x576. Do not expect a
+larger output shape or concurrency greater than one to fit.
+
+Because a single 950PR holds the whole partition in HBM, this route does not
+need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
+GPU recipes does not apply. The validated run completed inside a container
+limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
+actively harmful: it stages weights in non-reclaimable host memory and the
+container's OOM killer terminates the server with exit code -9.
+
+
+## Environment
+
+- Host architecture: x86_64
+- driver: 25.7.rc1.6 (hal 7.35.23)
+- firmware: 9.0.0.105.229
+- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
+- npu-smi: 25.7.rc1.6
+- Python: 3.12.13
+- PyTorch: 2.10.0+cpu
+- torch_npu: 2.10.0.post2
+- vLLM: 0.26.0
+- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
+
+Install vLLM-Omni from a checkout with MiniMax-H3 support:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv pip install -e .
+```
+
+## Start a server
+
+```bash
+export MODEL=/path/to/MiniMax-H3/FL2VA
+export PORT=8000
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
+
+vllm serve "${MODEL}" \
+ --omni \
+ --host 0.0.0.0 \
+ --port "${PORT}" \
+ --trust-remote-code \
+ --num-gpus 1 \
+ --tensor-parallel-size 1 \
+ --usp 1 \
+ --ring 1 \
+ --text-encoder-tp-size 1 \
+ --vae-patch-parallel-size 1 \
+ --vae-parallel-mode tile \
+ --vae-use-tiling \
+ --init-timeout 14400 \
+ --stage-init-timeout 14400 \
+ --quantization int8
+```
+
+Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
+NPU implementation and fails at the first denoise step. Leave the backend
+unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
+installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
+
+H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
+
+## Validated evidence
+
+Measured on one NPU 950PR with the server command above plus
+`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
+the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
+requested 60 denoise steps and executed 59 denoise updates, so per-step latency
+is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
+124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
+
+Five requests were issued back to back with no idle time between them. The card
+is measurably faster on the first request than in steady state, so both regimes
+are reported.
+
+| Stage | First request, cold card | Thermal steady state |
+| --- | ---: | ---: |
+| End-to-end request | 462.13 s | 506.38 s |
+| Text encode | 0.55 s | 0.05 s |
+| Denoise (59 updates) | 429.29 s | 471.74 s |
+| Per denoise update | 7,276 ms | 7,996 ms |
+| VAE decode | 8.43 s | 8.14 s |
+| Worker-to-server handoff | 24 s | 27 s |
+| MP4 muxing | 0.57 s | 0.50 s |
+| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
+| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
+| Average NPU power | 566.9 W | 556.9 W |
+| Peak NPU temperature | 104 C | 104 C |
+
+Sample counts differ per row. The cold column is a single request. In the steady
+column, text encode, denoise and VAE decode are means over the four subsequent
+requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
+end-to-end, handoff and muxing are means over the two of those four that
+survived the handoff timeout described above.
+
+**Report the steady-state column.** The cold card completes denoise 9.9% faster
+while drawing only 1.8% more average power, and both regimes reach the same
+104 C ceiling, so the first request is buying a short window of higher clocks
+before the die saturates. Any benchmark that issues a single request against an
+idle card will overstate throughput by about 10%.
+
+**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
+the entire stage wall time by the 60 *requested* steps, so it absorbs text
+encoding, VAE decode and the 837 MiB handoff. All three successful requests
+satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
+per-update figure in the table is `diffuse / 59` from the pipeline profiler.
+
+Denoise accounts for 93% of the steady-state request and the handoff accounts
+for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
+to roughly 480 s.
+
+Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
+reserved high-water mark rather than live allocation: it reads 111,272 MiB after
+weight loading and before the first request, and does not fall back after
+requests complete. Re-measure for longer outputs, a different output shape, or
+concurrency greater than one.
+
+### Startup
+
+Weight loading took **31 min 50 s** on the validated host (first log line
+08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
+shared GlusterFS network mount. This is why the server command sets
+`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
+1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
+Reduce both timeouts when the partition is staged on local disk.
+
+## T2VA request example
+
+```bash
+export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
+
+curl -sS --max-time 1800 -X POST "${API_URL}" \
+ -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
+ -F 'width=1024' \
+ -F 'height=576' \
+ -F 'aspect_ratio=16:9' \
+ -F 'fps=24' \
+ -F 'num_inference_steps=60' \
+ -F 'flow_shift=12' \
+ -F 'seed=1101' \
+ -F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
+ -o t2va.mp4
+```
+
+## Known limitations
+
+- INT8 online quantization is validated for T2VA on this card. Use BF16 for
+ FL2VA and Ref2VA, or re-measure before relying on it.
+- Single-card serving loads one task partition at a time. For Ref2VA, stop the
+ server and restart it against the `Ref2VA` directory.
+- Sustained load drives the die to 104 C and costs about 10% of denoise
+ throughput relative to a cold card. Warm the server with at least one full
+ request before measuring, and treat steady state as the reportable number.
+- Requests can fail during result handoff on hosts where the IPC spill path is
+ slow. See [§ Result handoff and shared
+ memory](#result-handoff-and-shared-memory).
+- Loading the checkpoint from a network mount takes over 30 minutes and forces
+ very large init timeouts. Stage the partition on local disk when possible.
+- The configuration measured here is 1024x576 at 60 steps, which differs from
+ the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
+ cost scales superlinearly with token count, so the numbers are not directly
+ comparable across recipes.
+- The image ships a `triton` package whose backend is not built
+ (`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
+ at startup and disables Triton; the diffusion path does not need it and the
+ run completes normally.
+
+## Additional resources
+
+- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
+- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
+- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
+- [Supported models](../../docs/models/supported_models.md)
+- [Video API](../../docs/serving/videos_api.md)
diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md
index ffe8eee09..b4563c1f7 100755
--- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md
+++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md
@@ -118,6 +118,21 @@ On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at
above) or HSDP — see
[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3).
+### CPU MP4 response encoding (Atlas A2)
+
+Non-streaming MP4 responses use one public automatic encoder. It checks the
+runtime frame shape, common dtype, and RGB channel-plane contiguity for every
+request; compatible inputs use direct planar PyAV frames, while unsupported
+inputs fall back to the legacy muxer before opening the PyAV container. No CLI
+flag, model declaration, or user configuration is required. Streaming fMP4
+output remains on its existing incremental path.
+
+Current performance and correctness validation is limited to one Atlas A2 host
+with 8x Ascend 910B4-1 NPUs, the `FL2VA`/`t2va` partition, one request at a
+time, 1344x768 at 24 fps, and 5, 8.7, and 15 second requests. This optimization
+only changes CPU MP4 response encoding; it does not change DiT execution or
+stage 0.
+
### Optional optimizations
Two independent optimizations may be enabled on top of the configuration
diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md
index 114352343..71de893fd 100644
--- a/recipes/MiniMaxAI/MiniMax-H3.md
+++ b/recipes/MiniMaxAI/MiniMax-H3.md
@@ -65,6 +65,32 @@ On AMD ROCm, install without the `[fa4]` extra (FA4 is CUDA-only) and use
`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
reference-video preparation and MP4 output.
+## CPU MP4 response encoding
+
+For CUDA and ROCm deployments, non-streaming MP4 responses are encoded on the
+host CPU through PyAV/libx264 after generation. The response encoder selects the
+path automatically at runtime. The server-owned parallel converter accepts
+supported frame shapes and dtypes with either per-channel-contiguous or strided
+RGB planes, including interleaved arrays materialized by output transport.
+Standalone callers without a parallel converter retain the legacy fallback for
+strided planes. No CLI flag, model declaration, or user configuration is
+required. Streaming fMP4 output is unchanged.
+
+A community benchmark on 2x Xeon 8480C reported the following comparison
+between the legacy and direct planar paths
+([full result](https://github.com/vllm-project/vllm-omni/pull/6288#issuecomment-5337546499)):
+
+| Metric | Legacy | Direct planar | Change |
+| --- | ---: | ---: | ---: |
+| Median wall time | 1.805 s | 1.394 s | -22.8% |
+| Median process CPU time | 3.613 s | 3.207 s | -11.2% |
+| Peak RSS | 3182 MiB | 2794 MiB | -387 MiB (-12.2%) |
+
+Across the 1.0-8.7 s sweep, wall-time improvement was approximately 21.8-22.6%;
+outputs were byte-identical and full decode passed. This is evidence for host
+CPU response encoding. Actual gains depend on the CPU and runtime, and should
+not be interpreted as GPU, DiT, or stage 0 speedups.
+
## Start a server
Pass the repository ID directly. The pipeline uses `FL2VA` for model discovery
@@ -158,7 +184,7 @@ vllm serve "${MODEL}" \
Use the profile that matches the per-GPU memory capacity:
| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
-|---|---:|---:|---:|---|---|---|
+| --- | ---: | ---: | ---: | --- | --- | --- |
| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
@@ -177,8 +203,8 @@ increasing it on a different request shape.
At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
2 x RTX 5090 without OOM:
-| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
-|---:|---:|---:|---:|---|
+| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
+| -------: | ------------: | ---------: | ---------------------: | -----------------------------------------------------------: |
| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
This is a single end-to-end validation run, not a warmed multi-run latency
@@ -247,7 +273,7 @@ vllm serve "${MODEL}" \
Do not add `--enforce-eager` to this performance configuration. The first
request includes regional compilation; warm the server once before measuring
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
-1. The H3 VAE supports its native `tile` mode, not
+`1`. The H3 VAE supports its native `tile` mode, not
`spatial_shard_height` or `spatial_shard_width`.
### Attention Backends
@@ -355,14 +381,78 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while
`task=ref2va` routes to
`Ref2VA/transformer`. T2VA uses the FL2VA DiT.
+### Step execution and continuous batching
+
+H3 implements the step-wise execution contract, so the scheduler can admit and
+retire requests between denoise steps instead of running one request end to
+end. Add the feature gate, then raise `--max-num-seqs` to co-batch:
+
+```bash
+--step-execution --max-num-seqs 4
+```
+
+Co-batched requests are packed into a single sequence that keeps one attention
+document per request, so a batch costs one DiT forward. That packing requires
+`--diffusion-attention-backend FLASH_ATTN`; other backends fall back to one
+forward per request. `--max-num-seqs 1` keeps the conservative single-request
+step path. Cache acceleration (`--cache-backend`) is not available in step mode.
+
+!!! warning "Co-batching does not improve H3 throughput for large simultaneous requests"
+ Keep `--max-num-seqs 1` unless you specifically need scheduler-level control
+ (admitting and retiring requests between denoise steps). Measured on two
+ H100s (TP2, BF16, 672x384, 209 frames, 30 steps, 4 requests at concurrency
+ 4 with `--request-rate inf`, i.e. all four submitted at time 0; one packed
+ request is 16384 rows):
+
+ | Configuration | Wall time | Mean latency | Peak memory |
+ |---------------|-----------|--------------|-------------|
+ | request mode | 174.8 s | 111.5 s | 72.4 GB |
+ | `--step-execution --max-num-seqs 1` | 179.0 s | 113.8 s | 72.4 GB |
+ | `--step-execution --max-num-seqs 4` | 182.1 s | 175.7 s | 78.3 GB |
+
+ A single H3 denoise step is a compute-bound dense GEMM over an already long
+ packed sequence, so fusing N requests costs N times the FLOPs and buys almost
+ no amortization — unlike LLM decoding, which is memory-bandwidth bound.
+ Going from one request per step to four cuts the per-request denoise cost
+ only from 1.323 s to 1.291 s (2.4%), which the step-mode bookkeeping then
+ spends. Mean latency degrades further because co-batched requests finish
+ together instead of staggered. Quantization moves the absolute numbers
+ without changing this: with online `int8` the same workload runs in 153.3 s
+ at 56.9 GB (request mode), and `--max-num-seqs 4` is still 5.0% slower than
+ request mode.
+
+ Two workloads outside this table are unmeasured and may behave differently:
+
+ - **Staggered arrivals (admission latency).** A single 16384-row H3 request
+ already saturates the GPUs, so submitting all requests at time 0 is the
+ one arrival pattern where co-batching cannot win: it can only bunch
+ completions. In request mode a new request queues behind the whole
+ in-flight generation (~45 s at this size); step mode admits at the next
+ denoise-step boundary (~1.3 s). Whether that shows up as a wall-clock
+ benefit under Poisson arrivals is not yet measured for H3. To reproduce,
+ run request mode vs `--step-execution --max-num-seqs 4` with
+ `diffusion_benchmark_serving.py --request-rate 0.05` (roughly one request
+ every 20 s) and report mean / p95 latency, which then includes queueing.
+ - **Small requests.** The numbers above are drawn from 672x384 / 209-frame
+ requests. A short clip at lower resolution packs a few thousand rows and
+ may not saturate the hardware; co-batching may amortize better there.
+ This is also unmeasured.
+
### Online FP8 quantization
-MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
-remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
-uses dynamic activation scaling during inference. By default, attention and
-MLP linears in the token refiner and main DiT blocks, the condition
-projection, and all AdaLN projections use FP8. Patch, timestep, and final
-projections remain FP32; the text encoder and VAEs are unchanged.
+MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
+text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
+weights at runtime and uses dynamic activation scaling during inference. By
+default, `--quantization fp8` quantizes eligible attention and MLP linears in
+the text decoder, token refiner, and main DiT blocks, as well as the DiT
+condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
+RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
+keep checkpoint precision.
+
+To select a component, use `--diffusion-quantization-config` with
+`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
+`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
+entries can be combined. The shorthand below enables both components.
Add this option to an existing H3 server command:
@@ -370,6 +460,28 @@ Add this option to an existing H3 server command:
--quantization fp8
```
+#### Single 96 GB GPU, no-offload capacity check
+
+Use the FL2VA-only partition for this capacity test. Loading the combined
+service would also load the Ref2VA DiT and would test a different memory
+budget. A no-offload capacity check should contain none of
+`--enable-cpu-offload`, `--enable-layerwise-offload`, or
+`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
+but does not offload model weights to the CPU.
+
+The run passes the capacity check when the server initializes, the request
+finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
+card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
+32 kHz stereo AAC audio. Report the measured headroom rather than assuming
+that every nominal 96 GB SKU exposes the same MiB total.
+
+As a capacity proxy only, the same five-second case on one B300 measured a
+92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
+encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
+6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
+they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
+reported device capacity differ across GPUs.
+
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
resolves the `transformer` component before constructing the DiT, so names do
not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -386,10 +498,11 @@ For example, keep the first main block's attention projections in BF16 with:
'{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
```
-The structured option replaces `--quantization fp8`. Online FP8 is currently
-incompatible with H3 layerwise offload because the offload path produces a
-weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
-parallelism and VAE tiling instead.
+The structured option replaces `--quantization fp8`. Online FP8 can be used
+with H3 layerwise offload and with both DLO transfer paths. The default
+AllGather path uses the ordinary loader to finalize FP8 weights and scales
+before sharding them across ranks. `--dlo-no-use-allgather` instead retains
+complete rank-local tensors and avoids the synchronized request-wave contract.
## AMD ROCm (gfx942 / gfx950)
@@ -475,7 +588,7 @@ vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
`0.26.0+rocm723` wheel (HIP 7.2).
| Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
+| ---------- | --------------- | ----------------- |
| T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
| FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
| T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
@@ -657,7 +770,7 @@ at most 15 seconds combined.
## Official input matrix and limits
| Task | Supported references | Limits |
-|------|----------------------|--------|
+| ------ | ---------------------- | -------- |
| T2VA | text only | prompt must be non-empty |
| FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
| Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
@@ -688,13 +801,17 @@ Omitting `quality` preserves the startup default: it uses the reference path
normally, or the server-configured profile when the server was started with
`--cache-backend cache_dit`.
+Turbo is independent of this quality switch: `quality` selects a Cache-DiT
+policy, while Turbo changes the active LoRA weights and sampling schedule. See
+[LoRA](#lora) below.
+
The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
`lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
balanced switch order.
| `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
-|---|---:|---:|---:|---:|---|
+| --- | ---: | ---: | ---: | ---: | --- |
| `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
| `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
@@ -704,10 +821,173 @@ balanced switch order.
> workload. The values above apply to this deployment and are not universal
> guarantees. `lossless` remains the exact reference path.
+## LoRA
+
+### Turbo LoRA
+
+Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
+
+```text
+lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+```
+
+Download only that file:
+
+```bash
+export TURBO_DIR=/path/to/minimax-h3-turbo
+export TURBO_FILE=minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
+hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
+export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
+```
+
+Start from a non-offloaded or DLO FL2VA server command and add
+`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
+`--lora-path` preloads the adapter; each request still activates it and uses
+the published sampling settings:
+
+```bash
+-F 'num_inference_steps=5' \
+-F 'flow_shift=6' \
+-F 'extra_params={"task":"t2va","duration":4.4,"audio_flow_shift":3.0}' \
+-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
+```
+
+For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
+ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
+dynamic-only and does not support prefusion or LoRA composition. DLO is
+supported by keeping the request-switchable LoRA A/B buffers resident on the
+accelerator while DLO streams only the base blocks; budget for this additional
+fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
+The five requested sigma points produce the four denoiser evaluations expected
+by the Turbo artifact.
+
+### FlashGen native LoRA
+
+The FlashGen 4-step T2VA artifact uses the native MiniMax-H3 module layout and
+declares its distilled sigma schedule in safetensors metadata. It is published on
+[ModelScope](https://modelscope.cn/models/FlashGen/Minimax-H3-4step-lora-flashgen):
+
+```text
+FlashGen/Minimax-H3-4step-lora-flashgen/minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
+```
+
+Download only that file:
+
+```bash
+python -m pip install modelscope
+export FLASHGEN_DIR=/path/to/minimax-h3-flashgen-lora
+export FLASHGEN_FILE=minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors
+modelscope download FlashGen/Minimax-H3-4step-lora-flashgen \
+ --local_dir "${FLASHGEN_DIR}" \
+ --include "${FLASHGEN_FILE}"
+export FLASHGEN_LORA="${FLASHGEN_DIR}/${FLASHGEN_FILE}"
+```
+
+Start from a non-offloaded or DLO FL2VA server command and add
+`--task-type fl2va --lora-backend peft --lora-path "${FLASHGEN_LORA}"`.
+Each request must use T2VA and the distilled interval-count contract:
+
+```bash
+-F 'num_inference_steps=4' \
+-F 'extra_params={"task":"t2va","duration":5.2}' \
+-F "lora={\"name\":\"h3-flashgen-v1.0\",\"path\":\"${FLASHGEN_LORA}\",\"scale\":1.0}"
+```
+
+This path rejects Ref2VA and checkpoints that already pin `base_schedule` in
+`model_index.json`. The adapter metadata carries
+`base_schedule=1.0,0.7,0.4,0.15,0.0`, so `num_inference_steps=4` means four
+denoiser evaluations, not five sigma points. Request-mode generation may omit
+the field and take the count from the adapter schedule; `--step-execution`
+requires it explicitly, because the step scheduler reads the total step count
+off the request at admission, before the adapter schedule is known.
+
+DLO is supported in request-mode generation on the same terms as the Turbo
+adapter: the request-switchable LoRA A/B buffers stay resident on the
+accelerator while DLO streams only the base blocks, so budget for that
+additional fixed HBM usage. The native artifact is rank 64 over 259 target
+modules, and its packed `qkv_proj` and `fc1` layers reuse the full-input A
+tensor per slice while B carries slice-local output rows, so the resident
+footprint exceeds the on-disk payload; measure it for your parallel layout
+rather than assuming the checkpoint size. Pure Ulysses replicates the adapter
+on every rank, while DiT tensor parallelism shards the B buffers. Model-level
+and standard layerwise offload remain unsupported, and `--step-execution`
+cannot be combined with `--enable-distributed-layerwise-offload`.
+
+To validate a deployment, post the same fixed-seed T2VA request twice with the
+adapter and twice without it, then compare the four output digests. The adapter
+is bound and deterministic when each pair matches internally and the two pairs
+differ from each other.
+
+### FastH3 adapter
+
+[FastH3](https://haoailab.com/blogs/fasth3-preview/) is FastVideo's four-step
+DMD2 student of H3-Base. It reuses H3's text encoder, VAEs, tokenizers, and
+schedulers unchanged, replacing 49 denoiser evaluations with four. It is fused
+into the checkpoint at load time rather than switched per request, because it
+carries full-rank deltas that no LoRA layer can express.
+
+The bundle publishes four variants, so download one and point `--lora-path` at
+it; the repository root is refused rather than guessed at:
+
+```bash
+export FASTH3_DIR=/path/to/fasth3
+hf download FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA \
+ dense-datafree/adapter_model.safetensors --local-dir "${FASTH3_DIR}"
+export FASTH3_LORA="${FASTH3_DIR}/dense-datafree/adapter_model.safetensors"
+```
+
+Add `--task-type fl2va --lora-path "${FASTH3_LORA}"` to a non-offloaded server
+command. T2VA is served by the FL2VA partition, so `--task-type fl2va` is
+correct even though FastH3 preview v1 distills T2VA only. Because the adapter is
+fused, `--lora-backend` does not apply and a request carrying a `lora=` field is
+rejected rather than served without the adapter it asked for.
+
+```bash
+-F 'num_inference_steps=4' \
+-F 'extra_params={"task":"t2va","duration":4.4}'
+```
+
+Requests must ask for `num_inference_steps=4` and `task=t2va`: the release's five
+sigma points bound four denoiser evaluations, and that count is what the step
+scheduler admits a request on. The server denoises on the release's own ladder,
+keeping H3's per-modality shifts at the checkpoint values, so a request that
+overrides `flow_shift` or `audio_flow_shift` is rejected - it would sample the
+student at noise levels it was never distilled at.
+
+Only a release that identifies itself as FastH3 is fused; any other
+`fastvideo-lora-v2` adapter stays on the dynamic LoRA route. A claimed artifact
+is then held to its own metadata: one that misdeclares its tensor counts or
+leaves a transformer block unedited is refused at startup instead of serving
+mostly base H3 weights on a four-step schedule. Offload is refused for the same
+reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
+`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
+
+> [!NOTE]
+> Only the **Dense / Data-Free** variant is supported today. The three VSA
+> variants carry `.set_weight` compression gates for modules that do not exist
+> in vLLM-Omni's dense H3 attention, and the release states dense attention is
+> not a supported substitute for them, so they are refused at startup rather
+> than silently run wrong. Video Sparse Attention for H3's packed
+> `[text | cond | audio | video]` sequence is follow-up work.
+
+Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
+1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
+
+| Adapter | Steps | End-to-end | Diffusion engine |
+| --- | ---: | ---: | ---: |
+| none (base H3) | 50 | 25.8 / 26.4 s | 16.22 / 16.28 s |
+| FastH3 Dense | 4 (5 sigma points) | 11.7 / 11.8 s | 2.37 / 2.36 s |
+
+The denoising speedup is 6.9x. End-to-end is 2.2x because text encoding, VAE
+decoding and muxing are a fixed cost that dominates a clip this short; longer
+generations move the end-to-end figure toward the denoising one. Fusing the
+adapter does not measurably change startup: weight loading took 77.3 s with it
+against 85.8 s without.
+
## Key parameters
| Parameter | Recommended value | Notes |
-|-----------|-------------------|-------|
+| ----------- | ------------------- | ------- |
| `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
| `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
@@ -732,9 +1012,9 @@ Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 servi
The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
excluded warmup followed by three requests.
-| Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
-| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
+| Workload | Configuration | Observed result |
+| -------------------------------------- | ------------------------------------------- | ------------------------------------ |
+| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
These measurements describe the validated shapes rather than a general
@@ -750,6 +1030,13 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
reduction.
+For direct human inspection, see the external
+[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
+which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
+[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
+also record per-task fidelity metrics and provenance without storing generated
+media in this repository.
+
## TeaCache acceleration
TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -819,7 +1106,12 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
- TeaCache is calibrated for FL2VA only; Ref2VA requests run uncached.
- Combined serving requires sibling `FL2VA` and `Ref2VA` directories, loads
both task-specific DiTs, and loads shared components once from `FL2VA`.
-- H3 currently executes one generation request per diffusion batch.
+- Request mode executes one generation request per diffusion batch. Use
+ `--step-execution` with `--max-num-seqs N` to admit several requests at once
+ (see
+ [Execution modes](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/execution_modes.md));
+ step
+ mode does not support `cache_backend`.
- The first regional-compile request is a warmup and should not be included in
steady-state performance measurements.
- The serving path accepts fewer references than the model supports. H3 documents up
@@ -834,7 +1126,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
H3 native `tile` mode only.
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
-- Online FP8 is not compatible with layerwise offload.
+- Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
+ in host memory on every rank during startup before retaining only each rank's
+ shard. Size startup host memory for that transient peak.
- TeaCache and Cache-DiT cannot be enabled on the same server.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
diff --git a/recipes/MiniMaxAI/MiniMax-Music3.md b/recipes/MiniMaxAI/MiniMax-Music3.md
new file mode 100644
index 000000000..3d078f645
--- /dev/null
+++ b/recipes/MiniMaxAI/MiniMax-Music3.md
@@ -0,0 +1,173 @@
+# MiniMax Music 3
+
+> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out
+
+## Summary
+
+- Vendor: MiniMax
+- Model: `MiniMaxAI/MiniMax-Music3`
+- Task: Text-to-music generation (vocals and instrumental together)
+- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: @linyueqian
+
+## When to use this recipe
+
+Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six
+minutes. This is not a TTS model with a musical voice: there is no reference
+speaker, no `voice`, and no `temperature`. What you control is the lyrics, the
+caption, the seed, and the length.
+
+## References
+
+- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3)
+- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/)
+- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU),
+ `vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two)
+
+## Architecture
+
+| Component | Spec |
+|---|---|
+| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k |
+| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 |
+| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder |
+| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 |
+| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling |
+| Solver | Euler, 30 steps, acoustic guidance scale 1.7 |
+| Window | 200 frames per acoustic chunk, 100-frame hop |
+| Frame rate | 25 frames per second |
+| Context length | 10,240 tokens |
+| Output | 32 kHz stereo WAV |
+
+Generation is two-stage. The backbone predicts one audio frame per decode step,
+each frame eight RVQ codebooks deep. Every 200 frames are handed to the
+flow-matching transformer, which solves a latent that the decoder turns into
+waveform.
+
+**Guidance doubles the rows.** Every request decodes twice: one row sees the
+real prompt, the other sees the same prompt with the caption-and-lyrics span
+replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole
+song, so a request costs twice the KV its length suggests. Size the pool for
+rows, not requests.
+
+## Hardware Support
+
+## GPU
+
+### 1x H200 141GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- CUDA: 12.9
+- vLLM version: 0.27.0
+- vLLM-Omni version or commit: branch `feat/minimax-music3`
+
+#### Command
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+ --host 0.0.0.0 --port 8000 \
+ --trust-remote-code --omni
+```
+
+The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered
+automatically from the checkpoint's `model_type=minimax_music3`. It colocates
+both stages on device 0.
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8000/v1/audio/speech \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "MiniMaxAI/MiniMax-Music3",
+ "input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me",
+ "instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.",
+ "seed": 1,
+ "max_new_tokens": 750
+ }' \
+ --output song.wav
+```
+
+Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames
+at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a
+target: the model ends the song itself when it emits the audio-end token, and a
+response shorter than the cap is the model finishing rather than a truncation.
+
+#### Notes
+
+- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone
+ and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights,
+ which sit outside any KV budget.
+- **Key flags:** none required. Guidance, the solver step count and the window
+ geometry are fixed by the checkpoint and are not request parameters.
+- **Known limitations:**
+ - The external API is non-streaming; `stream: true` is rejected.
+ - `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a
+ maximum-length song (9,000 frames) cannot both fit. Long captions shorten
+ the maximum song. This is a property of the checkpoint.
+ - The acoustic stage runs in float32. bfloat16 is accepted but measurably
+ degrades the solver.
+
+### 2x H200 141GB
+
+Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device:
+
+```bash
+vllm serve MiniMaxAI/MiniMax-Music3 \
+ --host 0.0.0.0 --port 8000 \
+ --trust-remote-code --omni \
+ --deploy-config minimax_music3_2gpu.yaml
+```
+
+Only the placement differs; both layouts run the acoustic stage in float32.
+
+## Writing the request
+
+Two fields carry the request, and both are required.
+
+- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`,
+ `[Outro]`) steer the arrangement and are part of the prompt contract.
+ **Put a tag on its own line.** Normalization keeps only the tags on a line
+ that starts with one and drops the rest of that line, so lyrics written next
+ to a tag are silently lost:
+
+ ```text
+ "[Verse]\nWalking down the street" -> [start] [verse] Walking down the street
+ "[Verse] Walking down the street" -> [start] [verse]
+ ```
+
+- `instructions` is the **caption** describing genre, instrumentation, tempo,
+ mood and production. It is the strongest control available: vague captions
+ give generic arrangements. Long structured captions work well, since the
+ model was trained on descriptions covering global attributes, emotional
+ progression, vocal detail and arrangement.
+
+A request is deterministic in its seed: the same lyrics, caption, seed and
+length return the same audio. Changing only the seed gives another take of the
+same song.
+
+## Request parameters
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `input` | string | required | Lyrics, non-empty. Tags on their own lines |
+| `instructions` | string | required | Caption, non-empty |
+| `seed` | int | `0` | Non-negative. Fixes the output for a given request |
+| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 |
+| `response_format` | string | `"wav"` | Output container |
+
+## Parameters the model rejects
+
+The request is refused rather than silently ignored, so a mistake is visible
+immediately.
+
+| Parameter | Reason |
+|---|---|
+| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw |
+| `voice` | There is no speaker to select; the vocal comes from the caption |
+| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag |
+| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" |
+| `stream: true` | The external API is non-streaming |
diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md
index 0a5270222..826309835 100644
--- a/recipes/NVIDIA/GR00T-N1.7.md
+++ b/recipes/NVIDIA/GR00T-N1.7.md
@@ -44,7 +44,7 @@ vllm serve nvidia/GR00T-N1.7-3B \
--host 127.0.0.1 \
--port 8000 \
--served-model-name gr00t-n1d7 \
- --stage-configs-path vllm_omni/deploy/Gr00tN1d7.yaml
+ --deploy-config vllm_omni/deploy/Gr00tN1d7.yaml
```
Notes:
diff --git a/recipes/NVIDIA/Nemotron-Labs-Audex.md b/recipes/NVIDIA/Nemotron-Labs-Audex.md
index 21ee7e655..13be59f4c 100644
--- a/recipes/NVIDIA/Nemotron-Labs-Audex.md
+++ b/recipes/NVIDIA/Nemotron-Labs-Audex.md
@@ -51,10 +51,10 @@ NemotronH (~3B active parameters). Pick the pipeline by task:
`model_type` auto-resolves `audex_tts.yaml`):
```bash
-vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
+vllm serve nvidia/Nemotron-Labs-Audex-2B \
--host 0.0.0.0 --port 8097 \
--trust-remote-code --omni
-# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
+# other modes: add --deploy-config vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
# or use the launcher: MODE=s2s PORT=8098 examples/online_serving/audex/run_server.sh
```
@@ -109,10 +109,10 @@ The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
auto-resolves to the 2B-tuned config otherwise):
```bash
-vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+vllm serve nvidia/Nemotron-Labs-Audex-30B-A3B \
--host 0.0.0.0 --port 8097 \
--trust-remote-code --omni \
- --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
+ --deploy-config vllm_omni/deploy/audex_tts_30b.yaml
# or: SIZE=30b MODE=tts examples/online_serving/audex/run_server.sh
```
diff --git a/recipes/NVIDIA/NemotronLabs-VoiceChat.md b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
index 8faa268d2..dbe7d5431 100644
--- a/recipes/NVIDIA/NemotronLabs-VoiceChat.md
+++ b/recipes/NVIDIA/NemotronLabs-VoiceChat.md
@@ -1,7 +1,7 @@
# NemotronLabs VoiceChat 11B
-> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
-> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
+> Speech-to-speech on a 3-stage vLLM-Omni pipeline: offline single-turn
+> inference and experimental model-native, frame-locked Realtime serving.
## Summary
@@ -11,8 +11,8 @@
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
-- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
- function calling, and batch>1 are documented follow-ups.
+- Mode: Offline single-turn inference (batch=1), plus experimental native
+ duplex Realtime serving at one 80 ms microphone frame per scheduler wake.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
@@ -33,6 +33,11 @@ NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
+- Native duplex config and E2E probe:
+ [`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml),
+ [`tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py`](../../tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py)
+- Nightly model-level gate:
+ [`tests/e2e/online_serving/test_nemotron_voicechat_duplex.py`](../../tests/e2e/online_serving/test_nemotron_voicechat_duplex.py)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
@@ -60,7 +65,7 @@ output stayed within one word of the reference in testing.
- OS: Linux
- Python: 3.12
-- vLLM version: 0.26.0
+- vLLM version: 0.27.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
@@ -88,6 +93,55 @@ The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
+#### Native duplex serving
+
+The checkpoint does not contain the underlying Nemotron text tokenizer. Set
+`NEMOTRON_VOICECHAT_LLM_PATH` to a local snapshot of
+`nvidia/NVIDIA-Nemotron-Nano-9B-v2` before starting an offline deployment.
+
+```bash
+export HF_HUB_OFFLINE=1
+export TRANSFORMERS_OFFLINE=1
+export NEMOTRON_VOICECHAT_LLM_PATH=/path/to/NVIDIA-Nemotron-Nano-9B-v2
+
+vllm-omni serve /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
+ --omni \
+ --served-model-name nemotron-voicechat \
+ --deploy-config vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
+```
+
+The single-GPU profile runs three eager fp32 engine processes on one device.
+
+In another shell, stream only the user channel of one of NVIDIA's bundled
+stereo conversation fixtures:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+ --model nemotron-voicechat \
+ --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/turn_taking.wav \
+ --input-channel 0 \
+ --output-dir results/nemotron_voicechat_duplex \
+ --timeout-s 600
+```
+
+The probe requires a completed response, non-silent 22.05 kHz output, and the
+advertised native 80 ms append capabilities. Add `--expect-function-call` and
+use `tool_call.wav` to validate the function-call channel. To validate the
+Realtime tool round trip, also provide the expected arguments and return a
+tool result:
+
+```bash
+python tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py \
+ --model nemotron-voicechat \
+ --input-wav /path/to/NVIDIA-NemotronLabs-VoiceChat-11B/tool_call.wav \
+ --output-dir results/nemotron_voicechat_tool_call \
+ --expect-function-call \
+ --expected-function-name generate_random_number \
+ --expected-function-arguments '{"min":1,"max":50}' \
+ --function-output 20 \
+ --expected-post-tool-text "random number"
+```
+
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
@@ -109,5 +163,14 @@ with any ASR model.
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
-- Known limitations: batch=1 only; offline single turn (no online/duplex
- serving yet); the function-call channel is decoded but not yet acted on.
+- The 80 ms frame period is a model/protocol contract, not a throughput
+ guarantee. In an eager fp32 H200 profile, Stage 2 emitted 80 ms packets at
+ roughly 0.87--0.88 s intervals (RTF about 10.8--10.9); placing the three
+ stages on separate GPUs did not remove the serial 8-step EAR-TTS bottleneck.
+ The current native duplex path is functionally streaming but not wall-clock
+ realtime.
+- Known limitations: batch=1 only. The native duplex deployment allows one
+ active session and does not support barge-in. Tool execution remains
+ client-owned: the server emits function-call events, accepts a validated
+ `function_call_output`, and resumes the live model with the returned result;
+ it does not execute arbitrary tools itself.
diff --git a/recipes/NVIDIA/SANA-Video-2B.md b/recipes/NVIDIA/SANA-Video-2B.md
new file mode 100644
index 000000000..c85d8ca40
--- /dev/null
+++ b/recipes/NVIDIA/SANA-Video-2B.md
@@ -0,0 +1,245 @@
+# SANA-Video 2B
+
+> Native text-to-video and image-to-video generation at 480p and 720p,
+> plus a validated Diffusers-adapter compatibility path
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `Efficient-Large-Model/SANA-Video_2B_480p_diffusers`,
+ `Efficient-Large-Model/SANA-Video_2B_720p_diffusers`
+- Task: Text-to-video and image-to-video
+- Mode: Offline inference and OpenAI-compatible online serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use the native `SanaVideoPipeline` for T2V and
+`SanaImageToVideoPipeline` for I2V. Both native pipelines support the 480p
+and 720p checkpoints. Use `--diffusion-load-format diffusers` when you need
+the black-box Diffusers compatibility baseline; adapter T2V and I2V are
+validated at both resolutions.
+
+The native pipeline loads the 480p checkpoint through
+`DistributedAutoencoderKLWan` and the 720p checkpoint through
+`DistributedAutoencoderKLLTX2Video`. These are vLLM-Omni distributed wrappers
+around the corresponding Diffusers autoencoders, not independent VAE
+implementations. The denoising loop also intentionally loads Diffusers'
+`DPMSolverMultistepScheduler` from the checkpoint to preserve its scheduler
+configuration and numerical behavior.
+
+## References
+
+- Upstream project:
+- Model cards:
+ [480p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p_diffusers),
+ [720p](https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_720p_diffusers)
+- Diffusers documentation:
+- Online serving guides:
+ [Text-to-Video](../../docs/user_guide/examples/online_serving/text_to_video.md),
+ [Image-to-Video](../../docs/user_guide/examples/online_serving/image_to_video.md)
+- Offline T2V example:
+ [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
+- Offline I2V example:
+ [`examples/offline_inference/image_to_video/image_to_video.py`](../../examples/offline_inference/image_to_video/image_to_video.py)
+- Support discussion: [vLLM-Omni issue #5432](https://github.com/vllm-project/vllm-omni/issues/5432)
+
+## Hardware Support
+
+## GPU
+
+### 1x RTX 5090 32GB
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS
+- Python: 3.12.3
+- Driver / runtime: NVIDIA driver 580.95.05; PyTorch 2.11.0+cu130
+- Diffusers: 0.38.0
+- vLLM version: 0.26.0
+- vLLM-Omni version or commit: PR #5508, commit `22037901`
+
+#### Command
+
+##### Native text-to-video inference
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+ --model Efficient-Large-Model/SANA-Video_2B_720p_diffusers \
+ --model-class-name SanaVideoPipeline \
+ --prompt "A cat walking on the grass, facing the camera." \
+ --negative-prompt "blurry, low quality, temporal artifacts" \
+ --height 704 --width 1280 --num-frames 81 \
+ --num-inference-steps 50 --guidance-scale 6 \
+ --extra-body '{"motion_score": 30}' \
+ --fps 16 --seed 42 --output sana_video_720p.mp4
+```
+
+For 480p, select `SANA-Video_2B_480p_diffusers` and use
+`--height 480 --width 832`.
+
+##### Native image-to-video inference
+
+SANA checkpoints declare `SanaVideoPipeline` in `model_index.json`, so I2V
+must be selected explicitly with `--model-class-name
+SanaImageToVideoPipeline`.
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+ --model Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ --model-class-name SanaImageToVideoPipeline \
+ --image input.jpg \
+ --prompt "A cat turns toward the camera with smooth, natural motion." \
+ --negative-prompt "blurry, low quality, temporal artifacts" \
+ --height 480 --width 832 --num-frames 81 \
+ --num-inference-steps 50 --guidance-scale 6 \
+ --extra-body '{"motion_score": 30}' \
+ --fps 16 --seed 42 --output sana_video_i2v_480p.mp4
+```
+
+The same pipeline class supports the 720p checkpoint through vLLM-Omni's
+distributed LTX-2 VAE wrapper; use `--height 704 --width 1280`.
+
+For online I2V serving:
+
+```bash
+MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ bash examples/online_serving/image_to_video/run_server_sana_video.sh
+
+INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v.mp4 \
+ bash examples/online_serving/image_to_video/run_curl_sana_video.sh
+```
+
+##### Native online serving
+
+```bash
+MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ bash examples/online_serving/text_to_video/run_server_sana_video.sh
+
+bash examples/online_serving/text_to_video/run_curl_sana_video.sh
+```
+
+To run the black-box compatibility backend for T2V, replace the server script
+with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request
+works; `num_frames` is adapted to Diffusers' `frames` argument. The script
+selects `TORCH_SDPA` because SANA-Video uses an attention mask that the
+AITER-backed Diffusers attention path does not accept.
+
+##### Diffusers-adapter image-to-video serving
+
+The validated I2V adapter commands are:
+
+```bash
+# 480p
+MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+ bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh
+
+INPUT_IMAGE=input.jpg OUTPUT_PATH=sana_video_i2v_adapter.mp4 \
+ bash examples/online_serving/image_to_video/run_curl_sana_video.sh
+
+# 720p
+MODEL=Efficient-Large-Model/SANA-Video_2B_720p_diffusers \
+ bash examples/online_serving/image_to_video/run_server_sana_video_diffusers.sh
+
+INPUT_IMAGE=input.jpg WIDTH=1280 HEIGHT=704 \
+ OUTPUT_PATH=sana_video_i2v_adapter_720p.mp4 \
+ bash examples/online_serving/image_to_video/run_curl_sana_video.sh
+```
+
+#### Verification
+
+Check the encoded 720p output metadata after running a generation command:
+
+```bash
+ffprobe -v error -select_streams v:0 \
+ -show_entries stream=width,height,r_frame_rate,nb_frames \
+ -of default=noprint_wrappers=1 sana_video_720p.mp4
+```
+
+The standard 720p request above should report:
+
+```text
+width=1280
+height=704
+r_frame_rate=16/1
+nb_frames=81
+```
+
+For 480p, expect `width=832` and `height=480`.
+
+The automated serving matrix covers both checkpoint variants:
+
+| Backend | 480p T2V | 720p T2V | 480p I2V | 720p I2V |
+| --- | --- | --- | --- | --- |
+| Native vLLM-Omni | Validated | Validated | Validated | Validated |
+| Diffusers adapter | Validated | Validated | Validated | Validated |
+
+Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the
+primary SANA execution paths. The Diffusers adapter is retained as a
+validated compatibility/reference backend.
+
+#### Notes
+
+- Key flags: select I2V explicitly with `--model-class-name
+ SanaImageToVideoPipeline`; SANA checkpoint `model_index.json` files declare
+ the T2V class. Pass `motion_score` through `--extra-body`. The supplied
+ Diffusers-adapter server scripts select `TORCH_SDPA` because the AITER-backed
+ path does not accept SANA's attention mask.
+- Output profile: 81 frames at 16 FPS is the standard checkpoint profile and
+ produces approximately five seconds of video. Minute-scale generation
+ requires the separate LongSANA/LongLive block-autoregressive workflow.
+- Backend boundary: the native pipelines and Transformer are owned by
+ vLLM-Omni. The 480p Wan VAE and 720p LTX-2 VAE run through vLLM-Omni
+ distributed wrappers derived from the corresponding Diffusers VAE classes.
+ The denoising loop intentionally retains the checkpoint-compatible
+ Diffusers `DPMSolverMultistepScheduler`.
+- Known limitations:
+ - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution
+ are not validated for the native pipeline.
+ - The Diffusers backend is a compatibility path and does not provide native
+ vLLM-Omni parallelism or continuous batching.
+ - Native describes pipeline and Transformer ownership, not a zero-Diffusers
+ dependency guarantee.
+
+### 1x NVIDIA L20X
+
+#### Environment
+
+- Driver: 570.133.20
+- PyTorch: 2.13.0+cu130
+- Diffusers: 0.38.0
+- vLLM: 0.27.0
+- Nsight Systems: 2026.1
+- vLLM-Omni commit: `54fcc6d`
+
+#### Same-workload T2V comparison
+
+Both backends used
+`Efficient-Large-Model/SANA-Video_2B_480p_diffusers`, the same T2V prompt
+and seeds, 832×480 output, 81 frames, and four inference steps. One warmup
+request was excluded, followed by three measured requests.
+
+| Metric | Native vLLM-Omni | Diffusers adapter | Native delta |
+| --- | ---: | ---: | ---: |
+| E2E latency | 7.1600 ± 0.0043 s | 6.9761 ± 0.0517 s | +2.636% |
+| Throughput | 0.1397 req/s | 0.1433 req/s | -2.568% |
+| Peak reserved VRAM | 23,936 MiB | 19,060 MiB | +25.582% |
+| Initialization | 11.669 s | 20.236 s | -42.336% |
+
+The peak-memory result is not a Transformer-only comparison. For the 480p
+checkpoint, the native pipeline deliberately loads and decodes with the Wan
+VAE in FP32. This matches the upstream SANA-Video reference and the recorded
+accuracy-golden setup for this integration. The adapter launch uses
+`--dtype bfloat16` as the Diffusers pipeline-wide dtype, so its Wan VAE also
+runs in BF16.
+
+Peak VRAM is the CUDA allocator's maximum reserved memory over the complete
+request. It includes VAE decode activations and allocator-reserved blocks, not
+only live Transformer tensors. The different VAE precision policies are a
+material code-level difference and a likely contributor to the higher native
+high-water mark; the available full-request measurements do not isolate how
+much of the delta they account for. The 25.582% result should therefore not be
+interpreted as an isolated native Transformer memory regression. Exact
+attribution requires both paths to use the same VAE dtype. Changing the native
+480p VAE to BF16 would instead change the recorded golden/reference setup and
+requires a new accuracy study.
diff --git a/recipes/NVIDIA/SANA-WM.md b/recipes/NVIDIA/SANA-WM.md
new file mode 100644
index 000000000..2926cc46f
--- /dev/null
+++ b/recipes/NVIDIA/SANA-WM.md
@@ -0,0 +1,193 @@
+# SANA-WM
+
+> Camera-controllable first-frame image-to-video world model.
+
+## Summary
+
+- Vendor: Efficient-Large-Model / NVlabs SANA
+- Model: `BBBBruce/SANA-WM_bidirectional-stage1-diffusers` (standard diffusers layout, converted offline from the NVlabs release; Stage-1 transformer + VAE only)
+- Task: First-frame image-to-video generation with camera control
+- Mode: Online serving with the OpenAI-compatible video API
+- Model weights: about 13 GB for the Stage-1 transformer (10 GB) and VAE (2.3 GB)
+- Local disk: reserve about 40 GB for the Hugging Face cache and runtime artifacts
+- Recommended GPU: 24 GB or larger CUDA GPU
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to serve SANA-WM through `/v1/videos` or
+`/v1/videos/sync`. The model takes a text prompt, a first-frame image, and
+either an action DSL string or explicit camera poses. vLLM-Omni serves the
+SANA-WM Stage-1 DiT, decoded through the SANA VAE. The optional LTX-2 refiner
+stage is not supported by this integration; it is a planned follow-up.
+
+## References
+
+- Upstream model card:
+- Video API: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA RTX PRO 6000 Blackwell 96GB
+
+#### Capacity
+
+- Model storage: the Stage-1 transformer is about 10 GB and the VAE about
+ 2.3 GB. The Gemma text encoder is a separate 4.9 GB repo. The engine prefetches
+ the whole model repo at startup (`allow_patterns=["*"]`), which is why the
+ Stage-1 weights live in their own repo — the two-stage one carries an
+ additional 84 GB `refiner/` that this path never loads.
+- Text encoder: the pipeline tries `google/gemma-2-2b-it` first, then falls back
+ to the ungated mirror `Efficient-Large-Model/gemma-2-2b-it`. The first repo is
+ gated, so without an accepted licence and `hf auth login` you will see one
+ failed load before the fallback succeeds — that is expected, not an error to
+ chase. Set `VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER` to pin a specific repo or a
+ local path and skip both.
+- Disk sizing: provision about 40 GB of local disk or Hugging Face cache volume
+ so the model, temporary downloads, and generated artifacts fit without cache
+ eviction.
+- GPU sizing: the default 1280x704, 161-frame, 60-step serving profile peaks at
+ 22.6 GB of device memory and takes about 133 s to generate on one RTX PRO
+ 6000 Blackwell. The peak lands in the VAE decode, which is why the pipeline
+ forces VAE tiling on regardless of `vae_use_tiling` — without it the same
+ request costs about 9 GB more, and a 321-frame one OOMs outright. On smaller
+ GPUs, lower `width`, `height`, or `num_frames` before serving production
+ requests.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
+ build
+- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
+ Python environment
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+The repo ships the standard Diffusers layout (`model_index.json` +
+`transformer/`, `vae/`), and its `model_index.json` names `SanaWmPipeline`, so
+the pipeline class resolves on its own.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve BBBBruce/SANA-WM_bidirectional-stage1-diffusers \
+ --omni \
+ --host 0.0.0.0 \
+ --port 8091
+```
+
+No deploy config: single-stage diffusion models are deliberately absent from
+`OMNI_PIPELINES` (`vllm_omni/config/pipeline_registry.py`), so stage resolution
+falls back to the default stage config and a YAML's stage settings — including
+`default_sampling_params` — would not be applied. The production generation
+settings therefore live in the model (`num_inference_steps=60`,
+`guidance_scale=5.0`); a request that omits a field gets them. The examples
+below still pass every field explicitly so the numbers are visible.
+
+If you point this at the older two-stage repo
+(`BBBBruce/SANA-WM_bidirectional-diffusers`), startup fails with `Model class
+SanaWmTwoStagesPipeline not found in diffusion model registry`, because that
+repo's `model_index.json` names a class this build does not register. Add
+`--model-class-name SanaWmPipeline` to override it.
+
+#### Verification
+
+Use a short smoke request first:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+ -H "Accept: video/mp4" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=9" \
+ -F "fps=16" \
+ -F "num_inference_steps=2" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-8","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+ -o sana_wm_smoke.mp4
+```
+
+For a production-length request, note that the action durations must sum to
+`num_frames - 1` — the rollout includes the identity start pose — and a
+mismatch is rejected rather than padded or truncated:
+
+```bash
+curl -sS -X POST http://localhost:8091/v1/videos/sync \
+ -H "Accept: video/mp4" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=161" \
+ -F "fps=16" \
+ -F "num_inference_steps=60" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}' \
+ -o sana_wm_output.mp4
+```
+
+Use `POST /v1/videos` instead when you want job storage and polling rather than
+inline MP4 bytes. It accepts the same form fields as `/v1/videos/sync`.
+
+```bash
+create_response=$(curl -sS -X POST http://localhost:8091/v1/videos \
+ -H "Accept: application/json" \
+ -F "prompt=A slow forward camera move through a quiet city street." \
+ -F "negative_prompt=blurry, low quality, distorted, watermark" \
+ -F "input_reference=@/path/to/first_frame.png;type=image/png" \
+ -F "width=1280" \
+ -F "height=704" \
+ -F "num_frames=161" \
+ -F "fps=16" \
+ -F "num_inference_steps=60" \
+ -F "guidance_scale=5.0" \
+ -F "seed=42" \
+ --form-string 'extra_params={"sana_wm":{"action":"w-160","translation_speed":0.055,"rotation_speed_deg":1.2,"intrinsics":{"fx":640,"fy":640,"cx":640,"cy":352}}}')
+
+video_id=$(echo "$create_response" | jq -r '.id')
+curl -sS "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o sana_wm_output.mp4
+```
+
+#### Notes
+
+- Sequence parallelism is not supported. The bidirectional gated delta
+ recurrence carries state across frames, so a rank cannot denoise a slice
+ of the token sequence in isolation; supporting it needs a distributed scan
+ or an all-gather before the GDN blocks.
+
+- `input_reference` is required for the first frame. Use `image_reference` only
+ when you need a JSON-safe image URL or data URL instead of a multipart upload.
+- `sana_wm` must provide exactly one of `action` or `camera`.
+- Action strings use comma-separated `-` segments. Supported
+ keys are `w`, `a`, `s`, `d` for translation and `i`, `j`, `k`, `l` for
+ pitch/yaw rotation. The durations must sum to `num_frames - 1`.
+- Explicit camera control (alternative to `action`): pass
+ `"camera": {"poses": [...]}` where `poses` is a list of `num_frames`
+ camera-to-world 4x4 matrices (row-major, OpenCV `+X right, +Y down, +Z forward`
+ convention), e.g.
+ `extra_params={"sana_wm":{"camera":{"poses":[[[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], ...]},"intrinsics":{...}}}`.
+ Most callers should prefer `action`; explicit poses exist for callers that
+ already have a per-frame trajectory.
+- Explicit `intrinsics` are recommended and take the mapping form
+ `{"fx":640,"fy":640,"cx":640,"cy":352}` (for 1280x704). This `{fx,fy,cx,cy}`
+ mapping is the only accepted intrinsics form; omit `intrinsics` to derive them
+ from the output resolution. All four values must be finite, and `fx`/`fy`
+ must be positive — the ray map divides by them.
+- The video API returns decoded MP4 bytes and has no `output_type` field, so
+ raw Stage-1 latents are reachable only from the offline API
+ (`OmniDiffusionSamplingParams(output_type="latent")`); see
+ [`tests/e2e/offline_inference/test_sana_wm.py`](../../tests/e2e/offline_inference/test_sana_wm.py).
+ Putting `output_type` in `extra_params` does not work: it lands in
+ `sampling_params.extra_args`, while the pipeline reads the top-level field.
diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md
index 99b064310..86f0883f7 100644
--- a/recipes/OpenBMB/MiniCPM-o-4_5.md
+++ b/recipes/OpenBMB/MiniCPM-o-4_5.md
@@ -71,6 +71,35 @@ single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
two-GPU profile. The removed fused two-stage implementation is not retained as
a fallback because it would duplicate state machines and correctness paths.
+### Graph execution
+
+Graph boundaries follow each stage's state model:
+
+| Stage | Graph mode on Ascend | Eager boundary |
+| --- | --- | --- |
+| 0 Thinker | vLLM `PIECEWISE` | multimodal preprocessing and output routing |
+| 1 Talker | vLLM `PIECEWISE` | conditioning preprocessing, codec sampling, request-state updates |
+| 2 Code2Wav | inner exact-shape NPUGraph for the CFM DiT estimator | encoder, timestep embedding, HiFT/RNG, request parsing, stream-state commit |
+
+Stage 2 keeps `enforce_eager: true` for the generation runner because its
+Python request metadata and per-request cache dictionaries are not valid
+outer-graph inputs. The backend captures only the deterministic CFM DiT
+estimator, with an exact graph key for each tensor shape. Timestep embedding
+stays eager because the upstream implementation creates a CPU frequency tensor;
+HiFT stays eager because it generates random phase. The graph cache stores up
+to 32 entries by default, after which unseen shapes run eagerly.
+
+An ACL graph capture failure is fatal to that Stage-2 process because older
+torch-npu releases can leave allocator and RNG capture state invalid. Restart
+the service after a capture failure. To run without capture, set
+`code2wav_enable_npu_graph: false` under the Stage-2
+`platforms.npu.stages[].additional_config` block before startup. Tune the cache
+limit there with `code2wav_max_npu_graphs`. Graph mode also requires
+`ASCEND_LAUNCH_BLOCKING` to be unset or set to `0`.
+
+The inner NPUGraph is independent of the outer runner setting, so do not use a
+global `--enforce-eager` override when Stage 0/1 `PIECEWISE` replay is desired.
+
## GPU
### 1 x GPU (default — single command)
@@ -211,9 +240,9 @@ speech output (TTS)"** checkbox on / off.
model processes still share one CUDA device.
- `--trust-remote-code` is required — the HF repo ships a custom
`MiniCPMO` config / model class.
-- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
- eager because its request-owned Flow/HiFT caches and variable chunk/cache
- shapes are not yet exposed through a static exact-shape graph wrapper.
+- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 keeps its
+ request orchestration eager while using inner CFM and HiFT CUDA Graphs. On
+ Ascend, the exact-shape CFM DiT estimator instead uses inner NPUGraph replay.
- All default stages use `max_num_seqs: 4` to reduce cross-process GPU
contention. Talker AR
state and Code2Wav caches are request-owned; Code2Wav batches only
@@ -224,8 +253,8 @@ speech output (TTS)"** checkbox on / off.
- `StageRequestStats.batch_size` is a request-scoped placeholder, not the
scheduler's execution batch.
- Single-GPU co-location trades throughput for hardware density: Stage 0/1
- CUDA Graph replay and eager Stage 2 vocoder kernels compete across three
- CUDA contexts. Use the 8x4090 config or a custom multi-GPU mapping for
+ CUDA Graph replay and Stage 2 vocoder kernels compete across three CUDA
+ contexts. Use the 8x4090 config or a custom multi-GPU mapping for
throughput-sensitive serving.
### 8 x RTX 4090 24GB (consumer-GPU layout)
diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md
index b6d2c9293..ede83d47d 100644
--- a/recipes/Qwen/Qwen3-TTS.md
+++ b/recipes/Qwen/Qwen3-TTS.md
@@ -18,11 +18,11 @@ examples in this repository.
Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
-| Task Type | Model | Description |
-| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
-| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
-| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
-| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
+| Task Type | Model | Description |
+| ------------- | ---------------------------------------- | ------------------------------------------------------------- |
+| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speaker voices with optional style/emotion control |
+| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Generate speech from a natural language voice description |
+| `Base` | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from reference audio + transcript |
Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
@@ -166,6 +166,15 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
+- Base codec termination: Base requests without an explicit `max_new_tokens`
+ use a text-scaled safety ceiling (at least 192 codec frames and no more than
+ the configured model limit). If the Talker reaches that ceiling without
+ codec EOS, non-streaming serving discards the incomplete audio and retries
+ once with a fresh seed; an explicit `seed` or `max_new_tokens` disables the
+ retry. SSE and WebSocket clients receive structured errors and must discard
+ previously emitted audio when the error contains `"action":"discard"`;
+ raw PCM streams terminate with a connection error after any partial bytes
+ already sent.
- Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.
## Hardware Support
@@ -252,3 +261,53 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- stage_id: 1
gpu_memory_utilization: 0.15
```
+
+### 1x AMD MI300X, 1.7B checkpoints
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `3fecb6953ca8dc51210cc0421ef24552267a41ef`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+- ONNX Runtime: onnxruntime-rocm 1.22.2.post3 with `ROCMExecutionProvider`
+- transformers: 5.15.0
+
+#### Command
+
+The MI300X tests also worked without `enforce_eager`. To use that setup, remove `enforce_eager: true` from stage 1 under `platforms.rocm` in `vllm_omni/deploy/qwen3_tts.yaml`.
+
+Set this environment variable before you start Qwen3-TTS:
+
+```bash
+export MIOPEN_FIND_MODE=FAST
+```
+
+#### Verification
+
+Each row records one offline request with one prompt and batch size one. Request time is the logged `e2e_wall_time_ms`, and stage 1 time is the logged `e2e_stage_1_wall_time_ms`.
+
+| Checkpoint | Code2Wav configuration | `MIOPEN_FIND_MODE` | Request time | Stage 1 time | Output duration | Real time factor | Maximum sampled device memory |
+| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
+| CustomVoice | Eager | Unset | 1.683 s | 1.613 s | 5.68 s | 0.30 | 60.91 GiB |
+| CustomVoice | Graphs allowed | `FAST` | 1.326 s | 1.259 s | 5.68 s | 0.23 | 62.19 GiB |
+| VoiceDesign | Eager | Unset | 1.195 s | 1.180 s | 4.40 s | 0.27 | 60.91 GiB |
+| VoiceDesign | Graphs allowed | `FAST` | 1.127 s | 1.054 s | 4.40 s | 0.26 | 62.21 GiB |
+| Base | Eager | Unset | 13.330 s | 13.263 s | 4.32 s | 3.09 | 61.74 GiB |
+| Base | Graphs allowed | `FAST` | 14.314 s | 14.245 s | 4.80 s | 2.98 | 62.84 GiB |
+
+All six processes finished successfully, and each output passed the 24 kHz mono WAV checks.
+
+#### Notes
+
+- The current ROCm config keeps Code2Wav in eager mode.
+- The tested alternative allows graphs and sets [`MIOPEN_FIND_MODE=FAST`](https://rocm.docs.amd.com/projects/MIOpen/en/develop/reference/env_variables.html). MIOpen uses a saved kernel choice when available and uses its immediate fallback otherwise.
+- CustomVoice used a graph for one of four logged Code2Wav batches. VoiceDesign used a graph for two of four batches. Base used no graphs for its four logged batches.
+- The recorded request time with graphs allowed was lower for CustomVoice and VoiceDesign and higher for Base. The Base runs generated different amounts of audio, so their request times are not a direct comparison.
+- Each setup was tested with one request, so the timing is not a performance benchmark.
+- GPU memory was sampled once per second with `rocm-smi`.
diff --git a/recipes/README.md b/recipes/README.md
index e010ade67..f82a5955d 100644
--- a/recipes/README.md
+++ b/recipes/README.md
@@ -39,7 +39,6 @@ recipes/
| Recipe | Task | Hardware |
|--------|------|----------|
-| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
| [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ |
@@ -47,18 +46,23 @@ recipes/
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
+| [`HiDream-ai/HiDream-O1-Image.md`](./HiDream-ai/HiDream-O1-Image.md) | Text-to-image with shared offline example | 1x or 2x H100 80GB |
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
+| [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB |
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
+| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
+| [`NVIDIA/SANA-Video-2B.md`](./NVIDIA/SANA-Video-2B.md) | Native and Diffusers-adapter 480p/720p T2V+I2V | 1x RTX 5090 32GB |
| [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
+| [`Robbyant/LingBot-World-2.0.md`](./Robbyant/LingBot-World-2.0.md) | Offline and experimental realtime interactive world generation | 1x H200/B200 or 2x B200 for TP=2 |
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
@@ -67,21 +71,23 @@ recipes/
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
+| [`NVIDIA/SANA-WM.md`](./NVIDIA/SANA-WM.md) | First-frame image-to-video serving with camera control | 1x 24GB+ CUDA GPU; ~40GB disk |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
-| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
-| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
-| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB |
+| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
+| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
+| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
+| [`XiaomiMiMo/MiMo-Audio.md`](./XiaomiMiMo/MiMo-Audio.md) | Offline + online omni audio (TTS, ASR, dialogue) | 1x RTX 5090 / 5090D 32GB |
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
-| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
+| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB / 2x non-standard RTX 4090 48GB (49,140 MiB each) |
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ |
diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md
new file mode 100644
index 000000000..851aeb930
--- /dev/null
+++ b/recipes/Robbyant/LingBot-World-2.0.md
@@ -0,0 +1,122 @@
+# LingBot-World 2.0
+
+> Offline and experimental realtime interactive world generation
+
+## Summary
+
+- Vendor: Robbyant
+- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
+- Task: image-conditioned interactive world generation
+- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
+- Hardware validated: NVIDIA H200 and B200
+- Maintainer: Community
+
+The checkpoint is separately licensed under CC BY-NC-SA and restricted to
+non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
+
+## Offline generation
+
+The offline path consumes one source image and an action directory containing
+`poses.npy` with shape `[frames, 4, 4]` and `intrinsics.npy` with shape
+`[frames, 4]`.
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2.py \
+ --prompt "The camera moves slowly forward through the scene." \
+ --image /path/to/first_frame.png \
+ --action-dir /path/to/actions/forward \
+ --num-frames 81 \
+ --output lingbot_world_v2.mp4
+```
+
+The current causal-fast checkpoint uses four DMD sampling steps. Raw frame
+counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
+longer bounded source trajectory (the official example contains 269 frames);
+the runtime consumes only the prefix needed by the request.
+
+## Realtime in-process generation
+
+The realtime example keeps AR-Diffusion KV and model-owned state across
+requests. Each JSONL line describes the prompt and/or three latent-frame camera
+actions applied at the next chunk boundary:
+
+```json
+{"event_id":1,"prompt":"A road through a forest","frames":[["j"],[],[]]}
+{"event_id":2,"frames":[["w"],["w"],["w"]]}
+{"event_id":3,"prompt":"The road enters a snowy valley","frames":[[],[],[]]}
+```
+
+Run:
+
+```bash
+python examples/offline_inference/diffusion/lingbot_world_v2_realtime.py \
+ --image /path/to/first_frame.png \
+ --events /path/to/events.jsonl \
+ --output-dir /tmp/lingbot-realtime \
+ --gpu-memory-fraction 0.6
+```
+
+The example writes one latent tensor and one metadata JSON file per chunk. It
+exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
+-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
+transport.
+
+This PR intentionally does not define a public realtime HTTP/WebSocket schema.
+An online serving client should be added together with the public transport
+API rather than exposing LingBot-specific event fields from the generic model
+runtime. The structured camera-interaction frontend and transport work is
+tracked separately in
+[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
+
+## Realtime identity and controls
+
+- `session_id` identifies the persistent world and its worker-owned state.
+- `event_id` identifies a prompt/control update and remains monotonic across reset.
+- `chunk_index` is contiguous from zero and restarts from zero after reset.
+- `request_id` correlates one chunk snapshot with its output metadata.
+- AsyncOmni uses a separate UUID-suffixed internal engine routing ID.
+
+The generic runtime transports controls as opaque snapshots. LingBot's adapter
+accepts:
+
+- `lingbot.camera_actions.v1` for per-latent-frame key states such as `w`, `a`,
+ `s`, `d`, `i`, `j`, `k`, and `l`;
+- `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories.
+
+## Validation
+
+Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42.
+The exercised matrix includes:
+
+- TP=1 and TP=2 execution;
+- two interleaved resident sessions;
+- action input and prompt switching;
+- seven contiguous chunks crossing the sink plus recent rolling window;
+- direct versus paged replay;
+- CUDA Graph execution;
+- reset, close, failure cleanup, and exact metadata matching; and
+- nine-frame VAE decode.
+
+Official `generate.py` accuracy and performance numbers, together with
+generation artifacts, are recorded in the PR validation section for the exact
+tested commit.
+
+## Current limitations
+
+- Only the 14B causal-fast checkpoint is supported.
+- The realtime control plane is internal; there is no public server transport yet.
+- AR-Diffusion stages currently require one replica because session-affine
+ routing across replicas is not implemented.
+- One AR block is generated per request and `max_num_seqs` must be one.
+- Stateful streaming VAE decode is not implemented; the realtime example emits
+ latent chunks.
+- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
+ Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed.
+- No AMD GPU, Ascend NPU, or Intel GPU support is claimed.
+
+## References
+
+- Checkpoint:
+- Official implementation:
+- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
+- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)
diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md
index be25df570..8e9091d63 100644
--- a/recipes/SenseNova/SenseNova-U1.md
+++ b/recipes/SenseNova/SenseNova-U1.md
@@ -178,6 +178,48 @@ with the expected 1536×2720 resolution.
- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
for fused QKV and gate/up projections with TP support.
+### 1x AMD MI300X 192GB
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `a704c8759c96e123c0d7c89b11f120b1c0f120cf`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Command
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+ --model SenseNova/SenseNova-U1-8B-MoT \
+ --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look." \
+ --width 1536 \
+ --height 2720 \
+ --seed 42 \
+ --num-inference-steps 50 \
+ --cfg-scale 4.0 \
+ --extra-body '{"think": true, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
+ --enable-diffusion-pipeline-profiler \
+ --log-stats \
+ --output sensenova_u1_mi300x.png
+```
+
+#### Verification
+
+The command completed and wrote a valid 1536 by 2720 RGB PNG.
+
+#### Notes
+
+- Generation with 50 inference steps took 34.011 seconds.
+- Model loading used 32.774 GiB and took 13.987 seconds.
+- The internal profiler recorded 35.67 GB reserved and 35.10 GB allocated for the request.
+- The highest one second whole device memory sample was 37.93 GiB.
+
## Online Serving
SenseNova-U1 supports all four modalities via the OpenAI-compatible
diff --git a/recipes/StabilityAI/Stable-Audio-Open.md b/recipes/StabilityAI/Stable-Audio-Open.md
index a7f18128d..d24d516ca 100644
--- a/recipes/StabilityAI/Stable-Audio-Open.md
+++ b/recipes/StabilityAI/Stable-Audio-Open.md
@@ -23,8 +23,6 @@ offline validation sample with TeaCache and online serving through the
- Model:
- Offline example:
[`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio)
-- Online example:
- [`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio)
- Audio generation API:
[`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md)
- Related issue:
@@ -78,7 +76,7 @@ python examples/offline_inference/text_to_audio/text_to_audio.py \
Start the online serving endpoint:
```bash
-vllm-omni serve /path/to/stable-audio-open-1.0 \
+vllm serve /path/to/stable-audio-open-1.0 \
--host 0.0.0.0 \
--port 8091 \
--gpu-memory-utilization 0.9 \
@@ -103,7 +101,7 @@ curl -X POST http://localhost:8091/v1/audio/generate \
"seed": 42,
"response_format": "wav"
}' \
- --output examples/online_serving/stable_audio/piano_10s.wav
+ --output piano_10s.wav
```
##### Verification
@@ -137,12 +135,12 @@ PY
Validate the online output:
```bash
-ls -lh examples/online_serving/stable_audio/piano_10s.wav
+ls -lh piano_10s.wav
python - <<'PY'
import soundfile as sf
-path = "examples/online_serving/stable_audio/piano_10s.wav"
+path = "piano_10s.wav"
audio, sample_rate = sf.read(path)
print("sample_rate:", sample_rate)
print("shape:", audio.shape)
@@ -164,7 +162,50 @@ PY
`torchaudio==2.11.0+cu130`.
- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary
warnings observed during validation did not prevent successful generation.
-- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts,
- ROCm, XPU, and NPU setups are not covered here.
+- The RTX 4090 entry was validated on one 24 GB GPU. The MI300X entry below covers one 192 GB GPU.
- Long generations, higher inference-step counts, and non-WAV response formats
were not benchmarked in this recipe.
+
+#### 1x AMD MI300X 192GB
+
+##### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+##### Commands
+
+```bash
+python3 examples/offline_inference/text_to_audio/text_to_audio.py \
+ --model stabilityai/stable-audio-open-1.0 \
+ --prompt "A gentle piano melody with soft room ambience" \
+ --negative-prompt "Low quality, distorted, noisy" \
+ --seed 42 \
+ --guidance-scale 7.0 \
+ --audio-length 10.0 \
+ --num-inference-steps 50 \
+ --cache-backend tea_cache \
+ --enable-diffusion-pipeline-profiler \
+ --output stable_audio_10s.wav
+```
+
+##### Verification
+
+The command completed and wrote a valid 44.1 kHz stereo WAV with 10.00 seconds of audio.
+
+##### Notes
+
+- TeaCache ran with `rel_l1_thresh=0.2`.
+- Model loading used 2.7891 GiB and took 3.706 seconds.
+- Generation took 4.750 seconds, which gives a real time factor of 0.475 for the 10.00 second output.
+- The internal profiler recorded 15.65 GB reserved and 9.68 GB allocated for the request.
+- The highest one second whole device memory sample was 19.61 GiB.
+- The output RMS was 0.0887, and its peak absolute amplitude was 0.5761.
+- The full process took 384 seconds, including startup and compilation.
diff --git a/recipes/TEMPLATE.md b/recipes/TEMPLATE.md
index 6af66f721..f2a145660 100644
--- a/recipes/TEMPLATE.md
+++ b/recipes/TEMPLATE.md
@@ -39,6 +39,9 @@ platform section, document one or more tested hardware configurations.
#### Command
+Serve commands use `vllm serve --omni`. The `--omni` flag is what
+selects the Omni pipeline, so it is required and cannot be dropped.
+
```bash
# Add the exact command(s) here
```
diff --git a/recipes/XiaomiMiMo/MiMo-Audio.md b/recipes/XiaomiMiMo/MiMo-Audio.md
new file mode 100644
index 000000000..2850cad2a
--- /dev/null
+++ b/recipes/XiaomiMiMo/MiMo-Audio.md
@@ -0,0 +1,161 @@
+# MiMo-Audio for omni audio generation and understanding
+
+> Offline + online TTS, ASR, audio understanding, and multi-turn dialogue
+
+## Summary
+
+- Vendor: Xiaomi MiMo
+- Model: `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- Task: Text-to-speech, speech-to-text, audio understanding, spoken dialogue
+- Mode: Offline end-to-end scripts and online OpenAI-compatible chat serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for running
+MiMo-Audio on a single **32 GB Blackwell consumer GPU** (RTX 5090 / 5090D).
+The bundled deploy profile [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+splits VRAM between the AR thinker (stage 0) and code2wav (stage 1), avoids
+FlashAttention-2 on SM 12.x, and keeps the audio tokenizer on CPU to reduce
+Stage-1 GPU pressure.
+
+## References
+
+- Upstream model card:
+- Audio tokenizer:
+- Offline example:
+ [`examples/offline_inference/mimo_audio/end2end.py`](../../examples/offline_inference/mimo_audio/end2end.py)
+- Online example:
+ [`examples/online_serving/mimo_audio/README.md`](../../examples/online_serving/mimo_audio/README.md)
+- Default deploy config (generic single GPU):
+ [`vllm_omni/deploy/mimo_audio.yaml`](../../vllm_omni/deploy/mimo_audio.yaml)
+- 5090D-tuned deploy config:
+ [`vllm_omni/deploy/mimo_audio_5090d.yaml`](../../vllm_omni/deploy/mimo_audio_5090d.yaml)
+- Pipeline / deploy docs:
+ [Stage configs](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/)
+- Related issue or discussion:
+ [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+- Tokenizer device alignment fix (required for CPU tokenizer on stage 1):
+ [PR #6539](https://github.com/vllm-project/vllm-omni/pull/6539)
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA GeForce RTX 5090 / 5090D 32GB
+
+Two-stage async-chunk pipeline (thinker + code2wav) on a single GPU. Tested
+offline with `tts_sft`.
+
+#### Environment
+
+- OS: Ubuntu (Linux x86_64)
+- Python: 3.12
+- PyTorch: 2.11.0+cu129
+- Driver / runtime: NVIDIA driver 570.x, CUDA 12.9 runtime (Blackwell SM 12.x)
+- GPU: NVIDIA GeForce RTX 5090 D, 32 GB
+- vLLM: 0.26.0+cu129
+- vLLM-Omni: `main` (tested while [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) was open for CPU-tokenizer support)
+
+Install vLLM-Omni from source and align the vLLM wheel with your CUDA stack.
+On Blackwell, prefer the **cu129** vLLM build over cu13 defaults when using
+PyTorch 2.11+cu129.
+
+Download or cache:
+
+- `XiaomiMiMo/MiMo-Audio-7B-Instruct`
+- `XiaomiMiMo/MiMo-Audio-Tokenizer` (or a local checkout of the tokenizer repo)
+
+Export the runtime knobs below before offline or online runs:
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="${MIMO_AUDIO_TOKENIZER_PATH:-XiaomiMiMo/MiMo-Audio-Tokenizer}"
+export MIMO_AUDIO_TOKENIZER_DEVICE=cpu
+export MIMO_AUDIO_TOKENIZER_CUDA_GRAPH=0
+export VLLM_USE_FLASHINFER_SAMPLER=0
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+```
+
+Optional (mirror / download stability):
+
+```bash
+export HF_HUB_DISABLE_XET=1
+export HF_ENDPOINT=https://hf-mirror.com # if you use a Hugging Face mirror
+```
+
+Do **not** set `VLLM_ATTENTION_BACKEND` on vLLM 0.26.x; configure
+`attention_backend: TRITON_ATTN` in the deploy YAML instead.
+
+#### Command
+
+Offline TTS validation from the repository root:
+
+```bash
+cd examples/offline_inference/mimo_audio
+
+python3 -u end2end.py \
+ --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+ --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+ --query-type tts_sft \
+ --text "The weather is so nice today."
+```
+
+Online serving (requires MiMo chat template):
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+
+vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
+ --deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml \
+ --served-model-name "MiMo-Audio-7B-Instruct" \
+ --port 18091 \
+ --chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
+```
+
+Pass the same environment exports to the server process when using a CPU
+tokenizer or disabling FlashInfer sampling.
+
+#### Verification
+
+Offline:
+
+```bash
+ls -lh output_audio/tts_sft/*.wav
+```
+
+Expect a non-empty `.wav` under `output_audio/tts_sft/` and a companion `.txt`
+with the decoded text. A successful run completes both stages without a
+Stage-1 `decode_vq` device mismatch.
+
+Online:
+
+```bash
+curl -s -o /tmp/mimo_health.txt -w '%{http_code}' http://127.0.0.1:18091/health
+```
+
+Then use the client under
+[`examples/online_serving/mimo_audio/`](../../examples/online_serving/mimo_audio/).
+
+#### Notes
+
+- Memory usage (observed offline): Stage-0 weights ~16.6 GiB; KV cache ~4.9 GiB
+ with `gpu_memory_utilization: 0.78` on stage 0 and `0.12` on stage 1.
+- Key flags:
+ - `--deploy-config vllm_omni/deploy/mimo_audio_5090d.yaml` for 5090D memory
+ and `TRITON_ATTN`.
+ - `MIMO_AUDIO_TOKENIZER_PATH` is mandatory.
+ - `MIMO_AUDIO_TOKENIZER_DEVICE=cpu` saves VRAM; requires
+ [#6539](https://github.com/vllm-project/vllm-omni/pull/6539) (or equivalent
+ fix) so code2wav decode uses the tokenizer device.
+ - `VLLM_USE_FLASHINFER_SAMPLER=0` avoids FlashInfer top-p/top-k sampler
+ issues on consumer GPUs.
+- Known limitations:
+ - Blackwell may log `SM 12.x requires CUDA >= 12.9` from some optional
+ tooling; the tested path uses TRITON attention and cu129 wheels.
+ - Without a working `flash-attn` build, audio quality may be slightly
+ metallic; install FlashAttention when a compatible wheel is available.
+ - First request triggers Triton JIT for rotary / attention kernels (latency
+ spike); subsequent requests are faster.
+ - For longer outputs or batching, you may need to raise `max_model_len` in
+ the deploy YAML and match `max_position_embeddings` in the model config
+ (see the offline example README).
diff --git a/recipes/audiox/AudioX.md b/recipes/audiox/AudioX.md
deleted file mode 100644
index a2e260736..000000000
--- a/recipes/audiox/AudioX.md
+++ /dev/null
@@ -1,130 +0,0 @@
-# AudioX
-
-> AudioX MMDiT for unified audio + music generation: t2a / t2m / v2a / v2m / tv2a / tv2m
-
-## Summary
-
-- Vendor: HKUSTAudio (project), `zhangj1an/AudioX` weight bundle
-- Model: `zhangj1an/AudioX`
-- Task: Text/video → audio or music. Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`.
-- Mode: Offline inference + online serving (pure diffusion)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe to run AudioX for sound-effect (`*2a`) or music (`*2m`) generation
-from a text prompt and/or video clip. AudioX is a unified diffusion transformer
-that produces stereo 44.1 kHz audio up to ~10 s per call.
-
-## References
-
-- Project page:
-- vLLM-Omni weight bundle:
-- Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline`
-- Input transforms: `vllm_omni.transformers_utils.processors.audiox`
-- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs)
-- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py)
-- Online: standard OpenAI chat-completions endpoint (see commands below)
-
-## Hardware Support
-
-## GPU
-
-### 1x L4 24GB
-
-#### Environment
-
-- OS: Ubuntu 22.04
-- Python: 3.10+
-- Driver / runtime: CUDA 12.4
-- vLLM version: 0.20.0
-- vLLM-Omni version: 0.1.x
-
-#### Command
-
-AudioX uses the **standard** `text_to_audio` example offline and the standard
-OpenAI chat-completions endpoint online. All model-specific knobs
-(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`,
-`cfg_rescale`, `video_path`, `audio_path`) are declared in
-`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) /
-`--extra-body` (offline). The offline example also exposes `--task`, `--video`,
-`--audio-start`, and `--audio-length` shortcuts for the most common knobs.
-
-Offline — text tasks (`t2a` / `t2m`):
-
-```bash
-huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
-python examples/offline_inference/text_to_audio/text_to_audio.py \
- --model ./audiox_weights --task t2a \
- --prompt "Fireworks burst twice, followed by a clock ticking." \
- --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
- --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \
- --output t2a.wav
-```
-
-Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`:
-
-```bash
-python examples/offline_inference/text_to_audio/text_to_audio.py \
- --model ./audiox_weights --task tv2a \
- --prompt "drum beating sound and human talking" \
- --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \
- --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \
- --output tv2a.wav
-```
-
-Online:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \
- vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099
-
-curl -sS -X POST http://localhost:8099/v1/chat/completions \
- -H 'Content-Type: application/json' \
- -d '{
- "model": "zhangj1an/AudioX",
- "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}],
- "extra_body": {
- "num_inference_steps": 250,
- "guidance_scale": 7.0,
- "seed": 42,
- "audiox_task": "t2m",
- "seconds_total": 10.0,
- "sigma_min": 0.3,
- "sigma_max": 500.0
- }
- }' > t2m.json
-```
-
-For `v2*` / `tv2*` online, attach the video as a `video_url` content item
-(data URI for local files) in the message `content`.
-
-#### Verification
-
-```bash
-# Health check
-curl http://localhost:8099/health
-
-# Listen to the saved file (stereo, 44.1 kHz, sigma_min=0.03, sigma_max=1000 — upstream defaults)
-ffprobe t2a.wav
-```
-
-#### Notes
-
-- Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio.
-- Output rate: 44.1 kHz stereo regardless of requested rate.
-- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the
- `audiox_task` knob (in `extra_body` online, or `--task` offline).
-- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video`
- (→ `video_path` knob), online via a `video_url` content item.
-- Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`).
-- Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with
- `QKVParallelLinear`); cross-attention K/V is also TP-sharded.
-
-### Known limitations
-
-- Inference uses an inlined DPM-Solver++(3M) SDE sampler (k-diffusion port). Replacing it with
- diffusers' `EDMDPMSolverMultistepScheduler` introduces a fixed ~861 Hz resonance and is not
- recommended.
-- Generation is fixed at 10 s (configured by the bundle's `sample_size`); longer outputs require
- a different bundle.
diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md
index 8eeea5cc2..2866e7a0c 100644
--- a/recipes/cosmos3/Cosmos3-DistOffload.md
+++ b/recipes/cosmos3/Cosmos3-DistOffload.md
@@ -189,8 +189,10 @@ wait
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
-- Online quantization (FP8) is incompatible with the DLO+AllGather mmap path;
- use `--dlo-no-use-allgather` or disable online quantization
+- Online per-tensor FP8 linears with DLO AllGather use the ordinary loader
+ instead of direct checkpoint mmap, then shard the finalized FP8 weights and
+ scales. Size host memory for a transient complete FP8 model on every rank at
+ startup
- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path. TP with
`--dlo-no-use-allgather` retains the standard TP loader and should be treated
as an experimental compatibility path
diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md
index b27d1b473..4aaebf0be 100644
--- a/recipes/cosmos3/Cosmos3-Nano.md
+++ b/recipes/cosmos3/Cosmos3-Nano.md
@@ -195,11 +195,13 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
# Transfer V2V with a precomputed depth control video. `control_path` can point
# to a local image/video; edge and blur can also be computed from `input_reference`
-# by passing `"edge":true` or `"blur":true`.
+# by passing `"edge":true` or `"blur":true`. The reference negative prompt is
+# optional; omit the `negative_prompt` form field to use an empty negative branch.
curl -sS -X POST http://localhost:8000/v1/videos/sync \
-H "Accept: video/mp4" \
-F "model=nvidia/Cosmos3-Nano" \
-F "prompt=Generate a realistic scene following the provided control video." \
+ --form-string "negative_prompt=$(jq -c . recipes/cosmos3/negative_prompt.json)" \
-F "size=1280x720" -F "num_frames=121" \
-F "num_inference_steps=50" -F "seed=125" \
-F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
@@ -322,15 +324,34 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \
`extra_params={"guardrails":false}` (per request) toggles safety. The
per-request flag only takes effect when the server was launched **with**
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
- `use_resolution_template` / `use_duration_template` are off by default and only
- needed when not using upsampled prompts that already encode resolution/duration.
+ Outside transfer mode, `use_resolution_template` / `use_duration_template`
+ are off by default and only needed when not using upsampled prompts that
+ already encode resolution/duration.
For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
API decodes the first or last needed reference frames (`"first"` by default).
- **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`,
`seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as
`{"control_path": "/path/to/control.mp4"}`; `edge` also accepts
- `preset_edge_threshold` and `blur` accepts `preset_blur_strength`.
+ `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. Every hint
+ accepts a non-negative `control_weight`; weights are normalized across active
+ controls and therefore only set their relative influence. A single positive
+ weight always normalizes to `1.0`; use `control_guidance` to change the
+ absolute strength of a single control. With two or more active controls, the
+ per-control attention passes run replicated on every sequence-parallel
+ (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for
+ multi-control transfer requests.
+ Transfer always uses Cosmos3's transfer-specific system prompt. By default it
+ also appends a directive naming every active hint and asking the model to
+ follow its shape, position, and motion precisely; set the request-level
+ `emphasize_control_in_prompt` option to `false` for prompt ablations. Transfer
+ does not add a negative prompt automatically. An optional reference prompt is
+ provided in [`negative_prompt.json`](negative_prompt.json); compact the JSON
+ with `jq -c` and pass it through `negative_prompt` as shown in the transfer
+ example. Transfer enables duration/FPS and resolution metadata on both CFG
+ branches. Set `use_duration_template` or `use_resolution_template` to
+ `false` to disable either template. `negative_metadata_mode` accepts `same`
+ (the transfer default), `inverse`, or `none`.
Transfer-level options include `control_guidance`,
`control_guidance_interval`, `num_video_frames_per_chunk` (default `93`,
`101` for WSM), `num_conditional_frames` (default `1`),
diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md
index 54d4d7e32..c8bfb6bbd 100644
--- a/recipes/cosmos3/Cosmos3-Super.md
+++ b/recipes/cosmos3/Cosmos3-Super.md
@@ -125,6 +125,23 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
under the top-level `action` field. Verified on the 64B Super under
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
and the rollout video reliably.
+- **Transfer controls:** same semantics as Nano — `extra_params` may include
+ `edge`, `blur`, `depth`, `seg`, or `wsm` hints; see the Cosmos3-Nano recipe
+ for the request shapes and full option list. Every hint accepts a
+ non-negative `control_weight`; weights are normalized across active controls
+ and therefore only set their relative influence. A single positive weight
+ always normalizes to `1.0`; use `control_guidance` to change the absolute
+ strength of a single control. With two or more active controls, the
+ per-control attention passes run replicated on every sequence-parallel
+ (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for
+ multi-control transfer requests. Transfer always uses Cosmos3's
+ transfer-specific system prompt, appends a control-adherence directive by
+ default (`emphasize_control_in_prompt: false` disables it), and enables
+ duration/FPS and resolution metadata on both CFG branches
+ (`use_duration_template` / `use_resolution_template` to disable;
+ `negative_metadata_mode` accepts `same`, `inverse`, or `none`, defaulting to
+ `same`). Transfer does not add a negative prompt automatically; an optional
+ reference prompt is provided in [`negative_prompt.json`](negative_prompt.json).
## NPU
diff --git a/recipes/cosmos3/negative_prompt.json b/recipes/cosmos3/negative_prompt.json
new file mode 100644
index 000000000..9cefcb4d5
--- /dev/null
+++ b/recipes/cosmos3/negative_prompt.json
@@ -0,0 +1,108 @@
+{
+ "subjects": [
+ {
+ "description": "Blurry, poorly defined subjects with inconsistent shapes and unrealistic proportions.",
+ "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+ "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+ "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+ "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+ "orientation": "Unnatural orientations that defy physics and spatial logic.",
+ "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.",
+ "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.",
+ "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.",
+ "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.",
+ "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+ "gender": "",
+ "age": "",
+ "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+ "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+ "number_of_subjects": 0,
+ "number_of_arms": 0,
+ "number_of_legs": 0
+ },
+ {
+ "description": "Extremely low-quality subjects with visible rendering artifacts, broken mesh geometry, and completely unrealistic proportions throughout.",
+ "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+ "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+ "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+ "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+ "orientation": "Unnatural orientations that defy physics and spatial logic.",
+ "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.",
+ "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.",
+ "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.",
+ "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.",
+ "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+ "gender": "",
+ "age": "",
+ "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+ "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+ "number_of_subjects": 0,
+ "number_of_arms": 0,
+ "number_of_legs": 0
+ },
+ {
+ "description": "Poorly generated subjects exhibiting all hallmarks of failed neural rendering \u2014 flickering edges, inconsistent depth, and uncanny spatial relationships.",
+ "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+ "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+ "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+ "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+ "orientation": "Unnatural orientations that defy physics and spatial logic.",
+ "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.",
+ "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.",
+ "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.",
+ "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.",
+ "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+ "gender": "",
+ "age": "",
+ "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+ "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+ "number_of_subjects": 0,
+ "number_of_arms": 0,
+ "number_of_legs": 0
+ }
+ ],
+ "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity \u2014 patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOAD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.",
+ "lighting": {
+ "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.",
+ "direction": "Inconsistent light sources \u2014 shadows point in multiple contradictory directions, breaking physical plausibility.",
+ "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable \u2014 even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.",
+ "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s."
+ },
+ "aesthetics": {
+ "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.",
+ "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.",
+ "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.",
+ "patterns": "Visible tiling artifacts in textures, moir\u00e9 patterns, and aliasing on edges."
+ },
+ "cinematography": {
+ "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow.",
+ "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space.",
+ "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content \u2014 panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles \u2014 instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.",
+ "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.",
+ "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.",
+ "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression."
+ },
+ "style_medium": "Low quality compressed digital video with visible encoding artifacts",
+ "artistic_style": "Amateur, unpolished with inconsistent visual style",
+ "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.",
+ "actions": [
+ {
+ "time": "0:00-0:08",
+ "description": "Subjects attempt to move but their motion is jerky, temporally inconsistent, and physically implausible. Background elements flicker and shift between frames."
+ }
+ ],
+ "text_and_signage_elements": [],
+ "segments": [
+ {
+ "segment_index": 0,
+ "time_range": "0:00-0:08",
+ "description": "A single continuous shot suffering from severe temporal inconsistencies \u2014 subjects that morph and deform between frames, backgrounds that shift and wobble, and rendering quality that fluctuates visibly over time. Motion blur is applied incorrectly, smearing in directions that don't match actual movement. Frame-to-frame coherence breaks down with individual pixels changing color randomly in flat areas. Texture detail level fluctuates between frames as if the rendering budget varied shot to shot. Color grading drifts over the duration with no creative motivation. Noise patterns change between frames in ways that draw attention rather than being invisible. Overall visual quality degrades progressively from start to finish.",
+ "key_changes": "No meaningful progression or narrative development. Visual quality degrades over time.",
+ "camera": "Unstable, poorly controlled camera work with visible mathematical interpolation artifacts."
+ }
+ ],
+ "transitions": [],
+ "temporal_caption": "The scene opens at 0.0 seconds with a poorly rendered establishing shot that immediately reveals low production quality. At 1.0 seconds, subjects begin to move but their motion is jerky and inconsistent, with limbs bending at unnatural angles and objects clipping through each other. From 2.0 to 4.0 seconds, the camera shakes violently while the scene exhibits visible compression artifacts, color banding in the sky, and flickering in the shadows. Between 4.0 and 6.0 seconds, temporal coherence breaks down as elements appear and disappear between frames, textures swim and morph unnaturally, and the lighting shifts abruptly without physical cause. In the final 2 seconds, the overall visual quality deteriorates further with increasing noise, blur, and a general loss of spatial coherence that makes the scene nearly unwatchable. Additionally, the frame rate appears inconsistent with visible judder and stuttering throughout. Color temperature shifts randomly between warm and cool tones with no motivation. The encode quality degrades in complex regions showing macro-blocking and mosquito noise around moving edges. Temporal noise patterns are spatially correlated, creating swimming artifacts on flat surfaces.",
+ "audio_description": "",
+ "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent \u2014 objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded."
+ }
diff --git a/recipes/inclusionAI/Ming-omni-tts.md b/recipes/inclusionAI/Ming-omni-tts.md
index e42959e8a..3666dce0f 100644
--- a/recipes/inclusionAI/Ming-omni-tts.md
+++ b/recipes/inclusionAI/Ming-omni-tts.md
@@ -5,20 +5,21 @@
## Summary
- Vendor: inclusionAI
-- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Model: `inclusionAI/Ming-omni-tts-16.8B-A3B`, `inclusionAI/Ming-omni-tts-0.5B`
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
- Maintainer: Community
## References
-- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- [Huggingface: inclusionAI/Ming-omni-tts-16.8B-A3B](https://huggingface.co/inclusionAI/Ming-omni-tts-16.8B-A3B)
+- [Huggingface: inclusionAI/Ming-omni-tts-0.5B](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
- Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
## Hardware Support
-This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B and the MoE 16.8B two-stage TTS pipeline deployment.
Other hardware is welcome as community validation lands.
## CUDA
@@ -39,11 +40,7 @@ Other hardware is welcome as community validation lands.
Launch the two-stage talker:
```bash
-vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
- --deploy-config vllm_omni/deploy/ming_tts.yaml \
- --omni \
- --port 8091 \
- --enforce-eager
+vllm serve inclusionAI/Ming-omni-tts-16.8B-A3B --omni --port 8091
```
#### Verification
@@ -54,7 +51,7 @@ Basic synthesis (save the WAV bytes):
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
- "model": "inclusionAI/Ming-omni-tts-0.5B",
+ "model": "inclusionAI/Ming-omni-tts-16.8B-A3B",
"input": "你好,这是 Ming 在线语音合成测试。",
"response_format": "wav"
}' --output ming_tts_basic.wav
@@ -74,7 +71,7 @@ Reference-audio zero-shot cloning:
REF_AUDIO="$BASE/10002287-00000094.wav"
jq -n --arg ref_audio "$REF_AUDIO" \
'{
- model: "inclusionAI/Ming-omni-tts-0.5B",
+ model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
input: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。",
ref_audio: $ref_audio,
ref_text: "在此奉劝大家别乱打美白针。",
@@ -107,7 +104,7 @@ jq -n \
--arg a "$REF_A" \
--arg b "$REF_B" \
'{
- model: "inclusionAI/Ming-omni-tts-0.5B",
+ model: "inclusionAI/Ming-omni-tts-16.8B-A3B",
input: $input,
ref_audio: [$a, $b],
ref_text: $ref_text,
@@ -170,6 +167,35 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
for the dialect case.
+**Offline smoke result on ROCm 7.2:**
+
+A later offline smoke run used the same MI300X model with the following software:
+
+- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
+- Kernel: Linux 6.8.0-134-generic, x86_64
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- ROCm or HIP: 7.2.53211
+- vLLM: 0.27.0+rocm723
+- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+
+```bash
+python3 examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+ --model inclusionAI/Ming-omni-tts-0.5B \
+ --deploy-config vllm_omni/deploy/ming_tts.yaml \
+ --case basic \
+ --text "你好,这是 AMD MI300X ROCm 配方验证。" \
+ --output-dir ./ming_tts_output \
+ --output-name ming_omni_tts_mi300x.wav \
+ --enforce-eager \
+ --log-stats \
+ --metadata-json ./ming_tts_manifest.json
+```
+
+The checked run produced a valid 44.1 kHz mono WAV with 4.80 seconds of audio, RMS 0.1449, and peak absolute amplitude 0.8621. The first request took 20.406 seconds, which gives a real time factor of 4.25. The stage weight loads used 1.31 GiB and 1.47 GiB and took 1.756 seconds and 1.264 seconds. The largest one second whole device memory sample was 91.05 GiB, including the reserved KV cache. Both stages used eager mode, and the AR stage selected `TRITON_ATTN`.
+
+The run used the official ROCm image built from `docker/Dockerfile.rocm`.
+
## Notes
- The official ROCm image includes the platform dependencies.
diff --git a/recipes/k2-fsa/OmniVoice.md b/recipes/k2-fsa/OmniVoice.md
new file mode 100644
index 000000000..74e1b12f5
--- /dev/null
+++ b/recipes/k2-fsa/OmniVoice.md
@@ -0,0 +1,98 @@
+# OmniVoice
+
+> Offline and online text to speech with automatic voice, voice design, and reference voice cloning
+
+## Summary
+
+- Vendor: k2-fsa
+- Model: `k2-fsa/OmniVoice`
+- Task: Multilingual text to speech
+- Mode: Offline inference for all modes and online serving for automatic voice
+- Maintainer: Community
+
+## References
+
+- [Hugging Face model](https://huggingface.co/k2-fsa/OmniVoice)
+- [Offline text to speech examples](../../docs/user_guide/examples/offline_inference/text_to_speech.md#omnivoice)
+- [Online text to speech examples](../../docs/user_guide/examples/online_serving/text_to_speech.md#omnivoice)
+- [Deploy config](../../vllm_omni/deploy/omnivoice.yaml)
+
+## Requirements
+
+Download the checkpoint before the first run:
+
+```bash
+huggingface-cli download k2-fsa/OmniVoice
+```
+
+Automatic voice and voice design require `transformers>=4.57.0`. Reference voice cloning requires `transformers>=5.3.0`.
+
+The vLLM Omni deploy config runs the OmniVoice generator and decoder inside one float32 diffusion stage.
+
+## ROCm
+
+### 1x AMD MI300X
+
+#### Environment
+
+- GPU: one AMD Instinct MI300X with 191.69 GiB visible HBM and gfx942
+- Kernel: Linux 6.8.0-134-generic, x86_64
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- ROCm or HIP: 7.2.53211
+- vLLM: 0.27.0+rocm723
+- vLLM Omni checkout: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- transformers: 5.15.0
+
+The run used the official ROCm image built from `docker/Dockerfile.rocm`.
+
+#### Automatic voice
+
+Run the offline example from the repository root:
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+ --model k2-fsa/OmniVoice \
+ --deploy-config vllm_omni/deploy/omnivoice.yaml \
+ --text "Hello, this is OmniVoice running on one AMD MI300X." \
+ --seed 42 \
+ --output omnivoice_mi300x.wav
+```
+
+The checked run produced a valid 24 kHz mono WAV with 3.60 seconds of audio, RMS 0.1194, and peak absolute amplitude 0.9129. The first request took 0.685 seconds, which gives a real time factor of 0.19. Model loading used 3.5739 GiB and took 9.564 seconds. The request used 4.56 GB reserved and 3.64 GB allocated according to the internal profiler, while the largest one second whole device sample was 8.33 GiB. The log states that the OmniVoice Triton kernels loaded.
+
+#### Reference voice cloning
+
+The offline interface also accepts a reference WAV and its transcript:
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+ --model k2-fsa/OmniVoice \
+ --text "Hello, this is a cloned voice." \
+ --ref-audio ref.wav \
+ --ref-text "This is the reference transcription." \
+ --output cloned.wav
+```
+
+#### Voice design
+
+```bash
+python3 examples/offline_inference/text_to_speech/omnivoice/end2end.py \
+ --model k2-fsa/OmniVoice \
+ --text "Hello, this is a designed voice." \
+ --instruct "female, low pitch, British accent" \
+ --output designed.wav
+```
+
+## Online serving
+
+Online serving currently supports automatic voice:
+
+```bash
+vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
+```
+
+```bash
+cd examples/online_serving/text_to_speech/omnivoice
+python speech_client.py --text "Hello, how are you?"
+```
diff --git a/recipes/lerobot/Pi0.md b/recipes/lerobot/Pi0.md
new file mode 100644
index 000000000..6341e57a0
--- /dev/null
+++ b/recipes/lerobot/Pi0.md
@@ -0,0 +1,108 @@
+# Pi0 (π0, Pi-Zero) VLA
+
+> Vision-Language-Action policy serving over the OpenPI realtime websocket
+
+## Summary
+
+- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
+- Model: `lerobot/pi0_base`
+- Task: Vision-Language-Action — multi-camera images + a language instruction +
+ robot proprioceptive state → a continuous action chunk
+ `[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
+ **not** emit text tokens.
+- Mode: Online serving via the OpenPI realtime robot API
+ (`/v1/realtime/robot/openpi` websocket)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
+client connects over the OpenPI websocket and, per control step, sends an
+observation (camera frames + state + instruction) and receives an action chunk.
+
+A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
+calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
+ignored.
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
+server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
+
+## Hardware Support
+
+### 1x H200 141GB (Online serving)
+
+#### Environment
+
+- OS: Ubuntu 22.04+
+- Python: 3.12+
+- Driver / runtime: NVIDIA CUDA environment
+- vLLM / vLLM-Omni: use the commit you are deploying from
+- OpenPI client extras (for the example client): `openpi-client`
+ (`pip install -e packages/openpi-client` from the
+ [openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
+ `msgpack`, `msgpack-numpy`
+
+#### Command
+
+```bash
+vllm serve lerobot/pi0_base \
+ --omni \
+ --host 0.0.0.0 --port 8000 \
+ --served-model-name pi0 \
+ --deploy-config vllm_omni/deploy/pi0.yaml \
+ --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
+handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
+32) that enables the OpenPI realtime serving layer.
+
+#### Verification
+
+Run the example client (sends a synthetic observation and prints the returned
+action chunk):
+
+```bash
+python examples/online_serving/pi0/openpi_client.py \
+ --host 127.0.0.1 --port 8000 \
+ --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
+is a flat dict per inference:
+
+```python
+{
+ "observation.images.base_0_rgb": np.uint8[H, W, 3],
+ "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
+ "observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
+ "state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
+ "prompt": "pick up the red block",
+ "session_id": "", # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
+reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
+
+## References
+
+- Model blog:
+- Checkpoint:
+- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
+- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
+- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
+- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
+ OpenPI websocket e2e in
+ [`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)
diff --git a/recipes/rednote-hilab/dots.tts.md b/recipes/rednote-hilab/dots.tts.md
new file mode 100644
index 000000000..d5c7868cb
--- /dev/null
+++ b/recipes/rednote-hilab/dots.tts.md
@@ -0,0 +1,152 @@
+# dots.tts
+
+> Offline continuous-AR TTS at 48 kHz (rednote-hilab)
+
+## Summary
+
+- Vendor: rednote-hilab
+- Model: `rednote-hilab/dots.tts-soar`
+- Task: Text-to-speech, zero-shot synthesis only
+- Mode: Offline end-to-end example only (not yet wired for online serving —
+ see [Known limitations](#known-limitations))
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe as a known-good starting point for running
+`rednote-hilab/dots.tts-soar` offline on vLLM-Omni on consumer-class GPUs.
+dots.tts is a ~1.7B-parameter continuous-AR TTS model (Qwen2.5-1.5B base LM
++ 344M DiT flow-matching head + 180M AudioVAE) that emits 48 kHz mono audio.
+It follows the same "vLLM-native base LM + side-path computation" pattern as
+VoxCPM2 — single-stage pipeline
+`Qwen2.5-1.5B base LM → DiT (10-step Euler flow matching) → patch_encoder AR
+loopback → AudioVAE (streaming decode)` — with a plain Qwen2 backbone
+instead of MiniCPM4, and no FSQ / residual-LM stage.
+
+This is an early integration (see [Known limitations](#known-limitations)
+below before depending on it for anything beyond offline experimentation).
+
+## References
+
+- Offline end-to-end script:
+ [`examples/offline_inference/text_to_speech/dots_tts/end2end.py`](../../examples/offline_inference/text_to_speech/dots_tts/end2end.py)
+- Example guide:
+ [`examples/offline_inference/text_to_speech/README.md`](../../examples/offline_inference/text_to_speech/README.md#dotstts)
+- Default deploy config:
+ [`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
+- Talker / pipeline source:
+ [`vllm_omni/model_executor/models/dots_tts/`](../../vllm_omni/model_executor/models/dots_tts/)
+- Upstream: [rednote-hilab/dots.tts](https://github.com/rednote-hilab/dots.tts)
+
+## Hardware Support
+
+This recipe documents one tested 16 GB consumer-GPU configuration. Other
+vendor sections (ROCm, NPU) and larger-VRAM configurations are welcome as
+community validation lands.
+
+## GPU
+
+### 1 x RTX 5080 16GB (Single GPU, Minimum Recommended)
+
+dots.tts (~1.7B params across the base LM + DiT + AudioVAE + CAM++
+speaker encoder, bfloat16) fits comfortably on a single 16 GB GPU. The
+bundled default config at
+[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
+(`gpu_memory_utilization: 0.8`, `max_num_seqs: 4`, `enforce_eager: true`,
+`enable_prefix_caching: false`) loads cleanly with ~5.1 GiB for model
+weights and ~0.3 GiB peak activation; the remainder of the configured
+budget is available for KV cache. Total resident footprint at idle is
+roughly **7-8 GiB / 16 GB** — the only tight spot in the full CUDA-Graph
+roadmap would be step 8's graph capture (not implemented yet; this
+release runs `enforce_eager: true`, so it doesn't apply today).
+
+#### Environment
+
+- OS: Linux (WSL2)
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 595.95
+- torch: 2.11.0+cu130
+- vLLM: 0.26.0
+- vLLM-Omni: 0.22.1.dev (current `main`)
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
+ --model rednote-hilab/dots.tts-soar \
+ --text "Hello, this is a test of dots TTS running on vLLM Omni."
+```
+
+The deploy config at
+[`vllm_omni/deploy/dots_tts.yaml`](../../vllm_omni/deploy/dots_tts.yaml)
+is loaded automatically by the model registry (HF `model_type=dots_tts`).
+Pass `--deploy-config ` to override.
+
+#### Verification
+
+**T1 — offline zero-shot synthesis**:
+
+```bash
+python examples/offline_inference/text_to_speech/dots_tts/end2end.py \
+ --model rednote-hilab/dots.tts-soar \
+ --text "Hello, this is a test of dots TTS running on vLLM Omni."
+```
+
+Observed: `output_audio/output.wav`, 3.52 s @ 48 kHz mono. Single
+one-shot process (init → one `generate()` → exit), so the reported
+numbers include engine init and first-request warmup, not just steady-state
+per-step throughput — same caveat as VoxCPM2's recipe. `Inference: 5.52s`,
+`RTF: 1.569`.
+
+Weight-loading breakdown from the same run (all tensors matched, no
+missing/extra keys): 951/951 AudioVAE, 244/244 DiT, 270/270
+patch_encoder, 198 Qwen2, 938/938 CAM++ speaker encoder.
+
+Whisper transcription of the output matched the input text with no
+dropped leading word (confirms the streaming-vocoder patch-boundary fix
+described in [Known limitations](#known-limitations)).
+
+#### Notes
+
+- Output: 48 kHz mono WAV.
+- Checkpoints: `rednote-hilab/dots.tts-soar` is the validated default
+ used throughout this recipe. `rednote-hilab/dots.tts-base` shares the
+ same architecture but is unvalidated in this repo. `rednote-hilab/dots.tts-mf`
+ (MeanFlow, 2-4 step) is not supported — see below.
+- `enforce_eager: true` and `enable_prefix_caching: false` in the deploy
+ config are load-bearing, not just conservative defaults: with prefix
+ caching enabled, vLLM-Omni's prefix-cache multimodal-output merge path
+ does not preserve the `sparse_audio` marker this model relies on to
+ route audio output correctly, and generation silently truncates to a
+ single ~160 ms patch. Do not override `enable_prefix_caching` for this
+ model until that framework-level gap is fixed.
+
+## Known limitations
+
+- **No online serving yet.** dots.tts is not registered in
+ `vllm_omni/entrypoints/openai/tts_adapters/` (the framework other TTS
+ models in this repo use for `/v1/audio/speech`), so `vllm serve
+ rednote-hilab/dots.tts-soar --omni` will start but the OpenAI-compatible
+ speech endpoint does not know how to build request params for it yet.
+ Offline `Omni()` (as used in `end2end.py`) is the only supported path
+ today.
+- **Voice cloning is not wired.** The CAM++ x-vector speaker encoder
+ weights load and are exercised by `load_weights()`, but nothing in the
+ prompt builder consumes a reference audio yet — generation is zero-shot
+ only. `end2end.py` has no `--ref-audio`/`--ref-text` flags.
+- **`dots.tts-mf` (MeanFlow, 2-4 step) checkpoint is not supported.** Only
+ the fixed 10-step Euler DiT sampler used by `dots.tts-soar` /
+ `dots.tts-base` is implemented.
+- **No CUDA graph capture.** The talker runs fully eager. voxcpm2's three
+ captured graphs (base LM decode, CFM solver, VAE decode) have no
+ dots.tts equivalent yet.
+- **Concurrent requests do not scale.** Each request's 10-step DiT Euler
+ integration runs serially in the side path (no cross-request batching,
+ unlike voxcpm2's `enable_batched_cfm`). A community review of this
+ integration measured no throughput gain at `c=4` concurrent requests
+ versus `c=1`.
+- **`SamplingParams.seed` does not control audio-generation randomness.**
+ The DiT's flow-matching noise is deterministically derived per-request
+ from a fixed internal seed (reproducible run-to-run, matching
+ voxcpm2's `deterministic_cfm_seed` convention) rather than from the
+ caller-supplied `seed` field — the same limitation voxcpm2 has today.
diff --git a/recipes/zai-org/GLM-TTS.md b/recipes/zai-org/GLM-TTS.md
index 2bf1a202a..79bb85c47 100644
--- a/recipes/zai-org/GLM-TTS.md
+++ b/recipes/zai-org/GLM-TTS.md
@@ -1,4 +1,4 @@
-# GLM-TTS for Chinese/English TTS on 1x GPU
+# GLM-TTS for Chinese/English TTS
## Summary
@@ -44,20 +44,20 @@ conditioned on reference audio and its transcript.
Start the server from the repository root:
```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
```
Async chunking is enabled by default in the bundled deployment config. For
the sync (non-streaming) path:
```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 --no-async-chunk
```
Use a custom deploy config for advanced cases:
```bash
-vllm-omni serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
+vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 \
--deploy-config /path/to/your_glm_tts_overrides.yaml
```
@@ -103,3 +103,137 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
- Key flags: `--omni` is required; `--trust-remote-code` is needed for the GLM-TTS phoneme tokenizer; the DiT stage enables bucketed model-internal CUDA graphs with eager fallback.
- Voice cloning: requires `ref_audio` + `ref_text` together. Reference audio should be 3-10 seconds. Feature extraction (WhisperVQ tokenizer, CampPlus ONNX, mel) runs on the model side.
- Known limitations: First request may be slow due to lazy model loading (WhisperVQ, CampPlus ONNX). Warm-cache RTF is approximately 0.6-0.7x on A40.
+
+### 2x non-standard RTX 4090 48GB
+
+This configuration was personally validated on two non-standard RTX 4090 cards
+reporting 49,140 MiB each in `nvidia-smi`. Stock 24GB RTX 4090 cards were not
+validated.
+
+#### Environment
+
+- OS: Ubuntu 22.04.1 LTS
+- NVIDIA driver: 580.76.05
+- Python: 3.12.13
+- PyTorch: 2.11.0+cu130
+- CUDA runtime and JIT toolkit: 13.0
+- vLLM: 0.26.0
+- vLLM-Omni: commit `66a9b0c8`, version
+ `0.26.0rc2.dev5+g66a9b0c84`
+- FlashInfer: 0.6.14
+
+The host's default `nvcc` was 11.8. The working CUDA toolkit came from the
+Python environment, so `CUDA_HOME` was set to
+`$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13`.
+
+#### Command
+
+From the repository root, start the async-chunk server with:
+
+```bash
+export CUDA_HOME="$VIRTUAL_ENV/lib/python3.12/site-packages/nvidia/cu13"
+export PATH="$CUDA_HOME/bin:$PATH"
+export LD_LIBRARY_PATH="$CUDA_HOME/lib:${LD_LIBRARY_PATH:-}"
+
+CUDA_VISIBLE_DEVICES=0,1 vllm-omni serve zai-org/GLM-TTS \
+ --omni \
+ --trust-remote-code \
+ --port 8091 \
+ --deploy-config vllm_omni/deploy/glm_tts.yaml \
+ --stage-overrides '{"1":{"devices":"1"}}' \
+ --allowed-local-media-path "${PWD}/tests/assets"
+```
+
+Stage 0 runs on GPU 0, while the override places stage 1 on GPU 1. For the sync
+variant, add `--no-async-chunk` to the same command.
+
+#### Verification
+
+Run the Chinese voice-cloning check:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+ --model zai-org/GLM-TTS \
+ --text "今天天气很好,我们一起学习语音模型推理。" \
+ --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
+ --ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
+ --response-format wav \
+ --output glm_tts_zh.wav
+```
+
+Run the English voice-cloning check with the same Chinese reference:
+
+```bash
+python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
+ --model zai-org/GLM-TTS \
+ --text "Hello, this is a voice cloning test with GLM-TTS." \
+ --ref-audio "file://${PWD}/tests/assets/glm_tts/jiayan_zh.wav" \
+ --ref-text "他当时还跟线下其他的站姐吵架,然后,打架进局子了。" \
+ --response-format wav \
+ --output glm_tts_en.wav
+```
+
+Inspect either output:
+
+```bash
+ffprobe -v error -select_streams a:0 \
+ -show_entries stream=codec_name,sample_rate,channels \
+ -of default=noprint_wrappers=1 glm_tts_zh.wav
+```
+
+Expected output:
+
+```text
+codec_name=pcm_s16le
+sample_rate=24000
+channels=1
+```
+
+The English smoke test produced a valid WAV with an actual duration of 5.020
+seconds.
+
+#### Observed warm-cache results
+
+These observations are from one controlled run, not a hardware or performance
+guarantee. Cold-start and JIT costs were excluded. Values are mean +/- standard
+deviation.
+
+##### Measurement procedure
+
+- For the single-GPU server, use the documented command with
+ `CUDA_VISIBLE_DEVICES=0` and omit `--stage-overrides`; both bundled stages
+ remain on GPU 0. For the dual-GPU server, use the command as shown.
+- Async chunking is the default. Add `--no-async-chunk` for sync mode.
+- For non-streaming measurements, use the exact Chinese client command above.
+ For streaming, add `--stream` to that command and change the output filename
+ to `glm_tts_zh.pcm`; streaming output is raw PCM, not WAV.
+- The fixed Chinese prompt and reference were exactly those shown in
+ Verification and yielded 3.900 seconds of audio. Each case used one warm-up
+ followed by three measured requests at concurrency 1.
+- Wall time starts immediately before the HTTP POST and stops after the complete
+ response body. TTFA runs from immediately before the POST to the first
+ non-empty audio bytes. RTF is wall time divided by audio duration.
+- Peak total device memory used was sampled every 0.2 seconds and reported by
+ NVML via `nvmlDeviceGetMemoryInfo.used`. The readings include background and
+ baseline use and are not baseline-subtracted.
+
+| GPUs | Mode | Response | Wall time (s) | TTFA (s) | RTF |
+|---:|---|---|---:|---:|---:|
+| 1 | Async | Non-streaming | 2.144 +/- 0.076 | -- | 0.550 +/- 0.019 |
+| 2 | Async | Non-streaming | 1.969 +/- 0.130 | -- | 0.505 +/- 0.033 |
+| 1 | Async | Streaming | 2.135 +/- 0.093 | 1.245 +/- 0.178 | 0.548 +/- 0.024 |
+| 2 | Async | Streaming | 1.891 +/- 0.006 | 1.170 +/- 0.042 | 0.485 +/- 0.002 |
+| 1 | Sync | Non-streaming | 1.822 +/- 0.016 | -- | 0.467 +/- 0.004 |
+| 2 | Sync | Non-streaming | 1.792 +/- 0.023 | -- | 0.459 +/- 0.006 |
+
+Moving stage 1 to GPU 1 reduced async non-streaming wall time by 8.2%. It
+reduced async streaming wall time by 11.4% and TTFA by 6.0%, and reduced sync
+wall time by 1.6%. Single- and dual-GPU outputs were bitwise identical only
+within matching execution and response modes.
+
+| Configuration | GPU 0 | GPU 1 |
+|---|---:|---:|
+| Single-GPU async | 31.52 GiB | -- |
+| Dual-GPU async | 29.23 GiB | 2.91 GiB |
+| Single-GPU sync | 31.47 GiB | -- |
+| Dual-GPU sync | 29.23 GiB | 2.91 GiB |
diff --git a/requirements/common.txt b/requirements/common.txt
index e0233bd8d..99d7412d2 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -1,18 +1,20 @@
# Common dependencies for all platforms
-# transformers: kept exactly aligned with the upstream vLLM constraint.
-# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
-# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
-# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
+# transformers: required unified multimodal model classes are unavailable
+# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
+# regression. transformers 5.9 also changed
+# _adjust_missing_and_unexpected_keys from
# list+list to set|set, breaking models that declare
# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
# That breakage is now handled in code instead of by pinning:
# MossTTSNanoForGeneration wraps weight loading in
# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 5.5.3
+transformers >= 5.10.1, < 5.15
av>=14.0.0
omegaconf>=2.3.0
-diffusers==0.38.0
+diffusers==0.40.0
+# Hub kernels, including LTX-2.5 NATTEN
+kernels==0.14.1
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
diff --git a/tests/assets/hidream_o1/hidream_o1_baseline.png b/tests/assets/hidream_o1/hidream_o1_baseline.png
new file mode 100644
index 000000000..38920f301
Binary files /dev/null and b/tests/assets/hidream_o1/hidream_o1_baseline.png differ
diff --git a/tests/assets/soulxsinger/music.json b/tests/assets/soulxsinger/music.json
deleted file mode 100644
index d9e0d8882..000000000
--- a/tests/assets/soulxsinger/music.json
+++ /dev/null
@@ -1,16 +0,0 @@
-[
- {
- "index": "vocal_240_51240",
- "language": "Mandarin",
- "time": [
- 240,
- 51240
- ],
- "duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
- "text": " 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 再 也 没 能 忘 掉 你 容 颜 梦 想 着 着 偶 偶 然 然 有 一 一 天 再 相 见 从 此 我 开 始 始 孤 孤 单 思 念 念 想 想 你 时 你 你 在 天 边 想 你 时 你 在 眼 前 前 想 你 时 你 你 在 脑 海 ",
- "phoneme": " zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 ",
- "note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0",
- "note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1",
- "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 309.0 320.1 329.6 334.4 334.4 334.5 338.2 346.0 355.1 361.4 366.1 369.5 372.7 373.3 372.9 375.2 375.8 374.7 367.6 360.1 0.0 0.0 0.0 0.0 0.0 0.0 283.6 280.4 277.2 274.3 271.8 271.0 273.3 277.2 278.8 278.1 277.2 277.4 278.8 279.3 280.2 281.4 279.2 274.9 268.3 254.1 239.9 0.0 0.0 267.7 267.7 270.1 271.8 273.1 276.3 278.3 269.1 256.1 0.0 0.0 0.0 0.0 286.0 281.1 275.7 272.7 271.7 271.7 272.5 274.0 275.5 277.9 281.1 287.1 314.0 342.2 363.0 376.8 375.4 365.8 351.8 323.3 0.0 0.0 0.0 0.0 299.6 322.7 336.5 336.9 335.4 332.9 330.0 326.6 323.4 322.3 324.4 326.7 328.5 328.5 326.4 324.0 322.5 323.2 325.5 327.7 328.6 328.1 325.4 320.9 316.5 316.1 320.8 327.7 333.7 333.2 325.6 315.5 306.9 305.5 315.1 331.0 343.3 341.6 331.4 319.5 308.3 307.6 318.1 329.5 338.2 342.0 337.4 329.4 321.0 316.6 319.8 330.8 340.6 344.5 342.3 331.4 319.3 315.7 322.9 329.3 336.5 344.0 341.1 328.3 318.7 320.0 328.2 333.5 337.9 339.7 338.5 327.3 323.3 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 285.3 284.6 284.3 288.2 295.1 305.6 312.2 315.5 287.8 0.0 0.0 0.0 0.0 0.0 309.8 312.4 313.8 312.8 310.8 310.7 312.5 313.2 310.7 305.7 305.9 308.1 309.9 311.3 311.0 309.8 309.7 310.6 311.1 311.1 311.0 311.9 313.7 316.4 319.5 311.6 281.3 283.4 306.6 0.0 0.0 0.0 334.4 326.7 322.6 321.6 321.7 324.5 332.0 336.7 333.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 332.5 331.7 332.7 334.8 334.7 335.9 339.7 345.6 353.9 362.3 370.2 374.0 373.8 371.4 370.0 368.2 367.0 368.0 370.4 371.2 371.3 370.0 360.3 0.0 0.0 0.0 319.1 326.8 331.5 332.0 331.6 327.2 328.2 334.1 343.1 350.7 358.0 363.5 368.4 370.8 371.8 371.1 370.2 368.5 368.0 371.1 375.5 376.3 369.5 0.0 0.0 0.0 280.4 267.2 263.6 265.8 271.2 275.0 276.4 276.9 277.1 280.9 287.5 284.7 280.7 0.0 0.0 0.0 0.0 0.0 0.0 231.6 226.0 222.9 224.1 226.2 228.5 234.3 240.9 246.6 249.8 249.6 246.3 245.2 245.8 247.2 248.3 249.3 249.0 245.9 242.7 235.8 226.3 217.8 213.1 209.3 208.0 207.0 205.9 205.8 203.4 201.5 205.3 207.5 209.1 210.4 208.9 203.2 199.2 199.3 201.2 205.6 206.3 204.0 202.7 202.2 203.3 206.1 205.6 201.9 198.8 195.5 195.6 198.7 204.2 214.5 218.3 214.3 207.0 199.1 192.4 189.9 193.6 201.0 210.6 212.4 209.2 202.0 195.2 190.0 190.4 196.0 204.2 209.2 208.9 203.0 195.8 190.7 189.8 194.6 200.8 208.7 213.5 210.5 202.0 193.9 187.7 187.8 193.1 199.8 202.8 204.1 203.8 200.4 197.0 193.9 191.6 192.5 198.4 204.8 203.9 203.5 201.3 200.2 198.4 198.5 201.0 204.0 204.6 205.4 202.0 199.2 194.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 426.6 432.4 436.6 437.2 434.6 433.0 434.0 453.1 473.8 492.8 503.2 501.8 491.3 479.6 478.1 489.7 503.6 508.3 503.2 489.5 482.3 484.2 495.7 505.4 506.6 501.6 495.4 492.7 497.1 500.4 496.1 490.1 480.4 453.8 412.6 373.6 357.7 354.5 356.8 361.0 365.0 367.8 369.5 370.1 371.1 371.3 370.8 370.1 369.7 369.0 369.5 370.4 371.1 370.9 370.2 370.8 371.3 372.7 380.2 386.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.0 427.8 417.0 411.6 412.6 411.2 410.7 412.5 413.7 414.9 414.8 415.7 415.4 414.0 413.8 414.6 415.2 416.3 415.2 413.6 413.4 414.9 416.8 416.4 414.7 413.8 413.2 414.7 417.5 420.7 428.2 443.1 459.8 481.7 502.3 510.8 509.1 497.8 490.6 493.5 500.8 508.3 511.8 504.6 493.5 487.3 492.3 507.2 518.8 514.8 496.4 489.6 490.4 491.8 493.3 0.0 0.0 0.0 0.0 0.0 391.9 368.8 360.2 358.3 360.0 362.0 361.8 360.9 360.3 361.8 365.7 368.7 369.3 367.7 364.7 365.6 368.4 370.4 371.7 369.0 366.8 367.5 367.9 369.4 374.6 384.6 396.8 388.2 382.8 0.0 0.0 0.0 0.0 0.0 352.6 339.5 328.2 326.0 326.5 326.9 328.2 329.4 329.4 329.3 337.3 349.1 350.9 343.6 333.4 326.5 327.1 331.5 333.4 326.8 318.4 0.0 0.0 0.0 282.0 283.6 281.6 279.5 273.2 267.9 268.8 273.2 277.5 279.8 278.7 276.8 273.3 268.7 265.0 266.1 272.4 282.3 284.9 279.4 268.5 255.3 253.0 259.9 270.2 281.5 284.3 279.7 270.1 260.2 256.4 260.1 266.6 273.0 279.6 283.6 282.3 272.3 262.1 255.8 259.8 269.6 274.9 278.7 278.4 271.5 261.7 255.8 256.8 263.9 272.1 279.7 278.4 269.9 259.0 255.9 262.1 269.2 274.5 279.6 282.1 280.7 274.3 269.9 270.9 270.7 273.8 276.9 274.9 271.1 266.5 266.7 269.4 276.9 282.9 282.0 279.5 274.9 271.2 267.9 263.1 270.2 281.3 285.5 283.9 280.6 271.1 262.8 263.9 263.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 362.9 364.0 366.8 366.6 363.6 359.8 357.7 359.9 367.6 371.6 365.4 359.3 358.6 359.5 365.5 375.7 378.1 371.5 359.5 350.8 358.6 374.7 381.7 376.8 363.6 353.3 356.9 367.8 378.3 380.8 375.7 372.3 372.6 374.7 372.9 365.2 346.4 314.8 285.6 277.3 275.4 275.3 277.5 279.2 279.2 278.8 277.6 276.4 276.5 277.6 278.8 279.9 281.8 282.7 281.0 279.2 278.6 279.4 279.9 279.2 278.7 281.7 285.8 289.9 288.7 284.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 428.0 417.9 410.0 409.5 411.2 414.0 417.5 417.9 418.1 419.3 419.0 417.8 417.1 417.2 416.6 415.8 415.0 415.1 414.5 410.9 401.5 388.2 371.5 359.2 363.0 371.7 375.7 370.6 359.6 353.4 358.7 367.0 374.5 378.2 373.7 362.8 358.0 361.6 368.5 373.7 377.7 375.8 368.8 361.7 360.5 366.7 372.5 375.6 380.1 384.9 383.5 376.6 0.0 0.0 0.0 0.0 0.0 331.7 320.0 316.8 320.7 325.1 327.0 327.4 326.6 326.4 325.5 325.1 326.2 327.6 328.5 328.4 327.7 328.7 329.9 329.8 329.7 327.7 326.5 327.8 328.9 329.3 329.8 329.2 328.8 329.4 330.4 331.0 331.0 330.5 329.2 328.2 328.3 328.7 329.4 327.5 322.9 314.8 299.3 290.8 292.3 299.0 305.4 309.2 312.1 313.9 316.8 320.4 327.6 333.1 338.0 341.2 342.1 338.4 331.0 0.0 0.0 0.0 0.0 0.0 0.0 311.7 292.5 281.8 276.1 275.1 276.6 277.6 277.4 274.0 264.8 249.6 239.1 238.8 243.0 245.7 245.0 241.9 237.9 234.6 237.1 242.7 250.4 252.5 248.3 240.7 231.6 228.3 234.0 240.6 247.0 250.6 248.5 242.1 233.6 226.8 227.2 234.0 240.7 246.9 248.9 244.4 237.6 230.8 228.8 236.4 246.9 250.2 249.2 245.4 239.9 232.5 224.7 226.5 238.5 252.0 257.4 255.9 248.4 237.4 230.5 231.6 238.8 247.6 252.6 253.9 253.0 248.9 241.8 236.4 234.1 236.5 246.2 257.7 255.4 248.2 238.5 234.8 237.8 244.4 250.5 253.8 251.7 246.7 240.0 238.3 244.8 251.6 256.4 256.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 484.0 484.5 487.5 490.6 487.2 479.7 478.9 483.8 494.1 505.7 510.3 501.1 485.4 472.4 472.3 495.6 515.8 517.1 503.8 486.5 475.1 479.9 491.4 507.9 509.9 505.0 496.2 485.7 484.5 490.8 495.4 497.7 496.6 489.9 469.1 424.3 387.2 373.0 367.3 364.4 364.6 366.2 366.6 368.8 371.9 373.1 371.4 368.9 367.1 367.0 367.7 368.8 370.1 369.8 369.2 369.8 370.0 369.0 369.0 371.0 375.2 385.3 391.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 425.2 418.5 409.4 409.5 408.1 406.3 408.4 410.3 412.8 413.8 412.6 411.2 410.8 410.4 410.2 410.5 412.2 413.6 413.4 412.6 412.6 413.4 415.6 417.9 418.5 417.0 415.0 415.0 416.8 421.5 437.9 455.0 472.6 493.0 502.2 496.7 484.6 480.6 487.2 498.6 506.3 506.2 499.4 489.7 489.5 496.2 503.8 507.2 501.8 491.1 489.9 501.1 520.5 533.7 525.8 0.0 0.0 0.0 0.0 0.0 421.8 365.2 351.0 352.0 358.0 362.6 365.6 365.4 363.2 362.5 362.2 362.4 364.1 365.2 364.9 363.5 363.4 366.3 368.5 370.0 369.7 368.5 366.9 365.7 366.3 368.6 370.8 372.2 370.2 373.1 377.0 376.5 371.0 346.8 326.1 325.2 328.4 331.4 334.0 332.2 327.2 324.4 331.3 349.7 363.2 362.0 346.1 327.4 321.8 329.8 340.7 348.4 349.9 346.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 269.5 275.7 282.6 282.9 278.8 275.2 274.3 273.5 273.9 275.3 276.4 277.3 276.1 272.3 268.0 268.3 272.5 279.1 281.7 280.1 271.3 261.0 257.3 262.7 273.2 280.8 283.3 279.3 269.1 258.3 258.0 266.3 278.2 286.7 287.8 283.4 274.9 264.7 257.7 260.0 272.0 286.3 294.7 289.8 274.5 263.8 263.4 268.8 277.3 284.1 286.5 285.2 281.7 272.7 264.3 260.3 267.7 281.7 289.2 289.0 281.2 266.2 256.0 254.4 261.5 276.4 286.8 288.8 286.7 273.8 260.9 260.8 270.2 283.4 291.3 292.8 286.1 273.8 265.2 264.0 271.3 283.8 290.3 289.7 277.7 266.9 260.5 263.0 267.8 281.5 286.7 286.1 285.5 280.1 279.0 283.0 284.1 284.5 285.5 282.2 0.0 0.0 280.0 276.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
- }
-]
diff --git a/tests/assets/soulxsinger/music.mp3 b/tests/assets/soulxsinger/music.mp3
deleted file mode 100644
index f3126e79e..000000000
Binary files a/tests/assets/soulxsinger/music.mp3 and /dev/null differ
diff --git a/tests/assets/soulxsinger/zh_prompt.json b/tests/assets/soulxsinger/zh_prompt.json
deleted file mode 100644
index 86ff3614e..000000000
--- a/tests/assets/soulxsinger/zh_prompt.json
+++ /dev/null
@@ -1,16 +0,0 @@
-[
- {
- "index": "vocal_320_10687",
- "language": "Mandarin",
- "time": [
- 320,
- 10687
- ],
- "duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
- "text": " 除 了 想 你 你 除 了 了 爱 你 你 我 什 么 什 么 都 愿 愿 意",
- "phoneme": " zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4",
- "note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0",
- "note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2",
- "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
- }
-]
diff --git a/tests/assets/soulxsinger/zh_prompt.mp3 b/tests/assets/soulxsinger/zh_prompt.mp3
deleted file mode 100644
index bc306410a..000000000
Binary files a/tests/assets/soulxsinger/zh_prompt.mp3 and /dev/null differ
diff --git a/tests/benchmarks/benchmark_video_response_encoding.py b/tests/benchmarks/benchmark_video_response_encoding.py
new file mode 100644
index 000000000..ec589e734
--- /dev/null
+++ b/tests/benchmarks/benchmark_video_response_encoding.py
@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Benchmark legacy and automatic MP4 response encoding on synthetic input."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import statistics
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import cast
+
+import numpy as np
+
+from vllm_omni.entrypoints.openai.video_api_utils import (
+ _encode_video_bytes,
+ _encode_video_bytes_legacy,
+ _PlanarFrameConverter,
+)
+
+
+@dataclass(frozen=True)
+class BenchmarkVariant:
+ label: str
+ encoder: Callable[..., bytes]
+
+
+def _positive_int(value: str) -> int:
+ parsed = int(value)
+ if parsed < 1:
+ raise argparse.ArgumentTypeError("value must be at least 1")
+ return parsed
+
+
+def _build_inputs(
+ *,
+ frames: int,
+ height: int,
+ width: int,
+ fps: int,
+ audio_sample_rate: int,
+ seed: int,
+ input_layout: str = "planar-backed",
+) -> tuple[np.ndarray, np.ndarray]:
+ rng = np.random.default_rng(seed)
+ planar = rng.random((3, frames, height, width), dtype=np.float32)
+ video = planar.transpose(1, 2, 3, 0)
+ if input_layout == "interleaved":
+ video = np.ascontiguousarray(video)
+ if not video.flags.c_contiguous:
+ raise RuntimeError("synthetic video is not C-contiguous")
+ elif input_layout == "planar-backed" and not all(video[..., channel].flags.c_contiguous for channel in range(3)):
+ raise RuntimeError("synthetic video does not expose contiguous channel planes")
+ elif input_layout != "planar-backed":
+ raise ValueError(f"unsupported input layout: {input_layout}")
+
+ sample_count = max(1, round(frames / fps * audio_sample_rate))
+ timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
+ left = np.sin(2 * np.pi * 220 * timeline).astype(np.float32, copy=False)
+ right = np.sin(2 * np.pi * 330 * timeline).astype(np.float32, copy=False)
+ return video, np.stack((left, right))
+
+
+def _measure(
+ variant: BenchmarkVariant,
+ *,
+ video: np.ndarray,
+ audio: np.ndarray,
+ fps: int,
+ audio_sample_rate: int,
+) -> tuple[bytes, dict[str, object]]:
+ cpu_start = time.process_time_ns()
+ wall_start = time.perf_counter_ns()
+ output = variant.encoder(
+ video,
+ fps=fps,
+ audio=audio,
+ audio_sample_rate=audio_sample_rate,
+ video_codec_options={"preset": "ultrafast", "threads": "0"},
+ )
+ wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000
+ process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000
+ return output, {
+ "label": variant.label,
+ "wall_ms": wall_ms,
+ "process_cpu_ms": process_cpu_ms,
+ "output_bytes": len(output),
+ "output_sha256": hashlib.sha256(output).hexdigest(),
+ }
+
+
+def _summarize(records: list[dict[str, object]], label: str) -> dict[str, object]:
+ selected = [record for record in records if record["label"] == label]
+ wall_values = [cast(float, record["wall_ms"]) for record in selected]
+ cpu_values = [cast(float, record["process_cpu_ms"]) for record in selected]
+ return {
+ "runs": len(selected),
+ "wall_ms": {
+ "median": statistics.median(wall_values),
+ "min": min(wall_values),
+ "max": max(wall_values),
+ },
+ "process_cpu_ms": {
+ "median": statistics.median(cpu_values),
+ "min": min(cpu_values),
+ "max": max(cpu_values),
+ },
+ }
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--frames", type=_positive_int, default=12)
+ parser.add_argument("--height", type=_positive_int, default=96)
+ parser.add_argument("--width", type=_positive_int, default=160)
+ parser.add_argument("--fps", type=_positive_int, default=24)
+ parser.add_argument("--audio-sample-rate", type=_positive_int, default=32000)
+ parser.add_argument("--warmup", type=int, default=1)
+ parser.add_argument("--rounds", type=_positive_int, default=3)
+ parser.add_argument("--seed", type=int, default=20260817)
+ parser.add_argument(
+ "--input-layout",
+ choices=("planar-backed", "interleaved"),
+ default="planar-backed",
+ )
+ parser.add_argument(
+ "--frame-conversion-workers",
+ type=_positive_int,
+ default=8,
+ )
+ parser.add_argument("--output-json", type=Path)
+ args = parser.parse_args()
+ if args.warmup < 0:
+ parser.error("--warmup must be non-negative")
+ return args
+
+
+def main() -> None:
+ args = _parse_args()
+ video, audio = _build_inputs(
+ frames=args.frames,
+ height=args.height,
+ width=args.width,
+ fps=args.fps,
+ audio_sample_rate=args.audio_sample_rate,
+ seed=args.seed,
+ input_layout=args.input_layout,
+ )
+ baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
+ converter = _PlanarFrameConverter(max_workers=args.frame_conversion_workers)
+ candidate = BenchmarkVariant(
+ "automatic",
+ partial(_encode_video_bytes, frame_converter=converter),
+ )
+
+ try:
+ for _ in range(args.warmup):
+ for variant in (baseline, candidate):
+ _measure(
+ variant,
+ video=video,
+ audio=audio,
+ fps=args.fps,
+ audio_sample_rate=args.audio_sample_rate,
+ )
+
+ records: list[dict[str, object]] = []
+ for round_index in range(args.rounds):
+ order = (baseline, candidate) if round_index % 2 == 0 else (candidate, baseline)
+ for variant in order:
+ _, record = _measure(
+ variant,
+ video=video,
+ audio=audio,
+ fps=args.fps,
+ audio_sample_rate=args.audio_sample_rate,
+ )
+ record["round"] = round_index + 1
+ records.append(record)
+ finally:
+ converter.shutdown()
+
+ output_hashes = {str(record["output_sha256"]) for record in records}
+ if len(output_hashes) != 1:
+ raise RuntimeError(f"benchmark variants produced different outputs: {sorted(output_hashes)}")
+
+ paired_savings_ms = []
+ paired_savings_percent = []
+ for round_index in range(args.rounds):
+ pair = [record for record in records if record["round"] == round_index + 1]
+ baseline_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
+ candidate_ms = cast(float, next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
+ paired_savings_ms.append(baseline_ms - candidate_ms)
+ paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
+
+ result = {
+ "config": {
+ "frames": args.frames,
+ "height": args.height,
+ "width": args.width,
+ "fps": args.fps,
+ "audio_sample_rate": args.audio_sample_rate,
+ "warmup": args.warmup,
+ "rounds": args.rounds,
+ "seed": args.seed,
+ "input_layout": args.input_layout,
+ "frame_conversion_workers": args.frame_conversion_workers,
+ "video_shape": list(video.shape),
+ "video_strides": list(video.strides),
+ "codec_options": {"preset": "ultrafast", "threads": "0"},
+ },
+ "records": records,
+ "summary": {
+ baseline.label: _summarize(records, baseline.label),
+ candidate.label: _summarize(records, candidate.label),
+ "paired_wall_saving_ms": {
+ "median": statistics.median(paired_savings_ms),
+ "min": min(paired_savings_ms),
+ "max": max(paired_savings_ms),
+ },
+ "paired_wall_saving_percent": {
+ "median": statistics.median(paired_savings_percent),
+ "min": min(paired_savings_percent),
+ "max": max(paired_savings_percent),
+ },
+ "output_sha256": output_hashes.pop(),
+ },
+ }
+ output_json = json.dumps(result, indent=2, sort_keys=True)
+ if args.output_json is not None:
+ args.output_json.write_text(output_json + "\n", encoding="utf-8")
+ print(output_json)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/benchmarks/duplex/test_omni_duplex_eval.py b/tests/benchmarks/duplex/test_omni_duplex_eval.py
new file mode 100644
index 000000000..5f77aacc0
--- /dev/null
+++ b/tests/benchmarks/duplex/test_omni_duplex_eval.py
@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
+ extract_timed_sentences,
+ normalize_response_items,
+ split_text,
+ validate_clock,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
+ DuplexSample,
+ canonical_task_type,
+ family_for_split,
+ load_samples,
+ task_type_for_split,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
+ PROTOCOL_PIN,
+ build_content_prompt,
+ build_reminder_prompt,
+ build_temporal_prompt,
+ parse_judge_json,
+ reminder_window,
+ summarize_pr_results,
+ temporal_window,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dataset_mapping_and_manifest(tmp_path):
+ manifest = tmp_path / "samples.json"
+ manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
+ sample = load_samples(manifest, media_root=tmp_path)[0]
+ assert sample.family == "rtd"
+ assert sample.video == str(tmp_path / "clip.mp4")
+ assert family_for_split("PR_event_reminder") == "pr"
+ assert canonical_task_type("post-event-reminder") == "post_event_reminder"
+ assert task_type_for_split("PR_correction") == "correction"
+ assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
+
+
+def test_hugging_face_subset_names_are_splits(monkeypatch):
+ calls = []
+
+ def fake_load_dataset(name, **kwargs):
+ calls.append((name, kwargs))
+ return [{"id": "pr", "question_text": "Correct this."}]
+
+ monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+ sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
+ assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
+ assert sample.split == "PR_correction"
+ assert sample.task_type == "correction"
+
+
+def test_response_aliases_and_clock_guard():
+ assert split_text("One. Two!") == ["One.", "Two!"]
+ assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
+ {"sentence": "x", "start": 800.0, "end": 800.0}
+ ]
+ with pytest.raises(ValueError, match="clock=invalid"):
+ validate_clock({"clock": "invalid"})
+ validate_clock({"clock": "invalid"}, allow_invalid=True)
+ timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
+ assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
+ with pytest.raises(ValueError, match="only clock=media"):
+ extract_timed_sentences([], clock="wall")
+
+
+def test_protocol_windows_and_parsing():
+ assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
+ assert temporal_window(4, 5, 10) == (2.0, 3.0)
+ assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
+ assert reminder_window(5) == (5.0, 15.0)
+ assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
+ "content_score": 2.35,
+ "is_relevant": 1,
+ }
+ assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
+ assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
+
+
+def test_protocol_prompt_wording_is_pinned():
+ prompts = [
+ build_temporal_prompt(1, 2, "response", "question"),
+ build_content_prompt("response", "question", ["one", "two"]),
+ build_reminder_prompt("instruction", "response", "proactive_reminder"),
+ build_reminder_prompt("instruction", "response", "post_event_reminder"),
+ build_reminder_prompt("instruction", "response", "correction", "answer"),
+ ]
+ assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
+ "5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
+ "1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
+ "d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
+ "08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
+ "3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
+ ]
+
+
+def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
+ response = tmp_path / "response.json"
+ response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
+ response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
+ video = tmp_path / "video.mp4"
+ video.write_bytes(b"video")
+ monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
+
+ class Judge:
+ def temporal(self, *args, **kwargs):
+ return '{"temporal_score": 3, "is_relevant": 1}'
+
+ def content(self, prompt, video, frames=None, *, mode="video_url"):
+ assert mode == "frame-sample"
+ assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
+ return '{"content_score": 3}'
+
+ sample = DuplexSample(
+ id="rtd",
+ split="RTD_OCR",
+ family="rtd",
+ task_type=None,
+ video=video,
+ video_duration=4,
+ )
+ score = eval_module.evaluate_sample(
+ sample,
+ response,
+ tmp_path / "score.json",
+ Judge(),
+ judge_video_mode="frame-sample",
+ )
+ assert score["content"]["frame_count"] == 2
diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py
index 32d9d7958..315685258 100644
--- a/tests/benchmarks/metrics/test_metrics.py
+++ b/tests/benchmarks/metrics/test_metrics.py
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for metrics.py
"""
+import math
+
import pytest
from vllm.benchmarks.serve import TaskType
@@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
return output
+def _calculate_test_metrics(outputs, goodput=None):
+ return calculate_metrics(
+ input_requests=[],
+ outputs=outputs,
+ dur_s=1.0,
+ tokenizer=None,
+ selected_percentiles=[50.0],
+ goodput_config_dict=goodput or {},
+ task_type=TaskType.GENERATION,
+ selected_percentile_metrics=[],
+ max_concurrency=None,
+ request_rate=float("inf"),
+ benchmark_duration=1.0,
+ )[0]
+
+
# ============================================================================
# total_input Tests
# ============================================================================
@@ -126,6 +144,228 @@ def test_audio_continuity_aggregation():
assert p99 is not None and p99 > 0.4
+def test_unmeasured_duplex_latency_does_not_add_zero_samples():
+ measured = _make_output(100, output_tokens=2)
+ measured.ttft = 0.1
+ measured.audio_ttfp = 0.2
+ measured.audio_rtf = 0.5
+ measured.duplex_session_metrics = {
+ "mean_ttft_ms": 100.0,
+ "mean_ttfp_ms": 200.0,
+ "mean_rtf": 0.5,
+ }
+ listen_only = _make_output(100, output_tokens=0)
+ listen_only.ttft = listen_only.audio_ttfp = listen_only.audio_rtf = 0.0
+ listen_only.duplex_session_metrics = {
+ "mean_ttft_ms": None,
+ "mean_ttfp_ms": None,
+ "mean_rtf": None,
+ }
+
+ metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0})
+
+ assert metrics.mean_ttft_ms == 100.0
+ assert metrics.mean_audio_ttfp_ms == 200.0
+ assert metrics.mean_audio_rtf == 0.5
+ assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1)
+ assert metrics.request_goodput == 1.0
+
+
+def test_all_unmeasured_duplex_latency_is_not_reported_as_zero():
+ listen_only = _make_output(100, output_tokens=0)
+ listen_only.duplex_session_metrics = {
+ "mean_ttft_ms": None,
+ "mean_ttfp_ms": None,
+ "mean_rtf": None,
+ }
+
+ metrics = _calculate_test_metrics(
+ [listen_only],
+ {"ttft": float("inf"), "audio_ttft": float("inf")},
+ )
+
+ assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0)
+ assert math.isnan(metrics.mean_ttft_ms)
+ assert math.isnan(metrics.mean_audio_ttfp_ms)
+ assert math.isnan(metrics.mean_audio_rtf)
+ assert metrics.request_goodput == 0.0
+
+
+def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput():
+ missing_tpot = _make_output(100, output_tokens=5)
+ missing_tpot.itl = []
+ missing_tpot.text_latency = missing_tpot.ttft = 0.1
+ missing_tpot.tpot_measured = False
+ missing_tpot.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+ slow_ttft = _make_output(100, output_tokens=5)
+ slow_ttft.ttft = 1.0
+ slow_ttft.text_latency = 1.4
+ slow_ttft.itl = [0.1] * 4
+ slow_ttft.duplex_session_metrics = {"mean_ttft_ms": 1000.0}
+
+ metrics = _calculate_test_metrics(
+ [missing_tpot, slow_ttft],
+ {"ttft": 500.0, "tpot": 200.0},
+ )
+
+ assert metrics.num_tpot_samples == 1
+ assert metrics.mean_tpot_ms == pytest.approx(100.0)
+ assert metrics.request_goodput == 0.0
+
+
+def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero():
+ output = _make_output(100, output_tokens=5)
+ output.itl = []
+ output.text_latency = output.ttft
+ output.tpot_measured = False
+ output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+ metrics = _calculate_test_metrics([output], {"tpot": float("inf")})
+
+ assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0)
+ assert math.isnan(metrics.mean_tpot_ms)
+ assert math.isnan(metrics.mean_itl_ms)
+ assert metrics.request_goodput == 0.0
+
+
+def test_duplex_response_timings_do_not_build_a_session_token_timeline():
+ output = _make_output(100, output_tokens=5)
+ output.latency = 101.0
+ output.duplex_request_metrics = [
+ {"response_id": "r1", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
+ {"response_id": "r2", "stage0_tokens": {"itls_ms": [100.0, 100.0]}},
+ ]
+ output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+ metrics = _calculate_test_metrics([output])
+
+ assert math.isnan(metrics.max_output_tokens_per_s)
+ assert metrics.max_concurrent_requests == 1
+ assert metrics.mean_tpot_ms == pytest.approx(100.0)
+ assert metrics.mean_itl_ms == pytest.approx(100.0)
+
+
+def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
+ output = _make_output(100, output_tokens=0)
+ output.generated_text = "timing metadata missing"
+ output.itl = []
+ output.text_latency = output.ttft
+ output.tpot_measured = False
+ output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+ def tokenizer(text, *, add_special_tokens):
+ assert text == output.generated_text
+ assert add_special_tokens is False
+ return type("Tokenized", (), {"input_ids": [1, 2, 3]})()
+
+ metrics, _ = calculate_metrics(
+ input_requests=[],
+ outputs=[output],
+ dur_s=1.0,
+ tokenizer=tokenizer,
+ selected_percentiles=[50.0],
+ goodput_config_dict={"tpot": float("inf")},
+ task_type=TaskType.GENERATION,
+ selected_percentile_metrics=[],
+ max_concurrency=None,
+ request_rate=float("inf"),
+ benchmark_duration=1.0,
+ )
+
+ assert metrics.total_output == 3
+ assert metrics.num_tpot_samples == 0
+ assert math.isnan(metrics.mean_tpot_ms)
+ assert metrics.request_goodput == 0.0
+
+
+def test_zero_itl_does_not_create_zero_tpot():
+ output = _make_output(100, output_tokens=3)
+ output.itl = [0.0, 0.0]
+ output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
+
+ metrics = _calculate_test_metrics([output], {"tpot": 1.0})
+
+ assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 2)
+ assert math.isnan(metrics.mean_tpot_ms)
+ assert metrics.mean_itl_ms == 0.0
+ assert metrics.request_goodput == 0.0
+
+
+def test_stage_output_tokens_without_client_timing_omit_tpot(capsys):
+ output = _make_output(100, output_tokens=5)
+ output.itl = []
+ output.text_latency = 0.0
+ output.ttft = 0.1
+
+ metrics, _ = calculate_metrics(
+ input_requests=[],
+ outputs=[output],
+ dur_s=1.0,
+ tokenizer=None,
+ selected_percentiles=[50.0, 99.0],
+ goodput_config_dict={},
+ task_type=TaskType.GENERATION,
+ selected_percentile_metrics=["tpot"],
+ max_concurrency=None,
+ request_rate=float("inf"),
+ benchmark_duration=1.0,
+ )
+
+ printed = capsys.readouterr().out
+ assert metrics.num_tpot_samples == 0
+ assert math.isnan(metrics.mean_tpot_ms)
+ assert "Time per Output Token" not in printed
+ assert "Mean TPOT" not in printed
+
+
+def test_consistent_client_latency_can_supply_tpot_fallback():
+ output = _make_output(100, output_tokens=10)
+ output.itl = []
+ output.ttft = 0.1
+ output.text_latency = 1.0
+
+ metrics = _calculate_test_metrics([output])
+
+ assert metrics.num_tpot_samples == 1
+ assert metrics.mean_tpot_ms == pytest.approx(100.0)
+
+
+def test_single_token_responses_do_not_report_zero_tpot(capsys):
+ output = _make_output(100, output_tokens=1)
+
+ metrics, _ = calculate_metrics(
+ input_requests=[],
+ outputs=[output],
+ dur_s=1.0,
+ tokenizer=None,
+ selected_percentiles=[50.0],
+ goodput_config_dict={},
+ task_type=TaskType.GENERATION,
+ selected_percentile_metrics=["tpot"],
+ max_concurrency=None,
+ request_rate=float("inf"),
+ benchmark_duration=1.0,
+ )
+
+ printed = capsys.readouterr().out
+ assert metrics.num_tpot_samples == 0
+ assert math.isnan(metrics.mean_tpot_ms)
+ assert "Time per Output Token" not in printed
+
+
+def test_duplex_goodput_does_not_pair_measurements_from_different_requests():
+ text_only, audio_only = _make_output(100), _make_output(100)
+ text_only.ttft, text_only.audio_ttfp = 0.1, 0.0
+ text_only.duplex_session_metrics = {"mean_ttft_ms": 100.0, "mean_ttfp_ms": None, "mean_rtf": None}
+ audio_only.ttft, audio_only.audio_ttfp = 0.0, 0.2
+ audio_only.duplex_session_metrics = {"mean_ttft_ms": None, "mean_ttfp_ms": 200.0, "mean_rtf": None}
+
+ metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0})
+
+ assert metrics.request_goodput == 0.0
+
+
# ============================================================================
# TTFT suppression for pure-audio (TTS) benchmarks
# ============================================================================
diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py
index e0a9ff49b..d5c714be8 100644
--- a/tests/benchmarks/patch/test_patch.py
+++ b/tests/benchmarks/patch/test_patch.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for patch.py
@@ -9,22 +9,98 @@ import asyncio
import base64
import json
import time
+from argparse import Namespace
from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
+from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
+ SeedTTSSampleRequest,
+ SeedTTSTextSampleRequest,
+)
from vllm_omni.benchmarks.patch.patch import (
MixRequestFuncOutput,
+ _apply_stage0_token_timings,
+ _attach_seed_tts_to_request_func_input,
async_request_openai_chat_omni_completions,
async_request_openai_realtime_duplex,
+ should_request_stage_metrics,
)
from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+def _seed_tts_request_func_input() -> RequestFuncInput:
+ return RequestFuncInput(
+ model="test-model",
+ model_name="test-model",
+ prompt="target text",
+ api_url="http://test.com/v1/chat/completions",
+ prompt_len=2,
+ output_len=20,
+ extra_body={"modalities": ["text", "audio"]},
+ )
+
+
+def test_seed_tts_chat_request_carries_reference_audio_once() -> None:
+ reference_audio_url = "data:audio/wav;base64,AAAA"
+ sample = SeedTTSSampleRequest(
+ prompt="target text",
+ prompt_len=2,
+ expected_output_len=20,
+ multi_modal_data=None,
+ seed_tts_system_prompt="Clone the supplied voice.",
+ seed_tts_speech_extra={
+ "ref_audio": reference_audio_url,
+ "ref_text": "reference text",
+ "task_type": "Base",
+ },
+ )
+ request_func_input = _seed_tts_request_func_input()
+
+ _attach_seed_tts_to_request_func_input(sample, request_func_input)
+
+ assert request_func_input.omni_chat_messages == [
+ {
+ "role": "system",
+ "content": [{"type": "text", "text": "Clone the supplied voice."}],
+ },
+ {"role": "user", "content": [{"type": "text", "text": "target text"}]},
+ ]
+ assert request_func_input.extra_body["ref_audio"] == reference_audio_url
+ serialized_request = json.dumps(
+ {"messages": request_func_input.omni_chat_messages, **request_func_input.extra_body},
+ )
+ assert serialized_request.count(reference_audio_url) == 1
+
+
+def test_seed_tts_text_chat_messages_do_not_add_reference_audio() -> None:
+ sample = SeedTTSTextSampleRequest(
+ prompt="target text",
+ prompt_len=2,
+ expected_output_len=20,
+ multi_modal_data=None,
+ seed_tts_system_prompt="Use the configured voice.",
+ seed_tts_speech_extra=None,
+ )
+ request_func_input = _seed_tts_request_func_input()
+
+ _attach_seed_tts_to_request_func_input(sample, request_func_input)
+
+ assert getattr(request_func_input, "seed_tts_row", False) is True
+ assert request_func_input.omni_chat_messages == [
+ {
+ "role": "system",
+ "content": [{"type": "text", "text": "Use the configured voice."}],
+ },
+ {"role": "user", "content": [{"type": "text", "text": "target text"}]},
+ ]
+ assert "ref_audio" not in request_func_input.extra_body
+
+
class MockResponse:
"""Mock aiohttp response for testing"""
@@ -102,7 +178,24 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
received_at_s=now + 0.03,
)
self.events.add(
- {"type": "response.done", "response": {"id": response_id}},
+ {
+ "type": "response.done",
+ "response": {
+ "id": response_id,
+ "metadata": {
+ "vllm_omni": {
+ "stage_metrics": {
+ "0": {
+ "num_tokens_out": 3,
+ "vllm_ttft_ms": 20.0,
+ "vllm_tpot_ms": 10.0,
+ "vllm_itls_ms": [10.0, 10.0],
+ }
+ }
+ }
+ },
+ },
+ },
received_at_s=now + 0.04,
)
@@ -143,6 +236,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
assert client.configure_kwargs["native_duplex"] is False
assert client.configure_kwargs["extra_body"] == {
"ref_audio": "data:audio/wav;base64,AAAA",
+ "return_stage_metrics": True,
}
assert client.sent == [
event
@@ -167,6 +261,10 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
assert output.audio_ttfp > output.ttft
assert output.audio_rtf > 0
assert output.latency > 0
+ assert output.output_tokens == 12
+ assert output.itl == [0.01, 0.01] * 4
+ assert output.text_latency == pytest.approx(output.ttft + 0.08)
+ assert output.tpot_measured is True
assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4
assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600
session_id = output.duplex_request_metrics[0]["session_id"]
@@ -198,11 +296,160 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
}
+def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete():
+ output = MixRequestFuncOutput(ttft=0.1)
+
+ measured = _apply_stage0_token_timings(
+ output,
+ [
+ {"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0},
+ {"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0},
+ ],
+ expected_output_tokens=5,
+ )
+
+ assert measured is True
+ assert output.itl == []
+ assert output.tpot_measured is True
+ assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02)
+
+
+def test_stage0_zero_itls_fall_back_to_positive_tpot():
+ output = MixRequestFuncOutput(ttft=0.1)
+
+ measured = _apply_stage0_token_timings(
+ output,
+ [{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}],
+ expected_output_tokens=3,
+ )
+
+ assert measured is True
+ assert output.itl == []
+ assert output.tpot_measured is True
+ assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025)
+
+
+def test_stage0_token_count_without_timing_is_not_measured():
+ output = MixRequestFuncOutput(ttft=0.1)
+
+ measured = _apply_stage0_token_timings(
+ output,
+ [{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}],
+ expected_output_tokens=3,
+ )
+
+ assert measured is False
+ assert output.itl == []
+ assert output.text_latency == output.ttft
+ assert output.tpot_measured is False
+
+
def create_sse_chunk(data_dict):
"""Helper to create SSE formatted chunk"""
return f"data: {json.dumps(data_dict)}\n\n".encode()
+def test_chat_text_timing_metrics_request_stage_metrics():
+ args = Namespace(
+ backend="openai-chat-omni",
+ percentile_metrics="ttft,tpot,itl,e2el",
+ print_stage=False,
+ extra_body={},
+ )
+
+ assert should_request_stage_metrics(args) is True
+
+
+@pytest.mark.asyncio
+async def test_bundled_first_text_chunk_uses_stage0_token_timings(mocker: MockerFixture):
+ """Engine timings recover TPOT when every text token arrives together."""
+ request_input = RequestFuncInput(
+ model="test-model",
+ model_name="test-model",
+ prompt="test prompt",
+ api_url="http://test.com/v1/chat/completions",
+ prompt_len=10,
+ output_len=20,
+ )
+ chunks = [
+ create_sse_chunk(
+ {
+ "choices": [{"delta": {"content": "ABCD"}}],
+ "modality": "text",
+ "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
+ "metrics": {
+ "stage_metrics": {
+ "0": {
+ "num_tokens_out": 4,
+ "vllm_itls_ms": [10.0, 11.0, 12.0],
+ "vllm_tpot_ms": 11.0,
+ }
+ }
+ },
+ }
+ ),
+ b"data: [DONE]\n\n",
+ ]
+ mock_response = MockResponse(200, chunks)
+ mock_session = mocker.AsyncMock()
+ mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+ output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+ assert output.success is True
+ assert output.output_tokens == 4
+ assert output.itl == pytest.approx([0.010, 0.011, 0.012])
+ assert output.text_latency - output.ttft == pytest.approx(0.033)
+ assert output.tpot_measured is True
+
+
+@pytest.mark.asyncio
+async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
+ request_input = RequestFuncInput(
+ model="test-model",
+ model_name="test-model",
+ prompt="test prompt",
+ api_url="http://test.com/v1/chat/completions",
+ prompt_len=10,
+ output_len=20,
+ )
+ chunks = [
+ create_sse_chunk(
+ {
+ "choices": [{"delta": {"content": "A"}}],
+ "modality": "text",
+ "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
+ }
+ ),
+ create_sse_chunk(
+ {
+ "choices": [{"delta": {"content": "B"}}],
+ "modality": "text",
+ "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
+ "metrics": {
+ "stage_metrics": {
+ "0": {
+ "num_tokens_out": 2,
+ "vllm_itls_ms": [1000.0],
+ "vllm_tpot_ms": 1000.0,
+ }
+ }
+ },
+ }
+ ),
+ b"data: [DONE]\n\n",
+ ]
+ mock_response = MockResponse(200, chunks, delay_between_chunks=0.01)
+ mock_session = mocker.AsyncMock()
+ mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+ output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+ assert output.success is True
+ assert len(output.itl) == 1
+ assert 0.0 < output.itl[0] < 0.1
+
+
# ============================================================================
# output_tokens Tests
# ============================================================================
diff --git a/tests/benchmarks/test_daily_omni_local_source.py b/tests/benchmarks/test_daily_omni_local_source.py
index 35045a04b..13ae0637a 100644
--- a/tests/benchmarks/test_daily_omni_local_source.py
+++ b/tests/benchmarks/test_daily_omni_local_source.py
@@ -9,13 +9,17 @@ from pathlib import Path
import pytest
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
+ DailyOmniDataset,
daily_omni_local_qa_json,
daily_omni_local_videos_dir,
+ ensure_daily_omni_hub_root,
resolve_daily_omni_local_root,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+_HUB_ID = "liarliar/Daily-Omni"
+
def _write_qa_json(root: Path) -> Path:
qa = root / "qa.json"
@@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path:
return qa
+def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
+ """Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
+
+ Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
+ only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
+ """
+ constants = pytest.importorskip("huggingface_hub").constants
+
+ cache = tmp_path / "hub"
+ repo_dir = cache / "datasets--liarliar--Daily-Omni"
+ snapshot = repo_dir / "snapshots" / "deadbeef"
+ snapshot.mkdir(parents=True)
+ if with_qa:
+ _write_qa_json(snapshot)
+ (repo_dir / "refs").mkdir()
+ (repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
+
+ monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
+ monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+ return snapshot
+
+
def test_resolve_local_root_returns_none_for_hub_id() -> None:
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
assert resolve_daily_omni_local_root("") is None
@@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
root = resolve_daily_omni_local_root(str(tmp_path))
assert daily_omni_local_qa_json(root) is None
assert daily_omni_local_videos_dir(root) is None
+
+
+def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+ snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+
+ root = ensure_daily_omni_hub_root(_HUB_ID)
+
+ assert root == snapshot.resolve()
+ assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
+
+
+def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+ """The regression this file exists for: a Hub id must not require ``load_dataset``."""
+ snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
+ # Any use of `datasets` here means we went back to the path that cannot read the hub cache.
+ monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+ dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
+
+ assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
+ assert [row["video_id"] for row in dataset.data] == ["vid0"]
+
+
+def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+ _stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
+ monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
+
+ with pytest.raises(RuntimeError) as excinfo:
+ DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
+
+ message = str(excinfo.value)
+ assert "someone-else/Daily-Omni" in message
+ assert "snapshot_download failed" in message
+ assert "--daily-omni-qa-json" in message
diff --git a/tests/benchmarks/test_daily_omni_pack_modes.py b/tests/benchmarks/test_daily_omni_pack_modes.py
index 0fd231857..34d2a5637 100644
--- a/tests/benchmarks/test_daily_omni_pack_modes.py
+++ b/tests/benchmarks/test_daily_omni_pack_modes.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
@@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
if _MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+ if _spec is None or _spec.loader is None:
+ raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}")
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod)
@@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
def fake_extract(video_path, *, audio_path, include_audio, **_kw):
counter.append(1)
frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
- segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+ segments: list[np.ndarray] = (
+ [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+ )
return frames, segments
monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py
new file mode 100644
index 000000000..c38e001bf
--- /dev/null
+++ b/tests/benchmarks/test_omniinteract.py
@@ -0,0 +1,782 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import io
+import json
+import subprocess
+import tarfile
+import threading
+import time
+import wave
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
+from vllm.benchmarks.serve import TaskType
+
+from vllm_omni.benchmarks import omniinteract as oi
+from vllm_omni.benchmarks import serve as benchmark_serve
+from vllm_omni.benchmarks.data_modules import omniinteract_dataset as data
+from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
+from vllm_omni.benchmarks.patch import patch as benchmark_patch
+from vllm_omni.entrypoints.cli.benchmark.cli_args import preprocess_serve_args
+from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
+from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
+from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+MODEL = "openbmb/MiniCPM-o-4_5"
+
+
+def _case(tmp_path: Path, *, subset: str = "1q1a", name: str = "video.mp4") -> data.OmniInteractCase:
+ video, annotation = tmp_path / name, tmp_path / f"{Path(name).stem}.json"
+ video.parent.mkdir(parents=True, exist_ok=True)
+ video.touch()
+ annotation.write_text("{}")
+ return data.OmniInteractCase(subset, name, video, annotation, "multi_turn")
+
+
+def _collector(*events: tuple[dict[str, object], float]) -> RealtimeEventCollector:
+ collector = RealtimeEventCollector()
+ for event, received_at in events:
+ collector.add(event, received_at_s=received_at)
+ return collector
+
+
+def _created(response_id: str = "r1") -> dict[str, object]:
+ return {"type": "response.created", "response": {"id": response_id}}
+
+
+def _done(response_id: str = "r1", status: str = "completed") -> dict[str, object]:
+ return {"type": "response.done", "response": {"id": response_id, "status": status}}
+
+
+def _audio(
+ response_id: str = "r1",
+ *,
+ samples: int = 2400,
+ rate: int | None = 24_000,
+ value: int = 1,
+) -> dict[str, object]:
+ event: dict[str, object] = {
+ "type": "response.audio.delta",
+ "response_id": response_id,
+ "format": "pcm16",
+ "delta": base64.b64encode(bytes((value, 0)) * samples).decode(),
+ }
+ if rate is not None:
+ event["sample_rate_hz"] = rate
+ return event
+
+
+def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]:
+ return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value}
+
+
+def _listen(*, buffering: bool = False) -> dict[str, object]:
+ return {
+ "type": "response.listen",
+ "response": {"metadata": {"model_listen": True, "buffering": buffering}},
+ }
+
+
+def _session_options(tmp_path: Path) -> data.OmniInteractSessionOptions:
+ return data.OmniInteractSessionOptions(tmp_path / "artifacts", 10.0, 5.0, "reference.wav", True)
+
+
+def _sample(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions):
+ return data.OmniInteractSampleRequest(
+ prompt="",
+ prompt_len=0,
+ expected_output_len=0,
+ multi_modal_data=None,
+ request_id="measured-0",
+ omniinteract_case=case,
+ omniinteract_options=options,
+ )
+
+
+def _write_dataset(root: Path) -> Path:
+ data_root = root / "data"
+ for subset in ("1q1a", "1q1a_math"):
+ subset_root = data_root / subset
+ (subset_root / "videos").mkdir(parents=True)
+ (subset_root / "annotations").mkdir()
+ (subset_root / "videos" / f"{subset}.mp4").touch()
+ (subset_root / "annotations" / f"{subset}.json").write_text("{}")
+ (subset_root / "video_json_map.json").write_text(
+ json.dumps(
+ {
+ "entries": [
+ {
+ "video": f"videos/{subset}.mp4",
+ "annotation": f"annotations/{subset}.json",
+ "scene_type": "multi_turn",
+ }
+ ]
+ }
+ )
+ )
+ one_to_many = data_root / "1qna"
+ (one_to_many / "videos_bench" / "nested").mkdir(parents=True)
+ (one_to_many / "annotations" / "nested").mkdir(parents=True)
+ (one_to_many / "videos_bench" / "nested" / "guide.mp4").touch()
+ (one_to_many / "annotations" / "nested" / "guide.json").write_text("{}")
+ return data_root
+
+
+def _successful_output(
+ tmp_path: Path,
+ *,
+ audio_time: float = 10.0,
+ status: str = "completed",
+) -> tuple[data.OmniInteractCase, RealtimeEventCollector, oi.OmniInteractCaseResult]:
+ case = _case(tmp_path)
+ collector = _collector(
+ (_created(), 9.9),
+ (_audio(), audio_time),
+ (_text(), audio_time + 0.01),
+ (_done(status=status), audio_time + 0.1),
+ )
+ result = oi.OmniInteractCaseResult(case.subset, str(case.video_path), "")
+ return case, collector, result
+
+
+def _write_success(
+ root: Path,
+ case: data.OmniInteractCase,
+ collector: RealtimeEventCollector,
+ result: oi.OmniInteractCaseResult,
+ **kwargs,
+):
+ require_response = kwargs.pop("require_response", True)
+ return oi.prepare_success_artifacts(
+ root,
+ case,
+ collector,
+ stream_start=10.0,
+ video_duration_s=1.0,
+ require_response=require_response,
+ result=result,
+ **kwargs,
+ )
+
+
+def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
+ root = _write_dataset(tmp_path)
+ all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
+ selected = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=2, disable_shuffle=True)
+ assert [case.subset for case in all_cases] == ["1q1a", "1q1a_math", "1qna"]
+ assert all_cases[-1].video_rel == "videos_bench/nested/guide.mp4"
+ assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
+
+
+def _archive(path: Path, member: str) -> None:
+ with tarfile.open(path, "w") as archive:
+ content = b"{}"
+ item = tarfile.TarInfo(member)
+ item.size = len(content)
+ archive.addfile(item, io.BytesIO(content))
+
+
+def test_archive_is_safe_and_atomically_shared(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ bad = tmp_path / "bad.tar"
+ _archive(bad, "../escape")
+ with pytest.raises(ValueError, match="Unsafe path"):
+ data._extract_archive(bad, tmp_path / "bad-cache")
+
+ good = tmp_path / "good.tar"
+ _archive(good, "data/1q1a/video_json_map.json")
+ target = tmp_path / "cache"
+ barrier, safe_extract = threading.Barrier(2), data._safe_extract
+ monkeypatch.setattr(data, "_safe_extract", lambda *args: (barrier.wait(), safe_extract(*args)))
+ with ThreadPoolExecutor(2) as executor:
+ roots = list(executor.map(lambda _: data._extract_archive(good, target), range(2)))
+ assert roots[0] == roots[1]
+ assert (roots[0] / "1q1a" / "video_json_map.json").is_file()
+ assert not list(target.glob(".tmp-*"))
+
+
+def test_hub_archive_uses_vllm_filesystem(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ source = tmp_path / "source.tar"
+ _archive(source, "data/1q1a/video_json_map.json")
+
+ class FS:
+ def get_file(self, remote: str, local: str) -> None:
+ assert remote == "datasets/org/repo/data.tar.gz"
+ Path(local).write_bytes(source.read_bytes())
+
+ monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
+ monkeypatch.setattr(data, "hf_fs", lambda: FS())
+ assert (data.resolve_omniinteract_root(None, "org/repo") / "1q1a").is_dir()
+
+
+def test_media_commands_are_bounded(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+ calls: list[float] = []
+ commands: list[list[str]] = []
+ bounded_run = oi._run_media_command
+
+ def run(command, *, timeout_s, text=False):
+ calls.append(timeout_s)
+ commands.append(command)
+ if command[0] == "ffprobe":
+ return subprocess.CompletedProcess(command, 0, "1.0", "")
+ return subprocess.CompletedProcess(command, 0, b"", b"")
+
+ monkeypatch.setattr(oi, "_run_media_command", run)
+ duration, pcm, frames = oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=1.0)
+ assert (duration, len(pcm), frames, calls) == (1.0, 32_000, [None], [3.0, 3.0, 3.0])
+ assert commands[1][commands[1].index("-t") + 1] == "1.0"
+ assert commands[1][commands[1].index("-fs") + 1] == "32000"
+
+ monkeypatch.setattr(subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(subprocess.TimeoutExpired("x", 3)))
+ with pytest.raises(TimeoutError, match="timed out"):
+ bounded_run(["ffmpeg"], timeout_s=3.0)
+
+
+def test_media_duration_is_bounded_before_decode(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
+ calls = 0
+
+ def run(command, *, timeout_s, text=False):
+ nonlocal calls
+ calls += 1
+ return subprocess.CompletedProcess(command, 0, "601.0", "")
+
+ monkeypatch.setattr(oi, "_run_media_command", run)
+ with pytest.raises(ValueError, match="Invalid video duration"):
+ oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=600.0)
+ assert calls == 1
+
+
+def test_final_commit_keeps_a_partial_model_unit():
+ """Exact units retain a sample for the server's correlated final flush."""
+
+ events = [{"session": {"capabilities": {"chunk_period_ms": 1000}}}]
+ exact_unit = b"\0\0" * 16_000
+ assert oi._ensure_final_commit_tail(exact_unit, events) == exact_unit[:-2]
+ assert oi._ensure_final_commit_tail(exact_unit + b"\0\0", events) == exact_unit + b"\0\0"
+
+
+@pytest.mark.parametrize(
+ ("events", "match"),
+ [
+ (((_done(), 0.0),), "without response.created"),
+ (((_created(), 0.0), (_created(), 0.1)), "duplicate response.created"),
+ (((_created(), 0.0), (_done(), 0.1), (_done(), 0.2)), "duplicate response.done"),
+ (
+ (
+ (_created(), 0.0),
+ ({"type": "response.done", "response": {"id": "r1", "status_details": {"type": "failed"}}}, 0.1),
+ ),
+ "reports failure",
+ ),
+ ],
+)
+def test_response_ledger_rejects_identity_errors(events, match: str):
+ with pytest.raises(ValueError, match=match):
+ oi.response_ledger(_collector(*events))
+
+
+class _CompletionClient:
+ def __init__(self, collector: RealtimeEventCollector):
+ self.events, self.acks = collector, []
+
+ def raise_if_reader_stopped(self) -> None:
+ return None
+
+ async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
+ self.acks.append((response_id, played_ms))
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("tail", "succeeds"),
+ [
+ (((_created(), 1.1), (_done(), 1.2)), True),
+ (((_listen(), 1.1),), True),
+ (((_listen(buffering=True), 1.1),), False),
+ ((({"type": "session.closed", "reason": "timeout"}, 1.1),), False),
+ ],
+)
+async def test_completion_requires_a_final_model_decision(tail, succeeds: bool):
+ events = _collector(({"type": "input_audio_buffer.committed"}, 1.0), *tail)
+ coroutine = oi.wait_for_session_completion(
+ _CompletionClient(events), oi._Playback(), commit_from=0, timeout_s=0.08, settle_s=0.0
+ )
+ if succeeds:
+ assert await coroutine == 0
+ else:
+ with pytest.raises((TimeoutError, RuntimeError)):
+ await coroutine
+
+
+@pytest.mark.asyncio
+async def test_deferred_commit_ignores_the_prior_response_terminal():
+ events = _collector(
+ (_created("old"), 0.0),
+ ({"type": "input_audio_buffer.committed", "event": {"overlap_deferred": True}}, 1.0),
+ (_done("old"), 1.1),
+ (_listen(), 1.2),
+ )
+ assert (
+ await oi.wait_for_session_completion(
+ _CompletionClient(events), oi._Playback(), commit_from=1, timeout_s=0.1, settle_s=0.0
+ )
+ == 1
+ )
+
+
+@pytest.mark.asyncio
+async def test_playback_acks_only_after_serial_audio_drain():
+ collector = _collector((_audio(samples=2400), 10.0), (_done(), 10.01))
+ client, playback = _CompletionClient(collector), oi._Playback()
+ await playback.acknowledge(client, now=10.05)
+ assert client.acks == []
+ await playback.acknowledge(client, now=10.1)
+ assert client.acks == [("r1", 100)]
+
+
+@pytest.mark.parametrize(
+ ("event", "match"),
+ [
+ (_audio(rate=16_000), "24000 Hz"),
+ ({**_audio(), "format": "opus"}, "pcm16"),
+ ({**_audio(), "delta": "!"}, "base64"),
+ ],
+)
+def test_playback_rejects_non_official_audio(event: dict[str, object], match: str):
+ playback = oi._Playback()
+ with pytest.raises(ValueError, match=match):
+ playback.ingest(_collector((event, 1.0)))
+
+
+def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: Path):
+ case, collector, result = _successful_output(tmp_path)
+ summary = _write_success(
+ tmp_path / "out",
+ case,
+ collector,
+ result,
+ )
+ context = result._artifact_context
+ assert result.success and result.eligible_for_official_eval
+ assert summary["scene_type"] == "multi_turn"
+ assert context and context.spans
+ assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta")
+ oi.publish_deferred_case_artifacts(tmp_path / "out", case, result)
+ directory = oi._output_dir(tmp_path / "out", case)
+ assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS)
+ with wave.open(str(directory / "output.wav")) as output:
+ assert (output.getframerate(), output.getnframes()) == (24_000, 24_000)
+ transcript = json.loads((directory / "wav_transcript.json").read_text())
+ assert transcript["text"] == "hello" and transcript["chunks"][0]["timestamp"] == [0.0, 0.1]
+ assert "serialized playback queue" in transcript["timestamp_semantics"]
+ assert result._artifact_context is None
+
+
+@pytest.mark.parametrize(
+ ("audio_time", "status", "reason"),
+ [(10.95, "completed", "audio_clipped"), (10.0, "cancelled", "cancelled_response")],
+)
+def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_time: float, status: str, reason: str):
+ case, collector, result = _successful_output(tmp_path, audio_time=audio_time, status=status)
+ _write_success(tmp_path / "out", case, collector, result, require_response=False)
+ assert reason in result.official_eval_ineligible_reasons
+ oi.write_batch_artifacts(tmp_path / "out", [case], [result])
+ assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
+
+
+def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):
+ case, _, result = _successful_output(tmp_path)
+ collector = _collector((_created(), 10.0), (_audio(), 10.0), (_done(), 10.1))
+ with pytest.raises(ValueError, match="audio and transcript"):
+ _write_success(tmp_path, case, collector, result)
+
+
+def test_publication_failure_removes_partial_success_bundle(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case, collector, result = _successful_output(tmp_path)
+ _write_success(tmp_path, case, collector, result)
+ original = oi._atomic_write_json
+
+ def fail(path: Path, value: object) -> None:
+ if path.name == "events.json":
+ raise OSError("disk full")
+ original(path, value)
+
+ monkeypatch.setattr(oi, "_atomic_write_json", fail)
+ with pytest.raises(OSError, match="disk full"):
+ oi.publish_deferred_case_artifacts(tmp_path, case, result)
+ directory = oi._output_dir(tmp_path, case)
+ assert not any((directory / name).exists() for name in oi.SUCCESS_ARTIFACTS)
+ assert not result.success and "artifact_write_failed" in result.official_eval_ineligible_reasons
+
+
+def test_atomic_write_preserves_destination_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ destination = tmp_path / "output.wav"
+ destination.write_bytes(b"old")
+ monkeypatch.setattr(oi, "write_pcm16_wav", lambda *a, **k: (_ for _ in ()).throw(OSError("write")))
+ with pytest.raises(OSError):
+ oi._atomic_write_wav(destination, b"\0\0", 24_000)
+ assert destination.read_bytes() == b"old"
+ assert not list(tmp_path.glob(".*.tmp"))
+
+
+class _RealtimeClient:
+ instances: list[_RealtimeClient] = []
+
+ def __init__(self, url: str, **kwargs):
+ self.url, self.events, self.acks, self.configure_kwargs = url, RealtimeEventCollector(), [], {}
+ self.instances.append(self)
+
+ async def __aenter__(self):
+ return self
+
+ async def __aexit__(self, *args):
+ return None
+
+ async def configure(self, model: str, **kwargs) -> None:
+ self.configure_kwargs = kwargs
+ self.events.add({"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 1000}}})
+
+ async def send(self, event: dict[str, object]) -> None:
+ return None
+
+ async def commit(self) -> None:
+ now = time.monotonic()
+ for event, offset in (
+ ({"type": "input_audio_buffer.committed"}, 0.0),
+ (_created(), 0.001),
+ (_audio(samples=24), 0.002),
+ (_text(value="answer"), 0.003),
+ (_done(), 0.004),
+ ):
+ self.events.add(event, received_at_s=now + offset)
+
+ async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
+ self.acks.append((response_id, played_ms))
+
+ async def close_session(self, **kwargs) -> None:
+ self.events.add({"type": "session.closed"})
+
+ def raise_if_reader_stopped(self) -> None:
+ return None
+
+
+@pytest.mark.asyncio
+async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case = _case(tmp_path)
+ config = oi.OmniInteractBenchmarkConfig(
+ model=MODEL,
+ output_root=tmp_path / "out",
+ ref_audio="ref.wav",
+ require_response=True,
+ extra_body={"custom": "value"},
+ )
+ prepared = data.OmniInteractPreparedInput(1.0, b"\0\0\0\0", ("frame",), "data:audio/wav;base64,ref")
+ monkeypatch.setattr(oi, "RealtimeDuplexClient", _RealtimeClient)
+ monkeypatch.setattr(oi, "_COMPLETION_SETTLE_S", 0.0)
+ result = await oi.run_omniinteract_case(
+ case,
+ config,
+ request_index=0,
+ prepared_input=prepared,
+ )
+ assert result.success and result.transcript == "answer" and result._artifact_context is not None
+ assert not config.output_root.exists()
+ assert "autostart=0" in _RealtimeClient.instances[-1].url
+ assert _RealtimeClient.instances[-1].configure_kwargs["extra_body"] == {"custom": "value"}
+ assert _RealtimeClient.instances[-1].acks == [("r1", 1)]
+
+
+def test_standard_sample_loading_prepares_media_before_timing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ _write_dataset(tmp_path)
+ ref = tmp_path / "ref.wav"
+ ref.touch()
+ parser = TrackingArgumentParser()
+ OmniBenchmarkServingSubcommand.add_cli_args(parser)
+ args = parser.parse_args(
+ [
+ "--backend",
+ "openai-realtime-duplex",
+ "--dataset-name",
+ "omniinteract",
+ "--dataset-path",
+ str(tmp_path),
+ "--model",
+ MODEL,
+ "--endpoint",
+ "/v1/realtime",
+ "--num-prompts",
+ "0",
+ "--omniinteract-ref-audio",
+ str(ref),
+ ]
+ )
+ preprocess_serve_args(args)
+ assert args.endpoint == "/v1/realtime"
+ monkeypatch.setattr(benchmark_patch, "reference_audio_data_url", lambda _: "data:audio/wav;base64,ref")
+ monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
+ samples = benchmark_patch.get_samples(args, None)
+ assert len(samples) == args.num_prompts == 3
+ assert all(sample.omniinteract_prepared_input.video_frames == ("frame",) for sample in samples)
+
+
+def _request(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions) -> RequestFuncInput:
+ request = RequestFuncInput(
+ model=MODEL,
+ model_name=MODEL,
+ prompt="",
+ api_url="http://server:8000/v1/realtime",
+ prompt_len=0,
+ output_len=0,
+ logprobs=None,
+ multi_modal_content=None,
+ ignore_eos=False,
+ request_id="measured-0",
+ )
+ request.omniinteract_case, request.omniinteract_options = case, options
+ return request
+
+
+@pytest.mark.asyncio
+async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case, options = _case(tmp_path), _session_options(tmp_path)
+ request = _request(case, options)
+ output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+ assert not output.success and "prepared before benchmark request timing" in output.error
+
+ prepared = data.OmniInteractPreparedInput(1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref")
+ request.omniinteract_prepared_input = prepared
+ request.extra_headers = {"X-Proxy": "test"}
+ request.extra_body = {"custom": "value"}
+ monkeypatch.setenv("OPENAI_API_KEY", "secret")
+ seen = {}
+
+ async def run(_case_arg, config, **kwargs):
+ seen["headers"] = config.extra_headers
+ seen["extra_body"] = config.extra_body
+ seen.update(kwargs)
+ return oi.OmniInteractCaseResult(
+ case.subset,
+ str(case.video_path),
+ str(options.output_root),
+ success=True,
+ transcript="timing metadata missing",
+ output_tokens=0,
+ duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}],
+ duplex_session_metrics={"mean_ttft_ms": 10.0},
+ )
+
+ monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
+ output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+ assert output.success and seen["prepared_input"] is prepared
+ assert seen["headers"] == {"Authorization": "Bearer secret", "X-Proxy": "test", "x-request-id": "measured-0"}
+ assert seen["extra_body"] == {"custom": "value"}
+ assert output.itl == []
+ assert output.tpot_measured is False
+ assert seen["capture_artifacts"] is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("request_metrics", "expected_itl", "expected_text_latency", "expected_tpot_measured"),
+ [
+ (
+ [
+ {"stage0_tokens": {"output_token_count": 3, "itls_ms": [10.0, 20.0]}},
+ {"stage0_tokens": {"output_token_count": 2, "itls_ms": [30.0]}},
+ ],
+ [0.01, 0.02, 0.03],
+ 0.16,
+ True,
+ ),
+ (
+ [
+ {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 10.0}},
+ {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 40.0}},
+ ],
+ [],
+ 0.18,
+ True,
+ ),
+ (
+ [
+ {"stage0_tokens": {"output_token_count": 3, "tpot_ms": 0.0}},
+ {"stage0_tokens": {"output_token_count": 2, "tpot_ms": 0.0}},
+ ],
+ [],
+ 0.1,
+ False,
+ ),
+ ],
+)
+async def test_adapter_reports_exact_or_weighted_token_timing(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ request_metrics: list[dict[str, object]],
+ expected_itl: list[float],
+ expected_text_latency: float,
+ expected_tpot_measured: bool,
+):
+ case, options = _case(tmp_path), _session_options(tmp_path)
+ request = _request(case, options)
+ request.omniinteract_prepared_input = data.OmniInteractPreparedInput(
+ 1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref"
+ )
+
+ async def run(*args, **kwargs):
+ return oi.OmniInteractCaseResult(
+ case.subset,
+ str(case.video_path),
+ str(options.output_root),
+ success=True,
+ output_tokens=5,
+ duplex_request_metrics=request_metrics,
+ duplex_session_metrics={"mean_ttft_ms": 100.0},
+ )
+
+ monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
+ output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
+
+ assert output.success
+ assert output.itl == expected_itl
+ assert output.text_latency == pytest.approx(expected_text_latency)
+ assert output.tpot_measured is expected_tpot_measured
+
+
+def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case, options = _case(tmp_path), _session_options(tmp_path)
+ options.output_root.mkdir()
+ sample = _sample(case, options)
+ result = oi.OmniInteractCaseResult(
+ case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+ )
+ output = benchmark_patch.MixRequestFuncOutput(success=True)
+ output.omniinteract_case_result = result
+ published = []
+ monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: published.append(args))
+ summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+ assert summary and summary["total"] == summary["success"] == summary["eligible_for_official_eval"] == 1
+ assert summary["failed"] == summary["successful_but_ineligible"] == summary["audio_clipped_bytes"] == 0
+ assert published == [(options.output_root, case, result)]
+
+
+def test_case_artifact_failures_do_not_discard_serving_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case, options = _case(tmp_path), _session_options(tmp_path)
+ sample = _sample(case, options)
+ result = oi.OmniInteractCaseResult(
+ case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+ )
+ output = benchmark_patch.MixRequestFuncOutput(success=True, error="inference warning")
+ output.omniinteract_case_result = result
+ monkeypatch.setattr(
+ benchmark_patch,
+ "publish_deferred_case_artifacts",
+ lambda *args: (_ for _ in ()).throw(OSError("case disk full")),
+ )
+ monkeypatch.setattr(
+ benchmark_patch,
+ "write_failure_artifacts",
+ lambda *args: (_ for _ in ()).throw(OSError("failure disk full")),
+ )
+ monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", lambda *args: None)
+
+ summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+ metrics, _ = calculate_metrics(
+ [],
+ [output],
+ 1.0,
+ None,
+ [50.0],
+ {},
+ TaskType.GENERATION,
+ [],
+ None,
+ float("inf"),
+ 1.0,
+ )
+
+ assert summary and summary["failed"] == 1
+ assert summary["artifacts_complete"] is False
+ assert len(summary["artifact_errors"]) == 2
+ assert output.success and metrics.completed == 1
+ assert output.error.startswith("inference warning\n")
+ assert "Artifact publication failed: case disk full" in output.error
+
+
+def test_batch_artifact_failure_removes_partial_batch_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ case, options = _case(tmp_path), _session_options(tmp_path)
+ sample = _sample(case, options)
+ result = oi.OmniInteractCaseResult(
+ case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
+ )
+ output = benchmark_patch.MixRequestFuncOutput(success=True)
+ output.omniinteract_case_result = result
+ monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: None)
+
+ def fail_batch(root, *args):
+ root.mkdir(parents=True, exist_ok=True)
+ (root / "batch_summary.json").write_text("partial")
+ raise OSError("batch disk full")
+
+ monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", fail_batch)
+
+ summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
+
+ assert summary and summary["success"] == 1 and summary["artifacts_complete"] is False
+ assert summary["artifact_errors"] == ["Batch artifact publication failed: batch disk full"]
+ assert output.success
+ assert not (options.output_root / "batch_summary.json").exists()
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "match"),
+ [
+ ("omniinteract_ref_audio", None, "ref-audio"),
+ ("endpoint", "/v1/completions", "v1/realtime"),
+ ("max_concurrency", 0, "positive"),
+ ("skip_tokenizer_init", True, "skip-tokenizer-init"),
+ ],
+)
+def test_cli_rejects_incompatible_omniinteract_options(field: str, value, match: str):
+ args = argparse.Namespace(
+ dataset_name="omniinteract",
+ backend="openai-realtime-duplex",
+ endpoint="/v1/realtime",
+ omniinteract_ref_audio="ref.wav",
+ max_concurrency=1,
+ skip_tokenizer_init=False,
+ )
+ setattr(args, field, value)
+ with pytest.raises(ValueError, match=match):
+ preprocess_serve_args(args)
+
+
+def test_whole_benchmark_lock_serializes_a_shared_output_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ entered, release = threading.Event(), threading.Event()
+ calls: list[int] = []
+
+ async def fake_main_async(args):
+ calls.append(len(calls))
+ if len(calls) == 1:
+ entered.set()
+ await asyncio.to_thread(release.wait)
+ return {"total": 1}
+
+ monkeypatch.setattr(benchmark_serve, "main_async", fake_main_async)
+ args = argparse.Namespace(dataset_name="omniinteract", omniinteract_output_dir=tmp_path, extra_body=None)
+ with ThreadPoolExecutor(2) as executor:
+ first = executor.submit(benchmark_serve.main, args)
+ assert entered.wait(1)
+ second = executor.submit(benchmark_serve.main, args)
+ time.sleep(0.05)
+ assert calls == [0]
+ release.set()
+ assert first.result() == second.result() == {"total": 1}
+ assert calls == [0, 1]
diff --git a/tests/benchmarks/test_seed_tts_dataset_variants.py b/tests/benchmarks/test_seed_tts_dataset_variants.py
index fe1e6b62f..31f66bd7d 100644
--- a/tests/benchmarks/test_seed_tts_dataset_variants.py
+++ b/tests/benchmarks/test_seed_tts_dataset_variants.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for SeedTTSTextDataset, SeedTTSTextSampleRequest, SeedTTSDesignDataset,
and SeedTTSDesignSampleRequest.
@@ -234,34 +237,6 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
assert req.seed_tts_utterance_id == "ok"
-def test_attach_sets_seed_tts_row_even_without_extra_body():
- """seed_tts_row=True must be set for SeedTTSTextSampleRequest (no extra body)."""
- from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSTextSampleRequest
-
- req = SeedTTSTextSampleRequest(
- prompt="hello world",
- prompt_len=2,
- expected_output_len=100,
- multi_modal_data=None,
- request_id="test-0",
- seed_tts_speech_extra=None,
- seed_tts_ref_wav_path="",
- )
- assert req.seed_tts_speech_extra is None
- assert req.seed_tts_ref_wav_path == ""
- # The fix ensures that even with speech_extra=None, the function
- # sets seed_tts_row=True. We verify the source code has the fix.
- import inspect
-
- import vllm_omni.benchmarks.patch.patch as patch_mod
-
- src = inspect.getsource(patch_mod._attach_seed_tts_to_request_func_input)
- # seed_tts_row must be set BEFORE the 'if not ex: return' check
- row_pos = src.index("seed_tts_row")
- not_ex_pos = src.index("if not ex:")
- assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
-
-
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
from vllm_omni.benchmarks.data_modules import seed_tts_eval
diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py
index 7dd12d000..5365d56ac 100644
--- a/tests/benchmarks/test_serve_cli.py
+++ b/tests/benchmarks/test_serve_cli.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import argparse
import json
import os
@@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
"dataset_name",
[
"daily-omni",
+ "omniinteract",
"seed-tts",
"seed-tts-text",
"seed-tts-design",
@@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
[
(["--print-stage"], "print_stage", True),
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
+ (["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
(["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
],
)
@@ -80,12 +85,14 @@ def test_add_omni_args_preserves_implicit_defaults() -> None:
args = parser.parse_args([])
assert args.print_stage is False
assert args.daily_omni_input_mode == "all"
+ assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
assert args.seed_tts_locale == "en"
assert args.explicit_keys == set()
def test_update_omni_help_updates_upstream_actions() -> None:
parser = TrackingArgumentParser()
+ parser.add_argument("--num-prompts", default=1000, help="Number of prompts.")
parser.add_argument("--percentile-metrics", help="Upstream percentile help.")
parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.")
parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.")
@@ -93,6 +100,7 @@ def test_update_omni_help_updates_upstream_actions() -> None:
update_omni_help(parser)
actions = {action.dest: action for action in parser._actions}
+ assert "OmniInteract uses 3" in actions["num_prompts"].help
assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf"))
assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help
assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help
@@ -120,6 +128,31 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
assert args.extra_body == expected
+@pytest.mark.parametrize(
+ ("argv", "expected"),
+ [
+ ([], 3),
+ (["--num-prompts", "0"], 0),
+ (["--num-prompts", "12"], 12),
+ ],
+)
+def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
+ argv: list[str],
+ expected: int,
+) -> None:
+ parser = TrackingArgumentParser()
+ parser.add_argument("--dataset-name", default="omniinteract")
+ parser.add_argument("--backend", default="openai-realtime-duplex")
+ parser.add_argument("--endpoint", default="/v1/realtime")
+ parser.add_argument("--omniinteract-ref-audio", default="reference.wav")
+ parser.add_argument("--num-prompts", type=int, default=1000)
+ args = parser.parse_args(argv)
+
+ preprocess_serve_args(args)
+
+ assert args.num_prompts == expected
+
+
@pytest.mark.parametrize(
("argv", "expected_extra_body", "expected_explicit"),
[
diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py
index 24a2a731e..c994f8dd2 100644
--- a/tests/buildkite/test_upload_pipeline.py
+++ b/tests/buildkite/test_upload_pipeline.py
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import sys
from pathlib import Path
import pytest
+import yaml
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts"))
@@ -38,8 +40,6 @@ def _render(changed_files: list[str]) -> str:
def test_bootstrap_if_injected_by_step_key() -> None:
- import yaml
-
rendered = _render_bootstrap_pipeline(
BOOTSTRAP_STEPS_TEMPLATE,
decision=resolve_ci_decision([]),
@@ -62,16 +62,23 @@ def test_bootstrap_steps_loaded_from_file() -> None:
def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
- """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 still gates scheduled CI."""
+ """skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 / NON_CRITICAL=1 still gates scheduled CI."""
rendered = _render(["docs/foo.md"])
assert "key: image-build" in rendered
assert "key: upload-nightly-pipeline" in rendered
assert "key: upload-weekly-pipeline" in rendered
- # Scheduled nightly also uploads L2/L3 with --e2e
- assert "key: upload-ready-pipeline" in rendered
- assert "key: upload-merge-pipeline" in rendered
- assert 'build.env("NIGHTLY") == "1"' in rendered
+ # Scheduled main+WEEKLY=1 uploads L2/L3 with --e2e; NIGHTLY still gates L4 only.
+ doc = yaml.safe_load(rendered)
+ by_key = {step["key"]: step for step in doc["steps"]}
+ assert "NIGHTLY" not in by_key["upload-ready-pipeline"]["if"]
+ assert 'build.branch == "main"' in by_key["upload-ready-pipeline"]["if"]
+ assert 'build.env("WEEKLY") == "1"' in by_key["upload-ready-pipeline"]["if"]
+ assert "NIGHTLY" not in by_key["upload-merge-pipeline"]["if"]
+ assert 'build.branch == "main"' in by_key["upload-merge-pipeline"]["if"]
+ assert 'build.env("WEEKLY") == "1"' in by_key["upload-merge-pipeline"]["if"]
+ assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
assert 'build.env("WEEKLY") == "1"' in rendered
+ assert 'build.env("NON_CRITICAL") == "1"' in rendered
assert "nightly-test" not in rendered
assert "weekly-test" not in rendered
assert "merge-test" not in rendered
@@ -79,6 +86,19 @@ def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
assert "if: false" not in rendered
+def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None:
+ """NPU skip_all: scheduled NIGHTLY still uploads L4, not L2 ready."""
+ rendered = _render_bootstrap_pipeline(
+ BOOTSTRAP_STEPS_TEMPLATE,
+ decision=resolve_ci_decision(["docs/foo.md"]),
+ path=Path(".buildkite/npu/bootstrap-upload-steps.yml"),
+ )
+ doc = yaml.safe_load(rendered)
+ by_key = {step["key"]: step for step in doc["steps"]}
+ assert "upload-ready-pipeline" not in by_key
+ assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
+
+
def test_yaml_gated_l45_only_does_not_unconditionally_build_image() -> None:
rendered = _render([".buildkite/cuda/test-nightly.yml"])
assert "if: true" not in rendered
@@ -139,3 +159,65 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [
{"name": "swr-secret"},
]
+
+
+def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
+ rendered = _render_test_pipeline(doc, changed_files=changed_files)
+ labels: set[str] = set()
+
+ def walk(steps: list | None) -> None:
+ for step in steps or []:
+ if not isinstance(step, dict):
+ continue
+ if "label" in step:
+ labels.add(step["label"])
+ walk(step.get("steps"))
+
+ walk(rendered.get("steps"))
+ return labels
+
+
+# Synthetic coverage-style job: shared inputs that change what the split measures.
+_COVERAGE_SHARED_INPUTS_DOC = {
+ "steps": [
+ {
+ "label": "Coverage Pilot",
+ "source_file_dependencies": [
+ "tests/e2e/online_serving/test_example.py",
+ ".buildkite/common/scripts/run_cov_split.sh",
+ "pyproject.toml",
+ ],
+ "commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
+ },
+ {
+ "label": "Unrelated Model Test",
+ "source_file_dependencies": [
+ "tests/e2e/online_serving/test_other.py",
+ ],
+ "commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
+ },
+ ],
+}
+
+
+@pytest.mark.parametrize(
+ "changed_file",
+ [
+ ".buildkite/common/scripts/run_cov_split.sh",
+ "pyproject.toml",
+ ],
+)
+def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
+ """Jobs that list coverage shared inputs must stay selected when those files change."""
+ labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
+ assert "Coverage Pilot" in labels
+ assert "Unrelated Model Test" not in labels
+
+
+def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
+ labels = _surviving_labels(
+ _COVERAGE_SHARED_INPUTS_DOC,
+ ["vllm_omni/entrypoints/openai/serving_chat.py"],
+ )
+ assert "Coverage Pilot" not in labels
+ assert "Unrelated Model Test" not in labels
diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py
index c3c3edbd1..e3c6c585f 100644
--- a/tests/config/test_config_factory.py
+++ b/tests/config/test_config_factory.py
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for StageConfigFactory and related classes.
"""
import importlib
import warnings
-from dataclasses import dataclass, fields
+from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
@@ -14,7 +14,8 @@ import pytest
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
-from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config import config_factory as config_factory_module
+from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
from vllm_omni.config.omni_config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -33,13 +34,18 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
+ normalize_pipeline_cli_overrides,
pipeline_cfg_resolver,
- strip_parent_engine_args,
)
-from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
+from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+# Validation strings to match against in post init
+NO_STAGES_MATCH_STR = "no stages"
+NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage"
+
@pytest.fixture(autouse=True)
def _stable_test_platform(monkeypatch):
@@ -56,6 +62,7 @@ def clear_config_factory_caches():
yield
StageConfigFactory.get_hf_config.cache_clear()
StageConfigFactory.try_infer_model_type.cache_clear()
+ _materialize_object_storage_configs.cache_clear()
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
@@ -428,6 +435,20 @@ class TestStageConfigFactory:
assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image"
+ @pytest.mark.parametrize(
+ "model_class_name",
+ [
+ "MiniMaxH3Pipeline",
+ "WanPipeline",
+ "LTX2Pipeline",
+ "Cosmos3OmniDiffusersPipeline",
+ ],
+ )
+ def test_default_video_diffusion_output_type(self, model_class_name):
+ configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
+
+ assert configs[0]["final_output_type"] == "video"
+
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
@@ -475,7 +496,7 @@ class TestStageConfigFactory:
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
- "stage_configs_path": "/path", # Internal
+ "deploy_config": "/path", # Internal
"batch_timeout": 10, # Internal
}
@@ -483,7 +504,7 @@ class TestStageConfigFactory:
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
- assert "stage_configs_path" not in overrides
+ assert "deploy_config" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
@@ -525,49 +546,6 @@ class TestStageResolutionHelpers:
assert "model" not in overrides
assert "parallel_config" not in overrides
- def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
- parent_fields = {f.name: f for f in fields(EngineArgs)}
- filtered, overridden = strip_parent_engine_args(
- {
- "model": "some/model",
- "stage_configs_path": "/tmp/stages.yaml",
- "tensor_parallel_size": 4,
- "worker_extension_cls": "some.Extension",
- "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
- },
- parent_fields=parent_fields,
- keep_keys={"worker_extension_cls"},
- strip_keys={"stage_configs_path"},
- no_warn_keys={"model"},
- )
-
- assert filtered == {
- "worker_extension_cls": "some.Extension",
- "custom_pipeline_args": {"pipeline_class": "demo.Pipeline"},
- }
- assert overridden == ["tensor_parallel_size"]
-
- def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
- parent_fields = {f.name: f for f in fields(EngineArgs)}
- filtered, overridden = strip_parent_engine_args(
- {
- "model": "some/model",
- "stage_configs_path": "/tmp/stages.yaml",
- "allowed_local_media_path": "/data/qwentts",
- "allowed_media_domains": ["example.com"],
- },
- parent_fields=parent_fields,
- keep_keys={"allowed_local_media_path", "allowed_media_domains"},
- strip_keys={"stage_configs_path"},
- no_warn_keys={"model"},
- )
-
- assert filtered == {
- "allowed_local_media_path": "/data/qwentts",
- "allowed_media_domains": ["example.com"],
- }
- assert overridden == []
-
class TestPipelineDiscovery:
"""Tests for the central pipeline registry (``OMNI_PIPELINES``)."""
@@ -576,6 +554,7 @@ class TestPipelineDiscovery:
"""Check that specific models are in OMNI_PIPELINES."""
assert "qwen2_5_omni" in OMNI_PIPELINES
assert "qwen3_omni_moe" in OMNI_PIPELINES
+ assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
assert "qwen3_tts" in OMNI_PIPELINES
def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -600,6 +579,18 @@ class TestPipelineDiscovery:
assert pipeline.default_deploy_config_name is None
assert len(pipeline.stages) == 1 # thinker only
+ def test_registry_qwen3_omni_thinker_only_key_is_static(self):
+ """Explicit registry key must not re-enter the HF-config resolver."""
+ pipeline = resolve_pipeline_config(
+ "qwen3_omni_moe_thinker_only",
+ Q3_OMNI_ALL_STAGES_HF_CONFIG,
+ )
+ assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
+ assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
+ assert len(pipeline.stages) == 1
+ assert pipeline.stages[0].engine_output_type == "text"
+ assert pipeline.stages[0].final_output_type == "text"
+
@pytest.mark.parametrize(
"pipeline",
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -632,6 +623,7 @@ class TestPipelineDiscovery:
p = PipelineConfig(
model_type="custom_collide",
hf_architectures=("SomeCollidingArch",),
+ stages=SINGLE_STAGE_PIPE_CFG,
)
assert p.hf_architectures == ("SomeCollidingArch",)
@@ -653,33 +645,63 @@ class TestStagePipelineConfig:
class TestPipelineConfigNew:
- def test_frozen(self):
- p = PipelineConfig(model_type="t", model_arch="A")
- with pytest.raises(AttributeError):
- p.model_type = "changed"
+ def test_simple_init(self):
+ # Ensure pipeline config validates in post init.
+ PipelineConfig(
+ model_type="t",
+ model_arch="A",
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
- def test_validate_valid(self):
+ def test_frozen(self):
p = PipelineConfig(
+ model_type="t",
+ model_arch="A",
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
+ with pytest.raises(AttributeError):
+ p.model_type = "changed"
+
+ def test_validate_no_stages(self):
+ with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR):
+ PipelineConfig(model_type="t", model_arch="A")
+
+ def test_validate_no_terminal_stage(self):
+ """A pipeline with no ``final_output`` stage can never emit a result."""
+ with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR):
+ PipelineConfig(
+ model_type="t",
+ model_arch="A",
+ stages=(
+ StagePipelineConfig(stage_id=0, model_stage="a"),
+ StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
+ ),
+ )
+
+ def test_pipeline_can_have_final_output_in_any_stage(self):
+ """Ensure any stage may carry ``final_output``."""
+ PipelineConfig(
model_type="t",
model_arch="A",
stages=(
- StagePipelineConfig(stage_id=0, model_stage="a"),
+ StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
),
)
- assert p.validate() == []
-
- def test_validate_no_stages(self):
- p = PipelineConfig(model_type="t", model_arch="A")
- assert any("no stages" in e.lower() for e in p.validate())
class TestPipelineRegistration:
def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
deploy_key = "deploy_selected_pipeline"
model_type_key = "hf_model_type_pipeline"
- deploy_pipe = PipelineConfig(model_type=deploy_key)
- model_type_pipe = PipelineConfig(model_type=model_type_key)
+ deploy_pipe = PipelineConfig(
+ model_type=deploy_key,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
+ model_type_pipe = PipelineConfig(
+ model_type=model_type_key,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
register_pipeline(deploy_pipe)
register_pipeline(model_type_pipe)
@@ -705,7 +727,11 @@ class TestPipelineRegistration:
tmp_path,
):
model_type_key = "registered_model_type_pipeline"
- register_pipeline(PipelineConfig(model_type=model_type_key))
+ pipeline_cfg = PipelineConfig(
+ model_type=model_type_key,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
+ register_pipeline(pipeline_cfg)
deploy_path = tmp_path / "deploy.yaml"
deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8")
@@ -740,6 +766,7 @@ class TestPipelineRegistration:
pipe_cfg = PipelineConfig(
model_type=pipeline_key,
hf_architectures=("ArchitectureFallbackForTest",),
+ stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -770,15 +797,18 @@ class TestPipelineRegistration:
model_type="rejected_by_predicate",
hf_architectures=(shared_arch,),
hf_config_predicate=rejecting_predicate,
+ stages=SINGLE_STAGE_PIPE_CFG,
)
raises_cfg = PipelineConfig(
model_type="raises_in_predicate",
hf_architectures=(shared_arch,),
hf_config_predicate=raising_predicate,
+ stages=SINGLE_STAGE_PIPE_CFG,
)
accept_cfg = PipelineConfig(
model_type="accepted_by_predicate",
hf_architectures=(shared_arch,),
+ stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(reject_cfg)
register_pipeline(raises_cfg)
@@ -812,6 +842,7 @@ class TestPipelineRegistration:
model_type="predicate_without_arch_match",
hf_architectures=("DifferentArchitectureForTest",),
hf_config_predicate=predicate,
+ stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -835,6 +866,7 @@ class TestPipelineRegistration:
resolved_cfg = PipelineConfig(
model_type="callable_resolved_pipeline",
hf_architectures=("CallableArchitectureForTest",),
+ stages=SINGLE_STAGE_PIPE_CFG,
)
seen_hf_configs = []
@@ -908,7 +940,10 @@ class TestPipelineRegistration:
def test_pipeline_registration(self, clean_pipeline_registry):
"""Ensure that we can register and create a custom pipeline config."""
new_model_type = "new_model_type"
- pipe_cfg = PipelineConfig(model_type=new_model_type)
+ pipe_cfg = PipelineConfig(
+ model_type=new_model_type,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
@@ -951,7 +986,10 @@ class TestPipelineRegistration:
def custom_resolver(
hf_config: FakeConfig,
) -> PipelineConfig:
- return PipelineConfig(model_type=resolved_type)
+ return PipelineConfig(
+ model_type=resolved_type,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
# Register the new PipelineConfig
assert new_model_type not in OMNI_PIPELINES
@@ -1008,7 +1046,7 @@ class TestPipelineRegistration:
pipe_cfg = PipelineConfig(
model_type=pipeline_key,
model_arch="DeployOnlyArch",
- stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),),
+ stages=SINGLE_STAGE_PIPE_CFG,
)
register_pipeline(pipe_cfg)
@@ -1041,7 +1079,11 @@ class TestPipelineRegistration:
def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path):
pipeline_key = "single_load_pipeline"
- register_pipeline(PipelineConfig(model_type=pipeline_key))
+ pipeline_cfg = PipelineConfig(
+ model_type=pipeline_key,
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
+ register_pipeline(pipeline_cfg)
deploy_path = tmp_path / "single_load.yaml"
deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8")
@@ -1071,6 +1113,7 @@ class TestPipelineRegistration:
pipeline = PipelineConfig(
model_type="pipeline_with_default",
default_deploy_config_name=default_name,
+ stages=SINGLE_STAGE_PIPE_CFG,
)
with patch(
@@ -1090,8 +1133,16 @@ class TestPipelineRegistration:
OmniServingCapability.COMPLETIONS,
"pipeline_a blocks completions",
)
- pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
- pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
+ pipe_a = PipelineConfig(
+ model_type="detect_type",
+ endpoint_restrictions=(restriction,),
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
+ pipe_b = PipelineConfig(
+ model_type="override_type",
+ endpoint_restrictions=(),
+ stages=SINGLE_STAGE_PIPE_CFG,
+ )
register_pipeline(pipe_a)
register_pipeline(pipe_b)
@@ -1144,8 +1195,42 @@ class TestResolveScheduler:
class TestDeployConfigLoading:
- def test_load_minicpmo_duplex_deploy_config(self):
- deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
+ def test_rejects_legacy_stage_args_schema(self, tmp_path):
+ deploy_path = tmp_path / "legacy.yaml"
+ deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
+
+ with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
+ load_deploy_config(deploy_path)
+
+ @pytest.mark.parametrize(
+ ("filename", "max_sessions"),
+ [
+ ("minicpmo_4_5.yaml", 4),
+ ("minicpmo_4_5_2gpu.yaml", 4),
+ ("minicpmo_4_5_3gpu.yaml", 4),
+ ("minicpmo_4_5_8x4090.yaml", 1),
+ ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
+ ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
+ ("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
+ ],
+ )
+ def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
+ """Every MiniCPM-o profile serves /v1/realtime, so the retired
+ minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
+ """
+ deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
+ pipeline = resolve_pipeline_config("minicpmo_4_5")
+ assert isinstance(pipeline, PipelineConfig)
+ stages = merge_pipeline_deploy(pipeline, deploy)
+
+ assert deploy.session_mode == "duplex"
+ assert deploy.active_stream_window == max_sessions
+ assert deploy.duplex_session.max_sessions == max_sessions
+ assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
+ assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
+
+ def test_load_minicpmo_default_deploy_config(self):
+ deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
deploy = load_deploy_config(deploy_path)
pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1153,18 +1238,18 @@ class TestDeployConfigLoading:
stages = merge_pipeline_deploy(pipeline, deploy)
- assert deploy.session_mode == "duplex"
assert deploy.async_chunk is True
- assert deploy.active_stream_window == 1
- assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
- assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
- assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
- assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
+ assert stages[0].yaml_engine_args["async_scheduling"] is True
+ assert stages[1].yaml_engine_args["async_scheduling"] is True
+ assert "Async" in (stages[0].scheduler_cls or "")
+ assert "Async" in (stages[1].scheduler_cls or "")
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
assert deploy.stages[1].enforce_eager is False
- assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048
- assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
- assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1]
+ assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
+ assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
+ assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
+ assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
+ assert "codec_sampling_params" not in stages[1].yaml_engine_args
@pytest.mark.parametrize(
("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"),
@@ -1233,14 +1318,22 @@ stages:
assert deploy.connectors is not None
assert deploy.platforms is not None
- def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
+ @pytest.mark.parametrize(
+ ("hf_config", "model"),
+ [
+ (Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
+ (Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
+ ],
+ ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
+ )
+ def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
with patch(
"vllm_omni.config.config_factory.get_config",
- return_value=Q3_OMNI_THINKER_HF_CONFIG,
+ return_value=hf_config,
):
pipeline = StageConfigFactory.get_pipeline_config(
- "Qwen/Qwen3-Omni-30B-A3B-Thinking",
+ model,
trust_remote_code=False,
deploy_config_path=str(deploy_path),
)
@@ -1249,6 +1342,7 @@ stages:
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
deploy = load_deploy_config(deploy_path)
+ assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
stages = merge_pipeline_deploy(pipeline, deploy)
assert len(stages) == 1
@@ -1310,6 +1404,13 @@ stages:
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
+ (
+ "minimax_h3_disaggregated.yaml",
+ "minimax_h3_disaggregated",
+ 2,
+ "video",
+ "minimax_h3_disaggregated",
+ ),
],
)
def test_load_new_registry_backed_deploy_configs(
@@ -1330,15 +1431,98 @@ stages:
assert len(stages) == stage_count
assert stages[-1].final_output is True
assert stages[-1].final_output_type == final_output_type
+ if pipeline_name == "minimax_h3_disaggregated":
+ assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
+ assert (
+ stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
+ )
+ assert "model_path_resolver" not in stages[0].yaml_extras
+ assert "model_path_resolver" not in stages[1].yaml_extras
if deploy.trust_remote_code is not None:
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
+ def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
+ from vllm_omni.engine import stage_init_utils
+
+ model_root = tmp_path / "MiniMax-H3"
+ (model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
+ (model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
+ monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
+
+ deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
+ assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
+ stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
+ resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
+
+ assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
+ assert resolved[1]["model"] == str(model_root / "FL2VA")
+ assert all("model_path_resolver" not in engine_args for engine_args in resolved)
+
+ def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+ stages, _ = StageConfigFactory._create_legacy_from_registry(
+ pipeline,
+ {"text_encoder_tp_size": 4},
+ )
+
+ assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
+ assert "text_encoder_tp_size" not in stages[0].runtime_overrides
+ assert "text_encoder_tp_size" not in stages[1].runtime_overrides
+ assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
+
+ def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+ deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
+ stages = merge_pipeline_deploy(pipeline, deploy)
+
+ assert stages[0].yaml_engine_args["max_num_seqs"] == 1
+ assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
+ assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
+ parallel = stages[1].yaml_engine_args["parallel_config"]
+ assert parallel["tensor_parallel_size"] == 1
+ assert parallel["ulysses_degree"] == 4
+ assert parallel["vae_patch_parallel_size"] == 4
+ assert stages[1].yaml_engine_args["inline_diffusion"] is True
+ assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
+
+ turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
+ turbo_stages = merge_pipeline_deploy(pipeline, turbo)
+ turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
+ assert turbo_sampling["num_inference_steps"] == 5
+ assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
+
+ def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+ assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
+ assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
+ "stage_0_tensor_parallel_size": 4
+ }
+
+ def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+ with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
+ normalized = normalize_pipeline_cli_overrides(
+ pipeline,
+ {"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
+ )
+
+ assert normalized == {"stage_0_tensor_parallel_size": 2}
+
+ def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+
+ with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
+ normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
+
@pytest.mark.parametrize(
("config_json", "model_index", "expected_pipeline"),
[
({"model_type": "step_audio_2"}, None, "step_audio_2"),
(None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
(None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
+ (None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"),
],
)
def test_migrated_models_are_discovered_without_explicit_deploy(
@@ -1367,6 +1551,29 @@ stages:
assert pipeline is not None
assert pipeline.model_type == expected_pipeline
+ def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
+ def get_model_file(filename, _model, revision=None):
+ del revision
+ if filename == "config.json":
+ return None
+ if filename == "model_index.json":
+ return {"_class_name": "MiniMaxH3Pipeline"}
+ return None
+
+ with (
+ patch.object(StageConfigFactory, "get_hf_config", return_value=None),
+ patch(
+ "vllm_omni.config.config_factory.get_hf_file_to_dict",
+ side_effect=get_model_file,
+ ),
+ ):
+ pipeline = StageConfigFactory.get_pipeline_config(
+ model="/models/unrelated-checkpoint-name",
+ trust_remote_code=False,
+ )
+
+ assert pipeline is None
+
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
@@ -1419,7 +1626,8 @@ stages:
assert async_stages[1].custom_process_input_func is None
def test_no_bundled_legacy_stage_config_yamls(self):
- stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
+ repo_root = Path(__file__).resolve().parents[2]
+ stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
assert not list(stage_config_dir.glob("*.yaml"))
def test_merge_pipeline_deploy(self):
@@ -1474,6 +1682,7 @@ stages:
model_stage="ar",
execution_type=StageExecutionType.LLM_AR,
requires_multimodal_data=True,
+ final_output=True,
),
),
)
@@ -1494,6 +1703,7 @@ stages:
model_stage="ar",
execution_type=StageExecutionType.LLM_AR,
scheduler_cls=scheduler_cls,
+ final_output=True,
),
),
)
@@ -1693,6 +1903,7 @@ stages:
]
}
original_foo = fake_config["stages"][0]["foo"]
+ assert isinstance(original_foo, dict)
original_b = original_foo["b"]
with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
@@ -1738,7 +1949,6 @@ class TestQwen3OmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
assert len(p.stages) == 3
- assert p.validate() == []
def test_thinker(self):
p = resolve_pipeline_config(
@@ -1792,7 +2002,6 @@ class TestQwen2_5OmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
assert len(p.stages) == 3
- assert p.validate() == []
def test_thinker(self):
p = resolve_pipeline_config("qwen2_5_omni")
@@ -1841,7 +2050,6 @@ class TestQwen3TTSPipeline:
assert p is not None
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
assert len(p.stages) == 2
- assert p.validate() == []
def test_talker_stage(self):
p = resolve_pipeline_config("qwen3_tts")
@@ -1921,7 +2129,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
- assert p.validate() == []
def test_thinker_stage(self):
p = resolve_pipeline_config("ming_flash_omni")
@@ -1978,7 +2185,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
assert len(p.stages) == 1
- assert p.validate() == []
def test_tts_stage(self):
p = resolve_pipeline_config("ming_flash_omni_tts")
@@ -2035,7 +2241,6 @@ class TestMingFlashOmniPipeline:
assert isinstance(p, PipelineConfig)
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 1
- assert p.validate() == []
def test_thinker_only_stage(self):
p = resolve_pipeline_config("ming_flash_omni_thinker_only")
@@ -2074,7 +2279,6 @@ class TestMingFlashOmniPipeline:
assert p is not None
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
assert len(p.stages) == 2
- assert p.validate() == []
def test_image_thinker_stage(self):
s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
@@ -2145,6 +2349,26 @@ class TestMingFlashOmniPipeline:
class TestBaseConfigInheritance:
"""Test deploy YAML base_config inheritance."""
+ def test_minicpmo_overlays_inherit_talker_sampling_params(self):
+ """Overlays must keep the base Talker codec Sampler knobs."""
+ for filename in (
+ "minicpmo_4_5_2gpu.yaml",
+ "minicpmo_4_5_3gpu.yaml",
+ "minicpmo_4_5_3gpu_stage1_replicas.yaml",
+ "minicpmo_4_5_4gpu_stage1_replicas.yaml",
+ "minicpmo_4_5_8x4090_stage1_replicas.yaml",
+ ):
+ path = Path(get_deploy_config_path(filename))
+ if not path.exists():
+ pytest.skip(f"{filename} not found")
+ deploy = load_deploy_config(path)
+ sampling = deploy.stages[1].default_sampling_params
+ assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
+ assert sampling["temperature"] == 0.8, filename
+ assert sampling["top_k"] == 25, filename
+ assert sampling["top_p"] == 0.85, filename
+ assert sampling["repetition_penalty"] == 1.05, filename
+
def test_ci_inherits_from_main(self):
ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
if not ci_path.exists():
@@ -2235,6 +2459,42 @@ class TestPlatformOverrides:
assert rocm.stages[0].enforce_eager is None
assert rocm.stages[1].enforce_eager is True
+ def test_higgs_audio_v3_rocm_uses_triton_attention(self):
+ deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
+
+ base = load_deploy_config(deploy_path)
+ assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER"
+
+ rocm = _apply_platform_overrides(base, platform="rocm")
+ assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN"
+
+ pipeline = resolve_pipeline_config("higgs_multimodal_qwen3")
+ assert isinstance(pipeline, PipelineConfig)
+ stages = merge_pipeline_deploy(pipeline, rocm)
+ assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN"
+
+ def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
+ deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
+ pipeline = resolve_pipeline_config(
+ "qwen3_omni_moe",
+ Q3_OMNI_ALL_STAGES_HF_CONFIG,
+ )
+ assert isinstance(pipeline, PipelineConfig)
+
+ cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
+ cuda_stages = merge_pipeline_deploy(pipeline, cuda)
+ expected = {"custom_ops": ["+rotary_embedding"]}
+ assert cuda_stages[0].yaml_engine_args["compilation_config"] == expected
+ assert all("compilation_config" not in stage.yaml_engine_args for stage in cuda_stages[1:])
+
+ for platform in ("cpu", "musa", "npu", "rocm", "xpu"):
+ deploy = _apply_platform_overrides(
+ load_deploy_config(deploy_path),
+ platform=platform,
+ )
+ config = deploy.stages[0].compilation_config or {}
+ assert "+rotary_embedding" not in config.get("custom_ops", [])
+
def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
pipeline = resolve_pipeline_config("minicpmo_4_5")
assert isinstance(pipeline, PipelineConfig)
@@ -2741,22 +3001,23 @@ class TestSentinelDefaultPrecedence:
)
def test_ming_flash_omni_topology(self):
- """Guard ming_flash_omni's SIP cleanup: stage 0 has no full-payload
- producer hook (arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and
- stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func).
+ """Guard ming_flash_omni's SIP cleanup: stage 0 has no
+ ``custom_process_next_stage_input_func``, and stage 1 does not set
+ ``requires_full_payload_input=True``. Stage 1 uses only
+ ``thinker2talker_token_only`` (sync_process_input_func).
The dead ``thinker2talker`` (custom_process_input_func) was removed
because ming_flash_omni has no async_chunk support and both functions
called ``_build_talker_inputs`` identically.
- Merge under either async_chunk mode must not re-introduce a
- stage-0 full-payload hook."""
+ Merge under either async_chunk mode must not re-introduce either
+ full-payload transport gate."""
pipeline = resolve_pipeline_config("ming_flash_omni")
assert isinstance(pipeline, PipelineConfig)
stage0, stage1 = pipeline.stages
assert stage0.custom_process_next_stage_input_func is None, (
- "ming_flash_omni stage 0 must not declare a full-payload producer "
- "(connector path is not active for this arch)."
+ "ming_flash_omni stage 0 must not declare a full-payload producer hook."
)
+ assert stage1.requires_full_payload_input is False
assert stage1.custom_process_input_func is None
assert stage1.sync_process_input_func is not None
assert stage1.sync_process_input_func.endswith("thinker2talker_token_only")
@@ -2809,3 +3070,90 @@ class TestPipelineConfigResolvers:
pass
assert resolver(NotTheRightHfConfig()) is None
+
+
+class TestObjectStorageConfigResolution:
+ """Object-storage URIs must be materialized locally before any HF-style reads.
+
+ Regression coverage for review on the Run:AI PR: ``get_config`` and the
+ config.json/model_index.json fallbacks crashed with ``HFValidationError``
+ for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
+ a bucket named after another pipeline could hijack pipeline selection.
+ """
+
+ @pytest.fixture
+ def fake_object_storage(self, monkeypatch, tmp_path):
+ """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
+
+ Returns a recorder: calls to the returned callable write files into
+ the materialized directory, and ``recorder.pulls`` records every
+ ``pull_files`` invocation.
+ """
+ pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
+ materialized_files: list[tuple[str, str]] = []
+
+ class FakeObjectStorageModel:
+ def __init__(self, url: str):
+ self.dir = str(tmp_path)
+
+ def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
+ pulls.append((model_path, allow_pattern, ignore_pattern))
+ for name, content in materialized_files:
+ (tmp_path / name).write_text(content)
+
+ monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
+
+ class Recorder:
+ def __init__(self) -> None:
+ self.pulls = pulls
+
+ def set_files(self, files: list[tuple[str, str]]) -> None:
+ materialized_files.clear()
+ materialized_files.extend(files)
+
+ return Recorder()
+
+ def test_passthrough_for_non_uri(self, fake_object_storage):
+ assert _materialize_object_storage_configs("org/model") == "org/model"
+ assert _materialize_object_storage_configs("/local/model") == "/local/model"
+ assert fake_object_storage.pulls == []
+
+ def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
+ uri = "s3://bucket/model"
+
+ first = _materialize_object_storage_configs(uri)
+ second = _materialize_object_storage_configs(uri)
+
+ assert first == second
+ assert len(fake_object_storage.pulls) == 1
+ pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
+ assert pulled_uri == uri
+ assert "*.json" in allow_pattern
+ assert ignore_pattern is None
+
+ def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
+ """A misleading bucket name must not beat the real config.json content."""
+ fake_object_storage.set_files(
+ [
+ (
+ "config.json",
+ '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
+ )
+ ]
+ )
+ uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
+
+ model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
+
+ assert model == "qwen3_omni_moe"
+ assert fake_object_storage.pulls # materialization happened
+
+ def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
+ # Empty config.json (CosyVoice3 style) forces the name-match fallback.
+ fake_object_storage.set_files([("config.json", "{}")])
+
+ deceptive = "s3://qwen3-tts-models/plain-checkpoint"
+ assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
+
+ matching = "s3://any-bucket/my-cosyvoice3-model"
+ assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
diff --git a/tests/config/test_config_import_cycle.py b/tests/config/test_config_import_cycle.py
new file mode 100644
index 000000000..e149b63ee
--- /dev/null
+++ b/tests/config/test_config_import_cycle.py
@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Guard the NPU pytest circular import through DiffusionOutput."""
+
+from __future__ import annotations
+
+import ast
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _top_level_import_modules(path: Path) -> set[str]:
+ tree = ast.parse(path.read_text(encoding="utf-8"))
+ modules: set[str] = set()
+ for node in tree.body:
+ if isinstance(node, ast.ImportFrom) and node.module:
+ modules.add(node.module)
+ elif isinstance(node, ast.Import):
+ modules.update(alias.name for alias in node.names)
+ return modules
+
+
+def _assert_isolated_import_succeeds(script: str) -> None:
+ env = os.environ.copy()
+ env.pop("CUDA_VISIBLE_DEVICES", None)
+ env.pop("HIP_VISIBLE_DEVICES", None)
+ result = subprocess.run(
+ [sys.executable, "-c", script],
+ cwd=REPO_ROOT,
+ env=env,
+ capture_output=True,
+ text=True,
+ timeout=120,
+ check=False,
+ )
+ assert result.returncode == 0, result.stderr
+
+
+def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
+ imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
+ assert "vllm_omni.config.config_factory" not in imported
+ assert "vllm_omni.config.pipeline_registry" not in imported
+
+
+def test_config_package_lazy_exports_still_resolve() -> None:
+ from vllm_omni.config import StageConfigFactory, register_pipeline
+ from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
+ from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
+
+ assert StageConfigFactory is DirectStageConfigFactory
+ assert register_pipeline is direct_register_pipeline
+
+
+def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
+ src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
+ init_start = src.index("def __init__(self) -> None:")
+ runtime_start = src.index("def init_diffusion_model_runner_runtime")
+ init_src = src[init_start:runtime_start]
+ runtime_src = src[runtime_start:]
+ assert "apply_minimax_h3_qwen3vl_patch" not in init_src
+ assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
+ assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
+ assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
+ assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
+
+
+def test_lora_config_import_does_not_load_pipeline_registry() -> None:
+ """Import LoRAConfig and diffusion.data in a fresh interpreter.
+
+ The NPU failure was ``data.py`` importing LoRAConfig while still loading,
+ then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
+ DiffusionOutput. This child process must not preload those exports.
+ """
+ _assert_isolated_import_succeeds(
+ """
+import sys
+
+from vllm_omni.config.lora import LoRAConfig
+
+if "vllm_omni.config.pipeline_registry" in sys.modules:
+ raise SystemExit("LoRAConfig import loaded pipeline_registry")
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+ raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
+
+from vllm_omni.diffusion.data import DiffusionOutput
+
+if DiffusionOutput is None or LoRAConfig is None:
+ raise SystemExit("expected LoRAConfig and DiffusionOutput")
+"""
+ )
+
+
+def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
+ """Resolving the pi0 topology must not import its runtime pipeline.
+
+ ``pipeline_pi0`` imports ``diffusion.data``. Importing it while
+ ``diffusion.data`` is only partially initialized closes the Ascend plugin's
+ circular-import loop, so the registry must depend only on a lightweight
+ pipeline-config module.
+ """
+ _assert_isolated_import_succeeds(
+ """
+import sys
+
+from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+
+if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
+ raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
+
+pipeline = resolve_pipeline_config("pi0")
+if pipeline is None or pipeline.model_type != "pi0":
+ raise SystemExit("expected the registered pi0 pipeline config")
+"""
+ )
diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py
new file mode 100644
index 000000000..00e0dffdb
--- /dev/null
+++ b/tests/config/test_environment_variables.py
@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Drift checks for the reviewed environment-variable inventory."""
+
+import ast
+import builtins
+from collections import Counter
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.config.environment_variable_inventory import (
+ ENVIRONMENT_VARIABLE_INVENTORY,
+ EnvironmentVariableCategory,
+ ModelEnvironmentVariableDisposition,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).parents[2]
+_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni"
+_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py"
+_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md"
+
+
+def _module_string_constants(tree: ast.Module) -> dict[str, str]:
+ constants: dict[str, str] = {}
+ for statement in tree.body:
+ targets: list[ast.expr] = []
+ value: ast.expr | None = None
+ if isinstance(statement, ast.Assign):
+ targets = statement.targets
+ value = statement.value
+ elif isinstance(statement, ast.AnnAssign):
+ targets = [statement.target]
+ value = statement.value
+ if isinstance(value, ast.Constant) and isinstance(value.value, str):
+ for target in targets:
+ if isinstance(target, ast.Name):
+ constants[target.id] = value.value
+ return constants
+
+
+def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]:
+ os_names: set[str] = set()
+ getenv_names: set[str] = set()
+ environ_names: set[str] = set()
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Import):
+ for imported in node.names:
+ if imported.name == "os":
+ os_names.add(imported.asname or "os")
+ elif isinstance(node, ast.ImportFrom) and node.module == "os":
+ for imported in node.names:
+ if imported.name == "getenv":
+ getenv_names.add(imported.asname or imported.name)
+ elif imported.name == "environ":
+ environ_names.add(imported.asname or imported.name)
+ return os_names, getenv_names, environ_names
+
+
+def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool:
+ return (
+ isinstance(node, ast.Attribute)
+ and isinstance(node.value, ast.Name)
+ and node.value.id in os_names
+ and node.attr == "environ"
+ ) or (isinstance(node, ast.Name) and node.id in environ_names)
+
+
+def _environment_key_expressions(
+ node: ast.AST,
+ os_names: set[str],
+ getenv_names: set[str],
+ environ_names: set[str],
+) -> list[ast.expr]:
+ """Return key expressions used by one direct environment access."""
+ if isinstance(node, ast.Call) and node.args:
+ func = node.func
+ if (
+ isinstance(func, ast.Attribute)
+ and isinstance(func.value, ast.Name)
+ and func.value.id in os_names
+ and func.attr == "getenv"
+ ) or (isinstance(func, ast.Name) and func.id in getenv_names):
+ return [node.args[0]]
+ if (
+ isinstance(func, ast.Attribute)
+ and _is_environ(func.value, os_names, environ_names)
+ and func.attr in {"get", "pop", "setdefault"}
+ ):
+ return [node.args[0]]
+ elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names):
+ return [node.slice]
+ elif isinstance(node, ast.Compare):
+ expressions: list[ast.expr] = []
+ left_values = [node.left, *node.comparators[:-1]]
+ for left, operator, right in zip(left_values, node.ops, node.comparators):
+ if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names):
+ expressions.append(left)
+ return expressions
+ return []
+
+
+def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None:
+ value: str | None = None
+ if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
+ value = expression.value
+ elif isinstance(expression, ast.Name):
+ value = constants.get(expression.id)
+ return value or None
+
+
+def _environment_wrapper_parameters(
+ tree: ast.Module,
+ os_names: set[str],
+ getenv_names: set[str],
+ environ_names: set[str],
+) -> dict[str, set[tuple[str, int | None]]]:
+ """Find local helpers that forward one parameter as an environment key."""
+ wrappers: dict[str, set[tuple[str, int | None]]] = {}
+ for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
+ positional = [*function.args.posonlyargs, *function.args.args]
+ parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)}
+ parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs})
+ forwarded: set[tuple[str, int | None]] = set()
+ for node in ast.walk(function):
+ for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
+ if isinstance(expression, ast.Name) and expression.id in parameter_indexes:
+ forwarded.add((expression.id, parameter_indexes[expression.id]))
+ if forwarded:
+ wrappers[function.name] = forwarded
+ return wrappers
+
+
+def _environment_accesses(path: Path) -> set[str]:
+ """Return statically resolvable names accessed through ``os``."""
+ tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+ constants = _module_string_constants(tree)
+ os_names, getenv_names, environ_names = _os_import_aliases(tree)
+ wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names)
+ names: set[str] = set()
+
+ for node in ast.walk(tree):
+ for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
+ if name := _resolve_environment_name(expression, constants):
+ names.add(name)
+
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+ for parameter_name, position in wrappers.get(node.func.id, set()):
+ argument = node.args[position] if position is not None and position < len(node.args) else None
+ if argument is None:
+ argument = next(
+ (keyword.value for keyword in node.keywords if keyword.arg == parameter_name),
+ None,
+ )
+ if argument is not None and (name := _resolve_environment_name(argument, constants)):
+ names.add(name)
+
+ return names
+
+
+def test_inventory_matches_reviewed_snapshot_counts():
+ """Make an inventory expansion an explicit review decision."""
+ category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
+ assert category_counts == {
+ EnvironmentVariableCategory.PUBLIC_OMNI: 24,
+ EnvironmentVariableCategory.INHERITED_VLLM: 20,
+ EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
+ EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
+ EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
+ EnvironmentVariableCategory.INTERNAL: 2,
+ }
+
+ disposition_counts = Counter(
+ item.model_disposition
+ for item in ENVIRONMENT_VARIABLE_INVENTORY.values()
+ if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
+ )
+ assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
+ ModelEnvironmentVariableDisposition.PROMOTE: 33,
+ ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
+ ModelEnvironmentVariableDisposition.EXTERNAL: 0,
+ ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
+ ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
+ }
+
+
+def test_new_public_omni_names_use_project_prefix():
+ """Grandfather legacy names without allowing more prefix exceptions."""
+ legacy_public_names = {
+ "DIFFUSION_ATTENTION_BACKEND",
+ "DIFFUSION_CACHE_ADAPTER",
+ "DIFFUSION_CACHE_BACKEND",
+ "OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
+ "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
+ "OMNI_DIFFUSION_SESSION_STATE_MANAGER",
+ "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
+ "SPEAKER_MAX_UPLOADED",
+ "SPEAKER_SAMPLES_DIR",
+ "VLLM_VIDEO_ASYNC_CHUNK",
+ "VLLM_VIDEO_AUDIO_DELTA_MODE",
+ }
+ public_names = {
+ name
+ for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+ if item.category is EnvironmentVariableCategory.PUBLIC_OMNI
+ }
+
+ assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names
+
+
+def test_statically_resolvable_environment_accesses_are_classified():
+ discovered: set[str] = set()
+ for path in _PACKAGE_ROOT.rglob("*.py"):
+ discovered.update(_environment_accesses(path))
+
+ assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set()
+
+
+def test_model_and_benchmark_inventory_entries_are_still_referenced():
+ """Reject transitional rows after their implementation has disappeared."""
+ package_source = "\n".join(
+ path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE
+ )
+ transitional_categories = {
+ EnvironmentVariableCategory.MODEL_SPECIFIC,
+ EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL,
+ }
+
+ stale_names = {
+ name
+ for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+ if item.category in transitional_categories and name not in package_source
+ }
+
+ assert stale_names == set()
+
+
+def test_environment_scanner_covers_indirection_aliases_membership_and_casing():
+ expected_by_path = {
+ "metrics/definitions.py": {
+ "VLLM_OMNI_BENCH_AUDIO_CHANNELS",
+ "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
+ },
+ "entrypoints/openai/video_stream_envs.py": {
+ "VLLM_VIDEO_ASYNC_CHUNK",
+ "VLLM_VIDEO_AUDIO_DELTA_MODE",
+ },
+ "model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"},
+ "model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"},
+ "distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"},
+ }
+ for relative_path, expected in expected_by_path.items():
+ assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path)
+
+
+def test_generated_server_storage_environment_names_are_classified():
+ from vllm_omni.config.server_settings import FileBackend
+
+ generated_names = {
+ f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}"
+ for field_name in FileBackend.model_fields
+ if field_name != "type"
+ }
+ assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys()
+
+
+def test_public_omni_variables_are_in_the_reference_page():
+ reference = _REFERENCE_PAGE.read_text(encoding="utf-8")
+ missing = {
+ name
+ for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
+ if item.is_public_omni and f"`{name}`" not in reference
+ }
+ assert missing == set()
+
+
+def test_secret_values_are_marked_for_redaction():
+ assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == {
+ "HF_TOKEN",
+ "HUGGINGFACE_HUB_TOKEN",
+ "OPENAI_API_KEY",
+ }
+
+
+def test_collect_env_reports_only_safe_public_omni_values(monkeypatch):
+ from collect_env import get_env_vars
+
+ monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
+ monkeypatch.setenv("HF_TOKEN", "hf-secret")
+ monkeypatch.setenv("OPENAI_API_KEY", "openai-secret")
+ monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica")
+
+ report = get_env_vars()
+
+ assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines()
+ assert "hf-secret" not in report
+ assert "openai-secret" not in report
+ assert "VLLM_OMNI_REPLICA_ID" not in report
+
+
+@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError])
+def test_collect_env_survives_inventory_import_error(monkeypatch, error_type):
+ from collect_env import get_env_vars
+
+ real_import = builtins.__import__
+
+ def fail_inventory_import(name, *args, **kwargs):
+ if name == "vllm_omni.config.environment_variable_inventory":
+ raise error_type("broken vllm-omni installation")
+ return real_import(name, *args, **kwargs)
+
+ monkeypatch.setattr(builtins, "__import__", fail_inventory_import)
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
+
+ assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines()
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 1fb32239d..126b51c18 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -1,15 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the additive structured Omni config."""
from __future__ import annotations
from dataclasses import fields
+from inspect import Parameter, signature
from pathlib import Path
+import msgspec
import pytest
from pydantic import ValidationError
+from pydantic.fields import FieldInfo
from transformers import Qwen3OmniMoeConfig
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.config import omni_config as omni_config_module
@@ -36,10 +45,12 @@ from vllm_omni.config.stage_config import (
PipelineConfig,
StageDeployConfig,
StageExecutionType,
+ StagePipelineConfig,
load_deploy_config,
merge_pipeline_deploy,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -82,10 +93,23 @@ def _from_pipeline_key(
)
+def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
+ config = _from_pipeline_key(
+ "minimax_h3_disaggregated",
+ cli_overrides={"text_encoder_tp_size": 4},
+ )
+
+ assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
+ diffusion = config.stage_by_id(1)
+ assert diffusion.parallel_config.tensor_parallel_size == 1
+ assert diffusion.parallel_config.text_encoder_tp_size == 1
+
+
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
omni_config = _from_pipeline_key("personaplex")
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
+ assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
@@ -95,6 +119,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
+ assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
@@ -122,6 +147,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
engine_args = legacy_stage.yaml_engine_args
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
+ assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
@@ -234,6 +260,36 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
+@pytest.mark.parametrize(
+ ("engine_extras", "cli_overrides"),
+ [
+ ({"requires_full_payload_input": False}, {}),
+ ({}, {"stage_1_requires_full_payload_input": False}),
+ ],
+ ids=["deploy", "stage-cli"],
+)
+def test_from_pipeline_config_rejects_full_payload_input_capability_overrides(
+ engine_extras,
+ cli_overrides,
+):
+ pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+ deploy = DeployConfig(
+ stages=[
+ StageDeployConfig(
+ stage_id=1,
+ engine_extras=engine_extras,
+ )
+ ]
+ )
+
+ with pytest.raises(ValueError, match=r"no structured config owner: requires_full_payload_input"):
+ VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ user_deploy_config=deploy,
+ cli_overrides=cli_overrides,
+ )
+
+
@pytest.mark.parametrize(
("cli_overrides", "stage_id"),
[
@@ -243,9 +299,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
],
ids=["global", "ar-stage", "generation-stage"],
)
-def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
- with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
- _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
+def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
+ stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
+
+ assert stage.cache_config.cache_dtype == "fp8"
+ assert "cache_dtype" in stage.cache_config._omni_explicit_fields
+
+
+@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
+def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
+ pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+ deploy = DeployConfig(
+ stages=[
+ StageDeployConfig(
+ stage_id=0,
+ engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
+ )
+ ]
+ )
+
+ stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+
+ assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
+ assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
@@ -413,10 +489,12 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating
async_config = _from_pipeline_key("qwen3_tts")
assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
assert async_config.stage_by_id(1).custom_process_input_func is None
+ assert async_config.stage_by_id(1).model_config.requires_full_payload_input is True
sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False})
assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
+ assert sync_config.stage_by_id(1).model_config.requires_full_payload_input is True
assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
assert pipeline.get_stage(1).custom_process_input_func is None
@@ -437,6 +515,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
"connector_config",
"runtime_config",
"parallel_config",
+ "compilation_config",
+ "profiler_config",
"quantization_config",
}
assert "diffusion_config" not in public_fields
@@ -455,10 +535,47 @@ def test_runtime_config_fields_match_structured_runtime_scope():
"num_gpus",
"log_level",
"log_stats",
- "profiler_config",
}
+def test_upstream_config_compatible_fields_materialize_mapping_inputs():
+ compilation_config = {"backend": "eager"}
+ profiler_config = {"profiler": "cuda"}
+
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+ compilation_config=compilation_config,
+ profiler_config=profiler_config,
+ )
+
+ assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+ assert stage_config.compilation_config.backend == "eager"
+ assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+ assert stage_config.profiler_config.profiler == "cuda"
+
+
+def test_upstream_config_compatible_fields_validate_mapping_inputs():
+ with pytest.raises(ValidationError):
+ VllmOmniARStageConfig(
+ stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+ profiler_config={"delay_iterations": -1},
+ )
+
+
+def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
+ compilation_config = VllmCompilationConfig(backend="eager")
+ profiler_config = VllmProfilerConfig(profiler="cuda")
+
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
+ compilation_config=compilation_config,
+ profiler_config=profiler_config,
+ )
+
+ assert stage_config.compilation_config is compilation_config
+ assert stage_config.profiler_config is profiler_config
+
+
def test_sub_config_fields_match_structured_scopes():
assert {f.name for f in fields(OmniStageModelConfig)} == {
"model",
@@ -471,16 +588,19 @@ def test_sub_config_fields_match_structured_scopes():
"trust_remote_code",
"dtype",
"attention_backend",
+ "attention_config",
"moe_backend",
"hf_overrides",
"limit_mm_per_prompt",
"interleave_mm_strings",
"media_io_kwargs",
"active_stream_window",
+ "session_mode",
"duplex_max_sessions",
"enable_sleep_mode",
"default_sampling_params",
"subtalker_sampling_params",
+ "silence_ban_frames",
"has_sampling_extra_args",
"custom_voice_dir",
"task_type",
@@ -488,7 +608,6 @@ def test_sub_config_fields_match_structured_scopes():
"enforce_eager",
"max_cudagraph_capture_size",
"enable_flashinfer_autotune",
- "compilation_config",
"enable_multithread_weight_load",
"num_weight_load_threads",
"disable_autocast",
@@ -497,30 +616,27 @@ def test_sub_config_fields_match_structured_scopes():
# legacy engine-args path.
"model_subdir",
"tokenizer_subdir",
+ "requires_full_payload_input",
}
- assert {f.name for f in fields(OmniStageLoadConfig)} == {
+ vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
+ assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
+ assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
"tokenizer",
- "download_dir",
"skip_tokenizer_init",
- "load_format",
"tokenizer_mode",
"config_format",
"skip_mm_profiling",
}
- assert {f.name for f in fields(OmniStageCacheConfig)} == {
- "kv_cache_memory_bytes",
- "gpu_memory_utilization",
- "enable_prefix_caching",
+ assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
+ assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
+ assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
"mamba_ssm_cache_dtype",
}
- assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
- "max_num_seqs",
- "max_num_batched_tokens",
+ assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
+ assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
"max_model_len",
- "enable_chunked_prefill",
- "async_scheduling",
}
assert {f.name for f in fields(OmniStageConnectorConfig)} == {
"async_chunk",
@@ -529,44 +645,104 @@ def test_sub_config_fields_match_structured_scopes():
"output_connectors",
"input_connectors",
}
- assert {f.name for f in fields(OmniStageParallelConfig)} == {
- "pipeline_parallel_size",
- "data_parallel_size",
- "tensor_parallel_size",
- "enable_expert_parallel",
- "world_size",
- }
- assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
- "pipeline_parallel_size",
- "data_parallel_size",
- "tensor_parallel_size",
+ vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
+ assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
+ assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
+ assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
"sequence_parallel_size",
"ulysses_degree",
"ring_degree",
"allgather_degree",
"ulysses_mode",
+ "ulysses_a2a_permute",
"cfg_parallel_size",
"vae_patch_parallel_size",
- "text_encoder_tp_size",
"vae_parallel_mode",
+ "text_encoder_tp_size",
"use_hsdp",
"mask_sp_padding",
"hsdp_shard_size",
"hsdp_replicate_size",
- "enable_expert_parallel",
- "world_size",
}
+def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
+ scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
+ assert scheduler_config.max_num_encoder_input_tokens == 4096
+ assert scheduler_config.encoder_cache_size == 4096
+
+ unresolved_scheduler_config = OmniStageSchedulerConfig()
+ assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
+ assert unresolved_scheduler_config.encoder_cache_size is None
+
+ parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+ assert parallel_config.data_parallel_index == 2
+
+ diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
+ assert diffusion_parallel_config.data_parallel_index == 2
+
+
+@pytest.mark.parametrize(
+ "config_cls",
+ [
+ OmniStageLoadConfig,
+ OmniStageCacheConfig,
+ OmniStageSchedulerConfig,
+ OmniStageParallelConfig,
+ OmniStageDiffusionParallelConfig,
+ ],
+)
+def test_inherited_sub_configs_are_keyword_only(config_cls):
+ assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
+
+ with pytest.raises(TypeError):
+ config_cls(1)
+
+
+@pytest.mark.parametrize(
+ "config_cls",
+ [
+ OmniStageLoadConfig,
+ OmniStageCacheConfig,
+ OmniStageSchedulerConfig,
+ OmniStageParallelConfig,
+ OmniStageDiffusionParallelConfig,
+ ],
+)
+def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
+ config = config_cls()
+
+ uninitialized_fields = [
+ config_field.name
+ for config_field in fields(config)
+ if isinstance(getattr(config, config_field.name), FieldInfo)
+ ]
+ assert uninitialized_fields == []
+
+ msgspec.msgpack.encode(_serialize_stage_config(config))
+
+
+def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
+ omni_config = _from_pipeline_key("qwen3_tts")
+
+ for stage_config in omni_config.stage_configs:
+ payload = {
+ "stage_id": stage_config.stage_id,
+ "stage_config": _serialize_stage_config(stage_config),
+ }
+ msgspec.msgpack.encode(payload)
+
+
def test_diffusion_parallel_config_fields_cover_legacy_surface():
from vllm_omni.diffusion.data import DiffusionParallelConfig
legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
expected_upstream_fields = {"mask_sp_padding"}
+ vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
assert legacy_fields | expected_upstream_fields <= structured_fields
- assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
+ assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
@@ -609,8 +785,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
cfg = OmniStageDiffusionParallelConfig(
- pipeline_parallel_size=2,
- ulysses_degree=2,
+ ulysses_degree=4,
use_hsdp=True,
hsdp_shard_size=-1,
hsdp_replicate_size=2,
@@ -621,10 +796,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
- with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+ with pytest.raises(ValueError, match="not compatible with TP"):
OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
- with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+ with pytest.raises(ValueError, match="not compatible with DP"):
OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
@@ -685,6 +860,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
assert stage.parallel_config.world_size == 2
+def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
+ pipeline = _resolve_pipeline_or_skip("dreamzero")
+ deploy = DeployConfig(
+ stages=[
+ StageDeployConfig(
+ stage_id=0,
+ ulysses_a2a_permute=True,
+ )
+ ]
+ )
+
+ stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+
+ assert isinstance(stage, VllmOmniDiffusionStageConfig)
+ assert stage.parallel_config.ulysses_a2a_permute is True
+
+
def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
@@ -779,6 +971,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
+def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
+ deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
+ deploy_path.write_text(
+ "\n".join(
+ [
+ "pipeline: dreamzero",
+ "async_chunk: false",
+ "stages:",
+ " - stage_id: 0",
+ " extras:",
+ " ltx2_use_conv_vae: true",
+ ]
+ )
+ )
+
+ stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+ assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
+def test_stage_override_routes_ltx2_conv_vae_extra():
+ stage = _from_pipeline_key(
+ "dreamzero",
+ cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
+ ).stage_by_id(0)
+
+ assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
@@ -819,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
def test_from_pipeline_config_accepts_pre_resolved_pipeline():
- resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant")
+ resolved_pipeline = PipelineConfig(
+ model_type="callable_resolved_variant",
+ stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
+ )
omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
@@ -852,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p
pipeline = PipelineConfig(
model_type="pipeline_with_default",
default_deploy_config_name=default_name,
+ stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
)
loaded_paths = []
@@ -966,7 +1191,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
+ from vllm_omni.platforms import current_omni_platform
+
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
+ monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
diffusion_attention_backend="flash_attn",
@@ -976,6 +1204,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
kv_cache_skip_layers=[2],
static_lora_scale=0.25,
diffusion_kv_mode="paged_scheduler",
+ diffusion_kv_max_rows_per_request=2,
diffusers_load_kwargs=None,
diffusers_call_kwargs=None,
)
@@ -992,7 +1221,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
assert cfg.diffusers_call_kwargs == {}
-def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
+def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
+ from vllm_omni.platforms import current_omni_platform
+
+ monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
deploy_path.write_text(
"\n".join(
@@ -1004,6 +1236,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
" diffusion_attention_backend: flash_attn",
" fa_deterministic: true",
" diffusion_kv_mode: paged_scheduler",
+ " diffusion_kv_max_rows_per_request: 2",
]
)
)
@@ -1017,6 +1250,36 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
assert stage.diffusion_config.fa_deterministic is True
assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+ assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
+
+
+def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
+ from vllm_omni.platforms import current_omni_platform
+
+ monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+ deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
+ deploy_path.write_text(
+ "\n".join(
+ [
+ "pipeline: dreamzero",
+ "async_chunk: false",
+ "stages:",
+ " - stage_id: 0",
+ " diffusion_attention_backend: FASTVIDEO_VSA",
+ " fastvideo_vsa_topk: 96",
+ ]
+ )
+ )
+
+ stage = _from_pipeline_key(
+ "dreamzero",
+ deploy_config_path=str(deploy_path),
+ ).stage_by_id(0)
+
+ attention_config = stage.diffusion_config.diffusion_attention_config
+ assert attention_config.default is not None
+ assert attention_config.default.backend == "FASTVIDEO_VSA"
+ assert attention_config.default.fastvideo_vsa_topk == 96
def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
@@ -1059,6 +1322,7 @@ def test_diffusion_config_field_classification_covers_current_fields():
"prompt_embed_cache_size",
"diffusion_kv_cache_dtype",
"diffusion_kv_mode",
+ "diffusion_kv_max_rows_per_request",
} <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
assert {
"revision",
@@ -1077,3 +1341,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
assert cfg.quantization_config == quantization_config
+
+
+def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+ quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
+ cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
+ quantization_config=quantization_config,
+ )
+
+ # Exercise the terminal construction performed by the future typed startup
+ # path without probing ports or loading remote model metadata.
+ monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+ monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
+ cfg.enrich_config()
+
+ assert cfg.quantization_config is not None
+ assert cfg.quantization_config.get_name() == "int8"
diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py
index 4bc19b116..88109e052 100644
--- a/tests/config/test_pipeline_registry.py
+++ b/tests/config/test_pipeline_registry.py
@@ -1,16 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for out of tree registration to OMNI_PIPELINES."""
import pytest
from transformers import PretrainedConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
-from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
+from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
+ return PipelineConfig(
+ model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+ )
+
+
@pytest.fixture
def custom_resolver():
"""Build a reusable custom resolver for PipelineConfigs."""
@@ -22,7 +28,7 @@ def custom_resolver():
def custom_resolver(
hf_config: CustomConfigType,
) -> PipelineConfig:
- return PipelineConfig(model_type="resolved_type")
+ return build_fake_pipeline_config("resolved_type")
return custom_resolver
@@ -30,7 +36,7 @@ def custom_resolver():
def test_register_pipeline_config(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
new_model_type = "new_model_type"
- pipe_cfg = PipelineConfig(model_type=new_model_type)
+ pipe_cfg = build_fake_pipeline_config(new_model_type)
assert new_model_type not in OMNI_PIPELINES
register_pipeline(pipe_cfg)
assert new_model_type in OMNI_PIPELINES
@@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
new_model_type = "new_model_type"
unused_model_type = "foo"
- pipe_cfg = PipelineConfig(model_type=unused_model_type)
+ pipe_cfg = build_fake_pipeline_config(unused_model_type)
assert new_model_type not in OMNI_PIPELINES
assert unused_model_type not in OMNI_PIPELINES
@@ -66,3 +72,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
with pytest.raises(ValueError):
register_pipeline(custom_resolver)
+
+
+def test_minimax_h3_disaggregation_is_explicit_opt_in():
+ assert "minimax_h3" not in OMNI_PIPELINES
+ pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
+ assert isinstance(pipeline, PipelineConfig)
+ assert pipeline.model_type == "minimax_h3_disaggregated"
diff --git a/tests/conftest.py b/tests/conftest.py
index b453bbcbf..fa55a23da 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,7 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Root pytest entrypoint for the vLLM-Omni test suite.
-- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports.
+- `tests/conftest.py` stays thin: plugin registration only.
- Importable utilities live under `tests/helpers/`.
- Helper unit tests live under `tests/helpers/tests/`.
- Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`.
@@ -19,57 +22,15 @@ bootstrap_vllm_layer_custom_op_modules()
pytest_plugins = (
"tests.helpers.fixtures.config",
- "tests.helpers.fixtures.env",
+ "tests.helpers.fixtures.clean",
"tests.helpers.fixtures.log",
- "tests.helpers.fixtures.run_args",
+ "tests.helpers.fixtures.media",
+ "tests.helpers.fixtures.pytest_collection",
+ "tests.helpers.fixtures.pytest_run_args",
"tests.helpers.fixtures.runtime",
- "tests.helpers.fixtures.speaker_cache",
)
def pytest_terminal_summary(terminalreporter, exitstatus, config):
# Marker for Buildkite log folding before pytest summary lines.
terminalreporter.write_line("--- Running Summary")
-
-
-# Backward-compatible lazy re-exports.
-# (Many tests still import from `tests.conftest`; migrate these imports to `tests.helpers.*` over time.)
-# Keep these lazy so conftest import does not trigger heavy helper dependencies.
-_ASSERTION_EXPORT_NAMES = (
- "assert_audio_speech_response",
- "assert_diffusion_response",
- "assert_http_error",
- "assert_image_diffusion_response",
- "assert_image_valid",
- "assert_omni_response",
- "assert_video_diffusion_response",
- "assert_video_valid",
-)
-_MEDIA_EXPORT_NAMES = (
- "convert_audio_bytes_to_text",
- "convert_audio_file_to_text",
- "cosine_similarity_text",
- "decode_b64_image",
- "generate_synthetic_audio",
- "generate_synthetic_image",
- "generate_synthetic_video",
-)
-_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",)
-_RUNTIME_EXPORT_NAMES = (
- "DiffusionResponse",
- "HttpResponse",
- "OmniResponse",
- "OmniRunner",
- "OmniRunnerHandler",
- "OmniServer",
- "OmniServerParams",
- "OmniServerStageCli",
- "OpenAIClientHandler",
- "dummy_messages_from_mix_data",
-)
-_LAZY_EXPORT_MODULES = {
- **{name: "tests.helpers.assertions" for name in _ASSERTION_EXPORT_NAMES},
- **{name: "tests.helpers.media" for name in _MEDIA_EXPORT_NAMES},
- **{name: "tests.helpers.stage_config" for name in _STAGE_CONFIG_EXPORT_NAMES},
- **{name: "tests.helpers.runtime" for name in _RUNTIME_EXPORT_NAMES},
-}
diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py
new file mode 100644
index 000000000..d9fbc452b
--- /dev/null
+++ b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py
@@ -0,0 +1,89 @@
+"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
+
+``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
+handing control to upstream, because upstream ``Scheduler.schedule()`` raises
+``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
+admits in a finished state -- which kills the stage's engine core, not just
+the request.
+
+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.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_POLICIES = [
+ pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
+ pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
+]
+
+
+class _StubRequest:
+ """Minimal ``Request`` stub with the surface the sweep exercises.
+
+ Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
+ relies on so the same stub works under both scheduling policies.
+ """
+
+ def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
+ self.request_id = request_id
+ self.status = status
+ self.priority = 0
+ self.arrival_time = arrival_time
+
+ def __lt__(self, other: _StubRequest) -> bool:
+ return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
+
+
+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.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_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, skipped_waiting=[aborted])
+ scheduler.chunk_transfer_adapter = None
+ scheduler.input_coordinator = None
+ scheduler.num_waiting_for_streaming_input = 0
+ scheduler._process_pending_omni_inputs = lambda model_mode: None
+ scheduler._resync_streaming_input_counter = lambda: None
+ scheduler._should_defer_waiting_admission = lambda: False
+ scheduler._restore_omni_wait_queues = lambda: None
+ scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
+ scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
+ scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
+
+ sentinel = object()
+
+ def upstream_schedule(_self, _throttle_prefills: bool = False):
+ assert list(scheduler.skipped_waiting) == []
+ return sentinel
+
+ monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
+
+ assert OmniARScheduler.schedule(scheduler) is sentinel
diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py
index fb92414ef..5c1290fc4 100644
--- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py
+++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py
@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
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
return sched
@@ -141,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None:
session.num_computed_tokens = 6
session.num_output_placeholders = 1
session.num_in_flight_tokens = 1
+ session.async_tokens_to_discard = 1
_replace_streaming_session(session)
assert session.num_stale_output_tokens == 1
+ assert session.async_tokens_to_discard == 1
sched = _make_drain_sched(session)
assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped
assert session.num_stale_output_tokens == 0
+ assert session.async_tokens_to_discard == 0
assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives
diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py
index ca80323a7..a1aa7351f 100644
--- a/tests/core/sched/test_omni_ar_scheduler_streaming.py
+++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for Omni AR streaming-session async placeholder handling."""
@@ -18,21 +18,26 @@ from vllm.sampling_params import SamplingParams
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 vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
# isort: on
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
+def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
sched = OmniARScheduler.__new__(OmniARScheduler)
sched._new_prompt_len_snapshot = {}
- sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+ sched.vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode),
+ )
sched.num_waiting_for_streaming_input = 0
sched.log_stats = False
sched.chunk_transfer_adapter = None
sched.skipped_waiting = set()
+ sched._free_request_blocks = MagicMock()
+ sched.encoder_cache_manager = MagicMock()
+ sched._inflight_prefills = set()
return sched
@@ -61,6 +66,9 @@ def _run_resumable_segment_stop(
session: Request,
*,
session_finished: bool = False,
+ handle_stopped=None,
+ chunk_transfer_adapter=None,
+ inter_stage_output=None,
):
sched = MagicMock()
sched.requests = {session.request_id: session}
@@ -72,11 +80,14 @@ def _run_resumable_segment_stop(
return [42], True
sched._update_request_with_output.side_effect = stop_request
- sched._handle_stopped_request.return_value = session_finished
# vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured
# MagicMock iterates empty and fails to unpack at the call site.
sched._free_request.return_value = (None, None)
- sched.chunk_transfer_adapter = None
+ if handle_stopped is None:
+ sched._handle_stopped_request.return_value = session_finished
+ else:
+ sched._handle_stopped_request.side_effect = handle_stopped
+ sched.chunk_transfer_adapter = chunk_transfer_adapter
sched.running = [session]
sched.waiting_for_transfer_free = set()
sched.transfer_triggered_requests = set()
@@ -84,6 +95,7 @@ def _run_resumable_segment_stop(
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
@@ -102,6 +114,7 @@ def _run_resumable_segment_stop(
model_runner_output.cudagraph_stats = None
model_runner_output.req_id_to_index = {session.request_id: 0}
model_runner_output.routed_experts = None
+ model_runner_output.inter_stage_outputs = [inter_stage_output] if inter_stage_output is not None else None
return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
@@ -157,6 +170,44 @@ def test_update_from_output_settles_in_flight_tokens() -> None:
assert session.num_in_flight_tokens == 0
+def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> None:
+ """The old segment's transfer must not observe the next segment's reset.
+
+ A queued streaming update can replace the same mutable Request while the
+ stop output is still being handled. The connector needs the confirmed
+ token count from before that replacement to avoid dropping the boundary.
+ """
+ session = _make_request()
+ session.status = RequestStatus.RUNNING
+ session.resumable = True
+ session.num_computed_tokens = 26
+ adapter = MagicMock()
+ adapter._confirmed_num_computed_tokens.return_value = 26
+ inter_stage_output = {"codes": {"audio": [7]}}
+
+ def replace_with_next_segment(request: Request) -> bool:
+ request.num_computed_tokens = 0
+ request._omni_segment_generation = 1
+ request.status = RequestStatus.WAITING
+ return False
+
+ _run_resumable_segment_stop(
+ session,
+ handle_stopped=replace_with_next_segment,
+ chunk_transfer_adapter=adapter,
+ inter_stage_output=inter_stage_output,
+ )
+
+ adapter.save_async.assert_called_once_with(
+ inter_stage_output,
+ session,
+ True,
+ new_token_ids=[42],
+ confirmed_num_computed_tokens=26,
+ segment_generation=0,
+ )
+
+
def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
"""A decode step that neither stops nor carries an inter-stage payload.
@@ -181,6 +232,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
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
@@ -207,6 +259,171 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
sched.chunk_transfer_adapter.save_async.assert_not_called()
+def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
+ """Native duplex + async: the next unit is usually already queued when the
+ current one stops, so _handle_stopped_request applies the update and the
+ stop site below it fences the same in-flight decode a second time. An
+ accumulated residue outlives the drain and swallows the next unit's
+ listen/speak, which the client sees as silence until its timeout.
+ """
+ session = _make_request()
+ session.status = RequestStatus.RUNNING
+ session.resumable = True
+ session.append_output_token_ids([7])
+ session.num_computed_tokens = 4
+ session.num_output_placeholders = 1
+ # This step's frame (settled to 0 by update_from_output) plus one extra
+ # async decode still in flight.
+ session.num_in_flight_tokens = 2
+
+ sched = _make_scheduler(stage_id=0)
+ sched._enqueue_waiting_request = MagicMock()
+
+ def handle_stopped(request: Request) -> bool:
+ sched._update_request_as_session(request, _make_update([10, 20]))
+ return False
+
+ _run_resumable_segment_stop(session, handle_stopped=handle_stopped)
+
+ # Exactly the one unreported decode, so the drain reaches zero before the
+ # new segment's first frame arrives.
+ assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
+
+
+def test_stale_async_frame_is_dropped_before_output_processing() -> None:
+ session = _make_request()
+ session.status = RequestStatus.RUNNING
+ session.num_in_flight_tokens = 2
+ session.num_computed_tokens = session.num_tokens + 1
+ session.num_output_placeholders = 1
+ session.async_tokens_to_discard = 1
+ session.sampling_params = SimpleNamespace(num_logprobs=1)
+ num_computed_tokens = session.num_computed_tokens
+ num_output_placeholders = session.num_output_placeholders
+
+ sched = MagicMock()
+ sched.requests = {session.request_id: session}
+ sched.perf_metrics = None
+ sched.structured_output_manager.should_advance.return_value = False
+
+ def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
+ request.async_tokens_to_discard = 0
+ return token_ids, False
+
+ sched._update_request_with_output.side_effect = discard_stale_output
+ sched._process_kv_transfer_trigger.return_value = False
+ sched.chunk_transfer_adapter = MagicMock()
+ sched.running = [session]
+ 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
+
+ scheduler_output = MagicMock(spec=SchedulerOutput)
+ scheduler_output.num_scheduled_tokens = {session.request_id: 1}
+ scheduler_output.scheduled_spec_decode_tokens = {}
+ scheduler_output.num_invalid_spec_tokens = 0
+
+ model_runner_output = MagicMock(spec=ModelRunnerOutput)
+ 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.multimodal_outputs = None
+ model_runner_output.inter_stage_outputs = [{"hidden": object()}]
+ 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 = {session.request_id: 0}
+ model_runner_output.routed_experts = None
+
+ OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+ assert session.async_tokens_to_discard == 0
+ assert session.status == RequestStatus.RUNNING
+ assert session.num_computed_tokens == num_computed_tokens
+ assert session.num_output_placeholders == num_output_placeholders
+ sched.chunk_transfer_adapter.save_async.assert_not_called()
+
+ session.sampling_params = SamplingParams(max_tokens=8)
+ next_payload = {"hidden": object()}
+ sched._update_request_with_output.side_effect = None
+ sched._update_request_with_output.return_value = ([43], False)
+ model_runner_output.sampled_token_ids = [[43]]
+ model_runner_output.inter_stage_outputs = [next_payload]
+
+ OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+ sched.chunk_transfer_adapter.save_async.assert_called_once_with(
+ next_payload,
+ session,
+ False,
+ new_token_ids=[43],
+ confirmed_num_computed_tokens=None,
+ )
+
+
+def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None:
+ request = _make_request()
+ request.status = RequestStatus.RUNNING
+ request.num_in_flight_tokens = 1
+ request.num_output_placeholders = 0
+ request.async_tokens_to_discard = 1
+ request.num_stale_output_tokens = 0
+
+ sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler)
+ sched.requests = {request.request_id: request}
+ sched.perf_metrics = None
+ sched.chunk_transfer_adapter = None
+ sched.connector = None
+ sched.ec_connector = None
+ sched.waiting_for_transfer_free = set()
+ sched.transfer_triggered_requests = set()
+ sched.active_kv_transfers = set()
+ sched.pending_stop_after_extraction = set()
+ sched.finished_req_ids_dict = {}
+ sched._new_prompt_len_snapshot = {}
+ sched.kv_cache_manager = MagicMock()
+ sched.kv_cache_manager.take_events.return_value = None
+ sched._remove_stopped_requests_from_queues = MagicMock()
+ sched._handle_failed_kv_load_outputs = MagicMock(return_value=[])
+ sched._cleanup_kv_tracking = MagicMock()
+ sched._aggregate_kv_connector_stats = MagicMock(return_value=None)
+ sched._publish_kv_cache_events = MagicMock()
+ sched._attach_finished_request_sets = MagicMock()
+ sched._attach_scheduler_stats = MagicMock()
+ sched._capture_omni_connector_output = MagicMock()
+
+ scheduler_output = MagicMock(spec=SchedulerOutput)
+ scheduler_output.num_scheduled_tokens = {request.request_id: 1}
+ scheduler_output.scheduled_spec_decode_tokens = {}
+ scheduler_output.num_invalid_spec_tokens = 0
+
+ model_runner_output = MagicMock(spec=ModelRunnerOutput)
+ 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.multimodal_outputs = None
+ model_runner_output.inter_stage_outputs = 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
+
+ OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
+
+ assert request.async_tokens_to_discard == 0
+ assert request.num_output_placeholders == 0
+ assert request._output_token_ids == []
+
+
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
sched = _make_scheduler(stage_id=0)
session = _make_request()
@@ -276,29 +493,33 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
"meta": {"turn_eos_token_id": 99},
}
assert session.status == RequestStatus.WAITING
+ sched._free_request_blocks.assert_called_once_with(session)
+ sched.encoder_cache_manager.free.assert_called_once_with(session)
-def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
+def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_watermark() -> None:
sched = _make_scheduler(stage_id=1)
+ session = _make_request()
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=False,
segment_finished_requests=set(),
+ requests_num_chunks_sent={session.external_req_id: 59},
)
- session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.prompt_token_ids = [0] * 59
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session.num_prompt_tokens = 59
session.num_computed_tokens = 59
+ session.num_in_flight_tokens = 2
update = _make_update([0] * 10)
update.additional_information = None
update.model_intermediate_buffer = {
"ids": {"tts": list(range(8))},
"hidden_states": {"tts": [[0.0]] * 8},
"meta": {
- "replace_streaming_prompt": True,
"next_stage_prompt_len": 10,
+ "replace_streaming_prompt": True,
},
}
@@ -308,6 +529,110 @@ def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None
assert list(session._all_token_ids) == [0] * 10
assert session.num_prompt_tokens == 10
assert session.num_computed_tokens == 0
+ assert session.num_stale_output_tokens == 2
assert session.additional_information is None
assert session.model_intermediate_buffer == update.model_intermediate_buffer
assert session.status == RequestStatus.WAITING
+ assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+ sched._free_request_blocks.assert_called_once_with(session)
+ sched.encoder_cache_manager.free.assert_called_once_with(session)
+
+
+def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
+ sched = _make_scheduler(stage_id=1)
+ session = _make_request()
+ session.external_req_id = "external-ar-streaming-test"
+ session.num_in_flight_tokens = 2
+ # _update_request_as_session() may have already fenced this frame before
+ # the connector marks the explicit replacement ready.
+ session.num_stale_output_tokens = 2
+ session.num_output_placeholders = 2
+ session.spec_token_ids = [-1, -1]
+ sched.requests = {session.request_id: session}
+ sched._inflight_prefills.add(session)
+ sched.chunk_transfer_adapter = SimpleNamespace(
+ replaced_streaming_prompt_ids={session.request_id},
+ requests_with_ready_chunks={session.request_id},
+ requests_num_chunks_sent={session.external_req_id: 4090},
+ )
+
+ sched._reset_ready_async_chunk_replacements()
+ sched._reset_ready_async_chunk_replacements()
+
+ sched._free_request_blocks.assert_called_once_with(session)
+ sched.encoder_cache_manager.free.assert_called_once_with(session)
+ assert session not in sched._inflight_prefills
+ assert session.num_stale_output_tokens == 2
+ assert session.num_output_placeholders == 0
+ assert session.spec_token_ids == []
+ assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
+ assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
+ assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+
+
+def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
+ """A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
+ session = _make_request()
+ session.resumable = True
+ session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+ def queue(*requests):
+ result = MagicMock()
+ result.requests = set(requests)
+ result.add_request.side_effect = result.requests.add
+ result.remove_requests.side_effect = result.requests.difference_update
+ return result
+
+ sched = OmniARScheduler.__new__(OmniARScheduler)
+ sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
+ sched.running = []
+ sched.waiting = queue()
+ sched.skipped_waiting = queue(session)
+ sched.num_waiting_for_streaming_input = 1
+ sched.chunk_transfer_adapter = SimpleNamespace(
+ receives_chunks=True,
+ segment_finished_requests={session.request_id},
+ )
+
+ sched._resume_downstream_chunk_receiver(session)
+ sched._remove_stopped_requests_from_queues(set(), {session})
+
+ assert session.status == RequestStatus.WAITING
+ assert session in sched.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
+
+
+@pytest.mark.parametrize(
+ ("receives_chunks", "session_mode"),
+ [
+ (False, "duplex"),
+ (True, "turn"),
+ ],
+)
+def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
+ receives_chunks: bool,
+ session_mode: str,
+) -> None:
+ """Only duplex connector-driven receivers resume without an update."""
+ session = _make_request()
+ session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+
+ sched = OmniARScheduler.__new__(OmniARScheduler)
+ sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
+ sched.num_waiting_for_streaming_input = 1
+ sched.chunk_transfer_adapter = SimpleNamespace(
+ receives_chunks=receives_chunks,
+ segment_finished_requests={session.request_id},
+ )
+ sched.skipped_waiting = MagicMock()
+ sched._enqueue_waiting_request = MagicMock()
+
+ sched._resume_downstream_chunk_receiver(session)
+
+ 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.skipped_waiting.remove_requests.assert_not_called()
+ sched._enqueue_waiting_request.assert_not_called()
diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py
index 27087d253..0e0ed04ad 100644
--- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py
+++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py
@@ -76,6 +76,9 @@ def _make_scheduler(
else:
hf_config.voxcpm2_runtime_config = runtime_config
sched.vllm_config = SimpleNamespace(model_config=model_config)
+ # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
+ # abort sweep at the top of ``schedule()`` reads it.
+ sched.skipped_waiting = _MockQueue()
return sched
diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py
index 0f654715a..3e1b53afe 100644
--- a/tests/core/sched/test_omni_generation_scheduler_update_session.py
+++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py
@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
sched.recompute_kv_load_failures = False
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
@@ -410,12 +411,13 @@ class TestRealignRequestStatusToQueues:
assert stale.status == RequestStatus.RUNNING
assert clean.status == RequestStatus.RUNNING
- def test_finished_request_is_skipped(self) -> None:
+ def test_non_resumable_finished_request_is_skipped(self) -> None:
"""Already-finished requests must not be touched -- they may
have legitimate finished statuses (FINISHED_STOPPED etc.) that
a status flip would corrupt.
"""
req = _make_request(request_id="req-finished")
+ req.resumable = False
req.status = RequestStatus.FINISHED_STOPPED
sched = _RealignSchedulerStub(
diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py
new file mode 100644
index 000000000..36b084f7c
--- /dev/null
+++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py
@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
+
+Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
+non-empty scheduled step and, at the top of ``update_from_output()``,
+advances ``processed_step_seq`` and drains the deferred frees. The Omni
+schedulers each carried only half of that contract:
+
+- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
+ advances) but its reimplemented ``update_from_output()`` never advanced
+ ``processed_step_seq`` nor drained -> deferred blocks accumulate without
+ bound under async/PP KV-consumer configurations.
+- ``OmniGenerationScheduler`` reimplements both: the fast-path
+ ``schedule()`` never advanced ``sched_step_seq`` (so finished requests
+ freed blocks an in-flight GPU step could still be writing) and its
+ ``update_from_output()`` never drained.
+
+These tests pin the restored contract at the call boundary; the drain and
+free logic itself is upstream-tested.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import SchedulingPolicy
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_update_sched(scheduler_cls) -> MagicMock:
+ """MagicMock scheduler driving the real ``update_from_output`` with an
+ empty batch; only the fence state carries real values."""
+ sched = MagicMock()
+ sched.requests = {}
+ sched.perf_metrics = None
+ sched.chunk_transfer_adapter = None
+ sched.defer_block_free = True
+ sched.processed_step_seq = 0
+ return sched
+
+
+def _empty_frame(*, total_num_scheduled_tokens: int):
+ scheduler_output = MagicMock(spec=SchedulerOutput)
+ scheduler_output.num_scheduled_tokens = {}
+ scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
+ model_runner_output = MagicMock(spec=ModelRunnerOutput)
+ model_runner_output.sampled_token_ids = []
+ 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
+ return scheduler_output, model_runner_output
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
+ sched = _make_update_sched(scheduler_cls)
+ scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+ scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+ assert sched.processed_step_seq == 1
+ sched._drain_deferred_frees.assert_called_once_with()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
+ """0-token steps do not advance the fence upstream (their seq was never
+ advanced in schedule()); draining on them would run ahead of the GPU."""
+ sched = _make_update_sched(scheduler_cls)
+ scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
+
+ scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+ assert sched.processed_step_seq == 0
+ sched._drain_deferred_frees.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
+ sched = _make_update_sched(scheduler_cls)
+ sched.defer_block_free = False
+ scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
+
+ scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
+
+ assert sched.processed_step_seq == 0
+ sched._drain_deferred_frees.assert_not_called()
+
+
+def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
+ """MagicMock scheduler that drives the real generation fast-path
+ ``schedule()`` over one running one-shot request."""
+ request = Request(
+ request_id="req-fence",
+ prompt_token_ids=[1, 2, 3],
+ sampling_params=SamplingParams(max_tokens=4),
+ pooling_params=None,
+ arrival_time=100.0,
+ block_hasher=None,
+ )
+ request.status = RequestStatus.RUNNING
+
+ sched = MagicMock()
+ sched.max_num_scheduled_tokens = 8
+ sched._pause_state = PauseState.UNPAUSED
+ sched.requests = {request.request_id: request}
+ sched.running = [request]
+ sched.waiting = [] # falsy: the waiting loop never runs
+ sched.policy = SchedulingPolicy.FCFS
+ sched.chunk_transfer_adapter = None
+ sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
+ sched.num_lookahead_tokens = 0
+ sched.log_stats = False
+ sched.kv_cache_config = MagicMock(kv_cache_groups=[])
+ sched.use_v2_model_runner = False
+ sched.prev_step_scheduled_req_ids = set()
+ sched.needs_kv_cache_zeroing = False
+ sched.connector = None
+ sched.ec_connector = None
+ sched.defer_block_free = defer_block_free
+ sched.sched_step_seq = 0
+ # last_sched_seq stamping order: _update_after_schedule must observe the
+ # already-advanced fence value, as upstream schedule() guarantees.
+ sched.seq_at_update_after_schedule = None
+
+ def _capture(_scheduler_output):
+ sched.seq_at_update_after_schedule = sched.sched_step_seq
+
+ sched._update_after_schedule.side_effect = _capture
+ return sched, request
+
+
+def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
+ sched, request = _make_fast_path_sched(defer_block_free=True)
+
+ OmniGenerationScheduler.schedule(sched)
+
+ assert sched.sched_step_seq == 1
+ assert sched.seq_at_update_after_schedule == 1
+
+
+def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
+ sched, request = _make_fast_path_sched(defer_block_free=False)
+
+ OmniGenerationScheduler.schedule(sched)
+
+ assert sched.sched_step_seq == 0
+
+
+def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None:
+ sched, request = _make_fast_path_sched(defer_block_free=False)
+ calls: list[str] = []
+ sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep")
+ sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode)
+ sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync")
+
+ OmniGenerationScheduler.schedule(sched)
+
+ assert calls == ["sweep", "generation", "sweep", "resync"]
diff --git a/tests/core/sched/test_omni_sched_ec_request_finish.py b/tests/core/sched/test_omni_sched_ec_request_finish.py
new file mode 100644
index 000000000..81b6577fd
--- /dev/null
+++ b/tests/core/sched/test_omni_sched_ec_request_finish.py
@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
+
+Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
+BEFORE the encoder cache is freed (so the connector can inspect per-request
+state), honors its delayed-free result, and returns the generated
+``ec_transfer_params``; ``update_from_output()`` then carries them into the
+``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
+always returned ``None``; both Omni schedulers discarded the second tuple
+item, so a remote encoder-cache handle could never reach the frontend.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
+
+
+def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
+ """Minimal real-`_free_request` harness; `order` records the EC hook
+ firing relative to the encoder-cache free."""
+ order = MagicMock()
+ sched = OmniARScheduler.__new__(OmniARScheduler)
+ sched._omits_kv_transfer_cache = {}
+ sched._connector_finished = lambda request: (False, None)
+ sched.ec_connector = ec_connector
+ if ec_connector is not None:
+ order.attach_mock(ec_connector.request_finished, "request_finished")
+ sched.encoder_cache_manager = MagicMock()
+ order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
+ sched.finished_req_ids = set()
+ sched._new_prompt_len_snapshot = {}
+ sched.finished_req_ids_dict = None
+ sched._should_transfer_kv_for_request = lambda req_id: False
+ sched._free_blocks = MagicMock()
+ sched._free_input_coordinator_request = MagicMock()
+ sched.chunk_transfer_adapter = None
+ return sched, order
+
+
+class _FakeFinishedRequest:
+ def __init__(self, request_id: str) -> None:
+ self.request_id = request_id
+
+ def is_finished(self) -> bool:
+ return True
+
+
+def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
+ ec_connector = MagicMock()
+ ec_connector.request_finished.return_value = (False, EC_PARAMS)
+ sched, order = _make_free_request_sched(ec_connector=ec_connector)
+ request = _FakeFinishedRequest("req-ec")
+
+ kv_params, ec_params = sched._free_request(request)
+
+ assert ec_params == EC_PARAMS
+ ec_connector.request_finished.assert_called_once_with(request)
+ hook_calls = [name for name, *_ in order.mock_calls]
+ assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
+ sched._free_blocks.assert_called_once_with(request)
+
+
+def test_free_request_honors_ec_delayed_free() -> None:
+ ec_connector = MagicMock()
+ ec_connector.request_finished.return_value = (True, EC_PARAMS)
+ sched, _ = _make_free_request_sched(ec_connector=ec_connector)
+ request = _FakeFinishedRequest("req-ec-delay")
+
+ kv_params, ec_params = sched._free_request(request)
+
+ assert ec_params == EC_PARAMS
+ sched._free_blocks.assert_not_called()
+
+
+def test_free_request_without_ec_connector_returns_none() -> None:
+ sched, _ = _make_free_request_sched(ec_connector=None)
+ request = _FakeFinishedRequest("req-no-ec")
+
+ kv_params, ec_params = sched._free_request(request)
+
+ assert ec_params is None
+ sched._free_blocks.assert_called_once_with(request)
+
+
+def _make_finishing_request() -> Request:
+ request = Request(
+ request_id="req-ec-finish",
+ prompt_token_ids=[1, 2, 3],
+ sampling_params=SamplingParams(max_tokens=4),
+ pooling_params=None,
+ arrival_time=100.0,
+ block_hasher=None,
+ )
+ request.status = RequestStatus.RUNNING
+ return request
+
+
+def _make_finish_sched(request: Request) -> MagicMock:
+ """MagicMock scheduler driving the real ``update_from_output`` through a
+ finishing request; ``_free_request`` reports EC transfer params."""
+ sched = 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._process_kv_transfer_trigger.return_value = False
+ sched._maybe_decode_pooling_output.return_value = None
+ sched._handle_stopped_request.return_value = True
+ sched._free_request.return_value = (None, EC_PARAMS)
+ sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+ sched.chunk_transfer_adapter = None
+ sched.finished_req_ids_dict = None
+ sched._new_prompt_len_snapshot = {}
+ return sched
+
+
+def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
+ request = _make_finishing_request()
+ sched = _make_finish_sched(request)
+ sched._update_request_with_output.return_value = ([42], True)
+
+ engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+ (eco,) = engine_core_outputs[request.client_index].outputs
+ assert eco.ec_transfer_params == EC_PARAMS
+ sched._free_request.assert_called_once_with(request)
+
+
+def _finish_frame_parts(request: Request):
+ scheduler_output = MagicMock(spec=SchedulerOutput)
+ 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 = MagicMock(spec=ModelRunnerOutput)
+ 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
+ return scheduler_output, model_runner_output
+
+
+def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
+ request = _make_finishing_request()
+ # One-shot generation finish: the whole prompt is computed.
+ request.num_computed_tokens = len(request.prompt_token_ids)
+ sched = _make_finish_sched(request)
+
+ engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
+
+ (eco,) = engine_core_outputs[request.client_index].outputs
+ assert eco.ec_transfer_params == EC_PARAMS
+ sched._free_request.assert_called_once_with(request)
diff --git a/tests/core/sched/test_omni_sched_prefill_stats_finalize.py b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py
new file mode 100644
index 000000000..d436367db
--- /dev/null
+++ b/tests/core/sched/test_omni_sched_prefill_stats_finalize.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
+
+Upstream v0.28 finalizes prefill statistics with
+``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
+before emitting them, which computes ``num_cache_creation_tokens`` (tokens
+computed and written to the prefix cache). Both Omni schedulers previously
+emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
+# before Request / RequestStatus are bound in this module.
+# isort: off
+import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_request() -> Request:
+ request = Request(
+ request_id="req-prefill-stats",
+ prompt_token_ids=[1, 2, 3],
+ sampling_params=SamplingParams(max_tokens=4),
+ pooling_params=None,
+ arrival_time=100.0,
+ block_hasher=None,
+ )
+ request.status = RequestStatus.RUNNING
+ # A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
+ request.prefill_stats.set(
+ num_prompt_tokens=3,
+ num_local_cached_tokens=1,
+ num_external_cached_tokens=0,
+ )
+ return request
+
+
+def _make_sched(request: Request) -> MagicMock:
+ sched = 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._process_kv_transfer_trigger.return_value = False
+ sched._maybe_decode_pooling_output.return_value = None
+ sched._handle_stopped_request.return_value = True
+ # The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
+ sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
+
+ def _free_and_invalidate(req):
+ # Lifecycle guard: freeing releases the KV blocks, after which the
+ # manager reports 0 cached tokens. If finalize ran after
+ # _free_request (the reviewed bug), creation count would compute
+ # against 0 and the assertions below would see 0, not 2.
+ sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
+ return (None, None)
+
+ sched._free_request.side_effect = _free_and_invalidate
+ sched.chunk_transfer_adapter = None
+ sched.finished_req_ids_dict = None
+ sched._new_prompt_len_snapshot = {}
+ return sched
+
+
+def _frame(request: Request):
+ scheduler_output = MagicMock(spec=SchedulerOutput)
+ 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 = MagicMock(spec=ModelRunnerOutput)
+ 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
+ return scheduler_output, model_runner_output
+
+
+def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+ request = _make_request()
+ sched = _make_sched(request)
+ sched._update_request_with_output.return_value = ([42], True)
+
+ engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
+
+ (eco,) = engine_core_outputs[request.client_index].outputs
+ assert eco.prefill_stats is not None
+ assert eco.prefill_stats.num_cached_tokens == 1
+ assert eco.prefill_stats.num_cache_creation_tokens == 2
+ sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
+
+
+def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
+ request = _make_request()
+ request.num_computed_tokens = len(request.prompt_token_ids)
+ sched = _make_sched(request)
+
+ engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
+
+ (eco,) = engine_core_outputs[request.client_index].outputs
+ assert eco.prefill_stats is not None
+ assert eco.prefill_stats.num_cached_tokens == 1
+ assert eco.prefill_stats.num_cache_creation_tokens == 2
+ sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py
index fbfeb384f..770c85578 100644
--- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py
+++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py
@@ -157,3 +157,71 @@ def test_finish_requests_leaves_healthy_running_intact(
)
assert scheduler.running == [alive]
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_preserves_unrelated_resumable_segment(
+ monkeypatch: pytest.MonkeyPatch,
+ scheduler_cls,
+ scheduler_mod,
+) -> None:
+ paused = _StubRequest("req-paused", RequestStatus.FINISHED_STOPPED)
+ paused.resumable = True
+ leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED)
+ scheduler = _make_scheduler(
+ scheduler_cls,
+ requests={paused.request_id: paused},
+ running=[paused],
+ waiting=[],
+ )
+
+ def fake_finish_requests(self, request_ids, finished_status):
+ return [leaving]
+
+ monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+ scheduler_cls.finish_requests(scheduler, [leaving.request_id], RequestStatus.FINISHED_ABORTED)
+
+ assert scheduler.running == [paused]
+ assert scheduler.requests == {paused.request_id: paused}
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_materializes_single_pass_request_ids(
+ monkeypatch: pytest.MonkeyPatch,
+ scheduler_cls,
+ scheduler_mod,
+) -> None:
+ request = _StubRequest("req-generator", RequestStatus.RUNNING)
+ scheduler = _make_scheduler(
+ scheduler_cls,
+ requests={request.request_id: request},
+ running=[request],
+ waiting=[],
+ )
+
+ def fake_finish_requests(self, request_ids, finished_status):
+ assert tuple(request_ids) == (request.request_id,)
+ return [request]
+
+ monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+ result = scheduler_cls.finish_requests(
+ scheduler,
+ (request_id for request_id in [request.request_id]),
+ RequestStatus.FINISHED_ABORTED,
+ )
+
+ assert result == [request]
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_finish_requests_does_not_reopen_off_queue_deferred_free_terminal(scheduler_cls) -> None:
+ terminal = _StubRequest("req-deferred-free", RequestStatus.FINISHED_STOPPED)
+ terminal.resumable = True
+ scheduler = _make_scheduler(scheduler_cls, requests={terminal.request_id: terminal}, running=[], waiting=[])
+ scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
+
+ assert scheduler_cls.finish_requests(scheduler, [terminal.request_id], RequestStatus.FINISHED_ABORTED) == []
+ assert scheduler.requests == {terminal.request_id: terminal}
+ assert terminal.status == RequestStatus.FINISHED_STOPPED
diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py
index 9aaa4c656..bb9fd8859 100644
--- a/tests/core/sched/test_omni_scheduler_mixin_shared.py
+++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py
@@ -4,6 +4,7 @@ from types import SimpleNamespace
import pytest
from vllm.v1.engine import FinishReason
+from vllm_omni.core.sched import omni_scheduler_mixin
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
from vllm_omni.core.sched.output import OmniChunkRecvHandle
@@ -14,6 +15,77 @@ class _Scheduler(OmniSchedulerMixin):
pass
+def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
+ created_adapters = []
+
+ def adapter_factory(config):
+ adapter = SimpleNamespace(vllm_config=config)
+ created_adapters.append(adapter)
+ return adapter
+
+ monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", adapter_factory)
+
+ def init_scheduler(**model_config):
+ scheduler = _Scheduler()
+ scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config))
+ scheduler._init_omni_io_scheduling_state()
+ return scheduler
+
+ producer = init_scheduler(
+ stage_id=0,
+ async_chunk=True,
+ requires_full_payload_input=False,
+ custom_process_next_stage_input_func="test.pipeline.produce_async_chunk",
+ )
+ receiver = init_scheduler(
+ stage_id=1,
+ async_chunk=True,
+ requires_full_payload_input=True,
+ custom_process_next_stage_input_func=None,
+ )
+ terminal_stage = init_scheduler(
+ stage_id=2,
+ async_chunk=True,
+ requires_full_payload_input=False,
+ custom_process_next_stage_input_func=None,
+ )
+
+ assert producer.chunk_transfer_adapter is not None
+ assert receiver.chunk_transfer_adapter is not None
+ assert terminal_stage.chunk_transfer_adapter is not None
+ assert receiver.input_coordinator is None
+ assert len(created_adapters) == 3
+
+
+@pytest.mark.parametrize(
+ ("stage_id", "async_chunk", "required", "enabled"),
+ [
+ (0, False, True, False),
+ (1, True, True, False),
+ (1, False, False, False),
+ (1, False, True, True),
+ ],
+)
+def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, async_chunk, required, enabled):
+ monkeypatch.setattr(
+ omni_scheduler_mixin,
+ "OmniChunkTransferAdapter",
+ lambda _config: SimpleNamespace(),
+ )
+ scheduler = _Scheduler()
+ scheduler.vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(
+ stage_id=stage_id,
+ async_chunk=async_chunk,
+ requires_full_payload_input=required,
+ )
+ )
+
+ scheduler._init_omni_io_scheduling_state()
+
+ assert (scheduler.input_coordinator is not None) is enabled
+
+
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
calls = []
scheduler = _Scheduler()
@@ -22,13 +94,25 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
scheduler.requests = {"request": object()}
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
+
+ def _collect_timed_out(timeout_s):
+ calls.append(("chunk-timeouts", timeout_s))
+ return set()
+
+ def _collect_failed_sends():
+ calls.append(("failed-sends",))
+ return {}
+
scheduler.chunk_transfer_adapter = SimpleNamespace(
+ receives_chunks=True,
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
("process", waiting, running, scheduler_requests)
),
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
("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: calls.append(("restore-full", waiting))
@@ -41,6 +125,10 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
("consume", "ar"),
("timeouts",),
("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", scheduler.waiting, scheduler.running, scheduler.requests),
("restore-full", scheduler.waiting),
]
diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
index 5fad4340c..e7d126cd9 100644
--- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
+++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py
@@ -1,6 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit coverage for _process_pending_input_timeouts.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit coverage for the two stage-input timeout hooks.
+
+``_process_pending_input_timeouts`` covers the full-payload transport;
+``_process_pending_chunk_timeouts`` covers async-chunk. Both share
+``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
Verifies that the mixin correctly *delegates* timed-out requests to the
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
@@ -12,6 +16,10 @@ this file only asserts the wiring from the mixin to that API.
from __future__ import annotations
+import importlib
+import logging
+import math
+import os
from types import SimpleNamespace
import pytest
@@ -102,5 +110,280 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
assert scheduler.finish_calls == []
+class _FakeChunkAdapter:
+ """Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
+
+ def __init__(self, timed_out_ids, *, receives_chunks=True):
+ self._timed_out_ids = set(timed_out_ids)
+ self.receives_chunks = receives_chunks
+ self.calls = []
+
+ def collect_timed_out_request_ids(self, timeout_s):
+ self.calls.append(timeout_s)
+ return set(self._timed_out_ids)
+
+
+class _FakeChunkScheduler(OmniSchedulerMixin):
+ def __init__(self, requests, adapter):
+ self.requests = requests
+ self.chunk_transfer_adapter = adapter
+ self.finish_calls = []
+
+ def finish_requests(self, req_ids, status):
+ self.finish_calls.append((set(req_ids), status))
+
+
+def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
+ """A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
+
+ This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
+ existed, the chunk path had no deadline at all and issue #3833 could park a
+ request forever.
+ """
+ req_id = "stalled-stream"
+ adapter = _FakeChunkAdapter(timed_out_ids={req_id})
+ scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
+
+ scheduler._process_pending_chunk_timeouts()
+
+ assert len(adapter.calls) == 1
+ assert adapter.calls[0] > 0
+ assert len(scheduler.finish_calls) == 1
+ finished_ids, status = scheduler.finish_calls[0]
+ assert finished_ids == {req_id}
+ assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
+
+
+def test_process_pending_chunk_timeouts_skips_already_freed_request():
+ adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
+ scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
+
+ scheduler._process_pending_chunk_timeouts()
+
+ assert scheduler.finish_calls == []
+
+
+def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
+ """Stage 0 does not receive chunks, so it has nothing to expire."""
+ adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
+ scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
+
+ scheduler._process_pending_chunk_timeouts()
+
+ assert adapter.calls == []
+ assert scheduler.finish_calls == []
+
+
+def test_process_pending_chunk_timeouts_noop_without_adapter():
+ class _NoAdapter(OmniSchedulerMixin):
+ def __init__(self):
+ self.requests = {}
+ self.chunk_transfer_adapter = None
+ self.finish_calls = []
+
+ def finish_requests(self, req_ids, status):
+ self.finish_calls.append((set(req_ids), status))
+
+ scheduler = _NoAdapter()
+ scheduler._process_pending_chunk_timeouts()
+ assert scheduler.finish_calls == []
+
+
+def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
+ """The shared knob disables both transports' nets together."""
+ from vllm_omni.core.sched import omni_scheduler_mixin
+
+ monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
+
+ adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
+ scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
+ scheduler._process_pending_chunk_timeouts()
+
+ assert adapter.calls == []
+ assert scheduler.finish_calls == []
+
+
+class _FakeSendFailureAdapter:
+ """Stands in for the adapter's drained send-failure map (R1.2)."""
+
+ def __init__(self, failures):
+ self._failures = dict(failures)
+ self.collect_calls = 0
+
+ def collect_failed_send_request_ids(self):
+ self.collect_calls += 1
+ failures, self._failures = self._failures, {}
+ return failures
+
+
+class _FakeSendFailureScheduler(OmniSchedulerMixin):
+ def __init__(self, requests, adapter):
+ self.requests = requests
+ self.chunk_transfer_adapter = adapter
+
+
+def test_log_failed_chunk_sends_reports_a_live_request(caplog):
+ adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
+ scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
+
+ with caplog.at_level(logging.ERROR):
+ scheduler._log_failed_chunk_sends()
+
+ assert "req-live" in caplog.text
+ assert "No space left on device" in caplog.text
+
+
+def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
+ """``collect_*`` drains the map, so a failure for a request that already left
+ this scheduler is the one place a drop could still go unrecorded."""
+ adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
+ scheduler = _FakeSendFailureScheduler({}, adapter)
+
+ with caplog.at_level(logging.WARNING):
+ scheduler._log_failed_chunk_sends()
+
+ assert adapter.collect_calls == 1
+ assert "req-gone" in caplog.text
+ assert "No space left on device" in caplog.text
+ assert "finished or aborted" in caplog.text
+
+
+def test_log_failed_chunk_sends_noop_without_adapter():
+ scheduler = _FakeSendFailureScheduler({}, None)
+ scheduler._log_failed_chunk_sends()
+
+
+class _FakeReceiveFailureAdapter:
+ def __init__(self, failures):
+ self._failures = dict(failures)
+
+ def collect_failed_receive_request_ids(self):
+ failures, self._failures = self._failures, {}
+ return failures
+
+
+class _FakeReceiveFailureScheduler(OmniSchedulerMixin):
+ def __init__(self, requests, adapter):
+ self.requests = requests
+ self.chunk_transfer_adapter = adapter
+ self.finish_calls = []
+
+ def finish_requests(self, req_ids, status):
+ self.finish_calls.append((set(req_ids), status))
+
+
+def test_process_chunk_receive_failures_delegates_to_finish_requests():
+ request_id = "invalid-chunk"
+ adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"})
+ scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter)
+
+ scheduler._process_chunk_receive_failures()
+
+ assert len(scheduler.finish_calls) == 1
+ finished_ids, status = scheduler.finish_calls[0]
+ assert finished_ids == {request_id}
+ assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
+
+
+# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
+# these reloads the module under a different environment. One knob now gates the
+# deadline on both transports, which is why a silent 0 matters more than it did
+# when only the full-payload path used it.
+
+
+@pytest.fixture
+def reload_with(monkeypatch):
+ """Reload the mixin module under a given env value, then put it back.
+
+ ``importlib.reload`` re-executes into the live module dict, so whatever the
+ last test parsed is what every later test in the session sees -- ``-k
+ zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
+ silently disable both deadlines. The rejection cases are worse: the
+ ValueError fires *after* the rebind, so ``nan`` lingers in the module.
+ ``monkeypatch`` restores the environment but not the module, so reload once
+ more on the way out.
+ """
+ import vllm_omni.core.sched.omni_scheduler_mixin as mod
+
+ def _reload(raw):
+ if raw is None:
+ monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
+ else:
+ monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
+ return importlib.reload(mod)
+
+ yield _reload
+
+ monkeypatch.undo()
+ importlib.reload(mod)
+
+
+@pytest.mark.parametrize(
+ "raw,expected",
+ [
+ (None, 600.0),
+ ("30", 30.0),
+ ("0.5", 0.5),
+ ("0", 0.0), # explicit opt-out is still honoured
+ ],
+)
+def test_timeout_accepts_valid_values(reload_with, raw, expected):
+ mod = reload_with(raw)
+ assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
+
+
+@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
+def test_negative_timeout_is_rejected(reload_with, raw):
+ """`-1` is a common "no limit" idiom; substituting 600 would give an operator
+ who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
+ with pytest.raises(ValueError, match="negative"):
+ reload_with(raw)
+
+
+@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
+def test_non_finite_timeout_is_rejected(reload_with, raw):
+ """nan compares false against every deadline check and inf is unreachable;
+ both disable the net while looking like a number."""
+ with pytest.raises(ValueError, match="not finite"):
+ reload_with(raw)
+
+
+def test_zero_warns_and_names_both_transports(reload_with, caplog):
+ with caplog.at_level(logging.WARNING):
+ mod = reload_with("0")
+ assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
+ text = caplog.text
+ assert "DISABLED" in text
+ assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
+
+
+def test_unparsable_timeout_falls_back(reload_with):
+ mod = reload_with("not-a-number")
+ assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
+
+
+def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
+ """The reason ``reload_with`` reloads on the way out, done here by hand.
+
+ ``nan`` is bound before the ValueError fires, so the reload that rejects it
+ still leaves it in the module -- and nan compares false against every
+ deadline check, which is exactly the silent-disable this item exists to
+ stop. Restoring the environment is not enough; the module has to be reloaded
+ too.
+ """
+ import vllm_omni.core.sched.omni_scheduler_mixin as mod
+
+ ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
+
+ monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
+ with pytest.raises(ValueError, match="not finite"):
+ importlib.reload(mod)
+ assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
+
+ monkeypatch.undo()
+ importlib.reload(mod)
+ assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
+
+
if __name__ == "__main__":
pytest.main([__file__, "-v"])
diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py
new file mode 100644
index 000000000..6f44aadac
--- /dev/null
+++ b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py
@@ -0,0 +1,168 @@
+"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
+``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
+and in ``OmniARScheduler.schedule``.
+
+Upstream keeps that counter incrementally (``+1`` when a resumable request
+parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
+next update arrives or the request is finished) and
+``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
+queues. Omni rewrites ``request.status`` outside both hooks -- the
+chunk-transfer adapter's park/restore and
+``_realign_request_status_to_queues`` -- so a counted request can leave
+``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
+counter for the lifetime of the process.
+
+An inflated counter does not degrade throughput, it hangs the stage:
+``EngineCore.has_work()`` reads false while a live request sits in
+``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
+``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
+server answered the first conversation after a boot and was silent from the
+second on while still reporting healthy.
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StubRequest:
+ def __init__(self, request_id: str, status: RequestStatus) -> None:
+ self.request_id = request_id
+ self.status = status
+
+ def is_finished(self) -> bool:
+ return RequestStatus.is_finished(self.status)
+
+
+_SCHEDULER_PARAMS = [
+ pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
+ pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
+]
+
+
+def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
+ scheduler = scheduler_cls.__new__(scheduler_cls)
+ scheduler.chunk_transfer_adapter = None
+ scheduler.input_coordinator = None
+ scheduler.requests = requests
+ scheduler.running = running
+ scheduler.waiting = waiting
+ scheduler.skipped_waiting = skipped if skipped is not None else []
+ scheduler.num_waiting_for_streaming_input = counter
+ return scheduler
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_clears_counter_left_by_a_stomped_status(
+ monkeypatch: pytest.MonkeyPatch,
+ scheduler_cls,
+ scheduler_mod,
+) -> None:
+ """The abort case that wedged the talker.
+
+ The request was counted while parked for streaming input, then had its
+ status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
+ ``finish_requests`` took the running branch and never decremented. With
+ no parked requests left, the counter must land at 0.
+ """
+ aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
+ scheduler = _make_scheduler(
+ scheduler_cls,
+ requests={}, # upstream already dropped it
+ running=[],
+ waiting=[],
+ counter=1,
+ )
+
+ monkeypatch.setattr(
+ scheduler_mod.VLLMScheduler,
+ "finish_requests",
+ lambda self, request_ids, finished_status: [aborted],
+ )
+
+ scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
+
+ assert scheduler.num_waiting_for_streaming_input == 0
+
+
+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)
+ scheduler = _make_scheduler(
+ OmniARScheduler,
+ requests={parked.request_id: parked},
+ running=[],
+ waiting=[],
+ counter=0,
+ skipped=[parked],
+ )
+
+ scheduler._resync_streaming_input_counter()
+
+ assert scheduler.num_waiting_for_streaming_input == 1
+
+
+def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Pins placement in ``schedule``: the chunk-transfer adapter rewrites
+ status during ``process_pending_chunks``, so the resync has to run after
+ that and before upstream schedules, while ``self.waiting`` is still the
+ real queue."""
+ parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
+ scheduler = OmniARScheduler.__new__(OmniARScheduler)
+ scheduler.chunk_transfer_adapter = None
+ scheduler.input_coordinator = None
+ scheduler.requests = {parked.request_id: parked}
+ scheduler.running = [parked]
+ scheduler.waiting = []
+ scheduler.skipped_waiting = []
+ scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite
+ scheduler.policy = "fcfs"
+
+ seen: dict[str, int] = {}
+
+ monkeypatch.setattr(
+ ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
+ )
+ monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
+ monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
+
+ def fake_schedule(self, throttle_prefills=False):
+ seen["counter"] = self.num_waiting_for_streaming_input
+ raise _StopScheduleError
+
+ class _StopScheduleError(Exception):
+ pass
+
+ monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
+
+ with pytest.raises(_StopScheduleError):
+ OmniARScheduler.schedule(scheduler)
+
+ assert seen["counter"] == 0
+
+
+def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None:
+ aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED)
+ scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+ scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+ scheduler.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.skipped_waiting) == []
+ assert scheduler.running == []
diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py
index 444d21c34..ae1b3924e 100644
--- a/tests/core/sched/test_omni_scheduling_coordinator.py
+++ b/tests/core/sched/test_omni_scheduling_coordinator.py
@@ -20,7 +20,6 @@ import torch
import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
from vllm_omni.core.sched.omni_scheduling_coordinator import (
OmniSchedulingCoordinator,
- uses_full_payload_input_coordinator,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -96,109 +95,6 @@ class MockQueue:
# ------------------------------------------------------------------ #
-class TestFullPayloadCoordinatorSelection(unittest.TestCase):
- """Tests for the (model_arch, model_stage) whitelist gate.
-
- The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
- is a superset of the stages registered here -- consumer-wait stages
- must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
- the init allowlist covers both producer- and consumer-side runners.
- These tests pin which ``(arch, stage)`` pairs the gate fires for today.
- """
-
- # Expected whitelist (model_arch, model_stage). Hardcoded to avoid the
- # tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
- # against itself; any drift between this matrix and the whitelist will
- # fail loudly here.
- EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
- {
- ("Qwen3OmniMoeForConditionalGeneration", "talker"),
- ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
- ("Qwen2_5OmniForConditionalGeneration", "talker"),
- ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
- ("CovoAudioForConditionalGeneration", "code2wav"),
- ("MiMoAudioModel", "code2wav"),
- ("Qwen3TTSCode2Wav", "code2wav"),
- ("MiniCPMO45Code2Wav", "code2wav"),
- # nemotron_voicechat: sync-path RVQ-VAE decode waits for the
- # talker's full [frames, 31] code-stack payload.
- ("NemotronVoiceChatCode2Wav", "code2wav"),
- ("CosyVoice3Model", "cosyvoice3_code2wav"),
- # audex: sync-path TTS decoder (async_chunk: false deploys) and
- # the TTA XCodec1 stage (always sync full-payload).
- ("AudexCode2Wav", "audex_code2wav"),
- ("AudexXCodec1", "audex_xcodec"),
- ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
- ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
- ("DyninOmniForConditionalGeneration", "token2image"),
- ("DyninOmniForConditionalGeneration", "token2audio"),
- }
- )
-
- def test_whitelist_matches_expected_matrix(self):
- """_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
-
- Catches both accidental additions (which would silently enable the
- consumer-wait gate for a new arch) and accidental removals (which
- would silently disable an enabled arch).
- """
- from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
-
- self.assertEqual(
- frozenset(_FULL_PAYLOAD_INPUT_STAGES),
- self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
- msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
- "update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
- )
-
- def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
- """Every (arch, stage) pair in the expected matrix must fire
- the gate when stage_id > 0 and async_chunk=False.
- """
- for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
- model_config = SimpleNamespace(
- stage_id=1,
- async_chunk=False,
- model_arch=arch,
- model_stage=stage,
- )
- self.assertTrue(
- uses_full_payload_input_coordinator(model_config),
- msg=f"expected gate to fire for {arch}/{stage}",
- )
-
- def test_other_arch_or_stage_or_mode_does_not_fire(self):
- cases = [
- SimpleNamespace(
- stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
- ),
- SimpleNamespace(
- stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
- ),
- SimpleNamespace(
- stage_id=1,
- async_chunk=False,
- model_arch="Qwen3OmniMoeForConditionalGeneration",
- model_stage="some_future_stage",
- ),
- SimpleNamespace(
- stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
- ),
- SimpleNamespace(
- stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
- ),
- SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
- SimpleNamespace(
- stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
- ),
- ]
- for model_config in cases:
- self.assertFalse(
- uses_full_payload_input_coordinator(model_config),
- msg=f"expected gate OFF for {model_config}",
- )
-
-
class TestCoordinatorUpdateRequestMetadata(unittest.TestCase):
"""Test update_request_metadata applies scheduling metadata to requests."""
diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py
index 7e3ab77c7..864fb7ab3 100644
--- a/tests/dfx/conftest.py
+++ b/tests/dfx/conftest.py
@@ -105,7 +105,7 @@ def create_unique_server_pytest_params(
configs: list[dict[str, Any]],
stage_configs_dir: Path,
) -> list[Any]:
- """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
+ """Like :func:`_create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
marks_by_name = _marks_by_test_name(configs)
return [
pytest.param(
@@ -113,7 +113,7 @@ def create_unique_server_pytest_params(
marks=marks_by_name.get(row[0], []),
id=row[0],
)
- for row in create_unique_server_params(configs, stage_configs_dir)
+ for row in _create_unique_server_params(configs, stage_configs_dir)
]
@@ -150,7 +150,7 @@ def create_paired_omni_benchmark_pytest_params(
"""Paired params for ``run_benchmark.py`` (omni/tts)."""
mapping = create_test_parameter_mapping(configs)
marks_by_name = _marks_by_test_name(configs)
- server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)]
+ server_entries = [(row, row[0]) for row in _create_unique_server_params(configs, stage_configs_dir)]
params_by_test_name = {
test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries
}
@@ -232,7 +232,7 @@ def _build_serve_args(serve_args: Any) -> list[str]:
return args
-def create_unique_server_params(
+def _create_unique_server_params(
configs: list[dict[str, Any]],
stage_configs_dir: Path,
) -> list[tuple[str, str, str | None, str | None, tuple[str, ...], bool]]:
@@ -311,7 +311,7 @@ def create_reliability_omni_server_params(
configs: list[dict[str, Any]], stage_configs_dir: Path
) -> list[OmniServerParams]:
adjusted_configs = configs_with_platform_stage_configs(configs)
- unique_params = create_unique_server_params(adjusted_configs, stage_configs_dir)
+ unique_params = _create_unique_server_params(adjusted_configs, stage_configs_dir)
server_args_by_name = extract_server_args_by_test_name(adjusted_configs)
return [
OmniServerParams(
@@ -653,6 +653,7 @@ def run_benchmark(
random_input_len: Any | None = None,
random_output_len: Any | None = None,
resource_label: str | None = None,
+ num_warmups: int = 2,
) -> dict[str, Any]:
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
@@ -677,7 +678,7 @@ def run_benchmark(
+ args
+ [
"--num-warmups",
- "2",
+ str(num_warmups),
"--save-result",
"--result-dir",
os.environ.get("BENCHMARK_DIR", "tests"),
diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py
index 03dfac1a0..64285dca3 100644
--- a/tests/dfx/perf/scripts/run_benchmark.py
+++ b/tests/dfx/perf/scripts/run_benchmark.py
@@ -1,4 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import json
+import math
import os
import threading
from collections.abc import Callable
@@ -159,6 +163,18 @@ def benchmark_params(request):
def assert_result(result, params, num_prompt) -> None:
assert result["completed"] == num_prompt, "Request failures exist"
+ baseline = params.get("baseline")
+ hardware = result.get("Hardware")
+ hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
+ if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline:
+ num_tpot_samples = result.get("num_tpot_samples")
+ mean_tpot_ms = result.get("mean_tpot_ms")
+ assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, (
+ "TPOT baseline is configured, but no measurable TPOT samples were produced"
+ )
+ assert (
+ isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms)
+ ), "TPOT baseline is configured, but mean_tpot_ms is not finite"
expected_audio_turns = params.get("expected_duplex_audio_turns_per_session")
if expected_audio_turns is not None:
session_metrics = result.get("duplex_session_metrics")
@@ -279,5 +295,6 @@ def test_performance_benchmark(omni_server, benchmark_params):
random_input_len=params.get("random_input_len"),
random_output_len=params.get("random_output_len"),
resource_label=resource_label,
+ num_warmups=max(2, int(concurrency)),
)
assert_result(result, params, num_prompt)
diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
index 9f6ff0241..1189ddd9d 100644
--- a/tests/dfx/perf/tests/test_bagel_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
@@ -43,8 +43,8 @@
"H100": {
"throughput_qps": 0.495,
"latency_mean": 2.0238,
- "peak_memory_mb_max": 29393.4286,
- "peak_memory_mb_mean": 29392.7143
+ "peak_memory_mb_max": 29392.0,
+ "peak_memory_mb_mean": 29392.0
}
}
}
@@ -61,7 +61,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
@@ -93,10 +93,10 @@
"num-input-images": 1,
"baseline": {
"H100": {
- "throughput_qps": 0.3451,
- "latency_mean": 2.9031,
- "peak_memory_mb_max": 30480.2857,
- "peak_memory_mb_mean": 30479.2857
+ "throughput_qps": 0.3499,
+ "latency_mean": 2.8701,
+ "peak_memory_mb_max": 30466.0,
+ "peak_memory_mb_mean": 30465.0
}
}
}
@@ -113,7 +113,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
@@ -164,7 +164,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
@@ -196,8 +196,8 @@
"num-input-images": 1,
"baseline": {
"H100": {
- "throughput_qps": 0.149,
- "latency_mean": 6.7157,
+ "throughput_qps": 0.1533,
+ "latency_mean": 6.5243,
"peak_memory_mb_max": 30574.0,
"peak_memory_mb_mean": 30574.0
}
diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
index c4d2305e6..bcd649de1 100644
--- a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
@@ -49,8 +49,8 @@
},
"baseline": {
"H100": {
- "throughput_qps": 1.107,
- "latency_mean": 0.8929,
+ "throughput_qps": 1.1958,
+ "latency_mean": 0.843,
"peak_memory_mb_mean": 80000.0
}
}
@@ -118,9 +118,9 @@
},
"baseline": {
"H100": {
- "throughput_qps": 0.027,
- "latency_mean": 36.6824,
- "peak_memory_mb_mean": 30155.4
+ "throughput_qps": 0.0307,
+ "latency_mean": 32.7665,
+ "peak_memory_mb_mean": 28762.4762
}
}
}
@@ -186,9 +186,9 @@
},
"baseline": {
"H100": {
- "throughput_qps": 0.027,
- "latency_mean": 36.968,
- "peak_memory_mb_mean": 30443.6
+ "throughput_qps": 0.0305,
+ "latency_mean": 33.1009,
+ "peak_memory_mb_mean": 29441.1429
}
}
}
@@ -258,9 +258,9 @@
},
"baseline": {
"H100": {
- "throughput_qps": 0.027,
- "latency_mean": 37.2117,
- "peak_memory_mb_mean": 30135.6
+ "throughput_qps": 0.0299,
+ "latency_mean": 33.925,
+ "peak_memory_mb_mean": 29443.1429
}
}
}
diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json
index 75e5983b0..1197c88f1 100644
--- a/tests/dfx/perf/tests/test_higgs_audio_v3.json
+++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"tts"
],
"server_params": {
diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
index c4ef34b96..098dcd5a0 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
@@ -137,9 +137,7 @@
"H200": {
"throughput_qps": 0.1874,
"latency_p99": 101.852,
- "latency_mean": 79.0385,
- "peak_memory_mb_max": 70000,
- "peak_memory_mb_mean": 70000
+ "latency_mean": 79.0385
}
}
}
diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json
new file mode 100644
index 000000000..8186edf2c
--- /dev/null
+++ b/tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json
@@ -0,0 +1,106 @@
+[
+ {
+ "test_name": "test_hunyuanvideo15_i2v_single_device",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H100"
+ },
+ "num_cards": 1
+ }
+ },
+ "full_model",
+ "diffusion"
+ ],
+ "description": "Single-device baseline",
+ "server_type": "vllm-omni",
+ "server_params": {
+ "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
+ "serve_args": {
+ "flow-shift": 5.0,
+ "enable-diffusion-pipeline-profiler": true
+ }
+ },
+ "benchmark_params": [
+ {
+ "name": "832x480_frames33_steps4",
+ "dataset": "random",
+ "task": "i2v",
+ "num-prompts": 10,
+ "max-concurrency": 1,
+ "num-input-images": 1,
+ "seed": 42,
+ "enable-negative-prompt": true,
+ "random-request-config": [
+ {
+ "width": 832,
+ "height": 480,
+ "num_inference_steps": 4,
+ "num_frames": 33,
+ "fps": 24,
+ "weight": 1
+ }
+ ],
+ "extra-body": {
+ "guidance_scale": 6.0,
+ "flow_shift": 5.0
+ }
+ }
+ ]
+ },
+ {
+ "test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H100"
+ },
+ "num_cards": 2
+ }
+ },
+ "full_model",
+ "diffusion"
+ ],
+ "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
+ "server_type": "vllm-omni",
+ "server_params": {
+ "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
+ "serve_args": {
+ "flow-shift": 5.0,
+ "cache-backend": "cache_dit",
+ "tensor-parallel-size": 2,
+ "vae-patch-parallel-size": 2,
+ "vae-use-tiling": true,
+ "enable-diffusion-pipeline-profiler": true
+ }
+ },
+ "benchmark_params": [
+ {
+ "name": "832x480_frames33_steps4",
+ "dataset": "random",
+ "task": "i2v",
+ "num-prompts": 10,
+ "max-concurrency": 1,
+ "num-input-images": 1,
+ "seed": 42,
+ "enable-negative-prompt": true,
+ "random-request-config": [
+ {
+ "width": 832,
+ "height": 480,
+ "num_inference_steps": 4,
+ "num_frames": 33,
+ "fps": 24,
+ "weight": 1
+ }
+ ],
+ "extra-body": {
+ "guidance_scale": 6.0,
+ "flow_shift": 5.0
+ }
+ }
+ ]
+ }
+]
diff --git a/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
new file mode 100644
index 000000000..3771c1c81
--- /dev/null
+++ b/tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json
@@ -0,0 +1,104 @@
+[
+ {
+ "test_name": "test_hunyuanvideo15_t2v_single_device",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H100"
+ },
+ "num_cards": 1
+ }
+ },
+ "full_model",
+ "diffusion"
+ ],
+ "description": "Single-device baseline",
+ "server_type": "vllm-omni",
+ "server_params": {
+ "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+ "serve_args": {
+ "flow-shift": 5.0,
+ "enable-diffusion-pipeline-profiler": true
+ }
+ },
+ "benchmark_params": [
+ {
+ "name": "832x480_frames33_steps4",
+ "dataset": "random",
+ "task": "t2v",
+ "num-prompts": 10,
+ "max-concurrency": 1,
+ "seed": 42,
+ "enable-negative-prompt": true,
+ "random-request-config": [
+ {
+ "width": 832,
+ "height": 480,
+ "num_inference_steps": 4,
+ "num_frames": 33,
+ "fps": 24,
+ "weight": 1
+ }
+ ],
+ "extra-body": {
+ "guidance_scale": 6.0,
+ "flow_shift": 5.0
+ }
+ }
+ ]
+ },
+ {
+ "test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H100"
+ },
+ "num_cards": 2
+ }
+ },
+ "full_model",
+ "diffusion"
+ ],
+ "description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
+ "server_type": "vllm-omni",
+ "server_params": {
+ "model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
+ "serve_args": {
+ "flow-shift": 5.0,
+ "cache-backend": "cache_dit",
+ "tensor-parallel-size": 2,
+ "vae-patch-parallel-size": 2,
+ "vae-use-tiling": true,
+ "enable-diffusion-pipeline-profiler": true
+ }
+ },
+ "benchmark_params": [
+ {
+ "name": "832x480_frames33_steps4",
+ "dataset": "random",
+ "task": "t2v",
+ "num-prompts": 10,
+ "max-concurrency": 1,
+ "seed": 42,
+ "enable-negative-prompt": true,
+ "random-request-config": [
+ {
+ "width": 832,
+ "height": 480,
+ "num_inference_steps": 4,
+ "num_frames": 33,
+ "fps": 24,
+ "weight": 1
+ }
+ ],
+ "extra-body": {
+ "guidance_scale": 6.0,
+ "flow_shift": 5.0
+ }
+ }
+ ]
+ }
+]
diff --git a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
index 1a0f31e98..9b0634a20 100644
--- a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-device baseline with enforce-eager (no torch.compile)",
@@ -64,7 +64,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-device with torch.compile (default, no enforce-eager)",
@@ -117,7 +117,7 @@
"num_cards": 2
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "CFG-parallel=2 with enforce-eager",
@@ -172,7 +172,7 @@
"num_cards": 2
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "CFG-parallel=2 with torch.compile",
@@ -226,7 +226,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "CacheDiT with enforce-eager",
diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json
index cfb57169b..cac1880c5 100644
--- a/tests/dfx/perf/tests/test_minicpmo_4_5.json
+++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json
@@ -39,6 +39,7 @@
8
],
"no_oversample": true,
+ "disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
@@ -55,33 +56,116 @@
"baseline": {
"H100": {
"request_throughput": [
- 0.7803,
- 1.1008,
- 1.2027
+ 0.9731,
+ 1.1986,
+ 1.686
],
"mean_ttft_ms": [
- 166.9764,
- 307.7133,
- 323.584
+ 152.5665,
+ 270.2559,
+ 301.455
],
"mean_e2el_ms": [
- 1361.1682,
- 3800.6332,
- 6961.8636
+ 1027.2953,
+ 3306.3735,
+ 4686.5366
],
"mean_audio_ttfp_ms": [
- 635.8099,
- 1615.6591,
- 2888.9831
+ 458.1553,
+ 2030.2335,
+ 2068.6233
],
"mean_audio_rtf": [
- 0.3151,
- 0.8568,
- 1.5425
+ 0.2357,
+ 0.791,
+ 1.0324
+ ]
+ },
+ "A3": {
+ "request_throughput": [
+ 0.5383,
+ 0.6042,
+ 0.7547
+ ],
+ "mean_ttft_ms": [
+ 333.2633,
+ 533.8656,
+ 547.3388
+ ],
+ "mean_e2el_ms": [
+ 1857.2154,
+ 6600.5101,
+ 10450.6811
+ ],
+ "mean_audio_ttfp_ms": [
+ 986.4666,
+ 3411.0754,
+ 3352.7515
+ ],
+ "mean_audio_rtf": [
+ 0.4423,
+ 1.5734,
+ 2.3024
]
}
}
}
]
+ },
+ {
+ "test_name": "test_minicpmo_4_5_a2",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "npu": "A2"
+ },
+ "num_cards": 1
+ }
+ },
+ "full_model",
+ "omni"
+ ],
+ "server_params": {
+ "model": "openbmb/MiniCPM-o-4_5",
+ "extra_cli_args": [
+ "--deploy-config",
+ "vllm_omni/deploy/minicpmo_4_5.yaml",
+ "--trust-remote-code"
+ ]
+ },
+ "benchmark_params": [
+ {
+ "dataset_name": "seed-tts",
+ "dataset_path": "zhaochenyang20/seed-tts-eval",
+ "backend": "openai-chat-omni",
+ "endpoint": "/v1/chat/completions",
+ "num_prompts": [
+ 32,
+ 64,
+ 128
+ ],
+ "max_concurrency": [
+ 1,
+ 4,
+ 8
+ ],
+ "no_oversample": true,
+ "disable_shuffle": true,
+ "trust_remote_code": true,
+ "seed_tts_locale": "en",
+ "extra_body": {
+ "modalities": [
+ "text",
+ "audio"
+ ],
+ "chat_template_kwargs": {
+ "enable_thinking": false,
+ "use_tts_template": true
+ }
+ },
+ "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
+ }
+ ]
}
]
diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json
index cfbcc3abb..801743134 100644
--- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json
+++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json
@@ -18,7 +18,7 @@
"model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [
"--deploy-config",
- "vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
+ "vllm_omni/deploy/minicpmo_4_5.yaml",
"--trust-remote-code"
]
},
@@ -35,6 +35,7 @@
1
],
"no_oversample": true,
+ "disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"seed_tts_turns_per_session": 4,
@@ -42,9 +43,26 @@
"extra_body": {
"save_duplex_request_metrics": true
},
- "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
+ "percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
+ "request_throughput": [
+ 0.3038
+ ],
+ "mean_ttft_ms": [
+ 139.997
+ ],
+ "mean_e2el_ms": [
+ 3269.8807
+ ],
+ "mean_audio_ttfp_ms": [
+ 456.073
+ ],
+ "mean_audio_rtf": [
+ 0.2254
+ ]
+ },
+ "A3": {
"request_throughput": [
0.1594
],
diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
index e327be510..26dc54aab 100644
--- a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
+++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
@@ -37,6 +37,7 @@
4
],
"no_oversample": true,
+ "disable_shuffle": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
@@ -53,24 +54,24 @@
"baseline": {
"H100": {
"request_throughput": [
- 0.7803,
- 1.1008
+ 0.9731,
+ 1.1986
],
"mean_ttft_ms": [
- 166.9764,
- 307.7133
+ 152.5665,
+ 270.2559
],
"mean_e2el_ms": [
- 1361.1682,
- 3800.6332
+ 1027.2953,
+ 3306.3735
],
"mean_audio_ttfp_ms": [
- 635.8099,
- 1615.6591
+ 458.1553,
+ 2030.2335
],
"mean_audio_rtf": [
- 0.3151,
- 0.8568
+ 0.2357,
+ 0.791
]
}
}
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
index a73abc5fb..6b0247aaf 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
@@ -21,54 +21,6 @@
]
},
"benchmark_params": [
- {
- "dataset_name": "random",
- "backend": "openai-chat-omni",
- "endpoint": "/v1/chat/completions",
- "num_prompts": [
- 4,
- 16,
- 32,
- 64,
- 128
- ],
- "max_concurrency": [
- 1,
- 4,
- 8,
- 16,
- 32
- ],
- "random_input_len": 2500,
- "random_output_len": 900,
- "ignore_eos": true,
- "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
- "baseline": {
- "H100": {
- "mean_ttft_ms": [
- 96.4061,
- 140.8147,
- 271.9025,
- 362.2916,
- 507.776
- ],
- "mean_e2el_ms": [
- 18507.4582,
- 28365.4956,
- 31907.9613,
- 48161.5829,
- 72630.8601
- ],
- "mean_audio_rtf": [
- 0.1175,
- 0.18,
- 0.2359,
- 0.3278,
- 0.6296
- ]
- }
- }
- },
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
@@ -94,14 +46,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 354.0951
- ],
- "mean_e2el_ms": [
- 3952.5615
- ],
"mean_audio_rtf": [
- 0.117
+ 0.1098
]
}
}
@@ -133,14 +79,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 259.6317
- ],
- "mean_e2el_ms": [
- 4253.3009
- ],
"mean_audio_rtf": [
- 0.1448
+ 0.1335
]
}
}
@@ -174,60 +114,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 457.4453
- ],
- "mean_e2el_ms": [
- 5912.5623
- ],
"mean_audio_rtf": [
- 0.1967
- ]
- }
- }
- },
- {
- "dataset_name": "random",
- "backend": "openai-chat-omni",
- "endpoint": "/v1/chat/completions",
- "num_prompts": [
- 4,
- 16,
- 32,
- 64,
- 128
- ],
- "max_concurrency": [
- 1,
- 4,
- 8,
- 16,
- 32
- ],
- "random_input_len": 2500,
- "random_output_len": 900,
- "ignore_eos": true,
- "extra_body": {
- "modalities": [
- "text"
- ]
- },
- "percentile-metrics": "ttft,tpot,itl,e2el",
- "baseline": {
- "H100": {
- "mean_ttft_ms": [
- 90.8033,
- 202.1779,
- 326.5864,
- 544.8914,
- 961.3929
- ],
- "mean_e2el_ms": [
- 4517.3411,
- 7857.9711,
- 9861.1687,
- 13283.3363,
- 19541.0736
+ 0.1798
]
}
}
@@ -262,11 +150,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 318.048
- ],
"mean_e2el_ms": [
- 807.6114
+ 756.8415
]
}
}
@@ -303,11 +188,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 186.168
- ],
"mean_e2el_ms": [
- 726.1137
+ 694.9777
]
}
}
@@ -346,9 +228,6 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 375.5282
- ],
"mean_e2el_ms": [
1006.1663
]
@@ -356,5 +235,61 @@
}
}
]
+ },
+ {
+ "test_name": "test_qwen3_omni_chunk_cuda",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H100"
+ },
+ "num_cards": 2
+ }
+ },
+ "omni",
+ "slow"
+ ],
+ "server_params": {
+ "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+ "extra_cli_args": [
+ "--async-chunk"
+ ]
+ },
+ "benchmark_params": [
+ {
+ "dataset_name": "random",
+ "backend": "openai-chat-omni",
+ "endpoint": "/v1/chat/completions",
+ "num_prompts": [
+ 4,
+ 32,
+ 128
+ ],
+ "max_concurrency": [
+ 1,
+ 8,
+ 32
+ ],
+ "random_input_len": 100,
+ "random_output_len": 100,
+ "ignore_eos": true,
+ "extra_body": {
+ "modalities": [
+ "text"
+ ]
+ },
+ "percentile-metrics": "ttft,tpot,itl,e2el",
+ "baseline": {
+ "H100": {
+ "mean_e2el_ms": [
+ 4354.8945,
+ 9861.1687,
+ 18716.709
+ ]
+ }
+ }
+ }
+ ]
}
]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
index a3d8aaea8..5453d6a90 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
@@ -10,7 +10,7 @@
"num_cards": 3
}
},
- "full_model",
+ "slow",
"omni"
],
"server_params": {
@@ -36,10 +36,9 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
- "mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
- "mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
- "mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
+ "mean_tpot_ms": [7.6955, 10.6457, 12.7363, 15.1608],
+ "mean_audio_ttfp_ms": [325.1624, 602.5192, 699.3891, 1024.088],
+ "mean_audio_rtf": [0.1689, 0.2088, 0.2397, 0.2791]
}
}
},
@@ -64,9 +63,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [340.5122],
- "mean_audio_ttfp_ms": [463.0752],
- "mean_audio_rtf": [0.1153]
+ "mean_audio_ttfp_ms": [436.4663],
+ "mean_audio_rtf": [0.1069]
}
}
},
@@ -93,9 +91,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [236.926],
- "mean_audio_ttfp_ms": [395.2182],
- "mean_audio_rtf": [0.1259]
+ "mean_audio_ttfp_ms": [376.4911],
+ "mean_audio_rtf": [0.1145]
}
}
},
@@ -124,9 +121,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [470.6126],
- "mean_audio_ttfp_ms": [665.5474],
- "mean_audio_rtf": [0.166]
+ "mean_audio_ttfp_ms": [638.2227],
+ "mean_audio_rtf": [0.1438]
}
}
}
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
index 38dbfb2e3..d73629666 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
@@ -45,26 +45,12 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 115.8348,
- 173.1823,
- 446.672,
- 464.5985,
- 3420.0634
- ],
- "mean_e2el_ms": [
- 25070.9523,
- 31992.0646,
- 37619.7669,
- 50393.8257,
- 136619.2743
- ],
"mean_audio_rtf": [
0.169,
- 0.2111,
- 0.2475,
- 0.3551,
- 0.9266
+ 0.2041,
+ 0.2382,
+ 0.3177,
+ 0.4595
]
}
}
@@ -94,14 +80,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 343.1377
- ],
- "mean_e2el_ms": [
- 4811.1123
- ],
"mean_audio_rtf": [
- 0.1463
+ 0.138
]
}
}
@@ -133,14 +113,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 245.5446
- ],
- "mean_e2el_ms": [
- 5019.1256
- ],
"mean_audio_rtf": [
- 0.1713
+ 0.1609
]
}
}
@@ -174,14 +148,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 506.1407
- ],
- "mean_e2el_ms": [
- 6814.9656
- ],
"mean_audio_rtf": [
- 0.2213
+ 0.2048
]
}
}
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
index 8775ee381..ef8d97a94 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
@@ -10,7 +10,7 @@
"num_cards": 2
}
},
- "full_model",
+ "slow",
"omni"
],
"server_params": {
@@ -29,16 +29,12 @@
"endpoint": "/v1/completions",
"num_prompts": [
4,
- 16,
32,
- 64,
128
],
"max_concurrency": [
1,
- 4,
8,
- 16,
32
],
"random_input_len": 2500,
@@ -52,19 +48,10 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 69.8081,
- 173.0651,
- 271.8108,
- 413.4079,
- 497.2546
- ],
"mean_e2el_ms": [
- 4720.8225,
- 7937.2532,
- 9962.3566,
- 13724.5479,
- 18324.091
+ 4619.0068,
+ 9760.6175,
+ 18007.6918
]
}
}
@@ -100,53 +87,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 25.9477
- ],
"mean_e2el_ms": [
- 517.3155
- ]
- }
- }
- },
- {
- "dataset_name": "random-mm",
- "backend": "vllm",
- "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
- "endpoint": "/v1/completions",
- "num_prompts": [
- 40
- ],
- "request_rate": [
- 0.5
- ],
- "random_input_len": 100,
- "random_output_len": 100,
- "random_range_ratio": 0.0,
- "ignore_eos": true,
- "extra_body": {
- "modalities": [
- "text"
- ]
- },
- "random_mm_base_items_per_request": 2,
- "random_mm_num_mm_items_range_ratio": 0.5,
- "random_mm_limit_mm_per_prompt": {
- "image": 1,
- "video": 1
- },
- "random_mm_bucket_config": {
- "(256, 256, 1)": 0.5,
- "(720, 1280, 2)": 0.5
- },
- "percentile-metrics": "ttft,tpot,itl,e2el",
- "baseline": {
- "H100": {
- "mean_ttft_ms": [
- 26.7056
- ],
- "mean_e2el_ms": [
- 551.308
+ 491.2534
]
}
}
@@ -186,11 +128,8 @@
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"H100": {
- "mean_ttft_ms": [
- 27.6906
- ],
"mean_e2el_ms": [
- 596.5284
+ 559.6674
]
}
}
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
index dcd2188c5..34d275ac9 100644
--- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-device baseline (two input images)",
@@ -35,10 +35,10 @@
"enable-negative-prompt": true,
"baseline": {
"H100": {
- "throughput_qps": 0.0721,
- "latency_mean": 13.8398,
- "peak_memory_mb_max": 57610.0,
- "peak_memory_mb_mean": 57610.0
+ "throughput_qps": 0.0684,
+ "latency_mean": 14.6249,
+ "peak_memory_mb_max": 57632.0,
+ "peak_memory_mb_mean": 57632.0
}
}
},
@@ -55,10 +55,10 @@
"enable-negative-prompt": true,
"baseline": {
"H100": {
- "throughput_qps": 0.0186,
- "latency_mean": 53.6942,
- "peak_memory_mb_max": 67790.0,
- "peak_memory_mb_mean": 67790.0
+ "throughput_qps": 0.0177,
+ "latency_mean": 56.5285,
+ "peak_memory_mb_max": 67840.0,
+ "peak_memory_mb_mean": 67840.0
}
}
}
@@ -76,7 +76,7 @@
"num_cards": 4
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
@@ -105,10 +105,10 @@
"enable-negative-prompt": true,
"baseline": {
"H100": {
- "throughput_qps": 0.0538,
- "latency_mean": 18.5591,
- "peak_memory_mb_max": 56449.4286,
- "peak_memory_mb_mean": 56449.4286
+ "throughput_qps": 0.0515,
+ "latency_mean": 19.4345,
+ "peak_memory_mb_max": 56348.0,
+ "peak_memory_mb_mean": 56348.0
}
}
}
@@ -125,7 +125,7 @@
"num_cards": 4
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
@@ -164,10 +164,10 @@
"enable-negative-prompt": true,
"baseline": {
"H100": {
- "throughput_qps": 0.2776,
- "latency_mean": 3.5938,
- "peak_memory_mb_max": 58113.7143,
- "peak_memory_mb_mean": 58075.1429
+ "throughput_qps": 0.2479,
+ "latency_mean": 4.0389,
+ "peak_memory_mb_max": 57750.0,
+ "peak_memory_mb_mean": 57750.0
}
}
},
@@ -184,10 +184,10 @@
"enable-negative-prompt": true,
"baseline": {
"H100": {
- "throughput_qps": 0.1051,
- "latency_mean": 9.4918,
- "peak_memory_mb_max": 68107.4286,
- "peak_memory_mb_mean": 68107.4286
+ "throughput_qps": 0.0948,
+ "latency_mean": 10.5462,
+ "peak_memory_mb_max": 67764.0,
+ "peak_memory_mb_mean": 67764.0
}
}
}
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
deleted file mode 100644
index 686bd6160..000000000
--- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
+++ /dev/null
@@ -1,190 +0,0 @@
-[
- {
- "test_name": "test_qwen_image_edit_single_device",
- "mark": [
- {
- "hardware_marks": {
- "res": {
- "cuda": "H100"
- },
- "num_cards": 1
- }
- },
- "full_model",
- "diffusion"
- ],
- "description": "Single-device baseline",
- "server_type": "vllm-omni",
- "server_params": {
- "model": "Qwen/Qwen-Image-Edit",
- "serve_args": {
- "enable-diffusion-pipeline-profiler": true
- }
- },
- "benchmark_params": [
- {
- "name": "512x512_steps20_i2i",
- "dataset": "random",
- "task": "i2i",
- "width": 512,
- "height": 512,
- "num-inference-steps": 20,
- "num-prompts": 10,
- "max-concurrency": 1,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": 0.1129,
- "latency_mean": 8.8352,
- "peak_memory_mb_max": 58572.0,
- "peak_memory_mb_mean": 58572.0
- }
- }
- },
- {
- "name": "1536x1536_steps35_i2i",
- "dataset": "random",
- "task": "i2i",
- "width": 1536,
- "height": 1536,
- "num-inference-steps": 35,
- "num-prompts": 10,
- "max-concurrency": 1,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": 0.024,
- "latency_mean": 41.5686,
- "peak_memory_mb_max": 66368.0,
- "peak_memory_mb_mean": 66368.0
- }
- }
- }
- ]
- },
- {
- "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
- "mark": [
- {
- "hardware_marks": {
- "res": {
- "cuda": "H100"
- },
- "num_cards": 4
- }
- },
- "full_model",
- "diffusion"
- ],
- "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
- "server_type": "vllm-omni",
- "server_params": {
- "model": "Qwen/Qwen-Image-Edit",
- "serve_args": {
- "ulysses-degree": 2,
- "cfg-parallel-size": 2,
- "vae-patch-parallel-size": 4,
- "vae-use-tiling": true,
- "enable-diffusion-pipeline-profiler": true
- }
- },
- "benchmark_params": [
- {
- "name": "1536x1536_steps35_i2i",
- "dataset": "random",
- "task": "i2i",
- "width": 1536,
- "height": 1536,
- "num-inference-steps": 35,
- "num-prompts": 10,
- "max-concurrency": 1,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": 0.0672,
- "latency_mean": 14.8517,
- "peak_memory_mb_max": 57076.0,
- "peak_memory_mb_mean": 57076.0
- }
- }
- }
- ]
- },
- {
- "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
- "mark": [
- {
- "hardware_marks": {
- "res": {
- "cuda": "H100"
- },
- "num_cards": 4
- }
- },
- "full_model",
- "diffusion"
- ],
- "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
- "server_type": "vllm-omni",
- "server_params": {
- "model": "Qwen/Qwen-Image-Edit",
- "serve_args": {
- "ulysses-degree": 2,
- "cfg-parallel-size": 2,
- "cache-backend": "cache_dit",
- "cache-config": {
- "Fn_compute_blocks": 1,
- "Bn_compute_blocks": 0,
- "max_warmup_steps": 4,
- "residual_diff_threshold": 0.24,
- "max_continuous_cached_steps": 3,
- "enable_taylorseer": false,
- "taylorseer_order": 1,
- "scm_steps_mask_policy": null,
- "scm_steps_policy": "dynamic"
- },
- "enable-diffusion-pipeline-profiler": true
- }
- },
- "benchmark_params": [
- {
- "name": "512x512_steps20_i2i",
- "dataset": "random",
- "task": "i2i",
- "width": 512,
- "height": 512,
- "num-inference-steps": 20,
- "num-prompts": 10,
- "max-concurrency": 1,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": 0.2879,
- "latency_mean": 3.4701,
- "peak_memory_mb_max": 58758.0,
- "peak_memory_mb_mean": 58758.0
- }
- }
- },
- {
- "name": "1536x1536_steps35_i2i",
- "dataset": "random",
- "task": "i2i",
- "width": 1536,
- "height": 1536,
- "num-inference-steps": 35,
- "num-prompts": 10,
- "max-concurrency": 1,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": 0.1113,
- "latency_mean": 8.9656,
- "peak_memory_mb_max": 67424.0,
- "peak_memory_mb_mean": 67424.0
- }
- }
- }
- ]
- }
-]
diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
index bc4d93a24..bc0f5c923 100644
--- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"diffusion"
],
"description": "Single-device baseline",
@@ -55,8 +55,8 @@
"H100": {
"throughput_qps": 0.0417,
"latency_mean": 23.914,
- "peak_memory_mb_max": 63225.1429,
- "peak_memory_mb_mean": 63225.1429
+ "peak_memory_mb_max": 62912.0,
+ "peak_memory_mb_mean": 62912.0
}
}
}
diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
index ebf9ccc29..4f64d6ec1 100644
--- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -36,7 +36,7 @@
"H100": {
"throughput_qps": 0.4153,
"latency_mean": 2.4065,
- "peak_memory_mb_mean": 55121.4286
+ "peak_memory_mb_mean": 55062.0
}
}
},
@@ -54,7 +54,7 @@
"H100": {
"throughput_qps": 0.0423,
"latency_mean": 23.5991,
- "peak_memory_mb_mean": 65118.0
+ "peak_memory_mb_mean": 64546.0
}
}
}
@@ -74,13 +74,14 @@
"full_model",
"diffusion"
],
- "description": "Single-device baseline (no parallelism) with step execution",
+ "description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).",
"server_type": "vllm-omni",
"server_params": {
"model": "Qwen/Qwen-Image",
"serve_args": {
"enable-diffusion-pipeline-profiler": true,
- "step-execution": true
+ "step-execution": true,
+ "max-num-seqs": 8
}
},
"benchmark_params": [
@@ -98,7 +99,7 @@
"H100": {
"throughput_qps": 0.4175,
"latency_mean": 2.3953,
- "peak_memory_mb_mean": 55088.5714
+ "peak_memory_mb_mean": 54952.0
}
}
},
@@ -116,7 +117,28 @@
"H100": {
"throughput_qps": 0.0423,
"latency_mean": 23.5768,
- "peak_memory_mb_mean": 66012.2857
+ "peak_memory_mb_mean": 64500.0
+ }
+ }
+ },
+ {
+ "name": "512x512_steps20_high_concurrency",
+ "dataset": "random",
+ "task": "t2i",
+ "width": 512,
+ "height": 512,
+ "num-inference-steps": 20,
+ "num-prompts": [20, 40, 80, 160],
+ "max-concurrency": [1, 2, 4, 8],
+ "warmup-requests": 8,
+ "warmup-concurrency": 8,
+ "warmup-num-inference-steps": 20,
+ "enable-negative-prompt": true,
+ "baseline": {
+ "H100": {
+ "throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
+ "latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
+ "peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
}
}
}
@@ -163,7 +185,7 @@
"H100": {
"throughput_qps": 0.1214,
"latency_mean": 8.2211,
- "peak_memory_mb_mean": 54501.1429
+ "peak_memory_mb_mean": 54474.0
}
}
}
@@ -220,7 +242,7 @@
"H100": {
"throughput_qps": 0.559,
"latency_mean": 1.794,
- "peak_memory_mb_mean": 55290.5714
+ "peak_memory_mb_mean": 55174.0
}
}
},
@@ -238,58 +260,10 @@
"H100": {
"throughput_qps": 0.1918,
"latency_mean": 5.2022,
- "peak_memory_mb_mean": 64838.8571
+ "peak_memory_mb_mean": 64750.0
}
}
}
]
- },
- {
- "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
- "mark": [
- {
- "hardware_marks": {
- "res": {
- "cuda": "H100"
- },
- "num_cards": 1
- }
- },
- "full_model",
- "diffusion"
- ],
- "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
- "server_type": "vllm-omni",
- "server_params": {
- "model": "Qwen/Qwen-Image",
- "serve_args": {
- "enable-diffusion-pipeline-profiler": true,
- "step-execution": true,
- "max-num-seqs": 8
- }
- },
- "benchmark_params": [
- {
- "name": "512x512_steps20_high_concurrency",
- "dataset": "random",
- "task": "t2i",
- "width": 512,
- "height": 512,
- "num-inference-steps": 20,
- "num-prompts": [20, 40, 80, 160],
- "max-concurrency": [1, 2, 4, 8],
- "warmup-requests": 8,
- "warmup-concurrency": 8,
- "warmup-num-inference-steps": 20,
- "enable-negative-prompt": true,
- "baseline": {
- "H100": {
- "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
- "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
- "peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
- }
- }
- }
- ]
}
]
diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py
index 74ad29cf6..2e98dcaf7 100644
--- a/tests/dfx/perf/tests/test_runner_metadata.py
+++ b/tests/dfx/perf/tests/test_runner_metadata.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for DFX runner metadata field exclusion."""
import json
@@ -105,7 +108,7 @@ def test_resolve_pytest_marks_hardware_dict_with_extra():
assert "H100" in names
assert "cuda" in names
assert "gpu" in names
- assert "distributed_cuda" in names
+ assert "cards_2" in names
assert "full_model" in names
assert "diffusion" in names
assert "local_model" in names
@@ -512,3 +515,41 @@ def test_omni_duplex_expected_audio_turns_rejects_incomplete_session():
{"expected_duplex_audio_turns_per_session": 4},
1,
)
+
+
+def test_omni_tpot_baseline_accepts_measured_finite_sample():
+ from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+ assert_result(
+ {
+ "completed": 1,
+ "Hardware": "H100",
+ "num_tpot_samples": 1,
+ "mean_tpot_ms": 10.0,
+ },
+ {"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
+ 1,
+ )
+
+
+@pytest.mark.parametrize(
+ ("num_tpot_samples", "mean_tpot_ms", "match"),
+ [
+ (0, float("nan"), "no measurable TPOT samples"),
+ (1, float("nan"), "mean_tpot_ms is not finite"),
+ ],
+)
+def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match):
+ from tests.dfx.perf.scripts.run_benchmark import assert_result
+
+ with pytest.raises(AssertionError, match=match):
+ assert_result(
+ {
+ "completed": 1,
+ "Hardware": "H100",
+ "num_tpot_samples": num_tpot_samples,
+ "mean_tpot_ms": mean_tpot_ms,
+ },
+ {"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
+ 1,
+ )
diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json
index f398e3271..46aa5e83f 100644
--- a/tests/dfx/perf/tests/test_tts.json
+++ b/tests/dfx/perf/tests/test_tts.json
@@ -424,10 +424,10 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
- 82.1886
+ 82.6344
],
"median_audio_rtf": [
- 0.1241
+ 0.1277
]
}
}
@@ -453,13 +453,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
- 160.9168
+ 170.3737
],
"median_audio_rtf": [
- 0.1907
+ 0.2076
],
"audio_throughput": [
- 39.2332
+ 37.0834
]
}
}
@@ -485,13 +485,13 @@
"baseline": {
"H100": {
"median_audio_ttfp_ms": [
- 100.1265
+ 101.0293
],
"median_audio_rtf": [
- 0.1437
+ 0.1498
],
"audio_throughput": [
- 11.7605
+ 10.1186
]
}
}
diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json
index e5163e5bb..1756bc206 100644
--- a/tests/dfx/perf/tests/test_voxcpm2.json
+++ b/tests/dfx/perf/tests/test_voxcpm2.json
@@ -10,7 +10,7 @@
"num_cards": 1
}
},
- "full_model",
+ "local_model",
"tts"
],
"server_params": {
diff --git a/tests/dfx/reliability/README.md b/tests/dfx/reliability/README.md
index 8dc69dc55..79d862025 100644
--- a/tests/dfx/reliability/README.md
+++ b/tests/dfx/reliability/README.md
@@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal
- Qwen3 tests: `test_reliability_qwen3_omni.py`
- Wan2.2 tests: `test_reliability_wan22.py`
- HunyuanImage tests: `test_reliability_hunyuan_image.py`
-- VoxCPM2 tests: `test_reliability_voxcpm2.py`
```bash
pytest --collect-only tests/dfx/reliability
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
-pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
```
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py
index ee1515fa6..f14dcf8c1 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py
@@ -1,13 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""``POST /v1/audio/generate`` validation (live Stable Audio Open server)."""
from __future__ import annotations
+from typing import Any
+
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
@@ -68,10 +70,14 @@ _PARAMS = [
pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
pytest.param(
- {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
+ {"num_inference_steps": 0},
+ ("num_inference_steps", "greater than or equal to 1"),
+ id="num_inference_steps_zero",
),
pytest.param(
- {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
+ {"num_inference_steps": -1},
+ ("num_inference_steps", "greater than or equal to 1"),
+ id="num_inference_steps_negative",
),
pytest.param(
{"num_inference_steps": 6000},
@@ -95,20 +101,20 @@ _PARAMS = [
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_audio_generate_invalid_field_values(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
- body = {"model": omni_server.model, "input": "ambient electronic pad"}
+ body: dict[str, Any] = {"model": omni_server.model, "input": "ambient electronic pad"}
body.update(overrides)
- openai_client.send_audio_generate_http_request(
+ online_client.send_audio_generate_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
-def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_audio_generate_http_request(
+def test_audio_generate_missing_input(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_audio_generate_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py
index 61cd30ec8..9861b609e 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""HTTP + WebSocket validation for Qwen3-TTS Base, Higgs Audio V3 and Audex TTA ``/v1/audio/speech`` (and related routes)."""
from __future__ import annotations
@@ -13,8 +13,8 @@ from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.media import load_test_audio_data_url
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
from tests.helpers.stage_config import get_deploy_config_path
# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni).
@@ -26,7 +26,7 @@ _SPEECH_API_MAX_NEW_TOKENS = 4096
_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024
_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192
_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size).
_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024
@@ -57,6 +57,7 @@ _QWEN3_TTS_SPEECH = [
),
]
+
_HIGGS_AUDIO_V3_SPEECH = [
pytest.param(
OmniServerParams(
@@ -84,7 +85,7 @@ _AUDEX_TTA_SPEECH = [
]
-def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
+def _apply_batch_overrides(body: dict[str, Any], *, loc: str, overrides: dict[str, object]) -> None:
if loc == "batch":
body.update(overrides)
return
@@ -112,8 +113,8 @@ def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) ->
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
-def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_audio_speech_http_request(
+def test_speech_malformed_json(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_audio_speech_http_request(
{
"raw_body": "{",
"timeout": 120,
@@ -124,8 +125,8 @@ def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAICli
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
-def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_audio_speech_http_request(
+def test_speech_missing_required_fields(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_audio_speech_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
@@ -156,8 +157,14 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
"ref_audio": None,
"ref_text": None,
},
- ("CustomVoice", "does not support"),
- id="customvoice_task_not_supported",
+ # NOTE: This is delegating down the TTS adapter, so the message is both specific
+ # to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
+ # (as opposed to having valida speakers, but getting an invalid speaker). it would
+ # be more ideal to validate that this is landing on `.validate()` on the adapter
+ # and ensuring that the returned message matches what we would get from calling the
+ # adapter validate directly.
+ ("no speakers configured",),
+ id="customvoice_invalid_voice",
),
pytest.param(
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
@@ -180,13 +187,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
("mutually exclusive", "speaker_embedding", "ref_audio"),
id="speaker_embedding_with_ref_audio_mutually_exclusive",
),
- pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+ pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
pytest.param(
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
("max_new_tokens", "exceed 4096"),
id="max_new_tokens_above_max",
),
- pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+ pytest.param(
+ {"seed": -(2**63) - 1},
+ ("seed", "greater_than_equal", "-9223372036854775808"),
+ id="seed_below_min",
+ ),
pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
pytest.param(
{"initial_codec_chunk_frames": -1},
@@ -242,13 +253,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_invalid_field_values(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
ov = dict(overrides)
embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False))
- body = {
+ body: dict[str, Any] = {
"model": omni_server.model,
"input": "Hello.",
"voice": "clone",
@@ -261,7 +272,7 @@ def test_speech_invalid_field_values(
body.pop("ref_audio", None)
body.pop("ref_text", None)
body.update(ov)
- openai_client.send_audio_speech_http_request(
+ online_client.send_audio_speech_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -274,7 +285,9 @@ def test_speech_invalid_field_values(
pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
pytest.param(
- {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
+ {"input": "Hello world.", "max_new_tokens": 0},
+ ("max_new_tokens", "greater_than_equal"),
+ id="max_new_tokens_below_min",
),
pytest.param(
{
@@ -328,12 +341,12 @@ def test_speech_invalid_field_values(
@pytest.mark.parametrize("omni_server", _HIGGS_AUDIO_V3_SPEECH, indirect=True)
def test_speech_higgs_invalid_field_values(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
extra_json: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
body = {"model": omni_server.model, **extra_json}
- openai_client.send_audio_speech_http_request(
+ online_client.send_audio_speech_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -359,12 +372,12 @@ def test_speech_higgs_invalid_field_values(
@pytest.mark.parametrize("omni_server", _AUDEX_TTA_SPEECH, indirect=True)
def test_speech_audex_tta_invalid_field_values(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
body = {"model": omni_server.model, "input": _AUDEX_TTA_CAPTION, **overrides}
- openai_client.send_audio_speech_http_request(
+ online_client.send_audio_speech_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -375,15 +388,15 @@ def test_speech_audex_tta_invalid_field_values(
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
-def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_audio_speech_batch_http_request(
+def test_speech_batch_empty_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_audio_speech_batch_http_request(
{"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")}
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
-def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_audio_speech_batch_http_request(
+def test_speech_batch_missing_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_audio_speech_batch_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
@@ -553,14 +566,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open
pytest.param(
"batch",
{"max_new_tokens": 0},
- ("max_new_tokens", "least 1"),
+ ("max_new_tokens", "greater_than_equal"),
id="batch_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"max_new_tokens": 0},
- ("max_new_tokens", "least 1"),
+ ("max_new_tokens", "greater_than_equal"),
id="item_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
@@ -618,14 +631,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_batch_invalid_field_values(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
loc: str,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
ov = dict(overrides)
strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False))
- body = {
+ body: dict[str, Any] = {
"model": omni_server.model,
"voice": "clone",
"ref_audio": REF_AUDIO_URL,
@@ -638,7 +651,7 @@ def test_speech_batch_invalid_field_values(
body.pop("ref_audio", None)
body.pop("ref_text", None)
_apply_batch_overrides(body, loc=loc, overrides=ov)
- openai_client.send_audio_speech_batch_http_request(
+ online_client.send_audio_speech_batch_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -670,7 +683,7 @@ _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object()
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_stream_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
send_frames_spec: Any,
err_message: str | tuple[str, ...],
) -> None:
@@ -687,7 +700,7 @@ def test_speech_stream_invalid_requests(
else:
send_frames = send_frames_spec
assert isinstance(send_frames, str)
- openai_client.send_audio_speech_stream_ws_request(
+ online_client.send_audio_speech_stream_ws_request(
{
"send_frames": send_frames,
"timeout": 120,
@@ -893,7 +906,7 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_voices_create_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
multipart_kind: str | None,
form_template: dict[str, Any],
err_message: str | tuple[str, ...],
@@ -909,7 +922,7 @@ def test_voices_create_invalid_requests(
files = _voices_upload_multipart_files(multipart_kind)
if files is not None:
cfg["files"] = files
- openai_client.send_audio_voices_create_http_request(cfg)
+ online_client.send_audio_voices_create_http_request(cfg)
@pytest.mark.parametrize(
@@ -962,13 +975,13 @@ def test_voices_create_invalid_requests(
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_voices_delete_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
voice_name: str,
err_code: int | tuple[int, ...],
err_message: str | tuple[str, ...],
) -> None:
"""DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405)."""
- openai_client.send_audio_voices_delete_http_request(
+ online_client.send_audio_voices_delete_http_request(
{
"name": voice_name,
"timeout": 120,
@@ -976,3 +989,68 @@ def test_voices_delete_invalid_requests(
"err_message": err_message,
}
)
+
+
+# ─── POST /v1/audio/speech · MiniMax Music 3 ───
+
+# Text-to-music on the speech endpoint. Most of the speech contract does not
+# apply: there is no speaker, no reference audio and no temperature, and the
+# lyrics alone are not a request. Each case below pins a rejection so a caller
+# who assumes the usual TTS shape finds out immediately.
+_MINIMAX_MUSIC3_SPEECH = [
+ pytest.param(
+ OmniServerParams(
+ model="MiniMaxAI/MiniMax-Music3",
+ stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+ server_args=_SPEECH_SERVER_ARGS,
+ ),
+ id="minimax_music3",
+ marks=hardware_marks(res={"cuda": "H100"}),
+ ),
+]
+
+_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
+_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
+# 25 frames per second, capped at six minutes of song.
+_MINIMAX_MUSIC3_MAX_FRAMES = 9000
+
+
+@pytest.mark.parametrize(
+ "overrides, err_message",
+ [
+ pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
+ pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"),
+ pytest.param({"input": " "}, ("input",), id="lyrics_blank"),
+ pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
+ pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
+ pytest.param(
+ {"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
+ ("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
+ id="max_new_tokens_over_cap",
+ ),
+ pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
+ ],
+)
+@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
+def test_minimax_music3_speech_invalid_field_values(
+ omni_server: OmniServer,
+ online_client: OnlineOmniClient,
+ overrides: dict[str, object],
+ err_message: str | tuple[str, ...],
+) -> None:
+ body: dict[str, Any] = {
+ "model": omni_server.model,
+ "input": _MINIMAX_MUSIC3_LYRICS,
+ "instructions": _MINIMAX_MUSIC3_CAPTION,
+ "seed": 1,
+ "max_new_tokens": 250,
+ "response_format": "wav",
+ }
+ for key, value in overrides.items():
+ if value is None:
+ body.pop(key, None)
+ else:
+ body[key] = value
+ online_client.send_audio_speech_http_request(
+ {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+ )
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
index 5ffd25632..9bc5e44fe 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``)."""
from __future__ import annotations
@@ -9,7 +9,7 @@ from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
@@ -72,7 +72,7 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_images_edits_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
tiny_png_bytes: bytes,
include_image: bool,
data_template: dict[str, Any],
@@ -87,4 +87,4 @@ def test_images_edits_invalid_requests(
}
if include_image:
cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")}
- openai_client.send_images_edits_http_request(cfg)
+ online_client.send_images_edits_http_request(cfg)
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
index 5e9424c88..daad00957 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``)."""
from __future__ import annotations
@@ -9,7 +9,7 @@ from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
@@ -126,7 +126,7 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_images_generations_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
body_spec: Any,
err_message: str | tuple[str, ...],
) -> None:
@@ -137,6 +137,6 @@ def test_images_generations_invalid_requests(
assert isinstance(body_spec, dict)
body = _minimal_images_gen_json(omni_server)
body.update(body_spec)
- openai_client.send_images_generations_http_request(
+ online_client.send_images_generations_http_request(
{"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
)
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py
index 9e39b059b..0d026a37b 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
from __future__ import annotations
@@ -10,7 +10,7 @@ from typing import Any
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
from tests.helpers.stage_config import get_deploy_config_path
pytestmark = [pytest.mark.slow, pytest.mark.omni]
@@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
body["top_logprobs"] = 5
elif case_id == "speaker_unknown":
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
+ elif case_id == "audio_format_unsupported":
+ body["modalities"] = ["text", "audio"]
+ body["audio"] = {"voice": "alloy", "format": "aac"}
else:
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
return {"json": body, "timeout": 120}
@@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
pytest.param(
"response_format_json_schema_incomplete",
400,
- ("response_format", "value_error", "json_schema"),
+ ("response_format", "BadRequestError", "json_schema"),
id="invalid_response_format_json_schema",
),
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
@@ -135,12 +138,18 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
("Invalid speaker", "Supported"),
id="speaker_unknown_preset",
),
+ pytest.param(
+ "audio_format_unsupported",
+ 400,
+ ("Invalid audio format", "aac", "Supported formats"),
+ id="unsupported_audio_format",
+ ),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_chat_completions_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
case_id: str,
err_code: int | tuple[int, ...],
err_message: str | tuple[str, ...],
@@ -148,7 +157,7 @@ def test_chat_completions_invalid_requests(
cfg = _chat_completions_request_without_expectations(omni_server, case_id)
cfg["err_code"] = err_code
cfg["err_message"] = err_message
- openai_client.send_chat_completions_http_request(cfg)[0]
+ online_client.send_chat_completions_http_request(cfg)[0]
# ─────────────────────────────────────────────────────────────────────────────
@@ -174,7 +183,7 @@ _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object()
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_video_chat_stream_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
send_frames_spec: Any,
err_message: str,
) -> None:
@@ -183,7 +192,7 @@ def test_video_chat_stream_invalid_requests(
else:
send_frames = send_frames_spec
assert isinstance(send_frames, str)
- openai_client.send_video_chat_stream_ws_request(
+ online_client.send_video_chat_stream_ws_request(
{
"send_frames": send_frames,
"timeout": 120,
@@ -243,11 +252,12 @@ _REALTIME_WS_INVALID_AUDIO_APPEND = object()
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_realtime_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
send_frames_spec: Any,
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(
{
@@ -262,7 +272,7 @@ def test_realtime_invalid_requests(
]
else:
send_frames = send_frames_spec
- openai_client.send_realtime_ws_request(
+ online_client.send_realtime_ws_request(
{
"send_frames": send_frames,
"timeout": 120,
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py
index ad3f659f1..4be838316 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``)."""
from __future__ import annotations
@@ -9,7 +9,7 @@ from typing import Any, Literal
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
@@ -86,7 +86,7 @@ _QWEN_IMAGE = [
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
def test_omni_sleep_wakeup_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
endpoint: Literal["sleep", "wakeup"],
json_body: dict[str, Any],
err_message: str | tuple[str, ...],
@@ -98,9 +98,9 @@ def test_omni_sleep_wakeup_invalid_requests(
"err_message": err_message,
}
if endpoint == "sleep":
- openai_client.send_omni_sleep_http_request(cfg)
+ online_client.send_omni_sleep_http_request(cfg)
else:
- openai_client.send_omni_wakeup_http_request(cfg)
+ online_client.send_omni_wakeup_http_request(cfg)
# Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the
@@ -108,7 +108,7 @@ def test_omni_sleep_wakeup_invalid_requests(
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
def test_wakeup_after_level2_sleep_fails(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
) -> None:
"""Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501.
@@ -116,7 +116,7 @@ def test_wakeup_after_level2_sleep_fails(
Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501.
"""
assert omni_server is not None
- sleep_resps = openai_client.send_omni_sleep_http_request(
+ sleep_resps = online_client.send_omni_sleep_http_request(
{
"json": {"stage_ids": [0], "level": 2},
"timeout": 180,
@@ -124,10 +124,11 @@ def test_wakeup_after_level2_sleep_fails(
)
sleep_resp = sleep_resps[0]
assert sleep_resp.status_code == 200, sleep_resp
- assert sleep_resp.json_body is not None
- assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body
+ json_body = sleep_resp.json_body
+ assert isinstance(json_body, dict)
+ assert json_body.get("status") == "SUCCESS", json_body
- openai_client.send_omni_wakeup_http_request(
+ online_client.send_omni_wakeup_http_request(
{
"json": {"stage_ids": [0]},
"timeout": 120,
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py
index c0dd133c9..75ff814df 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py
+++ b/tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``).
WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in
@@ -13,7 +13,7 @@ from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
@@ -54,12 +54,12 @@ def _minimal_video_form_data() -> dict[str, str]:
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_post_videos_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
form_extra: dict[str, str] | None,
err_message: str | tuple[str, ...],
) -> None:
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
- openai_client.send_videos_create_http_request(
+ online_client.send_videos_create_http_request(
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -85,12 +85,12 @@ def test_post_videos_invalid_requests(
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_post_videos_sync_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
form_extra: dict[str, str] | None,
err_message: str | tuple[str, ...],
) -> None:
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
- openai_client.send_videos_sync_http_request(
+ online_client.send_videos_sync_http_request(
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -111,11 +111,11 @@ def test_post_videos_sync_invalid_requests(
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_get_videos_list_invalid_requests(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
params: dict[str, Any],
err_message: str | tuple[str, ...],
) -> None:
- openai_client.send_videos_list_http_request(
+ online_client.send_videos_list_http_request(
{"params": params, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@@ -126,8 +126,8 @@ def test_get_videos_list_invalid_requests(
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
-def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_video_retrieve_http_request(
+def test_get_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_video_retrieve_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
@@ -143,8 +143,8 @@ def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: Open
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
-def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_video_delete_http_request(
+def test_delete_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_video_delete_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
@@ -160,8 +160,8 @@ def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: O
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
-def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- openai_client.send_video_content_http_request(
+def test_get_video_content_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
+ online_client.send_video_content_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
diff --git a/tests/dfx/reliability/test_reliability_hunyuan_image.py b/tests/dfx/reliability/test_reliability_hunyuan_image.py
index 43a684d1d..ae0133503 100644
--- a/tests/dfx/reliability/test_reliability_hunyuan_image.py
+++ b/tests/dfx/reliability/test_reliability_hunyuan_image.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
HunyuanImage-3.0-Instruct reliability integration tests (DiT-only deploy).
@@ -35,7 +38,7 @@ from tests.dfx.reliability.helpers import (
worker_residual_timeout_after_kill_signal,
)
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OpenAIClientHandler
+from tests.helpers.runtime import OnlineOmniClient
RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
{
@@ -133,8 +136,8 @@ def _image_request_config(omni_server: Any) -> dict[str, Any]:
}
-def _submit_image_generation(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
- responses = openai_client.send_images_generations_http_request(_image_request_config(omni_server))
+def _submit_image_generation(online_client: OnlineOmniClient, omni_server: Any) -> None:
+ responses = online_client.send_images_generations_http_request(_image_request_config(omni_server))
resp = responses[0]
if resp.status_code >= 400:
raise RuntimeError(
@@ -189,7 +192,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_gpu_oom_image_large_request_failure(
omni_server_function,
- openai_client_function,
+ online_client_function,
) -> None:
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
@@ -202,7 +205,7 @@ def test_reliability_fault_gpu_oom_image_large_request_failure(
)
try:
try:
- _submit_image_generation(openai_client_function, omni_server_function)
+ _submit_image_generation(online_client_function, omni_server_function)
except Exception as exc:
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
else:
@@ -222,10 +225,10 @@ def test_reliability_fault_gpu_oom_image_large_request_failure(
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_request_failure(
omni_server_after_fault_function,
- openai_client_function,
+ online_client_function,
) -> None:
try:
- _submit_image_generation(openai_client_function, omni_server_after_fault_function)
+ _submit_image_generation(online_client_function, omni_server_after_fault_function)
except Exception as exc:
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
else:
@@ -244,7 +247,7 @@ def test_reliability_fault_process_kill_request_failure(
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
omni_server_after_fault_function,
- openai_client_function,
+ online_client_function,
) -> None:
host = omni_server_after_fault_function.host
port = omni_server_after_fault_function.port
@@ -280,7 +283,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
start = time.monotonic()
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
futures = [
- executor.submit(_submit_image_generation, openai_client_function, omni_server_after_fault_function)
+ executor.submit(_submit_image_generation, online_client_function, omni_server_after_fault_function)
for _ in range(3)
]
done, pending = concurrent.futures.wait(
@@ -335,12 +338,12 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_worker_with_load_request_failure(
omni_server_function,
- openai_client_function,
+ online_client_function,
fault_injector: FaultInjector,
) -> None:
scenario = "kill_worker_with_load"
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+ submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
inject_fault=lambda: fault_injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -363,13 +366,13 @@ def test_reliability_fault_process_kill_worker_with_load_request_failure(
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+ submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -448,7 +451,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
@@ -458,7 +461,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
inter_kill_wait_seconds=0.1,
)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
+ submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py
index 32c660954..10331d88d 100644
--- a/tests/dfx/reliability/test_reliability_qwen3_omni.py
+++ b/tests/dfx/reliability/test_reliability_qwen3_omni.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Qwen3-Omni reliability integration tests.
"""
@@ -63,7 +66,7 @@ DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vll
def _default_oom_device_spec() -> str:
"""Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
- count = torch.accelerator.device_count()
+ count = torch.accelerator.device_count() or 0
if count <= 0:
return "0"
return ",".join(str(i) for i in range(count))
@@ -71,7 +74,7 @@ def _default_oom_device_spec() -> str:
OOM_INJECTION_CONFIG = {
"device": _default_oom_device_spec(),
- "target_mem_ratio": 0.92,
+ "target_mem_ratio": 0.95,
"hold_seconds": 0,
"startup_timeout_sec": 20,
"strict": False,
@@ -414,7 +417,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@pytest.mark.slow
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
+def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, online_client_function) -> None:
device_spec = resolve_oom_device_spec(
OOM_INJECTION_CONFIG,
_deploy_config_path_from_omni_server(omni_server_function),
@@ -444,7 +447,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
"key_words": {"audio": ["test"]},
}
try:
- openai_client_function.send_omni_request(request_config, request_num=1)
+ online_client_function.send_omni_request(request_config, request_num=1)
except Exception as exc:
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
else:
@@ -456,7 +459,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
@pytest.mark.slow
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
+def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, online_client_function) -> None:
device_spec = resolve_oom_device_spec(
OOM_INJECTION_CONFIG,
_deploy_config_path_from_omni_server(omni_server_function),
@@ -487,7 +490,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
"key_words": {"audio": ["test"]},
}
try:
- openai_client_function.send_omni_request(request_config, request_num=4)
+ online_client_function.send_omni_request(request_config, request_num=4)
except Exception as exc:
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
else:
@@ -505,7 +508,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
)
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_request_failure(
- omni_server_after_fault_function, openai_client_function
+ omni_server_after_fault_function, online_client_function
) -> None:
messages = dummy_messages_from_mix_data(
system_prompt=_get_system_prompt(),
@@ -519,7 +522,7 @@ def test_reliability_fault_process_kill_request_failure(
"key_words": {"text": ["beijing"]},
}
try:
- openai_client_function.send_omni_request(request_config, request_num=1)
+ online_client_function.send_omni_request(request_config, request_num=1)
except Exception as exc:
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
else:
@@ -567,7 +570,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
)
- assert ff_status >= 500, (
+ assert ff_status is not None and ff_status >= 500, (
f"[process_kill fast_fail] expected server-side failure after fault, "
f"got status={ff_status}, body={ff_body[:200]!r}"
)
@@ -652,13 +655,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_worker_with_load_request_failure(
omni_server_function,
- openai_client_function,
+ online_client_function,
fault_injector: FaultInjector,
) -> None:
request_config = _chat_request_config(omni_server_function)
scenario = "kill_worker_with_load"
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
inject_fault=lambda: fault_injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -704,14 +707,14 @@ def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
request_config = _chat_request_config(omni_server_function)
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -764,7 +767,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
request_config = _chat_request_config(omni_server_function)
@@ -775,7 +778,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
inter_kill_wait_seconds=0.1,
)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
diff --git a/tests/dfx/reliability/test_reliability_voxcpm2.py b/tests/dfx/reliability/test_reliability_voxcpm2.py
deleted file mode 100644
index 632c51081..000000000
--- a/tests/dfx/reliability/test_reliability_voxcpm2.py
+++ /dev/null
@@ -1,508 +0,0 @@
-"""
-VoxCPM2 reliability integration tests.
-
-Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
-"""
-
-from __future__ import annotations
-
-import concurrent.futures
-import os
-import time
-from pathlib import Path
-from typing import Any
-
-import pytest
-
-from tests.dfx.conftest import (
- assert_fault_exception,
- create_reliability_omni_server_params,
- resolve_oom_device_spec,
-)
-from tests.dfx.reliability.helpers import (
- FaultInjector,
- assert_no_server_tree_process_residual_and_gpu_release,
- extract_openai_error_contract_from_bytes,
- get_health_raw,
- inject_gpu_oom,
- make_process_kill_fault_injector,
- make_server_root_kill_fault_injector,
- make_server_tree_kill_fault_injector,
- post_json_raw,
- run_fault_injection_with_rate_load,
- stop_gpu_oom_hogs,
- worker_residual_timeout_after_kill_signal,
-)
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
-from tests.helpers.runtime import OpenAIClientHandler
-
-RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
- {
- "test_name": "voxcpm2_reliability_default",
- "server_params": {
- "model": "openbmb/VoxCPM2",
- "stage_config_name": "voxcpm2.yaml",
- # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
- # would otherwise override deploy YAML settings during server startup.
- "server_args": ["--trust-remote-code", "--disable-log-stats"],
- },
- },
-]
-
-DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
-OOM_INJECTION_CONFIG = {
- "target_mem_ratio": 0.96,
- "hold_seconds": 0,
- "startup_timeout_sec": 20,
- "strict": False,
-}
-FAULT_ERROR_KEYWORDS = (
- "the request failed",
- "oom",
- "out of memory",
- "cuda",
- "orchestrator",
- "timeout",
- "connection",
- "500",
- "503",
-)
-PROCESS_KILL_ERROR_KEYWORDS = (
- "timeout",
- "did not complete within",
- "connection",
- "engine",
- "orchestrator",
- "dead",
- "internal",
- "500",
- "503",
-)
-RUNTIME_WORKER_PATTERN = "VLLM::"
-SERVE_SIGNAL_PARAMS = [
- pytest.param("SIGTERM", id="sigterm"),
- pytest.param("SIGINT", id="sigint"),
- pytest.param("SIGKILL", id="sigkill"),
-]
-SERVE_WITH_LOAD_SIGNAL_PARAMS = [
- pytest.param("SIGTERM", id="sigterm"),
- pytest.param("SIGINT", id="sigint"),
- pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_SIGNAL_PARAMS = [
- pytest.param("SIGTERM", id="sigterm"),
- pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_WITH_LOAD_SIGNAL_PARAMS = [
- pytest.param("SIGTERM", id="sigterm"),
- pytest.param("SIGKILL", id="sigkill"),
-]
-WORKER_SIGNAL_FAULT_PARAMS = [
- pytest.param(
- make_process_kill_fault_injector(
- grep_patterns=RUNTIME_WORKER_PATTERN,
- signal_name="SIGTERM",
- limit=1,
- post_kill_wait_seconds=2.0,
- ),
- id="runtime_process_chain_sigterm",
- ),
- pytest.param(
- make_process_kill_fault_injector(
- grep_patterns=RUNTIME_WORKER_PATTERN,
- signal_name="SIGKILL",
- limit=1,
- post_kill_wait_seconds=2.0,
- ),
- id="runtime_process_chain_sigkill",
- ),
-]
-
-VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
-INFLIGHT_INJECTION_REQUEST_RATE = 0.3
-INFLIGHT_INJECTION_REQUEST_COUNT = 10
-_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
-_SPEECH_FF_MAX_ELAPSED_SEC = 30
-
-
-def _speech_request_config(omni_server: Any) -> dict[str, Any]:
- return {
- "model": omni_server.model,
- "input": "The weather is nice today, perfect for a walk in the park.",
- "stream": False,
- "timeout": 120.0,
- "response_format": "wav",
- "voice": "default",
- }
-
-
-def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
- openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
-
-
-def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
- base = "The weather is nice today, perfect for a walk in the park. "
- return {
- "model": omni_server.model,
- "input": base * 80,
- "max_new_tokens": 4096,
- "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
- "stream": False,
- "timeout": 300.0,
- "response_format": "wav",
- "voice": "default",
- }
-
-
-def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
- openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
-
-
-def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
- payload = {
- "model": model,
- "input": "post-fault fast-fail check",
- "voice": "default",
- }
- start = time.monotonic()
- try:
- status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
- elapsed = time.monotonic() - start
- assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
- assert status >= 500, (
- f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
- )
- except Exception: # noqa: BLE001
- elapsed = time.monotonic() - start
- assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
-
-
-def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
- deadline = time.monotonic() + 20.0
- last_observation = ""
- while time.monotonic() < deadline:
- try:
- status, body = get_health_raw(host, port, timeout_sec=5)
- last_observation = f"http={status}, body={body[:200]!r}"
- if status == 503:
- return
- except Exception as exc: # noqa: BLE001
- last_observation = f"exception={exc!r}"
- return
- time.sleep(0.5)
- pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_request_failure(
- omni_server_function,
- openai_client_function,
-) -> None:
- stage_config_path = getattr(omni_server_function, "stage_config_path", None)
- device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
- handle = inject_gpu_oom(
- device=device_spec,
- target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
- hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
- startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
- strict=OOM_INJECTION_CONFIG["strict"],
- )
- try:
- try:
- _submit_oom_speech_request(openai_client_function, omni_server_function)
- except Exception as exc:
- assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
- else:
- pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
- finally:
- stop_gpu_oom_hogs(handle)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skip(reason="issue#4285")
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_speech_error_contract(
- omni_server_function,
-) -> None:
- stage_config_path = getattr(omni_server_function, "stage_config_path", None)
- device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
- handle = inject_gpu_oom(
- device=device_spec,
- target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
- hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
- startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
- strict=OOM_INJECTION_CONFIG["strict"],
- )
- host = omni_server_function.host
- port = omni_server_function.port
- oom_config = _oom_speech_request_config(omni_server_function)
- payload = {
- "model": oom_config["model"],
- "input": oom_config["input"],
- "voice": oom_config["voice"],
- "max_new_tokens": oom_config["max_new_tokens"],
- "ref_audio": oom_config["ref_audio"],
- }
- try:
- status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
- finally:
- stop_gpu_oom_hogs(handle)
-
- assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
- speech_error = extract_openai_error_contract_from_bytes(body)
- assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
- assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
- "fault_injector",
- WORKER_SIGNAL_FAULT_PARAMS,
- indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_request_failure(
- omni_server_after_fault_function,
- openai_client_function,
-) -> None:
- try:
- _submit_speech_request(openai_client_function, omni_server_after_fault_function)
- except Exception as exc:
- assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
- else:
- pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
- "fault_injector",
- WORKER_SIGNAL_FAULT_PARAMS,
- indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_fast_fail_and_concurrent(
- omni_server_after_fault_function,
- openai_client_function,
-) -> None:
- host = omni_server_after_fault_function.host
- port = omni_server_after_fault_function.port
- model = omni_server_after_fault_function.model
- payload = {
- "model": model,
- "input": "Say hello in one short sentence.",
- "voice": "default",
- }
-
- start = time.monotonic()
- try:
- status, body = post_json_raw(
- host,
- port,
- "/v1/audio/speech",
- payload,
- timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
- )
- elapsed = time.monotonic() - start
- assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
- f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
- )
- assert status >= 500, (
- "[process_kill fast_fail] expected server-side failure after fault, "
- f"got status={status}, body={body[:200]!r}"
- )
- except Exception:
- elapsed = time.monotonic() - start
- assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
- f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
- )
-
- start = time.monotonic()
- with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
- futures = [
- executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
- for _ in range(3)
- ]
- done, pending = concurrent.futures.wait(
- futures,
- timeout=30,
- return_when=concurrent.futures.ALL_COMPLETED,
- )
-
- elapsed = time.monotonic() - start
- assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
- assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
-
- fault_observed = False
- for future in done:
- try:
- future.result()
- except Exception as exc:
- fault_observed = True
- assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
- assert fault_observed, (
- "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
- )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
- "fault_injector",
- WORKER_SIGNAL_FAULT_PARAMS,
- indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_worker_with_load_request_failure(
- omni_server_function,
- openai_client_function,
- fault_injector: FaultInjector,
-) -> None:
- scenario = "kill_worker_with_load"
- load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
- inject_fault=lambda: fault_injector(omni_server_function),
- num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
- request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
- completion_timeout_sec=120.0,
- )
- assert load_result["failure_observed"], (
- f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
- )
- host = omni_server_function.host
- port = omni_server_function.port
- _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
- _assert_post_fault_health_terminal(host, port, scenario=scenario)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
- omni_server_function,
- signal_name: str,
-) -> None:
- scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
- injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
- injector(omni_server_function)
- host = omni_server_function.host
- port = omni_server_function.port
- _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
- _assert_post_fault_health_terminal(host, port, scenario=scenario)
- assert_no_server_tree_process_residual_and_gpu_release(
- omni_server_function,
- scenario=scenario,
- timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
- )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
- omni_server_function,
- openai_client_function,
- signal_name: str,
-) -> None:
- scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
- injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
- load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
- inject_fault=lambda: injector(omni_server_function),
- num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
- request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
- completion_timeout_sec=300.0,
- )
- assert load_result["failure_observed"], (
- f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
- )
- host = omni_server_function.host
- port = omni_server_function.port
- _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
- _assert_post_fault_health_terminal(host, port, scenario=scenario)
- assert_no_server_tree_process_residual_and_gpu_release(
- omni_server_function,
- scenario=scenario,
- timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
- )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
- omni_server_function,
- signal_name: str,
-) -> None:
- scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
- injector = make_server_tree_kill_fault_injector(
- signal_name=signal_name,
- post_kill_wait_seconds=2.0,
- inter_kill_wait_seconds=0.1,
- )
- injector(omni_server_function)
- host = omni_server_function.host
- port = omni_server_function.port
- _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
- _assert_post_fault_health_terminal(host, port, scenario=scenario)
- assert_no_server_tree_process_residual_and_gpu_release(
- omni_server_function,
- scenario=scenario,
- timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
- )
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
- omni_server_function,
- openai_client_function,
- signal_name: str,
-) -> None:
- scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
- injector = make_server_tree_kill_fault_injector(
- signal_name=signal_name,
- post_kill_wait_seconds=2.0,
- inter_kill_wait_seconds=0.1,
- )
- load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
- inject_fault=lambda: injector(omni_server_function),
- num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
- request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
- completion_timeout_sec=120.0,
- )
- assert load_result["failure_observed"], (
- f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
- )
- host = omni_server_function.host
- port = omni_server_function.port
- _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
- _assert_post_fault_health_terminal(host, port, scenario=scenario)
- assert_no_server_tree_process_residual_and_gpu_release(
- omni_server_function,
- scenario=scenario,
- timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
- )
diff --git a/tests/dfx/reliability/test_reliability_wan22.py b/tests/dfx/reliability/test_reliability_wan22.py
index 84175429d..8a86c1c92 100644
--- a/tests/dfx/reliability/test_reliability_wan22.py
+++ b/tests/dfx/reliability/test_reliability_wan22.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Wan2.2 reliability integration tests.
"""
@@ -183,7 +186,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
@pytest.mark.slow
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, openai_client_function) -> None:
+def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, online_client_function) -> None:
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
handle = inject_gpu_oom(
@@ -211,7 +214,7 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
"stream": False,
}
try:
- openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+ online_client_function.send_video_diffusion_request(request_config, request_num=1)
except Exception as exc:
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
else:
@@ -230,11 +233,11 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
def test_reliability_fault_process_kill_request_failure(
omni_server_after_fault_function,
- openai_client_function,
+ online_client_function,
) -> None:
request_config = _video_request_config()
try:
- openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+ online_client_function.send_video_diffusion_request(request_config, request_num=1)
except Exception as exc:
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
else:
@@ -251,7 +254,7 @@ def test_reliability_fault_process_kill_request_failure(
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
omni_server_after_fault_function,
- openai_client_function,
+ online_client_function,
) -> None:
"""Black-box: after process kill, /v1/videos fails fast and concurrent calls don't hang.
@@ -308,7 +311,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
start = time.monotonic()
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
futures = [
- executor.submit(openai_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
+ executor.submit(online_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
]
done, pending = concurrent.futures.wait(
futures,
@@ -367,13 +370,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
def test_reliability_fault_process_kill_worker_with_load_request_failure(
omni_server_function,
- openai_client_function,
+ online_client_function,
fault_injector: FaultInjector,
) -> None:
request_config = _video_request_config()
scenario = "kill_worker_with_load"
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
inject_fault=lambda: fault_injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -516,7 +519,7 @@ def test_reliability_video_oom_recovers_after_fault_removed(
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
"""Black-box: during in-flight video requests, kill serve root and verify post-fault behavior/cleanup."""
@@ -524,7 +527,7 @@ def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_clean
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
@@ -601,7 +604,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
omni_server_function,
- openai_client_function,
+ online_client_function,
signal_name: str,
) -> None:
"""Black-box: during in-flight video requests, kill server process tree and verify post-fault behavior/cleanup."""
@@ -613,7 +616,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
inter_kill_wait_seconds=0.1,
)
load_result = run_fault_injection_with_rate_load(
- submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
+ submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
inject_fault=lambda: injector(omni_server_function),
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
diff --git a/tests/dfx/stability/conftest.py b/tests/dfx/stability/conftest.py
index eaa755d45..0889db885 100644
--- a/tests/dfx/stability/conftest.py
+++ b/tests/dfx/stability/conftest.py
@@ -37,7 +37,7 @@ def omni_server(request: pytest.FixtureRequest):
with _omni_server_lock:
# Same tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py``;
# ``serve_args`` from JSON are folded into ``extra_cli_args`` inside
- # ``create_unique_server_params``.
+ # ``create_unique_server_pytest_params`` (JSON ``mark`` → pytest marks).
test_name, model, deploy_path, stage_overrides, extra_cli_args, use_omni = request.param
print(f"Starting OmniServer with test: {test_name}, model: {model}")
diff --git a/tests/dfx/stability/scripts/test_stability_hunyuan_image.py b/tests/dfx/stability/scripts/test_stability_hunyuan_image.py
index 38858047f..be2aec805 100644
--- a/tests/dfx/stability/scripts/test_stability_hunyuan_image.py
+++ b/tests/dfx/stability/scripts/test_stability_hunyuan_image.py
@@ -13,7 +13,7 @@ import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
- create_unique_server_params,
+ create_unique_server_pytest_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
@@ -29,7 +29,7 @@ try:
except FileNotFoundError:
BENCHMARK_CONFIGS = []
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_omni.py b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py
index 5f487a27c..98ab5ee89 100644
--- a/tests/dfx/stability/scripts/test_stability_qwen3_omni.py
+++ b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py
@@ -13,7 +13,7 @@ import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
- create_unique_server_params,
+ create_unique_server_pytest_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
@@ -29,7 +29,7 @@ try:
except FileNotFoundError:
BENCHMARK_CONFIGS = []
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_tts.py b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py
index 0c8059ed9..293b723ca 100644
--- a/tests/dfx/stability/scripts/test_stability_qwen3_tts.py
+++ b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py
@@ -13,7 +13,7 @@ import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
- create_unique_server_params,
+ create_unique_server_pytest_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
@@ -29,7 +29,7 @@ try:
except FileNotFoundError:
BENCHMARK_CONFIGS = []
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
diff --git a/tests/dfx/stability/scripts/test_stability_qwen_image.py b/tests/dfx/stability/scripts/test_stability_qwen_image.py
index 7ff39b5ca..62b0cd86c 100644
--- a/tests/dfx/stability/scripts/test_stability_qwen_image.py
+++ b/tests/dfx/stability/scripts/test_stability_qwen_image.py
@@ -13,7 +13,7 @@ import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
- create_unique_server_params,
+ create_unique_server_pytest_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
@@ -29,7 +29,7 @@ try:
except FileNotFoundError:
BENCHMARK_CONFIGS = []
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
diff --git a/tests/dfx/stability/scripts/test_stability_voxcpm2.py b/tests/dfx/stability/scripts/test_stability_voxcpm2.py
deleted file mode 100644
index e26a2050c..000000000
--- a/tests/dfx/stability/scripts/test_stability_voxcpm2.py
+++ /dev/null
@@ -1,69 +0,0 @@
-"""
-VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
-
-Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from tests.dfx.conftest import (
- create_benchmark_indices,
- create_test_parameter_mapping,
- create_unique_server_params,
- load_configs,
-)
-from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
-
-STABILITY_DIR = Path(__file__).resolve().parent.parent
-DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
-CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
-DEFAULT_NUM_PROMPTS_PER_BATCH = 20
-STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
-
-try:
- BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-except FileNotFoundError:
- BENCHMARK_CONFIGS = []
-
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
-server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.mark.slow
-@pytest.mark.tts
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
-def test_stability_voxcpm2(omni_server, stability_benchmark_params):
- test_name = stability_benchmark_params["test_name"]
- params = stability_benchmark_params["params"]
- duration_sec = params.get("duration_sec", 300)
- num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
- request_rate = params.get("request_rate")
- max_concurrency = params.get("max_concurrency")
-
- bench_params = {
- k: v
- for k, v in params.items()
- if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
- }
-
- result = run_stability_benchmark_loop(
- host=omni_server.host,
- port=omni_server.port,
- model=omni_server.model,
- duration_sec=duration_sec,
- params=bench_params,
- request_rate=request_rate,
- max_concurrency=max_concurrency,
- result_dir=str(STABILITY_DIR),
- num_prompts_per_batch=num_prompts_per_batch,
- run_one_batch=_run_one_vllm_bench_batch,
- )
-
- assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
- assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"
diff --git a/tests/dfx/stability/scripts/test_stability_wan22.py b/tests/dfx/stability/scripts/test_stability_wan22.py
index a363133f5..aae0ae429 100644
--- a/tests/dfx/stability/scripts/test_stability_wan22.py
+++ b/tests/dfx/stability/scripts/test_stability_wan22.py
@@ -13,7 +13,7 @@ import pytest
from tests.dfx.conftest import (
create_benchmark_indices,
create_test_parameter_mapping,
- create_unique_server_params,
+ create_unique_server_pytest_params,
load_configs,
)
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
@@ -29,7 +29,7 @@ try:
except FileNotFoundError:
BENCHMARK_CONFIGS = []
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
+test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
diff --git a/tests/dfx/stability/tests/test_hunyuan_image.json b/tests/dfx/stability/tests/test_hunyuan_image.json
index b78fa812c..b468ceec3 100644
--- a/tests/dfx/stability/tests/test_hunyuan_image.json
+++ b/tests/dfx/stability/tests/test_hunyuan_image.json
@@ -1,6 +1,18 @@
[
{
"test_name": "test_hunyuan_image_stability_tp2_fp8_sp2",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H800"
+ },
+ "num_cards": 2
+ }
+ },
+ "local_model",
+ "diffusion"
+ ],
"server_params": {
"model": "tencent/HunyuanImage-3.0-Instruct",
"serve_args": {
diff --git a/tests/dfx/stability/tests/test_qwen3_omni.json b/tests/dfx/stability/tests/test_qwen3_omni.json
index cbdf25e67..816e981b6 100644
--- a/tests/dfx/stability/tests/test_qwen3_omni.json
+++ b/tests/dfx/stability/tests/test_qwen3_omni.json
@@ -1,6 +1,18 @@
[
{
"test_name": "test_qwen3_omni_stability_async_chunk",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H800"
+ },
+ "num_cards": 2
+ }
+ },
+ "local_model",
+ "omni"
+ ],
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"extra_cli_args": ["--async-chunk"]
diff --git a/tests/dfx/stability/tests/test_qwen3_tts.json b/tests/dfx/stability/tests/test_qwen3_tts.json
index 0208331bd..1c9356574 100644
--- a/tests/dfx/stability/tests/test_qwen3_tts.json
+++ b/tests/dfx/stability/tests/test_qwen3_tts.json
@@ -1,6 +1,19 @@
[
{
"test_name": "test_qwen3_tts_stability",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H800",
+ "npu": "A3"
+ },
+ "num_cards": 2
+ }
+ },
+ "local_model",
+ "tts"
+ ],
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
},
diff --git a/tests/dfx/stability/tests/test_qwen_image.json b/tests/dfx/stability/tests/test_qwen_image.json
index e28614973..0c21d0ac1 100644
--- a/tests/dfx/stability/tests/test_qwen_image.json
+++ b/tests/dfx/stability/tests/test_qwen_image.json
@@ -1,6 +1,18 @@
[
{
"test_name": "test_qwen_image_stability",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H800"
+ },
+ "num_cards": 1
+ }
+ },
+ "local_model",
+ "diffusion"
+ ],
"server_params": {
"model": "Qwen/Qwen-Image"
},
diff --git a/tests/dfx/stability/tests/test_voxcpm2.json b/tests/dfx/stability/tests/test_voxcpm2.json
deleted file mode 100644
index dd912da36..000000000
--- a/tests/dfx/stability/tests/test_voxcpm2.json
+++ /dev/null
@@ -1,35 +0,0 @@
-[
- {
- "test_name": "test_voxcpm2_stability",
- "server_params": {
- "model": "openbmb/VoxCPM2",
- "trust_remote_code": true
- },
- "benchmark_params": [
- {
- "dataset_name": "random",
- "backend": "openai-audio-speech",
- "endpoint": "/v1/audio/speech",
- "trust_remote_code": true,
- "duration_sec": 86400,
- "max_concurrency": 2,
- "num_prompts_per_batch": 50,
- "random_input_len": {
- "min": 0,
- "max": 1000
- },
- "random_output_len": {
- "min": 0,
- "max": 1000
- },
- "random_range_ratio": 0.0,
- "extra_body": {
- "voice": "default",
- "language": "English"
- },
- "ignore_eos": true,
- "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration"
- }
- ]
- }
-]
diff --git a/tests/dfx/stability/tests/test_wan22.json b/tests/dfx/stability/tests/test_wan22.json
index 7b3b54b65..e6ab0b5cf 100644
--- a/tests/dfx/stability/tests/test_wan22.json
+++ b/tests/dfx/stability/tests/test_wan22.json
@@ -1,6 +1,19 @@
[
{
- "test_name": "test_wan22_i2v_stability_v1_videos",
+ "test_name": "test_wan22_i2v_stability_ulysses2_vae_patch2_hsdp_slicing_tiling",
+ "description": "Ulysses=2 + VAE patch parallel=2 + HSDP + VAE slicing + tiling",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "cuda": "H800"
+ },
+ "num_cards": 2
+ }
+ },
+ "local_model",
+ "diffusion"
+ ],
"server_params": {
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
"serve_args": {
@@ -27,5 +40,45 @@
]
}
]
+ },
+ {
+ "test_name": "test_wan22_i2v_stability_usp2_vae_patch2_hsdp_slicing",
+ "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
+ "mark": [
+ {
+ "hardware_marks": {
+ "res": {
+ "npu": "A3"
+ },
+ "num_cards": 2
+ }
+ },
+ "local_model",
+ "diffusion"
+ ],
+ "server_params": {
+ "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+ "serve_args": {
+ "usp": 2,
+ "vae-patch-parallel-size": 2,
+ "use-hsdp": true,
+ "vae-use-slicing": true
+ }
+ },
+ "benchmark_params": [
+ {
+ "dataset": "random",
+ "task": "i2v",
+ "backend": "v1/videos",
+ "duration_sec": 86400,
+ "max_concurrency": 1,
+ "num_prompts_per_batch": 50,
+ "enable_negative_prompt": true,
+ "random_request_config": [
+ {"width": 832, "height": 480, "num_inference_steps": 4, "num_frames": 81, "fps": 16, "weight": 0.5},
+ {"width": 1280, "height": 720, "num_inference_steps": 4, "num_frames": 121, "fps": 16, "weight": 0.5}
+ ]
+ }
+ ]
}
]
diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py
index 7b94fe80f..10451c134 100644
--- a/tests/diffusion/ar_diffusion/test_capability_runner.py
+++ b/tests/diffusion/ar_diffusion/test_capability_runner.py
@@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
)
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
BLOCK = 16
POS = "positive"
@@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
)
+def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+ return ARDiffusionKVCacheSpec(
+ num_layers=1,
+ num_kv_heads=1,
+ head_size=1,
+ tokens_per_frame=1,
+ frames_per_block=3,
+ window_frames=3,
+ sink_frames=3,
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=capacity,
+ cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
+ )
+
+
class CapablePipeline:
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
self.spec = spec
@@ -103,6 +119,7 @@ def make_runner(
*,
available_bytes: int = 1 << 28,
step_execution: bool = False,
+ gpu_memory_fraction: float = 0.1,
) -> ARDiffusionModelRunner:
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(
@@ -113,7 +130,10 @@ def make_runner(
)
runner.device = torch.device("cpu")
runner.pipeline = pipeline
- runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+ runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
+ enable=True,
+ gpu_memory_fraction=gpu_memory_fraction,
+ )
runner.kv_cache = None
runner._ar_diffusion_capability = None
runner._ar_diffusion_kv_cache_spec = None
@@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
return state
+def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
+ state = runner._get_or_create_session(session_id)
+ assert runner.kv_cache is not None
+ seq_len = BLOCK * runner.kv_cache.frames_per_block
+ ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
+ ctx.ensure_video_slots(torch.device("cpu"))
+ state.commit_paged_context(kv_branch)
+ return state
+
+
def test_ar_runner_rejects_pipeline_without_capability():
runner = object.__new__(ARDiffusionModelRunner)
runner.od_config = SimpleNamespace(max_num_seqs=1)
@@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability():
runner._preallocate_kv_cache(available_bytes=1 << 20)
+def test_runner_uses_typed_tick_as_authoritative_session_contract():
+ tick = ARDiffusionTickRequest(
+ session_id="world-7",
+ request_id="request-3",
+ chunk_index=3,
+ reset=True,
+ )
+ req = SimpleNamespace(
+ request_id="request-3",
+ sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+ )
+
+ session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
+
+ assert session_id == "world-7"
+ assert extra_args == tick.to_extra_args()
+ assert parsed == tick
+
+
+def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
+ tick = ARDiffusionTickRequest(
+ session_id="world-7",
+ request_id="client-request-3",
+ chunk_index=3,
+ )
+ req = SimpleNamespace(
+ request_id="engine-request-uuid",
+ sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
+ )
+
+ session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
+
+ assert session_id == "world-7"
+ assert req.request_id == "engine-request-uuid"
+ assert parsed.request_id == "client-request-3"
+
+
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
pipeline = CapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline)
@@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
assert pipeline.closes == ["s1"]
+def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
+ runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
+ kv = runner.kv_cache
+ assert kv is not None
+
+ session_a = commit_one_block(runner, "world-a", "main")
+ a0_blocks = kv.window_block_ids(session_a.adapter("main"))
+ session_b = commit_one_block(runner, "world-b", "main")
+ b0_blocks = kv.window_block_ids(session_b.adapter("main"))
+ session_a_again = commit_one_block(runner, "world-a", "main")
+ a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
+
+ assert session_a_again is session_a
+ assert session_a.adapter("main").request_id == "ar::world-a::main"
+ assert session_b.adapter("main").request_id == "ar::world-b::main"
+ assert session_a.adapter("main").completed_chunks == 6
+ assert session_b.adapter("main").completed_chunks == 3
+ assert len(a0_blocks) == 3
+ assert len(b0_blocks) == 3
+ assert len(a1_blocks) == 6
+ assert set(a0_blocks) <= set(a1_blocks)
+ assert set(a1_blocks).isdisjoint(b0_blocks)
+ assert tuple(runner._sessions) == ("world-b", "world-a")
+
+
def test_lingbot_like_sink_survives_sliding_window_eviction():
runner = make_runner(CapablePipeline(lingbot_like_spec()))
kv = runner.kv_cache
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
- pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+ pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
runner = make_runner(pipeline)
kv = runner.kv_cache
assert kv is not None
- assert runner._session_capacity == 1
- assert kv.session_capacity == 1
+ assert runner._session_capacity == 2
+ assert kv.session_capacity == 2
free_total = kv.manager.block_pool.get_num_free_blocks()
old = commit_one_frame(runner, "old", "main")
k = torch.randn(1, 8, 4, 64)
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
assert kv.manager.block_pool.get_num_free_blocks() < free_total
runner._get_or_create_session("new")
+ assert tuple(runner._sessions) == ("old", "new")
+ runner._get_or_create_session("newest")
- assert tuple(runner._sessions) == ("new",)
+ assert tuple(runner._sessions) == ("new", "newest")
assert pipeline.closes == ["old"]
assert "old" not in kv._cross_sessions
assert kv.manager.block_pool.get_num_free_blocks() == free_total
+def test_budget_reduced_capacity_drives_runner_lru():
+ pipeline = CapablePipeline(tiny_spec(capacity=2))
+ # One tiny LingBot-like session requires 128 bytes; two require 192.
+ runner = make_runner(
+ pipeline,
+ available_bytes=128,
+ gpu_memory_fraction=1.0,
+ )
+ kv = runner.kv_cache
+ assert kv is not None
+ assert kv.requested_session_capacity == 2
+ assert kv.session_capacity == 1
+ assert runner._session_capacity == 1
+
+ runner._get_or_create_session("old")
+ runner._get_or_create_session("new")
+
+ assert tuple(runner._sessions) == ("new",)
+ assert pipeline.closes == ["old"]
+
+
+def test_dreamzero_like_requested_capacity_is_capped_by_budget():
+ spec = dreamzero_like_spec(capacity=64)
+ # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
+ # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
+ # cross-attention=1 page/session, for 44 pages total.
+ page_bytes = 65_536
+ pipeline = CapablePipeline(spec)
+ runner = make_runner(
+ pipeline,
+ available_bytes=44 * page_bytes,
+ gpu_memory_fraction=1.0,
+ )
+ kv = runner.kv_cache
+ assert kv is not None
+ assert kv.requested_session_capacity == 64
+ assert kv.session_capacity == 2
+ assert runner._session_capacity == 2
+ assert kv.cross_attention_reserved_bytes == 2 * page_bytes
+
+
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
pipeline = CapablePipeline(lingbot_like_spec())
runner = make_runner(pipeline)
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
raise RuntimeError("layer exploded")
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
- request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+ request = SimpleNamespace(
+ request_id="broken-request",
+ sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+ )
with pytest.raises(RuntimeError, match="layer exploded"):
runner.execute_model(request)
@@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
runner.device = torch.device("cuda")
- request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+ request = SimpleNamespace(
+ request_id="broken-request",
+ sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
+ )
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
runner.execute_model(request)
diff --git a/tests/diffusion/ar_diffusion/test_consumer.py b/tests/diffusion/ar_diffusion/test_consumer.py
new file mode 100644
index 000000000..d1197113c
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_consumer.py
@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Sequence
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSession,
+ ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+from vllm_omni.inputs.data import (
+ OmniDiffusionSamplingParams,
+ OmniPromptType,
+ OmniSamplingParams,
+)
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeStagePool:
+ def __init__(self, num_replicas: int = 1) -> None:
+ self.num_replicas = num_replicas
+
+
+class FakeEngine:
+ def __init__(self, num_replicas: int = 1) -> None:
+ self.stage_pools = [FakeStagePool(num_replicas)]
+
+
+class FakeGenerateClient:
+ def __init__(
+ self,
+ *,
+ include_metadata: bool = True,
+ applied_event_ids: list[int] | None = None,
+ output_error: str | None = None,
+ num_replicas: int = 1,
+ ) -> None:
+ self.engine = FakeEngine(num_replicas)
+ self.include_metadata = include_metadata
+ self.applied_event_ids = applied_event_ids
+ self.output_error = output_error
+ self.calls: list[dict[str, Any]] = []
+
+ async def generate(
+ self,
+ prompt: OmniPromptType,
+ *,
+ request_id: str,
+ sampling_params_list: Sequence[OmniSamplingParams],
+ output_modalities: list[str] | None = None,
+ ) -> AsyncIterator[OmniRequestOutput]:
+ self.calls.append(
+ {
+ "prompt": prompt,
+ "request_id": request_id,
+ "sampling_params_list": sampling_params_list,
+ "output_modalities": output_modalities,
+ }
+ )
+ params = sampling_params_list[0]
+ assert isinstance(params, OmniDiffusionSamplingParams)
+ tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
+ assert tick is not None
+ chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
+ if self.applied_event_ids is not None:
+ chunk_metadata["applied_event_ids"] = self.applied_event_ids
+ metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
+ yield OmniRequestOutput(
+ request_id=request_id,
+ stage_id=0,
+ final_output_type="video",
+ images=["chunk"], # type: ignore[list-item]
+ _multimodal_output=metadata,
+ finished=True,
+ error=self.output_error,
+ )
+
+
+class FakeLifecycle:
+ async def reset_session(self, session_id: str) -> None:
+ pass
+
+ async def close_session(self, session_id: str) -> None:
+ pass
+
+
+def make_tick() -> ARDiffusionTickRequest:
+ return ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="world-1-chunk-3",
+ chunk_index=3,
+ applied_event_ids=(7, 8),
+ prompt="turn left",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+
+
+def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
+ client = FakeGenerateClient(num_replicas=2)
+
+ with pytest.raises(
+ ValueError,
+ match=r"require exactly one replica.*num_replicas=2",
+ ):
+ ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ assert client.calls == []
+
+
+@pytest.mark.asyncio
+async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
+ client = FakeGenerateClient()
+ template = OmniDiffusionSamplingParams(
+ output_type="latent",
+ extra_args={"deployment_option": "kept"},
+ )
+ consumer = ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: {
+ "prompt": tick.prompt,
+ "multi_modal_data": {"image": "initial-image"},
+ },
+ sampling_params_list=[template],
+ diffusion_stage_id=0,
+ output_modalities=["video"],
+ )
+ tick = make_tick()
+
+ output = await consumer.execute_tick(tick)
+
+ assert output.images == ["chunk"]
+ assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
+ assert len(client.calls) == 1
+ call = client.calls[0]
+ assert call["request_id"] == tick.request_id
+ assert call["output_modalities"] == ["video"]
+ assert call["prompt"] == {
+ "prompt": "turn left",
+ "multi_modal_data": {"image": "initial-image"},
+ }
+ submitted_params = call["sampling_params_list"][0]
+ assert submitted_params.extra_args["deployment_option"] == "kept"
+ assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
+ assert "ar_diffusion_tick" not in template.extra_args
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_missing_chunk_metadata() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(include_metadata=False),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(ValueError, match="multimodal_output.metadata"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(applied_event_ids=[8, 7]),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(ValueError, match="unique and strictly increasing"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
+ consumer = ARDiffusionOmniTickConsumer(
+ FakeGenerateClient(include_metadata=False, output_error="model failed"),
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+
+ with pytest.raises(RuntimeError, match="model failed"):
+ await consumer.execute_tick(make_tick())
+
+
+@pytest.mark.asyncio
+async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
+ client = FakeGenerateClient()
+ consumer = ARDiffusionOmniTickConsumer(
+ client,
+ prompt_provider=lambda tick: tick.prompt or "",
+ sampling_params_list=[OmniDiffusionSamplingParams()],
+ diffusion_stage_id=0,
+ )
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=consumer,
+ lifecycle=FakeLifecycle(),
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
+ )
+ await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
+
+ output = await session.next_chunk()
+
+ assert output.request_id == "world-1-chunk-0"
+ assert session.chunk_index == 1
+ assert session.pending_event_count == 0
+ assert client.calls[0]["request_id"] == "world-1-chunk-0"
diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py
index d09c1c1b6..086666063 100644
--- a/tests/diffusion/ar_diffusion/test_kv_cache.py
+++ b/tests/diffusion/ar_diffusion/test_kv_cache.py
@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
BLOCK = 16
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
return ChunkWindowSpec(
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
def test_sink_chunks_protected():
# sink = 1 chunk (16 tokens) is never skipped.
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
- # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
assert skip == 32
@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
session_capacity=1,
cross_attention_lengths={"text": L},
- device=torch.device("cpu"),
)
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
- expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+ page_bytes = kv.spec.page_size_bytes * 2
+ scratch_bytes = kv.scratch_num_blocks * page_bytes
+ expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
assert kv.num_blocks == expected
- assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
+ assert kv.cross_attention_reserved_bytes == cross_bytes
+ assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
+ assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
+
+
+def _make_tiny_capacity_kv(
+ *,
+ requested_capacity: int,
+ available_bytes: int,
+ gpu_memory_fraction: float = 1.0,
+ model_owned_state_bytes_per_session: int = 0,
+) -> ARDiffusionKVCache:
+ return ARDiffusionKVCache(
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=1,
+ window_chunks=3,
+ sink_chunks=3,
+ gpu_memory_fraction=gpu_memory_fraction,
+ ),
+ num_layers=1,
+ num_kv_heads=1,
+ head_size=1,
+ dtype=torch.float32,
+ block_size=1,
+ max_model_len=128,
+ available_bytes=available_bytes,
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=requested_capacity,
+ cross_attention_lengths={"text": 2},
+ frames_per_block=3,
+ model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
+ device=torch.device("cpu"),
+ )
+
+
+def test_capacity_two_retains_both_windows_and_allocates_next_block():
+ # LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
+ # 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
+ # reservations: 192 bytes exactly.
+ kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
+ assert kv.requested_session_capacity == 2
+ assert kv.session_capacity == 2
+ assert kv.managed_num_blocks == 17
+
+ adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
+ for adapter in adapters:
+ for _ in range(6): # three sink + three window blocks
+ kv.allocate_token_slots(adapter, 1)
+ kv.commit_chunk(adapter)
+ assert len(kv.window_block_ids(adapter)) == 6
+
+ # Only one request is in flight, so its three-frame block is counted once.
+ kv.allocate_token_slots(adapters[0], 3)
+
+
+def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
+ kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
+
+ assert kv.requested_session_capacity == 64
+ assert kv.session_capacity == 2
+ assert kv.cross_attention_bytes_per_session == 16
+ assert kv.cross_attention_reserved_bytes == 32
+ assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
+
+
+def test_capacity_rejects_budget_that_cannot_fit_one_session():
+ # One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
+ # 14 * 8 + 16 = 128 bytes.
+ with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
+ _make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
+
+
+def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
+ kv = _make_tiny_capacity_kv(
+ requested_capacity=2,
+ available_bytes=1280,
+ gpu_memory_fraction=0.05,
+ )
+
+ assert kv.configured_memory_budget_bytes == 64
+ assert kv.memory_budget_bytes == 128
+ assert kv.session_capacity == 1
+
+
+def test_model_owned_state_reduces_effective_session_capacity():
+ kv = _make_tiny_capacity_kv(
+ requested_capacity=2,
+ available_bytes=200,
+ model_owned_state_bytes_per_session=16,
+ )
+
+ assert kv.session_capacity == 1
+ assert kv.model_owned_state_reserved_bytes == 16
+ assert (
+ kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
+ <= kv.memory_budget_bytes
+ )
+
+
+def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
+ return ARDiffusionKVCache(
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=1560,
+ window_chunks=9,
+ sink_chunks=9,
+ gpu_memory_fraction=gpu_memory_fraction,
+ ),
+ num_layers=40,
+ num_kv_heads=40,
+ head_size=128,
+ dtype=torch.bfloat16,
+ block_size=1560,
+ max_model_len=1 << 20,
+ available_bytes=100 * (1 << 30),
+ kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+ session_capacity=2,
+ cross_attention_lengths={"text": 512},
+ frames_per_block=3,
+ model_owned_state_bytes_per_session=7_488_000,
+ )
+
+
+def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
+ kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
+
+ assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
+ assert kv.memory_budget_bytes == 33_653_670_400
+ assert kv.session_capacity == 1
+
+
+def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
+ kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
+
+ assert kv.memory_budget_bytes == 60 * (1 << 30)
+ assert kv.session_capacity == 2
def _make_kv(
@@ -221,15 +359,24 @@ def _make_kv(
if local_branches == 1
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
)
+ page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
+ declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
+ scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
+ managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
return ARDiffusionKVCache(
- ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+ ARDiffusionKVConfig(
+ enable=True,
+ chunk_size=BLOCK,
+ window_chunks=window_chunks,
+ gpu_memory_fraction=1.0,
+ ),
num_layers=1,
num_kv_heads=4,
head_size=64,
dtype=torch.float32,
block_size=BLOCK,
max_model_len=4096,
- available_bytes=1 << 16, # tiny -> the floor binds
+ available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
device=torch.device("cpu"),
kv_branches=kv_branches,
session_capacity=1,
diff --git a/tests/diffusion/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
index 32603bf88..4482fa6fb 100644
--- a/tests/diffusion/ar_diffusion/test_kv_orchestration.py
+++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py
@@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction():
used = {int(s) // kv.block_size for s in slots}
assert kv.null_block_id not in used
assert used <= set(block_table)
- # Resident window stays bounded by window + the in-flight chunk.
+ # Allocation may transiently add one in-flight chunk.
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
- free_after.append(kv.manager.block_pool.get_num_free_blocks())
kv.commit_chunk(adapter)
+ # A successful commit immediately prunes physical ownership.
+ assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
+ free_after.append(kv.manager.block_pool.get_num_free_blocks())
# Pool memory plateaus once the window is full (eviction recycles blocks).
- assert free_after[-1] == free_after[kv.spec.window_chunks]
+ assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
kv.end_request(adapter)
assert kv.manager.block_pool.get_num_free_blocks() == free_total
diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py
index 0b66c01a0..603a2a175 100644
--- a/tests/diffusion/ar_diffusion/test_paged_attention.py
+++ b/tests/diffusion/ar_diffusion/test_paged_attention.py
@@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
return torch.einsum("bhqk,bkhd->bqhd", probs, value)
-def _cuda_flash_attn_usable() -> bool:
+def _gpu_flash_attn_usable() -> bool:
if not torch.cuda.is_available():
return False
+ if torch.version.hip is not None:
+ for module in ("aiter", "flash_attn"):
+ try:
+ imported = __import__(module, fromlist=["flash_attn_varlen_func"])
+ if getattr(imported, "flash_attn_varlen_func", None) is not None:
+ return True
+ except ImportError:
+ pass
+ return False
try:
spec = find_spec("vllm.vllm_flash_attn")
if spec is None or spec.origin is None:
@@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
-@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
@pytest.mark.parametrize("history_chunks", [1, 3])
@pytest.mark.parametrize("action_len", [0, 3])
@pytest.mark.parametrize("commit_current", [False, True])
def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
- pytest.importorskip("vllm.vllm_flash_attn")
torch.manual_seed(0)
device = torch.device("cuda")
dtype = torch.float16
@@ -324,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
def test_prepare_is_idempotent_and_layers_share_metadata():
device = torch.device("cpu")
- _, st = make_state(num_layers=2)
+ kv, st = make_state(num_layers=2)
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
fctx = contexts[0].forward_ctx
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
@@ -341,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
assert i0.block_table is i1.block_table
assert i0.seq_lens is i1.seq_lens
assert i0.video_slots is i1.video_slots
+ assert i0.key_pool is kv._k_pools[0]
+ assert i0.value_pool is kv._v_pools[0]
+ assert i1.key_pool is kv._k_pools[1]
+ assert i1.value_pool is kv._v_pools[1]
def test_layer_inputs_before_prepare_raises():
diff --git a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
index ec60290b9..ed5233800 100644
--- a/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
+++ b/tests/diffusion/ar_diffusion/test_pipeline_bridge.py
@@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache():
assert not st.is_cross_attention_populated(POS, "text")
+def test_clear_cross_attention_preserves_self_attention_history():
+ kv, st = make_state(num_layers=1, cross_attn_length=8)
+ _prepare_and_commit(st, POS, 2)
+ k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+ st.populate_cross_attention(POS, "text", [(k, k)])
+ resident_before = kv.window_block_ids(st.adapter(POS))
+
+ st.clear_cross_attention()
+
+ assert kv.window_block_ids(st.adapter(POS)) == resident_before
+ assert not st.is_cross_attention_populated(POS, "text")
+
+
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
assert st.kv_cache.num_local_kv_branches == 1
diff --git a/tests/diffusion/ar_diffusion/test_session.py b/tests/diffusion/ar_diffusion/test_session.py
new file mode 100644
index 000000000..a31be388b
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_session.py
@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionChunkMetadataError,
+ ARDiffusionEventAcceptanceStatus,
+ ARDiffusionEventQueueFullError,
+ ARDiffusionPreparedControls,
+ ARDiffusionSession,
+ ARDiffusionSessionError,
+ ARDiffusionSessionEvent,
+ ARDiffusionSessionManager,
+ ARDiffusionSessionStateError,
+ ARDiffusionSessionStatus,
+ ARDiffusionStaleEventError,
+ ARDiffusionWorkerLifecycle,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeTickConsumer:
+ def __init__(self) -> None:
+ self.ticks: list[ARDiffusionTickRequest] = []
+ self.failures_remaining = 0
+ self.started: asyncio.Event | None = None
+ self.release: asyncio.Event | None = None
+ self.metadata_override: ARDiffusionChunkMetadata | None = None
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+ self.ticks.append(tick)
+ if self.started is not None:
+ self.started.set()
+ if self.release is not None:
+ await self.release.wait()
+ if self.failures_remaining:
+ self.failures_remaining -= 1
+ raise RuntimeError("fake tick failed")
+ return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
+
+ def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+ return output
+
+
+class FakeLifecycle:
+ def __init__(self) -> None:
+ self.resets: list[str] = []
+ self.closes: list[str] = []
+ self.close_failures_remaining = 0
+
+ async def reset_session(self, session_id: str) -> None:
+ self.resets.append(session_id)
+
+ async def close_session(self, session_id: str) -> None:
+ self.closes.append(session_id)
+ if self.close_failures_remaining:
+ self.close_failures_remaining -= 1
+ raise RuntimeError("fake close failed")
+
+
+class FailingCommitReducer:
+ def __init__(self) -> None:
+ self.commit_calls = 0
+ self.reset_calls = 0
+
+ def prepare(
+ self,
+ *,
+ current_controls: dict[str, ARDiffusionControlInput],
+ events: tuple[ARDiffusionSessionEvent, ...],
+ chunk_index: int,
+ ) -> ARDiffusionPreparedControls:
+ del current_controls, events, chunk_index
+ return ARDiffusionPreparedControls(controls=())
+
+ def commit(self, prepared: ARDiffusionPreparedControls) -> None:
+ del prepared
+ self.commit_calls += 1
+ raise RuntimeError("fake reducer commit failed")
+
+ def reset(self) -> None:
+ self.reset_calls += 1
+
+
+def make_session(
+ *,
+ max_pending_events: int = 8,
+ consumer: FakeTickConsumer | None = None,
+ lifecycle: FakeLifecycle | None = None,
+) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
+ resolved_consumer = consumer or FakeTickConsumer()
+ resolved_lifecycle = lifecycle or FakeLifecycle()
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=resolved_consumer,
+ lifecycle=resolved_lifecycle,
+ max_pending_events=max_pending_events,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ )
+ return session, resolved_consumer, resolved_lifecycle
+
+
+@pytest.mark.asyncio
+async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
+ session, consumer, _ = make_session()
+ event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
+ event_one = ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+
+ accepted_two = await session.accept_event(event_two)
+ accepted_one = await session.accept_event(event_one)
+ duplicate = await session.accept_event(event_two)
+
+ assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+ assert consumer.ticks == []
+
+ metadata = await session.next_chunk()
+
+ assert metadata.applied_event_ids == (1, 2)
+ assert consumer.ticks[0] == ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="world-1-request-0",
+ chunk_index=0,
+ applied_event_ids=(1, 2),
+ prompt="turn",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"samples": [{"x": 1.0}]},
+ ),
+ ),
+ )
+ assert session.pending_event_count == 0
+
+ applied_duplicate = await session.accept_event(event_two)
+ assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+ next_metadata = await session.next_chunk()
+ assert next_metadata.chunk_index == 1
+ assert next_metadata.applied_event_ids == ()
+ assert consumer.ticks[1].prompt == "turn"
+ assert consumer.ticks[1].controls == consumer.ticks[0].controls
+
+
+@pytest.mark.asyncio
+async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
+ session, consumer, _ = make_session(max_pending_events=1)
+ payload = {"samples": [{"x": 1.0}]}
+ composite = ARDiffusionSessionEvent(
+ event_id=5,
+ prompt="new prompt",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data=payload,
+ ),
+ ARDiffusionControlInput(
+ track="action",
+ schema="example.action.v1",
+ data={"keys": ["left"]},
+ ),
+ ),
+ )
+ await session.accept_event(composite)
+ payload["samples"][0]["x"] = 99.0
+
+ duplicate = await session.accept_event(composite)
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+ with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
+ await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
+
+ await session.next_chunk()
+ assert consumer.ticks[0].applied_event_ids == (5,)
+ assert consumer.ticks[0].prompt == "new prompt"
+ assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
+ assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
+ assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
+ assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
+
+
+def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
+ with pytest.raises(ValueError, match="at most once"):
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ ARDiffusionControlInput(track="camera", schema="one", data={}),
+ ARDiffusionControlInput(track="camera", schema="two", data={}),
+ ),
+ )
+
+ with pytest.raises(ValueError, match="transport-safe"):
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data={"object": object()},
+ )
+
+
+@pytest.mark.asyncio
+async def test_failed_tick_closes_worker_and_requires_reset() -> None:
+ consumer = FakeTickConsumer()
+ consumer.failures_remaining = 1
+ lifecycle = FakeLifecycle()
+ session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
+
+ with pytest.raises(RuntimeError, match="fake tick failed"):
+ await session.next_chunk()
+
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert lifecycle.closes == ["world-1"]
+ with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+ await session.next_chunk()
+
+ await session.reset()
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.pending_event_count == 0
+ assert lifecycle.resets == ["world-1"]
+ await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
+ metadata = await session.next_chunk()
+ assert metadata.chunk_index == 0
+ assert metadata.applied_event_ids == (4,)
+ assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
+
+
+@pytest.mark.asyncio
+async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
+ consumer = FakeTickConsumer()
+ consumer.started = asyncio.Event()
+ consumer.release = asyncio.Event()
+ session, _, _ = make_session(consumer=consumer)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
+
+ first_chunk = asyncio.create_task(session.next_chunk())
+ await consumer.started.wait()
+ accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
+ assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+ assert consumer.ticks[0].prompt == "first"
+ assert consumer.ticks[0].applied_event_ids == (10,)
+
+ with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
+ await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
+
+ consumer.release.set()
+ first_metadata = await first_chunk
+ second_metadata = await session.next_chunk()
+ assert first_metadata.applied_event_ids == (10,)
+ assert second_metadata.applied_event_ids == (11,)
+ assert consumer.ticks[1].prompt == "second"
+
+
+@pytest.mark.asyncio
+async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
+ consumer = FakeTickConsumer()
+ consumer.metadata_override = ARDiffusionChunkMetadata(
+ session_id="another-world",
+ request_id="wrong",
+ chunk_index=0,
+ )
+ lifecycle = FakeLifecycle()
+ session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
+ await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+ with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
+ await session.next_chunk()
+
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert lifecycle.closes == ["world-1"]
+
+
+@pytest.mark.asyncio
+async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ reducer = FailingCommitReducer()
+ session = ARDiffusionSession(
+ "world-1",
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ control_reducer=reducer,
+ )
+ await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
+
+ with pytest.raises(RuntimeError, match="fake reducer commit failed"):
+ await session.next_chunk()
+
+ assert reducer.commit_calls == 1
+ assert session.status is ARDiffusionSessionStatus.FAILED
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 1
+ assert lifecycle.closes == ["world-1"]
+ with pytest.raises(ARDiffusionSessionStateError, match="failed"):
+ await session.next_chunk()
+
+ await session.reset()
+ assert reducer.reset_calls == 1
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.chunk_index == 0
+ assert session.pending_event_count == 0
+
+
+@pytest.mark.asyncio
+async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+ )
+ session = await manager.create_session("world-reset")
+ await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
+ await session.next_chunk()
+ await manager.reset_session("world-reset")
+
+ assert lifecycle.resets == ["world-reset"]
+ assert session.status is ARDiffusionSessionStatus.ACTIVE
+ assert session.chunk_index == 0
+ duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
+ assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
+
+ reset_tick = await session.next_chunk()
+ assert reset_tick.chunk_index == 0
+ assert consumer.ticks[-1].prompt is None
+ assert consumer.ticks[-1].controls == ()
+
+ await manager.close_session("world-reset")
+ assert lifecycle.closes == ["world-reset"]
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+ with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+ await manager.get_session("world-reset")
+
+ disconnected = await manager.create_session("world-disconnect")
+ await manager.disconnect("world-disconnect")
+ assert disconnected.status is ARDiffusionSessionStatus.CLOSED
+ assert lifecycle.closes[-1] == "world-disconnect"
+
+
+@pytest.mark.asyncio
+async def test_manager_keeps_interleaved_sessions_independent() -> None:
+ consumer = FakeTickConsumer()
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=FakeLifecycle(),
+ max_pending_events=4,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
+ )
+ session_a = await manager.create_session("world-a")
+ session_b = await manager.create_session("world-b")
+
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
+ await session_a.next_chunk()
+ await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
+ await session_b.next_chunk()
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
+ await session_a.next_chunk()
+
+ assert [
+ (
+ tick.session_id,
+ tick.request_id,
+ tick.chunk_index,
+ tick.applied_event_ids,
+ )
+ for tick in consumer.ticks
+ ] == [
+ ("world-a", "world-a-request-0", 0, (10,)),
+ ("world-b", "world-b-request-0", 0, (1,)),
+ ("world-a", "world-a-request-1", 1, (11,)),
+ ]
+ assert session_a.chunk_index == 2
+ assert session_b.chunk_index == 1
+ with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
+ await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
+ assert (
+ await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
+ ).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
+
+
+@pytest.mark.asyncio
+async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
+ consumer = FakeTickConsumer()
+ lifecycle = FakeLifecycle()
+ lifecycle.close_failures_remaining = 1
+ manager = ARDiffusionSessionManager(
+ tick_consumer=consumer,
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ )
+ session = await manager.create_session("world-close-retry")
+
+ with pytest.raises(RuntimeError, match="fake close failed"):
+ await manager.close_session("world-close-retry")
+
+ assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+ assert await manager.get_session("world-close-retry") is session
+ with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
+ await manager.create_session("world-close-retry")
+ with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
+ await manager.reset_session("world-close-retry")
+
+ await manager.close_session("world-close-retry")
+
+ assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+ with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
+ await manager.get_session("world-close-retry")
+ recreated = await manager.create_session("world-close-retry")
+ assert recreated is not session
+
+
+@pytest.mark.asyncio
+async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
+ lifecycle = FakeLifecycle()
+ lifecycle.close_failures_remaining = 1
+ manager = ARDiffusionSessionManager(
+ tick_consumer=FakeTickConsumer(),
+ lifecycle=lifecycle,
+ max_pending_events=4,
+ )
+ session = await manager.create_session("world-disconnect-retry")
+
+ with pytest.raises(RuntimeError, match="fake close failed"):
+ await manager.disconnect("world-disconnect-retry")
+
+ assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
+ assert await manager.get_session("world-disconnect-retry") is session
+
+ await manager.close_session("world-disconnect-retry")
+ assert session.status is ARDiffusionSessionStatus.CLOSED
+
+
+class FakeRPCClient:
+ def __init__(self, result: list[Any]) -> None:
+ self.result = result
+ self.calls: list[dict[str, Any]] = []
+
+ async def collective_rpc(self, **kwargs: Any) -> list[Any]:
+ self.calls.append(kwargs)
+ return self.result
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
+ client = FakeRPCClient([[True]])
+ lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
+
+ await lifecycle.reset_session("world-1")
+ await lifecycle.close_session("world-1")
+
+ assert client.calls == [
+ {
+ "method": "reset_ar_diffusion_session",
+ "timeout": 3.0,
+ "args": ("world-1",),
+ "stage_ids": [2],
+ },
+ {
+ "method": "close_ar_diffusion_session",
+ "timeout": 3.0,
+ "args": ("world-1",),
+ "stage_ids": [2],
+ },
+ ]
+
+
+@pytest.mark.asyncio
+async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
+ lifecycle = ARDiffusionWorkerLifecycle(
+ FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
+ stage_ids=[1],
+ )
+
+ with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
+ await lifecycle.close_session("world-1")
diff --git a/tests/diffusion/ar_diffusion/test_tick_protocol.py b/tests/diffusion/ar_diffusion/test_tick_protocol.py
new file mode 100644
index 000000000..70e71ba80
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_tick_protocol.py
@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="chunk-request-7",
+ chunk_index=7,
+ applied_event_ids=(10, 11),
+ prompt="turn left",
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
+ ),
+ ),
+ )
+
+ restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
+
+ assert restored == tick
+ assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
+ "session_id": "world-1",
+ "request_id": "chunk-request-7",
+ "chunk_index": 7,
+ "applied_event_ids": [10, 11],
+ }
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "message"),
+ [
+ ("session_id", "", "session_id"),
+ ("request_id", "", "request_id"),
+ ("chunk_index", -1, "chunk_index"),
+ ("applied_event_ids", (2, 1), "strictly increasing"),
+ ("applied_event_ids", (1, 1), "strictly increasing"),
+ ],
+)
+def test_tick_rejects_invalid_identity(
+ field: str,
+ value: object,
+ message: str,
+) -> None:
+ values = {
+ "session_id": "world-1",
+ "request_id": "request-1",
+ "chunk_index": 0,
+ "applied_event_ids": (),
+ }
+ values[field] = value
+
+ with pytest.raises(ValueError, match=message):
+ ARDiffusionTickRequest(**values)
+
+
+def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
+ payload = {"samples": [{"keys": ["w"]}]}
+ event_ids = [1, 2]
+ controls = [
+ ARDiffusionControlInput(
+ track="camera",
+ schema="example.camera.v1",
+ data=payload,
+ )
+ ]
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="request-a",
+ chunk_index=0,
+ applied_event_ids=event_ids, # type: ignore[arg-type]
+ controls=controls, # type: ignore[arg-type]
+ )
+ serialized = tick.to_extra_args()
+
+ payload["samples"][0]["keys"].append("j")
+ event_ids.append(3)
+ controls.clear()
+ serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
+
+ assert tick.applied_event_ids == (1, 2)
+ assert isinstance(tick.controls, tuple)
+ assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
+ {
+ "track": "camera",
+ "schema": "example.camera.v1",
+ "data": {"samples": [{"keys": ["w"]}]},
+ }
+ ]
+ with pytest.raises(TypeError):
+ tick.controls[0].data["samples"] = () # type: ignore[index]
+ with pytest.raises(AttributeError):
+ tick.controls[0].data["samples"][0]["keys"].append("j")
+
+
+def test_legacy_request_has_no_typed_tick() -> None:
+ assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
diff --git a/tests/diffusion/ar_diffusion/test_worker_lifecycle.py b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py
new file mode 100644
index 000000000..42204049c
--- /dev/null
+++ b/tests/diffusion/ar_diffusion/test_worker_lifecycle.py
@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeARDiffusionRunner:
+ def __init__(self) -> None:
+ self.resets: list[str] = []
+ self.closes: list[str] = []
+
+ def reset_session(self, session_id: str) -> None:
+ self.resets.append(session_id)
+
+ def close_session(self, session_id: str) -> None:
+ self.closes.append(session_id)
+
+
+class FakeRegularRunner:
+ pass
+
+
+def make_worker(runner: object) -> DiffusionWorker:
+ worker = object.__new__(DiffusionWorker)
+ worker.model_runner = runner
+ return worker
+
+
+def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
+ runner = FakeARDiffusionRunner()
+ worker = make_worker(runner)
+
+ assert worker.reset_ar_diffusion_session("world-1") is True
+ assert worker.close_ar_diffusion_session("world-1") is True
+ assert runner.resets == ["world-1"]
+ assert runner.closes == ["world-1"]
+
+
+def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
+ worker = make_worker(FakeRegularRunner())
+
+ assert worker.reset_ar_diffusion_session("world-1") is False
+ assert worker.close_ar_diffusion_session("world-1") is False
+
+
+def test_worker_rejects_empty_session_id() -> None:
+ worker = make_worker(FakeARDiffusionRunner())
+
+ with pytest.raises(ValueError, match="session_id"):
+ worker.close_ar_diffusion_session("")
diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py
index 4a2c65cbb..54ad2a32b 100644
--- a/tests/diffusion/attention/test_attention_config.py
+++ b/tests/diffusion/attention/test_attention_config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
@@ -11,6 +11,7 @@ Tests cover:
"""
from types import SimpleNamespace
+from unittest.mock import Mock
import pytest
import torch
@@ -90,10 +91,23 @@ class TestAttentionSpec:
with pytest.raises(ValueError, match=match):
AttentionSpec(**spec)
+ def test_fastvideo_vsa_topk_serialized(self):
+ spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96)
+ assert spec.backend_kwargs() == {"topk": 96}
+
+ def test_fastvideo_vsa_topk_rejected_for_other_backend(self):
+ with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"):
+ AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96)
+
def test_block_sparse_defaults_applied_when_backend_selected(self):
spec = AttentionSpec(backend="RAINFUSION_ATTN")
assert spec.block_sparse.sparsity == 0.8
- assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
+ assert spec.backend_kwargs() == {
+ "sparsity": 0.8,
+ "start_step": 0,
+ "end_step": 0,
+ "precision": "bf16",
+ }
def test_block_sparse_skip_layers_selector_expanded(self):
spec = AttentionSpec(
@@ -104,6 +118,8 @@ class TestAttentionSpec:
assert spec.backend_kwargs() == {
"sparsity": 0.9,
"start_step": 12,
+ "end_step": 0,
+ "precision": "bf16",
"skip_layers": [0, 1, 2, 38],
}
@@ -113,12 +129,44 @@ class TestAttentionSpec:
@pytest.mark.parametrize(
"block_sparse",
- [{"sparsity": 1.5}, {"start_step": -1}],
+ [{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}],
)
def test_block_sparse_invalid_values(self, block_sparse):
with pytest.raises(ValueError):
AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
+ def test_block_sparse_end_step_precision_forwarded(self):
+ spec = AttentionSpec(
+ backend="RAINFUSION_ATTN",
+ block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"},
+ )
+ assert spec.backend_kwargs() == {
+ "sparsity": 0.9,
+ "start_step": 5,
+ "end_step": 3,
+ "precision": "fp8",
+ }
+
+ def test_block_sparse_precision_enum_input_normalized(self):
+ """RainFusionPrecision enum members must normalize to their string value.
+
+ Regression: the str,Enum members passed the value check but str() later
+ produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
+ """
+ from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
+ from vllm_omni.diffusion.data import RainFusionPrecision
+
+ spec = AttentionSpec(
+ backend="RAINFUSION_ATTN",
+ block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX},
+ )
+ kwargs = spec.backend_kwargs()
+ assert kwargs["precision"] == "mix"
+ # The config path is where str(value) used to mangle enum members.
+ cfg = RainFusionConfig.from_backend_kwargs(kwargs)
+ assert cfg.precision == "mix"
+ assert type(cfg.precision) is str
+
class TestAttentionConfig:
def test_empty_config(self):
@@ -419,6 +467,120 @@ class TestCurrentDiffusionConfig:
class TestAttentionInitUsesCurrentDiffusionConfig:
+ @pytest.mark.parametrize("mindiesd_available", [False, True])
+ def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
+ attention = Attention.__new__(Attention)
+ attention._scheduler_paged_kv = True
+ attention.paged_kv_cache_role = "primary"
+ attention.backend_pref = "FLASH_ATTN"
+ attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+ flash_output = torch.ones(1)
+ sdpa_output = torch.zeros(1)
+ dense_flash = Mock(return_value=flash_output)
+ sdpa = Mock(return_value=sdpa_output)
+ attention.attention = SimpleNamespace(forward=dense_flash)
+ attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
+ monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+ monkeypatch.setattr(
+ layer_mod,
+ "get_forward_context",
+ lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
+ )
+ monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True)
+ monkeypatch.setattr(
+ layer_mod.current_omni_platform,
+ "supports_diffusion_dense_flash_attention",
+ lambda: mindiesd_available,
+ )
+
+ result = attention._run_local_attention(
+ torch.zeros(1, dtype=torch.bfloat16),
+ torch.zeros(1, dtype=torch.bfloat16),
+ torch.zeros(1, dtype=torch.bfloat16),
+ AttentionMetadata(),
+ )
+
+ if mindiesd_available:
+ assert result is flash_output
+ dense_flash.assert_called_once()
+ sdpa.assert_not_called()
+ else:
+ assert result is sdpa_output
+ dense_flash.assert_not_called()
+ sdpa.assert_called_once()
+
+ def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
+ attention = Attention.__new__(Attention)
+ attention._scheduler_paged_kv = False
+ attention.paged_kv_cache_role = "primary"
+ attention.backend_pref = "FLASH_ATTN"
+ attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+ attention.attention = SimpleNamespace(
+ forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd"))
+ )
+ attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back")))
+ monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+ with pytest.raises(ModuleNotFoundError, match="mindiesd"):
+ attention._run_local_attention(
+ torch.zeros(1, dtype=torch.bfloat16),
+ torch.zeros(1, dtype=torch.bfloat16),
+ torch.zeros(1, dtype=torch.bfloat16),
+ AttentionMetadata(),
+ )
+
+ def test_paged_marker_does_not_change_dense_backend(self, monkeypatch):
+ class _FakeAttentionImpl:
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+
+ def forward(self, query, key, value, attn_metadata=None):
+ return query
+
+ class _DenseBackend:
+ supports_paged_kv = False
+
+ @staticmethod
+ def get_name() -> str:
+ return "TORCH_SDPA"
+
+ @staticmethod
+ def get_impl_cls():
+ return _FakeAttentionImpl
+
+ calls = []
+
+ def _fake_get_attn_backend_for_role(*args, **kwargs):
+ calls.append((args, kwargs))
+ return _DenseBackend, None
+
+ monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+ monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+ monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+ monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+ od_config = SimpleNamespace(
+ diffusion_attention_config=AttentionConfig(),
+ diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY,
+ parallel_config=SimpleNamespace(ring_degree=1),
+ diffusion_kv_cache_dtype=None,
+ diffusion_kv_cache_skip_step_indices=None,
+ diffusion_kv_cache_skip_layer_indices=None,
+ )
+
+ with set_current_diffusion_config(od_config):
+ attention = Attention(
+ num_heads=4,
+ head_size=64,
+ causal=False,
+ softmax_scale=1.0,
+ paged_kv_cache_role="primary",
+ )
+
+ assert len(calls) == 1
+ assert attention.attn_backend is _DenseBackend
+ assert attention.backend_pref is None
+
def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
class _FakeAttentionImpl:
def __init__(self, **kwargs):
diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py
index de1adf8c7..1ff3baafb 100644
--- a/tests/diffusion/attention/test_cudnn_attn.py
+++ b/tests/diffusion/attention/test_cudnn_attn.py
@@ -5,6 +5,7 @@ from contextlib import nullcontext
import pytest
import torch
+from torch.nn.attention import SDPBackend
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+@pytest.mark.parametrize(
+ ("query_sequence_length", "key_sequence_length", "expected_backend"),
+ [
+ (1, 1, SDPBackend.MATH),
+ (8, 1, SDPBackend.MATH),
+ (1, 8, SDPBackend.MATH),
+ (8, 8, SDPBackend.CUDNN_ATTENTION),
+ ],
+)
+def test_cudnn_routes_single_token_attention_to_math(
+ monkeypatch, query_sequence_length, key_sequence_length, expected_backend
+):
+ selected_backends = []
+
+ def fake_sdpa_kernel(backends):
+ selected_backends.append(backends)
+ return nullcontext()
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
+ fake_sdpa_kernel,
+ )
+ monkeypatch.setattr(
+ torch.nn.functional,
+ "scaled_dot_product_attention",
+ lambda query, _key, _value, **_kwargs: query,
+ )
+ impl = CuDNNAttentionImpl(
+ num_heads=2,
+ head_size=4,
+ softmax_scale=0.5,
+ )
+ query = torch.randn(1, query_sequence_length, 2, 4)
+ key = torch.randn(1, key_sequence_length, 2, 4)
+ value = torch.randn_like(key)
+
+ output = impl.forward_cuda(query, key, value)
+
+ assert output.shape == query.shape
+ assert selected_backends == [[expected_backend]]
+
+
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
observed = {}
diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py
new file mode 100644
index 000000000..d6209bf42
--- /dev/null
+++ b/tests/diffusion/attention/test_fastvideo_vsa.py
@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+ FastVideoVSAImpl,
+)
+from vllm_omni.diffusion.attention.backends.registry import (
+ DiffusionAttentionBackendEnum,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_fastvideo_vsa_backend_is_registered():
+ assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
+
+
+def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
+ calls: dict[str, Any] = {}
+ fake_module = types.ModuleType("fastvideo_kernel")
+
+ def fake_video_sparse_attn_bshd(
+ q,
+ k,
+ v,
+ variable_block_sizes,
+ q_variable_block_sizes,
+ compress_attn_weight,
+ topk,
+ block_size,
+ ):
+ calls["q_shape"] = tuple(q.shape)
+ calls["vbs"] = variable_block_sizes.detach().cpu().tolist()
+ calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist()
+ calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+ calls["compress_shape"] = tuple(compress_attn_weight.shape)
+ calls["topk"] = topk
+ calls["block_size"] = block_size
+ return q + k + v
+
+ setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+ monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+ impl = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={
+ "topk": 1,
+ "block_size": (4, 8, 8),
+ "min_seq_len": 1,
+ "disable_when_sp_active": False,
+ },
+ )
+ query = torch.randn(1, 300, 2, 8)
+ key = torch.randn_like(query)
+ value = torch.randn_like(query)
+ attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)})
+ monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+ output = impl.forward_cuda(query, key, value, attn_metadata)
+
+ assert output.shape == query.shape
+ assert calls["q_shape"] == (1, 1024, 2, 8)
+ assert calls["vbs"] == [192, 48, 48, 12]
+ assert calls["q_vbs"] == [192, 48, 48, 12]
+ assert calls["compress_shape"] == (1, 1024, 2, 8)
+ assert calls["compress_sum"] == 0.0
+ assert calls["topk"] == 1
+ assert calls["block_size"] == (4, 8, 8)
+
+
+def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch):
+ calls: dict[str, Any] = {}
+ fake_module = types.ModuleType("fastvideo_kernel")
+
+ def fake_video_sparse_attn_bshd(
+ q,
+ k,
+ v,
+ variable_block_sizes,
+ q_variable_block_sizes,
+ compress_attn_weight,
+ topk,
+ block_size,
+ ):
+ calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
+ calls["compress_shape"] = tuple(compress_attn_weight.shape)
+ return q + k + v
+
+ setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
+ monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
+
+ impl = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={
+ "topk": 1,
+ "block_size": (4, 8, 8),
+ "min_seq_len": 1,
+ "disable_when_sp_active": False,
+ },
+ )
+ query = torch.randn(1, 300, 2, 8)
+ gate = torch.ones_like(query)
+ metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate})
+ monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
+
+ output = impl.forward_cuda(query, query, query, metadata)
+
+ assert output.shape == query.shape
+ assert calls["compress_shape"] == (1, 1024, 2, 8)
+ # Only the 300 real tokens carry ones; padded gate slots stay zero.
+ assert calls["compress_sum"] == gate.numel()
+
+
+def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch):
+ calls: dict[str, Any] = {}
+
+ impl = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={
+ "topk": 1,
+ "block_size": (4, 8, 8),
+ "min_seq_len": 1,
+ "disable_when_sp_active": False,
+ },
+ )
+
+ def fake_fallback(query, key, value, attn_metadata, reason):
+ calls["reason"] = reason
+ return torch.zeros_like(query)
+
+ monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+ query = torch.randn(1, 512, 2, 8)
+ output = impl.forward_cuda(query, query, query, AttentionMetadata())
+
+ assert output.shape == query.shape
+ assert calls["reason"] == "vsa_dit_seq_shape metadata is required"
+
+
+def test_fastvideo_vsa_falls_back_for_mask(monkeypatch):
+ calls: dict[str, Any] = {}
+
+ impl = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={"min_seq_len": 1},
+ )
+
+ def fake_fallback(query, key, value, attn_metadata, reason):
+ calls["reason"] = reason
+ return torch.zeros_like(query)
+
+ monkeypatch.setattr(impl, "_fallback", fake_fallback)
+
+ query = torch.randn(1, 512, 2, 8)
+ mask = torch.ones(1, 512, dtype=torch.bool)
+ output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask))
+
+ assert output.shape == query.shape
+ assert calls["reason"]
+
+
+def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
+ query = torch.randn(1, 512, 2, 8)
+ metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)})
+
+ all_blocks = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1},
+ )
+ too_many = FastVideoVSAImpl(
+ num_heads=2,
+ head_size=8,
+ softmax_scale=8**-0.5,
+ causal=False,
+ backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1},
+ )
+
+ # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
+ assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
+ assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py
index e3f5355bf..3b5c877cb 100644
--- a/tests/diffusion/attention/test_flash_attn.py
+++ b/tests/diffusion/attention/test_flash_attn.py
@@ -265,10 +265,15 @@ def test_fa_vs_sdpa():
print(f"Max absolute difference (valid region): {max_diff:.6f}")
print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
- # Assert that outputs are close
- # Using higher tolerance for bfloat16 and different implementations
- assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
- assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+ # AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
+ # about 1.3% relative error. An absolute-only bound misclassifies that
+ # expected rounding difference for larger output values (for example,
+ # -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
+ if current_omni_platform.is_rocm():
+ torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
+ else:
+ assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+ assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
@@ -303,6 +308,98 @@ def test_flash_attn_func_preferred_over_varlen():
print("✓ flash_attn_func forward works correctly!")
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
+@pytest.mark.parametrize("k_len", [40, 256])
+def test_cross_attn_key_padding_vs_sdpa(k_len):
+ """
+ Case 3: Cross-attention with a key-padding mask.
+
+ The mask covers only the key sequence; every query position must attend
+ to the valid keys. FA and SDPA outputs should be very close. k_len=256
+ covers equal Q/K lengths, which must not be mistaken for self-attention,
+ and valid key lengths differ per batch element.
+ """
+ device = torch.device(current_omni_platform.device_type)
+ dtype = torch.bfloat16
+
+ batch_size = 2
+ q_len = 256
+ valid_lens = (17, 29)
+ num_heads = 8
+ head_dim = 64
+
+ fa_impl = FlashAttentionImpl(
+ num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
+ )
+ sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+
+ torch.manual_seed(7)
+ query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
+ key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
+ value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
+
+ attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
+ for i, valid_len in enumerate(valid_lens):
+ attn_mask[i, :valid_len] = True
+ attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+
+ output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
+ output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
+
+ max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
+ mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
+
+ print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
+ print(f"Max absolute difference: {max_diff:.6f}")
+ print(f"Mean absolute difference: {mean_diff:.6f}")
+
+ # Match the platform-specific comparison used by test_fa_vs_sdpa above.
+ if current_omni_platform.is_rocm():
+ # AITER FlashAttention and PyTorch SDPA use different BF16 reduction
+ # implementations on ROCm, so individual values can differ by one BF16
+ # step even though the outputs agree within the expected relative error.
+ torch.testing.assert_close(output_fa, output_sdpa, rtol=1e-2, atol=1.5e-2)
+ else:
+ assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+ assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+
+ print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
+
+
+def test_varlen_masked_routing_by_role(monkeypatch):
+ """The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
+ calls = []
+
+ def fake_varlen_func(q, k, v, **kwargs):
+ calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
+ return q
+
+ monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
+ monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
+
+ batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
+ query = torch.randn(batch_size, seq_len, num_heads, head_dim)
+ metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
+
+ cross_impl = FlashAttentionImpl(
+ num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
+ )
+ output = cross_impl.forward_cuda(query, query, query, metadata)
+ assert output.shape == query.shape
+ num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
+ assert num_q_rows == batch_size * seq_len
+ assert cu_seqlens_q.tolist() == [0, 4, 8]
+ assert cu_seqlens_k.tolist() == [0, 2, 5]
+
+ self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
+ output = self_impl.forward_cuda(query, query, query, metadata)
+ assert output.shape == query.shape
+ num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
+ assert num_q_rows == 5
+ assert cu_seqlens_q.tolist() == [0, 2, 5]
+ assert cu_seqlens_k.tolist() == [0, 2, 5]
+
+
def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
calls = []
diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py
index 05967e28f..fc4746773 100644
--- a/tests/diffusion/attention/test_piecewise_attn.py
+++ b/tests/diffusion/attention/test_piecewise_attn.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""End-to-end test for ``piecewise_attn`` (CPU).
Verify that running attention in segments (causal outside full-attn spans,
@@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_
from vllm_omni.diffusion.attention.backends.abstract import QueryRange
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
Segment,
+ build_paged_piecewise_plan,
build_segments,
piecewise_attn,
+ run_paged_piecewise_plan,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -84,6 +86,234 @@ def test_build_segments_full_span_retains_global_kv_end():
assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
+def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
+ full_attn_spans = [[(2, 5)], [(4, 7)]]
+ plan = build_paged_piecewise_plan(
+ full_attn_spans,
+ query_offsets=[0, 2],
+ query_lens=[6, 6],
+ seq_lens=[6, 8],
+ device=DEVICE,
+ )
+ query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+ key = query + 100
+ value = query + 200
+ calls: list[object] = []
+
+ output_buffer = torch.empty_like(query)
+
+ def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer):
+ assert segment_output_buffer is None
+ calls.append(
+ (
+ [segment.mode for segment in segment_batch.row_segments],
+ [segment.kv_end for segment in segment_batch.row_segments],
+ segment_batch.query_indices.tolist(),
+ )
+ )
+ torch.testing.assert_close(segment_key, segment_query + 100)
+ torch.testing.assert_close(segment_value, segment_query + 200)
+ return segment_query
+
+ output = run_paged_piecewise_plan(
+ query,
+ key,
+ value,
+ plan,
+ plan.segments,
+ run_segment,
+ output_buffer=output_buffer,
+ )
+
+ assert output is output_buffer
+ torch.testing.assert_close(output, query)
+ assert all(segment.query_range is None for segment in plan.segments)
+ assert not plan.segments_cover_query_contiguously
+ assert calls == [
+ (["causal", "causal"], [2, 4], [0, 1, 6, 7]),
+ (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
+ (["causal", "causal"], [6, 8], [5, 11]),
+ ]
+
+
+def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch):
+ plan = build_paged_piecewise_plan(
+ [[(2, 5)], [(2, 5)]],
+ query_offsets=[0, 0],
+ query_lens=[6, 6],
+ seq_lens=[6, 6],
+ device=DEVICE,
+ )
+ assert plan.homogeneous_batch_shape == (2, 6)
+ query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+ calls: list[torch.Tensor] = []
+
+ def fail_index_copy(*_args, **_kwargs):
+ raise AssertionError("homogeneous piecewise path must not use index_copy_")
+
+ monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy)
+
+ def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer):
+ assert segment_key is None
+ assert segment_value is None
+ assert segment_output_buffer is None
+ calls.append(segment_query.clone())
+ return segment_query + 1
+
+ output_buffer = torch.empty_like(query)
+ output = run_paged_piecewise_plan(
+ query,
+ None,
+ None,
+ plan,
+ plan.segments,
+ run_segment,
+ output_buffer=output_buffer,
+ use_homogeneous_batch=True,
+ )
+
+ assert output is output_buffer
+ torch.testing.assert_close(output, query + 1)
+ assert [call.shape[0] for call in calls] == [4, 6, 2]
+ torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]])
+ torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]])
+ torch.testing.assert_close(calls[2], query[[5, 11]])
+
+
+def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout():
+ plan = build_paged_piecewise_plan(
+ [[(2, 5)], [(1, 4)]],
+ query_offsets=[0, 0],
+ query_lens=[6, 6],
+ seq_lens=[6, 6],
+ device=DEVICE,
+ )
+ assert plan.homogeneous_batch_shape is None
+ query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+ output = run_paged_piecewise_plan(
+ query,
+ None,
+ None,
+ plan,
+ plan.segments,
+ lambda segment_query, *_args: segment_query,
+ use_homogeneous_batch=True,
+ )
+ torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_slices_contiguous_segments():
+ plan = build_paged_piecewise_plan(
+ [[(2, 5)]],
+ query_offsets=[0],
+ query_lens=[6],
+ seq_lens=[6],
+ device=DEVICE,
+ )
+ query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+ key = query + 100
+ value = query + 200
+ expected_ranges = [(0, 2), (2, 5), (5, 6)]
+ calls: list[tuple[int, int]] = []
+
+ def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer):
+ assert segment_output_buffer is None
+ expected_range = expected_ranges[len(calls)]
+ start, end = expected_range
+ calls.append(expected_range)
+ assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr()
+ assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr()
+ assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr()
+ torch.testing.assert_close(segment_query, query[start:end])
+ return segment_query
+
+ output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment)
+
+ assert calls == expected_ranges
+ assert [segment.query_range for segment in plan.segments] == expected_ranges
+ assert plan.segments_cover_query_contiguously
+ torch.testing.assert_close(output, query)
+
+
+def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer():
+ plan = build_paged_piecewise_plan(
+ [[(2, 5)]],
+ query_offsets=[0],
+ query_lens=[6],
+ seq_lens=[6],
+ device=DEVICE,
+ )
+ query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+ output_buffer = torch.empty(6, 2, 3)
+ targets = []
+
+ def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer):
+ assert segment_output_buffer is not None
+ assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr()
+ targets.append(segment_output_buffer)
+ return segment_output_buffer.copy_(segment_query[..., :3])
+
+ output = run_paged_piecewise_plan(
+ query,
+ query,
+ query,
+ plan,
+ plan.segments,
+ run_segment,
+ output_buffer=output_buffer,
+ )
+
+ assert output is output_buffer
+ assert [target.shape[0] for target in targets] == [2, 3, 1]
+ torch.testing.assert_close(output, query[..., :3])
+
+
+def test_paged_piecewise_runner_returns_single_full_segment_directly():
+ plan = build_paged_piecewise_plan(
+ [[]],
+ query_offsets=[0],
+ query_lens=[6],
+ seq_lens=[6],
+ device=DEVICE,
+ )
+ query = torch.randn(6, 2, 4)
+ runner_output = query + 1
+
+ output = run_paged_piecewise_plan(
+ query,
+ query,
+ query,
+ plan,
+ plan.segments,
+ lambda *_args: runner_output,
+ )
+
+ assert output is runner_output
+
+
+def test_paged_piecewise_runner_uses_runner_output_shape():
+ plan = build_paged_piecewise_plan(
+ [[(2, 5)]],
+ query_offsets=[0],
+ query_lens=[6],
+ seq_lens=[6],
+ device=DEVICE,
+ )
+ query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+
+ output = run_paged_piecewise_plan(
+ query,
+ query,
+ query,
+ plan,
+ plan.segments,
+ lambda segment_query, *_args: segment_query[..., :3],
+ )
+
+ assert output.shape == (6, 2, 3)
+ torch.testing.assert_close(output, query[..., :3])
+
+
@pytest.mark.parametrize("global_spans", SPAN_CASES)
@pytest.mark.parametrize("q_range", Q_RANGE_CASES)
@pytest.mark.parametrize("batch_size", BATCH_CASES)
diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py
index f079a1758..5b5fcc405 100644
--- a/tests/diffusion/attention/test_rainfusion_plan.py
+++ b/tests/diffusion/attention/test_rainfusion_plan.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
@@ -11,18 +11,25 @@ to the always-kept prefix segment before generating the block mask.
"""
import dataclasses
+import sys
+import types
+from unittest import mock
import pytest
import torch
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends import rainfusion_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
_BLOCK_SIZE,
+ RainFusionAttentionBackend,
RainFusionAttentionImpl,
RainFusionPlan,
)
from vllm_omni.platforms import current_omni_platform
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
PREFIX_ROWS = 710 # 14 text rows + 696 audio rows
ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks
MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
@@ -97,6 +104,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
assert make_impl()._resolve_plan(metadata) is None
+def test_ref2va_multi_video_spans_resolve():
+ layout = VideoTokenLayout(
+ used_len=12000,
+ video_spans=(
+ VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
+ VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
+ ),
+ )
+ plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
+
+ assert plan is not None
+ assert plan.used_len == 12000
+ assert plan.video_spans == [
+ {"start": 128, "latent_shape": [4, 16, 64]},
+ {"start": 5000, "latent_shape": [4, 16, 64]},
+ ]
+
+
+def test_invalid_ref2va_video_spans_fall_back_to_dense():
+ layout = VideoTokenLayout(
+ used_len=12000,
+ video_spans=(
+ VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
+ VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
+ ),
+ )
+ assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
+
+
+def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
+ mindiesd = types.ModuleType("mindiesd")
+
+ def sparse_attention(query, key, value, **kwargs):
+ return query
+
+ mindiesd.sparse_attention = sparse_attention
+ monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
+ monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
+
+ with pytest.raises(ValueError, match="video_spans"):
+ RainFusionAttentionBackend.validate_available()
+
+
+def test_validate_available_accepts_new_mindiesd(monkeypatch):
+ mindiesd = types.ModuleType("mindiesd")
+
+ def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
+ return query
+
+ mindiesd.sparse_attention = sparse_attention
+ monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
+ monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
+
+ RainFusionAttentionBackend.validate_available()
+
+
@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
def test_short_video_stays_dense(grid):
assert make_impl()._resolve_plan(make_metadata(grid)) is None
@@ -169,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
heads, head_dim = 4, 128
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
- impl = make_impl()
+ # Explicit bf16: the dense comparison exercises the plain sparse kernel,
+ # and older MindIE-SD releases do not accept precision= (gate would raise).
+ impl = make_impl(precision="bf16")
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
out = impl._forward_sparse_npu(q, k, v, plan)
@@ -182,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
).transpose(1, 2)
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+
+
+# end_step tail fallback: the last ``end_step`` denoise steps must stay dense.
+
+
+def test_end_step_tail_window_stays_dense():
+ """step_idx inside the last end_step window must resolve no sparse plan."""
+ impl = make_impl(end_step=3)
+ fc = mock.Mock()
+ fc.denoise_step_idx = 47
+ fc.total_denoise_steps = 50
+ with (
+ mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
+ mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
+ ):
+ assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
+
+
+def test_end_step_before_tail_window_still_sparse():
+ """step_idx before the tail window must still resolve a sparse plan."""
+ impl = make_impl(end_step=3)
+ fc = mock.Mock()
+ fc.denoise_step_idx = 46
+ fc.total_denoise_steps = 50
+ with (
+ mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
+ mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
+ ):
+ assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
+
+
+def test_end_step_zero_never_triggers_tail_fallback():
+ """end_step=0 must not fall back even on the final denoise step."""
+ impl = make_impl(end_step=0)
+ fc = mock.Mock()
+ fc.denoise_step_idx = 49
+ fc.total_denoise_steps = 50
+ with (
+ mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
+ mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
+ ):
+ assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
+
+
+# precision capability gate: a mindiesd without sparse_attention(precision=)
+# must fail fast instead of silently running the BF16 path.
+
+
+def _fake_mindiesd_module():
+ """Install a minimal stand-in so the import inside _forward_sparse_npu succeeds."""
+ import types
+
+ fake = types.ModuleType("mindiesd")
+ fake.sparse_attention = lambda *args, **kwargs: None
+ return fake
+
+
+def test_precision_non_bf16_requires_mindiesd_support():
+ """precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
+ import sys
+
+ impl = make_impl(precision="mix")
+ sys.modules["mindiesd"] = _fake_mindiesd_module()
+ try:
+ with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
+ with pytest.raises(RuntimeError, match="requires MindIE-SD"):
+ impl._forward_sparse_npu(None, None, None, None)
+ finally:
+ sys.modules.pop("mindiesd", None)
+
+
+def test_precision_non_bf16_passes_gate_when_supported():
+ """precision != bf16 with a capable mindiesd must not raise from the gate."""
+ import sys
+
+ impl = make_impl(precision="mix")
+ sys.modules["mindiesd"] = _fake_mindiesd_module()
+ try:
+ with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
+ # q/k/v shapes: [B, S, N, D]; plan geometry must match S.
+ q = torch.randn(1, 8, 4, 128)
+ plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
+ # The gate must pass; the fake mindiesd returns None so no crash.
+ out = impl._forward_sparse_npu(q, q, q, plan)
+ assert out is None
+ finally:
+ sys.modules.pop("mindiesd", None)
diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py
index 76daf9d3b..336f7d037 100644
--- a/tests/diffusion/attention/test_trtllm_attn.py
+++ b/tests/diffusion/attention/test_trtllm_attn.py
@@ -1,14 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import math
+from unittest.mock import Mock
import pytest
import torch
import torch.nn.functional as F
from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import (
+ AttentionMetadata,
+ PackedPaddingMetadata,
+)
from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
@@ -35,6 +39,15 @@ def _impl(**bk):
return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk)
+def _packed_padding(cu_seqlens_q, cu_seqlens_k, q_length, kv_length):
+ return PackedPaddingMetadata(
+ q_length=q_length,
+ kv_length=kv_length,
+ cu_seqlens_q=cu_seqlens_q[:2],
+ cu_seqlens_k=cu_seqlens_k[:2],
+ )
+
+
def test_skip_config_pure_resolution():
from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
@@ -75,7 +88,7 @@ def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
def test_quant_quantize_calls_routine_and_shapes_sfs():
from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
- captured = {}
+ captured: dict[str, object] = {}
def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
@@ -209,62 +222,80 @@ def test_masked_layer_raises():
_impl().forward_cuda(q, k, v, _Meta())
-def test_packed_metadata_trims_padding_mask(monkeypatch):
- b, s, h, d = 1, 8, 8, 128
- q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
- mask = torch.tensor([[True, True, True, True, True, True, False, False]])
- cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+def test_generic_packed_metadata_preserves_ragged_batch(monkeypatch):
+ b, h, d = 1, 8, 128
+ q = torch.zeros(b, 5, h, d)
+ k, v = (torch.zeros(b, 7, h, d) for _ in range(2))
metadata = AttentionMetadata(
- attn_mask=mask,
extra={
- "cu_seqlens_q": cu_seqlens,
- "cu_seqlens_k": cu_seqlens,
- "max_seqlen_q": 6,
- "max_seqlen_k": 6,
+ "cu_seqlens_q": torch.tensor([0, 2, 5], dtype=torch.int32),
+ "cu_seqlens_k": torch.tensor([0, 3, 7], dtype=torch.int32),
+ "max_seqlen_q": 3,
+ "max_seqlen_k": 4,
+ }
+ )
+ fake_attention = Mock(return_value=torch.zeros_like(q.reshape(5, h, d)))
+
+ monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+ monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
+ monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
+
+ _impl().forward_cuda(q, k, v, metadata)
+ captured = fake_attention.call_args.kwargs
+
+ assert captured["query"].shape[0] == 5
+ assert captured["key"].shape[0] == 7
+ assert captured["value"].shape[0] == 7
+ assert captured["batch_size"] == 2
+ assert captured["seq_lens"].tolist() == [3, 4]
+ assert tg.TrtllmAttentionBackend.supports_multi_doc_packed_varlen()
+
+
+def test_packed_padding_metadata_trims_without_mask(monkeypatch):
+ q_len, kv_len = 8, 9
+ valid_q_tokens, valid_kv_tokens = 5, 6
+ b, h, d = 1, 8, 128
+ q = torch.zeros(b, q_len, h, d)
+ k, v = (torch.zeros(b, kv_len, h, d) for _ in range(2))
+ cu_seqlens_q = torch.tensor([0, valid_q_tokens, q_len], dtype=torch.int32)
+ cu_seqlens_k = torch.tensor([0, valid_kv_tokens, kv_len], dtype=torch.int32)
+ metadata = AttentionMetadata(
+ packed_padding=_packed_padding(
+ cu_seqlens_q,
+ cu_seqlens_k,
+ valid_q_tokens,
+ valid_kv_tokens,
+ ),
+ extra={
+ "cu_seqlens_q": cu_seqlens_q,
+ "cu_seqlens_k": cu_seqlens_k,
+ "max_seqlen_q": valid_q_tokens,
+ "max_seqlen_k": valid_kv_tokens,
+ "valid_kv_length": valid_kv_tokens,
},
)
- captured = {}
-
- def fake_attention(**kwargs):
- captured.update(kwargs)
- return torch.zeros_like(kwargs["query"])
+ fake_attention = Mock(return_value=torch.ones(valid_q_tokens, h, d))
monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
out = _impl().forward_cuda(q, k, v, metadata)
+ captured = fake_attention.call_args.kwargs
+ assert tg.TrtllmAttentionBackend.supports_packed_mask_free()
assert out.shape == q.shape
- assert captured["query"].shape[0] == 6
- assert captured["key"].shape[0] == 6
- assert captured["value"].shape[0] == 6
+ assert captured["query"].shape[0] == valid_q_tokens
+ assert captured["key"].shape[0] == valid_kv_tokens
+ assert captured["value"].shape[0] == valid_kv_tokens
assert captured["batch_size"] == 1
- assert captured["seq_lens"].tolist() == [6]
- assert captured["cum_seq_lens_q"].tolist() == [0, 6]
- assert captured["cum_seq_lens_kv"].tolist() == [0, 6]
- assert captured["max_q_len"] == 6
- assert captured["max_kv_len"] == 6
- assert torch.count_nonzero(out[:, 6:]) == 0
-
-
-def test_packed_metadata_rejects_non_prefix_mask(monkeypatch):
- b, s, h, d = 1, 8, 8, 128
- q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
- cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
- metadata = AttentionMetadata(
- attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]),
- extra={
- "cu_seqlens_q": cu_seqlens,
- "cu_seqlens_k": cu_seqlens,
- "max_seqlen_q": 6,
- "max_seqlen_k": 6,
- },
- )
- monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
-
- with pytest.raises(ValueError, match="prefix-valid"):
- _impl().forward_cuda(q, k, v, metadata)
+ assert captured["seq_lens"].tolist() == [valid_kv_tokens]
+ assert captured["cum_seq_lens_q"].tolist() == [0, valid_q_tokens]
+ assert captured["cum_seq_lens_kv"].tolist() == [0, valid_kv_tokens]
+ assert captured["max_q_len"] == valid_q_tokens
+ assert captured["max_kv_len"] == valid_kv_tokens
+ assert torch.count_nonzero(out[:, :valid_q_tokens] != 1) == 0
+ assert torch.count_nonzero(out[:, valid_q_tokens:]) == 0
def test_packed_metadata_must_be_complete():
@@ -393,12 +424,13 @@ def test_bf16_packed_padding_matches_sdpa():
q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
metadata = AttentionMetadata(
- attn_mask=torch.arange(s, device="cuda")[None] < used,
+ packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
extra={
"cu_seqlens_q": cu_seqlens,
"cu_seqlens_k": cu_seqlens,
"max_seqlen_q": used,
"max_seqlen_k": used,
+ "valid_kv_length": used,
},
)
@@ -443,12 +475,13 @@ def test_sage_packed_non_aligned_length_matches_sdpa():
q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
metadata = AttentionMetadata(
- attn_mask=torch.arange(s, device="cuda")[None] < used,
+ packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
extra={
"cu_seqlens_q": cu_seqlens,
"cu_seqlens_k": cu_seqlens,
"max_seqlen_q": used,
"max_seqlen_k": used,
+ "valid_kv_length": used,
},
)
diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py
index 184807d98..77e00edda 100644
--- a/tests/diffusion/attention/test_ulysses_uaa.py
+++ b/tests/diffusion/attention/test_ulysses_uaa.py
@@ -1,10 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import os
-import socket
import tempfile
from types import SimpleNamespace
@@ -12,6 +11,7 @@ import numpy as np
import pytest
import torch
+from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
from vllm_omni.diffusion.config import set_current_diffusion_config
@@ -27,20 +27,6 @@ from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-def _find_free_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
- os.environ["RANK"] = str(rank)
- os.environ["LOCAL_RANK"] = str(rank)
- os.environ["WORLD_SIZE"] = str(world_size)
- os.environ["MASTER_ADDR"] = "127.0.0.1"
- os.environ["MASTER_PORT"] = str(master_port)
-
-
@pytest.mark.core_model
@pytest.mark.cpu
def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
@@ -71,7 +57,7 @@ def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
def _run_attention_case(
local_rank: int,
world_size: int,
- master_port: int,
+ init_method: str,
input_file: str,
output_file: str,
num_heads: int,
@@ -85,10 +71,9 @@ def _run_attention_case(
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
- _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port)
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
- init_distributed_environment(world_size=world_size, rank=local_rank)
+ init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
initialize_model_parallel(
data_parallel_size=1,
cfg_parallel_size=1,
@@ -202,10 +187,8 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
batch_size = 2
head_size = 8
- base_port = _find_free_port()
- sp_port = _find_free_port()
- while sp_port == base_port:
- sp_port = _find_free_port()
+ base_init_method = get_distributed_init_method()
+ sp_init_method = get_distributed_init_method()
with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
input_file = f_in.name
@@ -224,7 +207,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
# Baseline (no SP)
torch.multiprocessing.spawn(
_run_attention_case,
- args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"),
+ args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
nprocs=1,
)
@@ -233,7 +216,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
_run_attention_case,
args=(
sp_world_size,
- sp_port,
+ sp_init_method,
input_file,
sp_file,
num_heads,
@@ -278,10 +261,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
split_sizes = [3, 2, 3, 2]
- base_port = _find_free_port()
- sp_port = _find_free_port()
- while sp_port == base_port:
- sp_port = _find_free_port()
+ base_init_method = get_distributed_init_method()
+ sp_init_method = get_distributed_init_method()
with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
input_file = f_in.name
@@ -300,7 +281,19 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
# Baseline (no SP)
torch.multiprocessing.spawn(
_run_attention_case,
- args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"),
+ args=(
+ 1,
+ base_init_method,
+ input_file,
+ baseline_file,
+ num_heads,
+ head_size,
+ 1,
+ "strict",
+ 1,
+ None,
+ "mem_efficient",
+ ),
nprocs=1,
)
@@ -309,7 +302,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
_run_attention_case,
args=(
sp_world_size,
- sp_port,
+ sp_init_method,
input_file,
sp_file,
num_heads,
diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py
index 60136ce4d..c9ef78383 100644
--- a/tests/diffusion/batching/test_diffusion_batching.py
+++ b/tests/diffusion/batching/test_diffusion_batching.py
@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
diffusion model and validates that every caller receives its correct
individual result. When the underlying diffusion stage is configured with
-``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
+``max_num_seqs > 1``, the
requests will be batched internally.
Even without explicit batching config this test is useful for verifying
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
) -> None:
"""Run the same prompts sequentially then in parallel and print a comparison."""
- omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+ omni = AsyncOmni(model=model, max_num_seqs=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
single_time = await run_single(omni, prompts)
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
return
- omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+ omni = AsyncOmni(model=model, max_num_seqs=batch_size)
try:
await warmup(omni, WARMUP_PROMPTS)
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
fired via asyncio.gather. Each request_id must appear in the results."""
async def _inner():
- omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
+ omni = AsyncOmni(model=model_name, max_num_seqs=1)
try:
prompts = TEST_PROMPTS[:4]
sp = _default_sampling_params()
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
"""
async def _inner():
- omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
+ omni = AsyncOmni(model=model_name, max_num_seqs=4)
try:
prompts = TEST_PROMPTS[:4]
sp = _default_sampling_params()
@@ -524,7 +524,12 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
parser.add_argument("--model", type=str, required=True, help="Model name or path")
parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
- parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ default=1,
+ help="Scheduler in-flight width (--max-num-seqs)",
+ )
parser.add_argument(
"--mode",
choices=["batch", "single", "compare", "validate"],
diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py
index fb912b08e..85bee131b 100644
--- a/tests/diffusion/cache/test_cache_backends.py
+++ b/tests/diffusion/cache/test_cache_backends.py
@@ -365,13 +365,10 @@ class TestCacheDiTBackend:
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
- def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
- """Test DreamID uses pipeline.transformer for cache enable/refresh.
-
- NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
- """
+ def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
+ """Generic cache path should pick up ``transformer.fused_blocks``."""
mock_pipeline = Mock()
- mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
+ mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
mock_pipeline.transformer = Mock()
mock_pipeline.transformer.fused_blocks = Mock()
mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py
index 2799f1607..165a98cc6 100644
--- a/tests/diffusion/cache/test_cache_dit.py
+++ b/tests/diffusion/cache/test_cache_dit.py
@@ -1,18 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Model specific tests for CacheDiT enablement.
"""
import ast
-import sys
from pathlib import Path
from unittest.mock import Mock, patch
import pytest
import torch
from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
+from vllm.distributed import parallel_state
import vllm_omni.diffusion.cache.cachedit as cd_backend
import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
@@ -24,19 +24,9 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import L
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.platforms import current_omni_platform
-# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect
-# the class even though the dependency may not be set up correctly; this is ok for these
-# tests because we just inspect it and never initialize the model.
-for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"):
- sys.modules.setdefault(mod, Mock())
-# isort: split
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402
-
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
SEPARATE_CFG_TRANSFORMERS = [
- DreamIdOmniModel,
- LTX2VideoTransformer3DModel,
HeliosTransformer3DModel,
LongCatImageTransformer2DModel,
Cosmos3VFMTransformer,
@@ -62,6 +52,18 @@ def test_custom_cache_dit_enablers_are_registered_explicitly():
assert cd_backend.CUSTOM_DIT_ENABLERS == expected_enablers
+@pytest.fixture()
+def init_fake_tp_group(mocker):
+ """Provide a fake TP group so vLLM linear layers can be instantiated."""
+ mock_tp = mocker.MagicMock()
+ mock_tp.world_size = 1
+ mock_tp.rank_in_group = 0
+ old = parallel_state._TP
+ parallel_state._TP = mock_tp
+ yield
+ parallel_state._TP = old
+
+
def test_wan22_vace_uses_wan22_custom_cache_dit_enabler():
assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_model_specific.enable_cache_for_wan22
@@ -111,7 +113,9 @@ def test_cachedit_consumers_use_package_api():
for module in modules:
if module == legacy_module or module.startswith(internal_prefix):
- invalid_imports.append(f"{source_path.relative_to(package_root)}:{node.lineno}: {module}")
+ invalid_imports.append(
+ f"{source_path.relative_to(package_root)}:{getattr(node, 'lineno', '?')}: {module}"
+ )
assert not invalid_imports, "Cache-DiT consumers must use the package API:\n" + "\n".join(invalid_imports)
@@ -124,6 +128,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
+def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
+ """LTX batches all guidance passes into one Transformer invocation."""
+ adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
+
+ assert adapter_config.has_separate_cfg is False
+
+
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py
index fdeca8e17..d821541d0 100644
--- a/tests/diffusion/cache/test_teacache_extractors.py
+++ b/tests/diffusion/cache/test_teacache_extractors.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for TeaCache extractor functions.
@@ -18,6 +18,7 @@ Currently implemented:
- TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor
"""
+import math
from abc import ABC, abstractmethod
from types import SimpleNamespace
from unittest.mock import MagicMock, Mock, patch
@@ -650,6 +651,29 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
torch.testing.assert_close(actual_video, expected_video)
torch.testing.assert_close(actual_audio, expected_audio)
+ def test_extractor_uses_prepared_rope_table(
+ self,
+ minimax_h3_module,
+ sample_inputs,
+ monkeypatch,
+ ):
+ """TeaCache must consume the denoise-loop RoPE table without rebuilding it."""
+ rope_table = minimax_h3_module.prepare_rope_table(
+ sample_inputs["img_position_ids"],
+ seq_len=sample_inputs["x"].shape[1],
+ )
+
+ def unexpected_rope_build(*args, **kwargs):
+ raise AssertionError("TeaCache extractor rebuilt the prepared RoPE table")
+
+ monkeypatch.setattr(minimax_h3_module, "prepare_rope_table", unexpected_rope_build)
+ monkeypatch.setattr(minimax_h3_module.rope, "forward", unexpected_rope_build)
+ context = extract_minimax_h3_context(
+ minimax_h3_module,
+ **{**sample_inputs, "rope_table": rope_table},
+ )
+ context.run_transformer_blocks()
+
def test_run_transformer_blocks_forwards_packed_layout(
self,
minimax_h3_module,
@@ -705,6 +729,222 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0]
assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0]
+ @pytest.mark.sp
+ def test_teacache_strict_sp_keeps_cache_state_rank_local(
+ self,
+ minimax_h3_module,
+ sample_inputs,
+ monkeypatch,
+ ):
+ """Strict SP must not mix local TeaCache state with gathered block rows."""
+ from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
+ from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+ from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
+ from vllm_omni.diffusion.distributed import parallel_state
+
+ seq_len = sample_inputs["x"].shape[1]
+ local_len = seq_len // 2
+ hidden_size = minimax_h3_module.hidden_size
+ video_width = minimax_h3_module.arch.latents_dim * math.prod(minimax_h3_module.arch.patch_size)
+ audio_width = minimax_h3_module.arch.audio_latents_dim
+ captured_modulation_indices: list[torch.Tensor] = []
+
+ class _StrictSPPrepare(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.global_indices: list[torch.Tensor] = []
+
+ def forward(self, hidden, rope_table, combined_indices):
+ assert hidden.shape[0] == local_len
+ assert rope_table.shape[0] == local_len
+ self.global_indices.append(combined_indices.clone())
+ return hidden, rope_table, combined_indices.narrow(0, 0, local_len)
+
+ class _StrictSPGather(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.inputs: list[torch.Tensor] = []
+ self.outputs: list[torch.Tensor] = []
+
+ def forward(self, tensor):
+ self.inputs.append(tensor.clone())
+ if tensor.shape[-1] == hidden_size:
+ raise AssertionError("TeaCache must not gather hidden states in strict SP")
+ assert tensor.shape == (local_len, video_width + audio_width)
+ # Model the all-gather of compact logits from every SP rank.
+ # postprocess() indexes the resulting global layout with
+ # infer_out_pos/audio_pos, including rows owned by another rank.
+ gathered = tensor.repeat(seq_len // local_len, 1)
+ assert gathered.shape == (seq_len, video_width + audio_width)
+ self.outputs.append(gathered)
+ return gathered
+
+ class _StrictSPGroup:
+ def __init__(self):
+ super().__init__()
+ self.local_decisions: list[int] = []
+ self.ops = []
+
+ def all_reduce(self, decision, op):
+ self.local_decisions.append(int(decision.item()))
+ self.ops.append(op)
+ # Simulate another rank requiring the second step to execute
+ # the collective block path even though this rank would cache.
+ if len(self.local_decisions) == 2:
+ decision.fill_(1)
+ return decision
+
+ class _StrictSPBlock(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.block_indices: list[torch.Tensor] = []
+
+ def adaln_proj(self, t_emb):
+ del t_emb
+ zeros = torch.zeros(3, hidden_size, dtype=torch.float32)
+ return zeros, zeros
+
+ def norm1(self, hidden):
+ return hidden.clone()
+
+ def forward(self, hidden, *, combined_indices, rope_table, **kwargs):
+ del kwargs
+ assert hidden.shape[0] == local_len
+ assert rope_table.shape[0] == local_len
+ self.block_indices.append(combined_indices.clone())
+ return hidden + 1
+
+ class _StrictSPFinalLayer(nn.Module):
+ def forward(self, hidden, *, t_emb, inverse_indices):
+ del t_emb, inverse_indices
+ assert hidden.shape == (local_len, hidden_size)
+ return hidden[:, :video_width].float(), hidden[:, :audio_width].float()
+
+ local_prepare = _StrictSPPrepare()
+ gather = _StrictSPGather()
+ sp_group = _StrictSPGroup()
+ block = _StrictSPBlock()
+ minimax_h3_module.local_sp_prepare = local_prepare
+ minimax_h3_module.sp_gather = gather
+ minimax_h3_module.blocks = nn.ModuleList([block])
+ minimax_h3_module.final_layer = _StrictSPFinalLayer()
+ monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, local_len))
+ monkeypatch.setattr(
+ minimax_h3_module,
+ "_embed",
+ lambda **kwargs: (
+ torch.arange(local_len * hidden_size, dtype=torch.float32).reshape(local_len, hidden_size),
+ torch.zeros(3, hidden_size, dtype=torch.float32),
+ ),
+ )
+ monkeypatch.setattr(
+ minimax_h3_module,
+ "prepare_rope_table",
+ lambda *args, **kwargs: torch.zeros(local_len, 6, dtype=torch.float32),
+ )
+ monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
+
+ original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
+
+ def capture_indexed_scale_shift(hidden, shift, scale, indices):
+ captured_modulation_indices.append(indices.clone())
+ return original_indexed_scale_shift(hidden, shift, scale, indices)
+
+ monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
+
+ local_img_pos = torch.tensor([0, 1])
+ local_audio_pos = torch.tensor([2, 3])
+ token_tags = torch.full((seq_len,), -1, dtype=torch.long)
+ token_tags[local_img_pos] = 0
+ token_tags[local_audio_pos] = 2
+ inverse_indices = torch.zeros(seq_len, dtype=torch.long)
+ inputs = {
+ **sample_inputs,
+ "inverse_indices": inverse_indices,
+ "token_tags": token_tags,
+ "update_mask": torch.ones(local_img_pos.numel(), dtype=torch.bool),
+ "prompt_embeds": torch.empty(0, 6),
+ "img_pos_info": {"position_ids": local_img_pos},
+ "audio_pos_info": {"position_ids": local_audio_pos},
+ "text_pos_info": {"position_ids": torch.empty(0, dtype=torch.long)},
+ "img_pos_for_infer_output_info": {"position_ids": local_img_pos},
+ "refiner_packed_seq_params": {
+ "cu_seqlens_q": torch.tensor([0, 0], dtype=torch.int32),
+ "max_seqlen_q": 0,
+ },
+ }
+ expected_local_indices = (inverse_indices * 3 + token_tags.clamp(min=0))[:local_len]
+
+ hook = TeaCacheHook(
+ TeaCacheConfig(
+ transformer_type="MiniMaxH3DiTModel",
+ rel_l1_thresh=0.3,
+ )
+ )
+ hook.initialize_hook(minimax_h3_module)
+
+ hook.new_forward(minimax_h3_module, **inputs)
+ hook.new_forward(minimax_h3_module, **inputs)
+ video_logits, audio_logits = hook.new_forward(minimax_h3_module, **inputs)
+ state = hook.state_manager.get_state()
+
+ assert state.previous_residual is not None
+ assert state.previous_residual.shape == (local_len, hidden_size)
+ assert len(captured_modulation_indices) == 3
+ assert all(torch.equal(indices, expected_local_indices) for indices in captured_modulation_indices)
+ assert sp_group.local_decisions == [1, 0, 0]
+ assert len(sp_group.ops) == 3
+ assert all(op == torch.distributed.ReduceOp.MAX for op in sp_group.ops)
+ assert len(local_prepare.global_indices) == 2
+ assert all(indices.shape == (seq_len,) for indices in local_prepare.global_indices)
+ assert len(block.block_indices) == 2
+ assert all(torch.equal(indices, expected_local_indices) for indices in block.block_indices)
+ assert len(gather.inputs) == 3
+ assert video_logits.shape == (local_img_pos.numel(), video_width)
+ assert audio_logits.shape == (local_audio_pos.numel(), audio_width)
+
+ @pytest.mark.cpu
+ def test_teacache_non_strict_sp_gathers_hidden_once(
+ self,
+ minimax_h3_module,
+ sample_inputs,
+ monkeypatch,
+ ):
+ """The non-strict SP cache path must not gather global hidden rows twice."""
+ from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+ from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
+
+ class _Prepare(nn.Module):
+ def forward(self, hidden, rope_table, combined_indices):
+ return hidden, rope_table, combined_indices
+
+ class _Gather(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.inputs: list[torch.Tensor] = []
+
+ def forward(self, tensor):
+ self.inputs.append(tensor.clone())
+ return tensor
+
+ seq_len = sample_inputs["x"].shape[1]
+ gather = _Gather()
+ minimax_h3_module.sp_prepare = _Prepare()
+ minimax_h3_module.sp_gather = gather
+ monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, seq_len))
+
+ hook = TeaCacheHook(
+ TeaCacheConfig(
+ transformer_type="MiniMaxH3DiTModel",
+ rel_l1_thresh=0.3,
+ )
+ )
+ hook.initialize_hook(minimax_h3_module)
+ hook.new_forward(minimax_h3_module, **sample_inputs)
+
+ assert len(gather.inputs) == 1
+ assert gather.inputs[0].shape == (seq_len, minimax_h3_module.hidden_size)
+
def test_invalid_module_raises_error(self):
"""Test that invalid module without blocks raises ValueError."""
invalid_module = Mock()
diff --git a/tests/diffusion/conftest.py b/tests/diffusion/conftest.py
new file mode 100644
index 000000000..e12d15503
--- /dev/null
+++ b/tests/diffusion/conftest.py
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def restore_torch_default_dtype():
+ """Keep process-global dtype changes from leaking between diffusion tests."""
+ default_dtype = torch.get_default_dtype()
+ try:
+ yield
+ finally:
+ torch.set_default_dtype(default_dtype)
diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py
index babc23b4a..3b094c37f 100644
--- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py
+++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections import namedtuple
from types import SimpleNamespace
@@ -9,8 +9,9 @@ import torch
from diffusers import DiffusionPipeline # pyright: ignore[reportPrivateImportUsage]
from PIL import Image
+import vllm_omni.diffusion.data as diffusion_data
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
from vllm_omni.diffusion.data import (
DiffusionOutput,
DiffusionParallelConfig,
@@ -629,6 +630,46 @@ class TestPipelineArgumentsHandling:
with pytest.raises(ValueError):
pipeline.forward(DiffusionRequestBatch(requests=[problematic_request]))
+ interpolation_request = _make_request(
+ sampling_params=OmniDiffusionSamplingParams(
+ enable_frame_interpolation=True,
+ ),
+ )
+ with pytest.raises(ValueError, match="Frame interpolation is not supported with the Diffusers backend"):
+ pipeline.forward(DiffusionRequestBatch(requests=[interpolation_request]))
+
+ def test_diffusers_config_preserves_checkpoint_class_for_metadata(self, mocker):
+ class MockWanImageToVideoPipeline:
+ pass
+
+ od_config = _make_od_config(model_class_name=None)
+ mocker.patch(
+ "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+ return_value={"_class_name": "WanImageToVideoPipeline"},
+ )
+ mocker.patch.object(
+ diffusion_data,
+ "diffusers",
+ SimpleNamespace(WanImageToVideoPipeline=MockWanImageToVideoPipeline),
+ )
+ mocker.patch.object(OmniDiffusionConfig, "update_multimodal_support")
+
+ od_config.enrich_config()
+
+ assert od_config.model_class_name == "WanImageToVideoPipeline"
+ assert od_config.diffusers_pipeline_cls is MockWanImageToVideoPipeline
+
+ def test_diffusers_config_falls_back_to_adapter_without_pipeline_index(self, mocker):
+ od_config = _make_od_config(model_class_name=None)
+ mocker.patch(
+ "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+ return_value=None,
+ )
+
+ od_config.enrich_config()
+
+ assert od_config.model_class_name == "DiffusersAdapterPipeline"
+
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@@ -651,7 +692,7 @@ class TestDiffusersBackendEndToEndExecution:
def test_t2i_random_weights(
self,
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
messages = dummy_messages_from_mix_data(content_text="a photo of an astronaut riding a horse on mars")
@@ -667,8 +708,10 @@ class TestDiffusersBackendEndToEndExecution:
},
}
- response = openai_client.send_diffusion_request(request_config)
- image: Image.Image = response[0].images[0] # pyright: ignore[reportOptionalSubscript]
+ response = online_client.send_diffusion_request(request_config)
+ images = response[0].images
+ assert images is not None
+ image: Image.Image = images[0]
# Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete.
assert image.size == (512, 512)
diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py
new file mode 100644
index 000000000..f2002f1c6
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_block_tables.py
@@ -0,0 +1,780 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+from contextlib import nullcontext
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.worker import block_table as native_block_table
+
+from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.metadata import (
+ DiffusionKVContextMetadata,
+ DiffusionKVMetadata,
+ DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _registration_backend(parallel_config: object):
+ model_config = SimpleNamespace(set_attention_geometry=Mock())
+ vllm_config = SimpleNamespace(
+ attention_config=SimpleNamespace(use_non_causal=False),
+ compilation_config=SimpleNamespace(static_forward_context={}),
+ model_config=model_config,
+ )
+ backend = DiffusionKVModelRunnerBackend(
+ vllm_config=vllm_config,
+ od_config=SimpleNamespace(parallel_config=parallel_config),
+ device=torch.device("cpu"),
+ )
+ return backend, vllm_config, model_config
+
+
+def _parallel_config(**overrides):
+ values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict")
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
+def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2))
+ layer = SimpleNamespace(
+ num_heads=8,
+ softmax_scale=0.125,
+ skip_sequence_parallel=False,
+ attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"),
+ )
+ unrelated = object()
+ vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated})
+ spec = FullAttentionSpec(
+ block_size=16,
+ num_kv_heads=4,
+ head_size=8,
+ dtype=torch.float16,
+ non_causal=True,
+ )
+ constructor_calls = []
+
+ class FakeLayerAdapter:
+ def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None:
+ constructor_calls.append((layer_name, layer, spec, ulysses_degree))
+ self.num_heads = layer.num_heads // ulysses_degree
+ self.num_kv_heads = spec.num_kv_heads // ulysses_degree
+ self.head_size = spec.head_size
+ self.spec = replace(spec, num_kv_heads=self.num_kv_heads)
+
+ monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter)
+ monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+ specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+ installed = vllm_config.compilation_config.static_forward_context["layer-0"]
+ assert isinstance(installed, FakeLayerAdapter)
+ assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated
+ assert specs["layer-0"].num_kv_heads == 2
+ assert constructor_calls == [("layer-0", layer, spec, 2)]
+ model_config.set_attention_geometry.assert_called_once_with(
+ num_heads=4,
+ num_kv_heads=2,
+ head_size=8,
+ )
+ assert vllm_config.attention_config.use_non_causal is True
+
+
+def test_cache_layer_registration_rejects_backend_without_paged_support() -> None:
+ backend, _, _ = _registration_backend(_parallel_config())
+ layer = SimpleNamespace(
+ attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"),
+ )
+ spec = FullAttentionSpec(
+ block_size=16,
+ num_kv_heads=4,
+ head_size=8,
+ dtype=torch.float16,
+ )
+
+ with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"):
+ backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+
+def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None:
+ backend, _, _ = _registration_backend(_parallel_config())
+
+ def unsupported_block_tables():
+ raise NotImplementedError("platform has no paged BlockTables")
+
+ monkeypatch.setattr(
+ model_runner_backend_module.current_omni_platform,
+ "get_diffusion_kv_block_tables_cls",
+ unsupported_block_tables,
+ )
+
+ with pytest.raises(NotImplementedError, match="no paged BlockTables"):
+ backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))})
+
+
+@pytest.mark.parametrize(
+ ("parallel_config", "message"),
+ [
+ (_parallel_config(ring_degree=2), "Ring"),
+ (_parallel_config(allgather_degree=2), "AllGather-KV"),
+ (_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"),
+ ],
+)
+def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None:
+ backend, _, _ = _registration_backend(parallel_config)
+
+ with pytest.raises(NotImplementedError, match=message):
+ backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False))
+
+
+def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None:
+ backend, _, _ = _registration_backend(
+ _parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa")
+ )
+
+ assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1
+
+
+class _FakeSpec:
+ def __init__(self, block_size: int) -> None:
+ self.block_size = block_size
+ self.capacity_calls: list[tuple[object, int]] = []
+
+ def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int:
+ self.capacity_calls.append((vllm_config, max_len))
+ return math.ceil(max_len / self.block_size)
+
+
+class _FakeStagedTable:
+ def __init__(self) -> None:
+ self.staged: list[tuple[int, tuple[int, ...]]] = []
+ self.clear_calls = 0
+
+ def clear_staged_writes(self) -> None:
+ self.staged.clear()
+ self.clear_calls += 1
+
+
+class _FakeNumBlocks:
+ def __init__(self, num_groups: int, max_num_reqs: int) -> None:
+ self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32)
+ self.copy_calls = 0
+
+ def copy_to_uva(self) -> None:
+ self.copy_calls += 1
+
+
+class _FakeBlockTables:
+ instances: list[_FakeBlockTables] = []
+
+ def __init__(self, **kwargs: object) -> None:
+ self.kwargs = kwargs
+ self.block_sizes = list(kwargs["block_sizes"])
+ self.max_num_reqs = int(kwargs["max_num_reqs"])
+ self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"])
+ self.cp_size = int(kwargs["cp_size"])
+ self.cp_rank = int(kwargs["cp_rank"])
+ self.cp_interleave = int(kwargs["cp_interleave"])
+ self.blocks_per_kv_block = [1] * len(self.block_sizes)
+ self.block_tables = [_FakeStagedTable() for _ in self.block_sizes]
+ self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs)
+ self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = []
+ self.apply_calls = 0
+ self.layout_refreshes = 0
+ self.fail_on_append_call: int | None = None
+ self.fail_apply = False
+ self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes]
+ self.__class__.instances.append(self)
+
+ def append_block_ids(
+ self,
+ req_index: int,
+ new_block_ids: tuple[list[int], ...],
+ overwrite: bool,
+ ) -> None:
+ call_number = len(self.append_calls) + 1
+ if self.fail_on_append_call == call_number:
+ raise RuntimeError("fake append failure")
+ copied_ids = tuple(list(group_ids) for group_ids in new_block_ids)
+ self.append_calls.append((req_index, copied_ids, overwrite))
+ for group_index, group_ids in enumerate(copied_ids):
+ self.num_blocks.np[group_index, req_index] = len(group_ids)
+ if group_ids:
+ self.block_tables[group_index].staged.append((req_index, tuple(group_ids)))
+
+ def apply_staged_writes(self) -> None:
+ self.apply_calls += 1
+ if self.fail_apply:
+ raise RuntimeError("fake apply failure")
+ for group_index, group_table in enumerate(self.block_tables):
+ for row, group_ids in group_table.staged:
+ self.committed[group_index][row] = group_ids
+ group_table.staged.clear()
+ self.num_blocks.copy_to_uva()
+
+ def init_block_table_layout_tensors(self) -> None:
+ self.layout_refreshes += 1
+
+
+def _metadata(
+ request_id: str = "req-0",
+ generation: int = 1,
+ *,
+ seq_len: int = 9,
+ sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]),
+ contexts: tuple[DiffusionKVContextMetadata, ...] | None = None,
+) -> DiffusionKVMetadata:
+ if contexts is None:
+ contexts = (
+ DiffusionKVContextMetadata(
+ context_id="text",
+ cache_role="cross_attention",
+ num_tokens=5,
+ block_ids=([6, 7], [8]),
+ ),
+ )
+ return DiffusionKVMetadata(
+ request_id=request_id,
+ allocation_generation=generation,
+ sequences=(
+ DiffusionKVSequenceMetadata(
+ sequence_id=0,
+ prefix_len=4,
+ target_len=5,
+ seq_len=seq_len,
+ block_ids=sequence_block_ids,
+ context_ids=tuple(context.context_id for context in contexts),
+ ),
+ ),
+ contexts=contexts,
+ )
+
+
+def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None:
+ _FakeBlockTables.instances.clear()
+
+ def get_block_table_width(max_num_blocks, block_size):
+ events.append(("width", max_num_blocks, block_size))
+ return max_num_blocks + 1
+
+ def init_backend(kv_cache_config, vllm_config, device):
+ events.append(("backend", kv_cache_config, vllm_config, device))
+ group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
+ return [[group], [group]], "cg-support", [4, 8]
+
+ def init_cache(
+ runner_kv_caches,
+ forward_context,
+ kv_cache_config,
+ attn_groups,
+ device,
+ cache_dtype,
+ kernel_block_sizes,
+ vllm_config,
+ ):
+ events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+ runner_kv_caches.extend(["cache-0", "cache-1"])
+ return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
+
+ monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend)
+ monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache)
+ monkeypatch.setattr(
+ model_runner_backend_module.current_omni_platform,
+ "get_diffusion_kv_block_tables_cls",
+ lambda: _FakeBlockTables,
+ )
+ monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False)
+
+
+def _runner(
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ max_rows_per_request: int = 3,
+ max_num_seqs: int = 2,
+ max_len: int = 16,
+ max_num_batched_tokens: int = 32,
+ cp_size: int = 1,
+ cp_rank: int = 0,
+ cp_interleave: int = 1,
+) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]:
+ events: list[tuple] = []
+ _install_native_modules(monkeypatch, events)
+ od_config = SimpleNamespace(
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ diffusion_kv_max_rows_per_request=max_rows_per_request,
+ max_num_seqs=max_num_seqs,
+ )
+ vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(max_model_len=max_len),
+ scheduler_config=SimpleNamespace(
+ max_num_batched_tokens=max_num_batched_tokens,
+ max_num_seqs=max_num_seqs,
+ ),
+ parallel_config=SimpleNamespace(
+ decode_context_parallel_size=cp_size,
+ cp_kv_cache_interleave_size=cp_interleave,
+ ),
+ compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+ cache_config=SimpleNamespace(cache_dtype="auto"),
+ )
+ monkeypatch.setattr(
+ model_runner_backend_module,
+ "get_dcp_group",
+ lambda: SimpleNamespace(rank_in_group=cp_rank),
+ )
+ backend = DiffusionKVModelRunnerBackend(
+ vllm_config=vllm_config,
+ od_config=od_config,
+ device="cuda:1",
+ )
+ backend._kv_cache_layer_adapters = {
+ "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+ "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+ }
+
+ config = SimpleNamespace(
+ num_blocks=32,
+ kv_cache_groups=[
+ SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+ SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+ ],
+ )
+ backend.initialize_kv_cache(config)
+ return backend, _FakeBlockTables.instances[-1], events
+
+
+def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, events = _runner(monkeypatch)
+
+ assert events[:2] == [("width", 4, 4), ("width", 2, 8)]
+ assert [event[0] for event in events[2:]] == ["backend", "cache"]
+ assert block_tables.kwargs == {
+ "block_sizes": [4, 8],
+ "max_num_reqs": 6,
+ "max_num_batched_tokens": 32,
+ "max_num_blocks_per_group": [5, 3],
+ "device": "cuda:1",
+ "kernel_block_sizes": [4, 8],
+ "cp_size": 1,
+ "cp_rank": 0,
+ "cp_interleave": 1,
+ }
+ assert runner.block_tables is block_tables
+ assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+ assert runner.kv_caches == ["cache-0", "cache-1"]
+ assert runner.vllm_config.scheduler_config.max_num_seqs == 6
+
+
+def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+ _, block_tables, _ = _runner(
+ monkeypatch,
+ cp_size=2,
+ cp_rank=1,
+ cp_interleave=4,
+ )
+
+ assert block_tables.kwargs["cp_size"] == 2
+ assert block_tables.kwargs["cp_rank"] == 1
+ assert block_tables.kwargs["cp_interleave"] == 4
+
+
+def test_initialize_failure_does_not_publish_partial_cache_state(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ events: list[tuple] = []
+ _install_native_modules(monkeypatch, events)
+ successful_init_cache = model_runner_backend_module.init_kv_cache
+ vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(max_model_len=16),
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
+ parallel_config=SimpleNamespace(
+ decode_context_parallel_size=1,
+ cp_kv_cache_interleave_size=1,
+ ),
+ compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+ cache_config=SimpleNamespace(cache_dtype="auto"),
+ )
+ backend = DiffusionKVModelRunnerBackend(
+ vllm_config=vllm_config,
+ od_config=SimpleNamespace(
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ diffusion_kv_max_rows_per_request=2,
+ ),
+ device="cuda:1",
+ )
+ backend._kv_cache_layer_adapters = {
+ "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"),
+ "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"),
+ }
+ config = SimpleNamespace(
+ num_blocks=32,
+ kv_cache_groups=[
+ SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+ SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+ ],
+ )
+
+ def fail_init_cache(runner_kv_caches, *_args, **_kwargs):
+ runner_kv_caches.append("partial-cache")
+ backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding"
+ raise RuntimeError("cache binding failed")
+
+ monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache)
+ with pytest.raises(RuntimeError, match="cache binding failed"):
+ backend.initialize_kv_cache(config)
+
+ assert backend.kv_cache_config is None
+ assert backend.kv_caches == []
+ assert backend.block_tables is None
+ assert backend.paged_attention_adapter is None
+ assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
+ assert backend.vllm_config.scheduler_config.max_num_seqs == 1
+
+ monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
+ backend.initialize_kv_cache(config)
+ assert backend.kv_cache_config is not None
+ assert backend.kv_caches == ["cache-0", "cache-1"]
+ assert backend.vllm_config.scheduler_config.max_num_seqs == 2
+
+
+def test_initialize_rejects_config_for_different_registered_layers() -> None:
+ backend = DiffusionKVModelRunnerBackend(
+ vllm_config=SimpleNamespace(),
+ od_config=SimpleNamespace(),
+ device="cuda:1",
+ )
+ backend._kv_cache_layer_adapters = {
+ "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None)
+ }
+ config = SimpleNamespace(
+ kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])],
+ )
+
+ with pytest.raises(ValueError, match="layer mismatch"):
+ backend.initialize_kv_cache(config)
+
+
+def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+
+ assert runner.install_diffusion_kv_metadata(_metadata()) is True
+
+ assert block_tables.append_calls == [
+ (0, ([1, 2, 3], [4, 5]), True),
+ (1, ([6, 7], [8]), True),
+ ]
+ assert block_tables.apply_calls == 1
+ assert runner.get_diffusion_kv_row("req-0", 0) == 0
+ assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
+ sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
+ context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
+ assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == (
+ 0,
+ 9,
+ ((1, 2, 3), (4, 5)),
+ )
+ assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
+ 1,
+ 5,
+ ((6, 7), (8,)),
+ )
+
+
+def test_block_table_mutations_invalidate_prepared_attention_batches(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ runner, _, _ = _runner(monkeypatch)
+ assert runner.paged_attention_adapter is not None
+ runner.paged_attention_adapter.invalidate_prepared_batches = Mock()
+
+ runner.install_diffusion_kv_metadata(_metadata())
+ runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with()
+
+ runner.remove_diffusion_kv_requests(["req-0"])
+ assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2
+
+ runner.refresh_block_table_layout()
+ assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3
+
+
+def test_cache_layer_registration_is_frozen_after_physical_initialization(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ runner, _, _ = _runner(monkeypatch)
+
+ with pytest.raises(RuntimeError, match="cannot be changed"):
+ runner.register_kv_cache_layers({})
+
+
+def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ first_sequence = _metadata().sequences[0]
+ metadata = replace(
+ _metadata(),
+ sequences=(
+ first_sequence,
+ replace(
+ first_sequence,
+ sequence_id=1,
+ block_ids=([9, 10, 11], [12, 13]),
+ ),
+ ),
+ )
+
+ assert runner.install_diffusion_kv_metadata(metadata) is True
+
+ assert block_tables.append_calls == [
+ (0, ([1, 2, 3], [4, 5]), True),
+ (1, ([9, 10, 11], [12, 13]), True),
+ (2, ([6, 7], [8]), True),
+ ]
+ assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2
+ assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2
+
+
+def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+ spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8)
+
+ metadata = _metadata(
+ sequence_block_ids=([1, 2], [4, 5]),
+ contexts=(),
+ )
+
+ assert runner.install_diffusion_kv_metadata(metadata) is True
+ assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)]
+
+
+def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ metadata = _metadata(
+ sequence_block_ids=([1, 1, 3], [4, 4]),
+ contexts=(),
+ )
+
+ assert runner.install_diffusion_kv_metadata(metadata) is True
+ assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)]
+
+
+def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ metadata = _metadata()
+ runner.install_diffusion_kv_metadata(metadata)
+
+ assert runner.install_diffusion_kv_metadata(metadata) is False
+ assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False
+ assert block_tables.apply_calls == 1
+
+ with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"):
+ runner.install_diffusion_kv_metadata(metadata)
+ with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"):
+ runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5])))
+ assert block_tables.apply_calls == 1
+
+
+@pytest.mark.parametrize(
+ ("metadata", "message"),
+ [
+ (_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"),
+ (_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"),
+ (_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"),
+ (
+ _metadata(
+ seq_len=17,
+ sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]),
+ contexts=(),
+ ),
+ "exceeds row capacity",
+ ),
+ ],
+)
+def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic(
+ monkeypatch: pytest.MonkeyPatch,
+ metadata: DiffusionKVMetadata,
+ message: str,
+) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+
+ with pytest.raises(ValueError, match=message):
+ runner.install_diffusion_kv_metadata(metadata)
+
+ assert block_tables.append_calls == []
+ assert block_tables.apply_calls == 0
+ assert runner._diffusion_kv_identity_to_row == {}
+ assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ bad_context = DiffusionKVContextMetadata(
+ context_id="text",
+ cache_role="cross_attention",
+ num_tokens=5,
+ block_ids=([6, 32], [8]),
+ )
+
+ with pytest.raises(ValueError, match="outside"):
+ runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,)))
+
+ assert block_tables.append_calls == []
+ assert all(not group_table.staged for group_table in block_tables.block_tables)
+
+
+def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ block_tables.fail_on_append_call = 2
+
+ with pytest.raises(RuntimeError, match="fake append failure"):
+ runner.install_diffusion_kv_metadata(_metadata())
+
+ assert np.count_nonzero(block_tables.num_blocks.np) == 0
+ assert all(not group_table.staged for group_table in block_tables.block_tables)
+ assert runner._diffusion_kv_identity_to_row == {}
+ assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ block_tables.fail_apply = True
+
+ with pytest.raises(RuntimeError, match="fake apply failure"):
+ runner.install_diffusion_kv_metadata(_metadata())
+
+ assert block_tables.apply_calls == 1
+ assert np.count_nonzero(block_tables.num_blocks.np) == 0
+ assert all(not group_table.staged for group_table in block_tables.block_tables)
+ assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables)
+ assert runner._diffusion_kv_identity_to_row == {}
+ assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+ second_context = DiffusionKVContextMetadata(
+ context_id="image",
+ cache_role="cross_attention",
+ num_tokens=1,
+ block_ids=([9], [10]),
+ )
+
+ with pytest.raises(ValueError, match="adapter limit is 2"):
+ runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context)))
+ assert block_tables.append_calls == []
+
+ runner.install_diffusion_kv_metadata(_metadata())
+ append_count = len(block_tables.append_calls)
+ with pytest.raises(ValueError, match="only 0 rows are free"):
+ runner.install_diffusion_kv_metadata(_metadata("req-1"))
+ assert len(block_tables.append_calls) == append_count
+
+
+def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=())
+ duplicate_sequence = _metadata().sequences[0]
+ duplicate = DiffusionKVMetadata(
+ request_id="req-0",
+ allocation_generation=0,
+ sequences=(duplicate_sequence, duplicate_sequence),
+ )
+
+ with pytest.raises(ValueError, match="at least one sequence"):
+ runner.install_diffusion_kv_metadata(empty)
+ with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"):
+ runner.install_diffusion_kv_metadata(duplicate)
+ assert block_tables.append_calls == []
+
+
+@pytest.mark.parametrize(
+ ("metadata", "message"),
+ [
+ (_metadata(generation=-1), "non-negative integer"),
+ (_metadata(request_id=""), "non-empty string"),
+ (_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"),
+ ],
+)
+def test_request_generation_and_block_id_types_are_validated(
+ monkeypatch: pytest.MonkeyPatch,
+ metadata: DiffusionKVMetadata,
+ message: str,
+) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+
+ with pytest.raises(ValueError, match=message):
+ runner.install_diffusion_kv_metadata(metadata)
+
+ assert block_tables.append_calls == []
+ assert runner._diffusion_kv_request_states == {}
+
+
+def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+ runner.install_diffusion_kv_metadata(_metadata())
+
+ assert runner.remove_diffusion_kv_requests(["missing"]) == 0
+ assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2
+ assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)]
+ assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]]
+ assert runner.remove_diffusion_kv_requests(["req-0"]) == 0
+ with pytest.raises(KeyError, match="No native diffusion KV row"):
+ runner.get_diffusion_kv_row("req-0", 0)
+
+ runner.install_diffusion_kv_metadata(_metadata("req-1"))
+ assert runner.get_diffusion_kv_row("req-1", 0) == 0
+ assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1
+
+
+def test_cached_step_does_not_reinstall_metadata() -> None:
+ runner = object.__new__(DiffusionModelRunner)
+ runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+ runner.pipeline = object()
+ runner.install_diffusion_kv_metadata = Mock()
+ runner._supports_step_mode = Mock(return_value=False)
+ scheduler_output = DiffusionSchedulerOutput(
+ step_id=1,
+ scheduled_new_reqs=[],
+ scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]),
+ finished_req_ids=set(),
+ num_running_reqs=1,
+ num_waiting_reqs=0,
+ )
+
+ with pytest.raises(ValueError, match="does not support step execution"):
+ runner.execute_stepwise(scheduler_output)
+ runner.install_diffusion_kv_metadata.assert_not_called()
+
+
+def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner, block_tables, _ = _runner(monkeypatch)
+ runner.refresh_block_table_layout()
+ assert block_tables.layout_refreshes == 1
+
+ dense_backend = DiffusionKVModelRunnerBackend(
+ vllm_config=SimpleNamespace(),
+ od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY),
+ device="cuda:1",
+ )
+ dense_backend.block_tables = Mock()
+ assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0
+ dense_backend.refresh_block_table_layout()
+ dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called()
+ with pytest.raises(ValueError, match="Dense diffusion execution"):
+ dense_backend.install_diffusion_kv_metadata(_metadata())
diff --git a/tests/diffusion/diffusion_kv/test_config.py b/tests/diffusion/diffusion_kv/test_config.py
index f03120784..61b1b3041 100644
--- a/tests/diffusion/diffusion_kv/test_config.py
+++ b/tests/diffusion/diffusion_kv/test_config.py
@@ -1,8 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
+from vllm_omni.config.omni_config import _DiffusionConfigProjection
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -35,10 +36,12 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None:
OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
-def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
+def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None:
with pytest.raises(ValueError, match="does not support imported AR KV"):
OmniDiffusionConfig.from_kwargs(
diffusion_kv_mode="paged_scheduler",
+ diffusion_kv_max_rows_per_request=1,
+ max_num_batched_tokens=1,
omni_kv_config={"need_recv_cache": True},
)
@@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
config = OmniDiffusionConfig.from_kwargs(
model_class_name="FutureDiffusionModel",
diffusion_kv_mode="paged_scheduler",
+ diffusion_kv_max_rows_per_request=1,
+ max_num_batched_tokens=1,
)
assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None:
+ config = config_cls.from_kwargs(
+ diffusion_kv_mode="paged_scheduler",
+ diffusion_kv_max_rows_per_request=1,
+ )
+
+ assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+def test_paged_scheduler_requires_a_worker_row_limit() -> None:
+ with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+ OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler")
+ with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+ _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler")
+
+
+@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5])
+def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None:
+ with pytest.raises((TypeError, ValueError)):
+ OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
+ with pytest.raises((TypeError, ValueError)):
+ _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py
new file mode 100644
index 000000000..a82dc01ff
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_initialization.py
@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+import vllm_omni.diffusion.diffusion_kv.initialization as initialization
+import vllm_omni.diffusion.vllm_config as diffusion_vllm_config
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Executor:
+ def __init__(self) -> None:
+ self.specs = [{"layer0": object()}, {"layer0": object()}]
+ self.memory = [1024, 1024]
+ self.configs = None
+ self.resolved_max_model_len = None
+ self.profile_requests = None
+
+ def get_kv_cache_specs(self):
+ return self.specs
+
+ def determine_available_kv_memory(self, profile_requests):
+ self.profile_requests = profile_requests
+ return self.memory
+
+ def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None:
+ self.configs = configs
+ self.resolved_max_model_len = resolved_max_model_len
+
+
+def _od_config(**overrides):
+ values = dict(
+ model="test-diffusion",
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ dtype=torch.bfloat16,
+ max_model_len=64,
+ quantization_config=None,
+ tf_model_config=None,
+ enforce_eager=True,
+ is_moe=False,
+ additional_config={},
+ profiler_config=None,
+ parallel_config=SimpleNamespace(
+ tensor_parallel_size=1,
+ pipeline_parallel_size=1,
+ data_parallel_size=1,
+ enable_expert_parallel=False,
+ ),
+ gpu_memory_utilization=0.9,
+ kv_cache_memory_bytes=None,
+ max_num_seqs=1,
+ max_num_batched_tokens=64,
+ num_gpus=1,
+ )
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
+def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
+ executor = _Executor()
+ od_config = _od_config(num_gpus=2, max_model_len=-1)
+ vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(max_model_len=256),
+ max_in_flight_tokens=128,
+ )
+ worker_configs = [object(), object()]
+ scheduler_kv_cache_config = object()
+ profile_requests = [object()]
+ monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: vllm_config)
+ build = Mock(return_value=(worker_configs, scheduler_kv_cache_config))
+ monkeypatch.setattr(
+ initialization,
+ "build_native_kv_cache_configs",
+ build,
+ )
+ monkeypatch.setattr(initialization, "resolve_kv_cache_block_sizes", lambda *args: (16, 16))
+
+ result = initialization.initialize_diffusion_kv_control_plane(
+ executor,
+ od_config,
+ profile_requests=profile_requests,
+ device=torch.device("cpu"),
+ )
+
+ build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
+ assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
+ assert executor.configs == worker_configs
+ assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len
+ assert executor.profile_requests is profile_requests
+
+
+def test_dense_mode_skips_worker_kv_initialization(monkeypatch) -> None:
+ executor = _Executor()
+ od_config = _od_config(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+ get_device = Mock(side_effect=AssertionError("dense mode must not resolve a KV device"))
+ monkeypatch.setattr(initialization.current_omni_platform, "get_torch_device", get_device)
+
+ result = initialization.initialize_diffusion_kv_control_plane(executor, od_config)
+
+ assert result is None
+ assert executor.configs is None
+ get_device.assert_not_called()
+
+
+def test_control_plane_rejects_rank_count_mismatch(monkeypatch) -> None:
+ executor = _Executor()
+ executor.memory = [1024]
+ od_config = _od_config(num_gpus=2)
+ monkeypatch.setattr(initialization, "create_diffusion_vllm_config", lambda *args, **kwargs: object())
+
+ with pytest.raises(ValueError, match="rank count mismatch"):
+ initialization.initialize_diffusion_kv_control_plane(
+ executor,
+ od_config,
+ profile_requests=[object()],
+ device=torch.device("cpu"),
+ )
+
+
+def test_paged_control_plane_requires_prepared_profile_request() -> None:
+ with pytest.raises(ValueError, match="prepared memory profile requests"):
+ initialization.initialize_diffusion_kv_control_plane(
+ _Executor(),
+ _od_config(),
+ device=torch.device("cpu"),
+ )
+
+
+def test_native_config_builder_produces_matching_worker_and_scheduler_pools() -> None:
+ od_config = _od_config()
+ vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+ spec = FullAttentionSpec(
+ block_size=vllm_config.cache_config.block_size,
+ num_kv_heads=2,
+ head_size=8,
+ dtype=torch.bfloat16,
+ non_causal=True,
+ )
+ num_blocks = 32
+
+ worker_configs, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+ vllm_config,
+ [{"layer0": spec}],
+ [spec.page_size_bytes * num_blocks],
+ )
+
+ assert len(worker_configs) == 1
+ assert worker_configs[0].num_blocks == num_blocks
+ assert scheduler_kv_cache_config.num_blocks == num_blocks
+ assert worker_configs[0].kv_cache_groups == scheduler_kv_cache_config.kv_cache_groups
+
+
+def test_native_control_plane_config_initializes_scheduler() -> None:
+ od_config = _od_config()
+ vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+ spec = FullAttentionSpec(
+ block_size=vllm_config.cache_config.block_size,
+ num_kv_heads=2,
+ head_size=8,
+ dtype=torch.bfloat16,
+ non_causal=True,
+ )
+ _, scheduler_kv_cache_config = initialization.build_native_kv_cache_configs(
+ vllm_config,
+ [{"layer0": spec}],
+ [spec.page_size_bytes * 32],
+ )
+
+ scheduler = RequestScheduler()
+ scheduler.initialize(
+ od_config,
+ kv_cache_config=scheduler_kv_cache_config,
+ scheduler_block_size=spec.block_size,
+ hash_block_size=spec.block_size,
+ kv_vllm_config=vllm_config,
+ )
+
+ assert scheduler._diffusion_kv_manager is not None
+ scheduler.close()
+
+
+def test_dense_config_keeps_new_paged_sizing_inputs_inactive() -> None:
+ od_config = _od_config(
+ diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+ parallel_config=SimpleNamespace(
+ tensor_parallel_size=1,
+ pipeline_parallel_size=4,
+ data_parallel_size=1,
+ enable_expert_parallel=False,
+ ),
+ gpu_memory_utilization=0.5,
+ kv_cache_memory_bytes=4096,
+ max_num_batched_tokens=64,
+ )
+ vllm_config = diffusion_vllm_config.create_base_diffusion_vllm_config(torch.device("cpu"), od_config)
+ original_pp_size = vllm_config.parallel_config.pipeline_parallel_size
+ original_memory_utilization = vllm_config.cache_config.gpu_memory_utilization
+ original_memory_bytes = vllm_config.cache_config.kv_cache_memory_bytes
+ original_token_budget = vllm_config.scheduler_config.max_num_batched_tokens
+
+ diffusion_vllm_config.configure_diffusion_vllm_config(vllm_config, od_config)
+
+ assert vllm_config.parallel_config.pipeline_parallel_size == original_pp_size
+ assert vllm_config.cache_config.gpu_memory_utilization == original_memory_utilization
+ assert vllm_config.cache_config.kv_cache_memory_bytes == original_memory_bytes
+ assert vllm_config.scheduler_config.max_num_batched_tokens == original_token_budget
+
+
+def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -> None:
+ vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+ torch.device("cpu"),
+ _od_config(gpu_memory_utilization=0.42),
+ )
+
+ assert vllm_config.cache_config.gpu_memory_utilization == 0.42
+
+
+def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
+ quantization_config = build_quant_config(
+ {
+ "quant_method": "modelopt",
+ "quant_algo": "FP8",
+ "ignore": [],
+ }
+ )
+ model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+ _od_config(
+ diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+ quantization_config=quantization_config,
+ )
+ )
+
+ assert model_config.dtype is torch.bfloat16
+ assert model_config.quantization == "modelopt"
+ assert model_config.quantization_config is quantization_config
+ assert model_config.is_quantized
+ assert model_config.is_moe is False
+ assert model_config.is_nvfp4_quantized() is False
+ assert model_config.original_max_model_len == model_config.max_model_len
+
+ unquantized_model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+ _od_config(
+ diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
+ quantization_config=None,
+ )
+ )
+ assert not unquantized_model_config.is_quantized
+
+
+def test_minus_one_max_model_len_uses_model_limit() -> None:
+ od_config = SimpleNamespace(
+ max_model_len=-1,
+ tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ )
+
+ assert diffusion_vllm_config.resolve_diffusion_max_model_len(od_config) == 4096
+
+ model_config = diffusion_vllm_config._make_diffusion_vllm_model_config(
+ _od_config(
+ max_model_len=-1,
+ tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+ )
+ )
+ assert model_config.max_model_len == 4096
+ assert model_config.original_max_model_len == -1
+
+
+def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None:
+ vllm_config = diffusion_vllm_config.create_diffusion_vllm_config(
+ torch.device("cpu"),
+ _od_config(
+ max_model_len=-1,
+ tf_model_config=SimpleNamespace(max_position_embeddings=4096),
+ ),
+ )
+ spec = FullAttentionSpec(
+ block_size=vllm_config.cache_config.block_size,
+ num_kv_heads=2,
+ head_size=8,
+ dtype=torch.bfloat16,
+ non_causal=True,
+ )
+ model_limit = vllm_config.model_config.max_model_len
+
+ initialization.build_native_kv_cache_configs(
+ vllm_config,
+ [{"layer0": spec}],
+ [spec.page_size_bytes * 16],
+ )
+
+ assert vllm_config.model_config.original_max_model_len == -1
+ # Sixteen physical pages provide ``16 * block_size`` tokens. The native
+ # backend chooses the block geometry, so keep this assertion backend
+ # independent instead of assuming the generic 16-token block size.
+ expected_max_model_len = min(model_limit, spec.block_size * 16)
+ assert vllm_config.model_config.max_model_len == expected_max_model_len
diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py
new file mode 100644
index 000000000..5152c2a67
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_manager.py
@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+BLOCK_SIZE = 4
+
+
+def _config(num_blocks: int) -> KVCacheConfig:
+ native_kv_managers.register_all_kvcache_specs(None)
+ spec = FullAttentionSpec(
+ block_size=BLOCK_SIZE,
+ num_kv_heads=2,
+ head_size=8,
+ dtype=torch.bfloat16,
+ )
+ return KVCacheConfig(
+ num_blocks=num_blocks,
+ kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+ kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+ )
+
+
+def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+ return DiffusionKVRequest(
+ f"{public_id}/diffusion-kv/{sequence_id}",
+ sequence_id=sequence_id,
+ prefix_len=4,
+ target_len=4,
+ seq_len=seq_len,
+ kv_contexts=kv_contexts,
+ )
+
+
+def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+ return DiffusionKVCacheManager(
+ _config(num_blocks),
+ max_model_len=max_model_len,
+ scheduler_block_size=BLOCK_SIZE,
+ hash_block_size=BLOCK_SIZE,
+ )
+
+
+def test_reserve_and_free_multi_cfg_request() -> None:
+ manager = _manager(8)
+ free_before = manager.native_manager.block_pool.get_num_free_blocks()
+ requests = (_request("public", 0), _request("public", 1))
+
+ metadata = manager.reserve_request("public", requests)
+
+ assert metadata is not None
+ assert metadata.request_id == "public"
+ assert metadata.allocation_generation == 1
+ assert [sequence.sequence_id for sequence in metadata.sequences] == [0, 1]
+ assert [sequence.prefix_len for sequence in metadata.sequences] == [4, 4]
+ assert [sequence.target_len for sequence in metadata.sequences] == [4, 4]
+ assert [sequence.seq_len for sequence in metadata.sequences] == [8, 8]
+ assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+ assert all(sequence.context_ids == () for sequence in metadata.sequences)
+ assert metadata.contexts == ()
+ assert metadata == manager.get_metadata("public")
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before - 4
+
+ manager.free_request("public")
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_rolls_back_and_fails_fast() -> None:
+ manager = _manager(3)
+ free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+ with pytest.raises(DiffusionKVAdmissionError, match="cannot fit even when the block pool is empty"):
+ manager.reserve_request(
+ "public",
+ (_request("public", 0), _request("public", 1)),
+ )
+
+ assert manager.has_request("public") is False
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_impossible_cfg_allocation_fails_fast_while_pool_is_busy() -> None:
+ manager = _manager(5)
+ assert manager.reserve_request("running", (_request("running", 0),)) is not None
+ free_before = manager.native_manager.block_pool.get_num_free_blocks()
+
+ with pytest.raises(DiffusionKVAdmissionError, match="required_blocks=6, available_blocks=4"):
+ manager.reserve_request(
+ "impossible",
+ tuple(_request("impossible", sequence_id) for sequence_id in range(3)),
+ )
+
+ assert manager.has_request("impossible") is False
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_temporary_capacity_pressure_returns_none() -> None:
+ manager = _manager(3)
+ assert manager.reserve_request("running", (_request("running", 0),)) is not None
+
+ allocation = manager.reserve_request("waiting", (_request("waiting", 0),))
+
+ assert allocation is None
+ assert manager.has_request("waiting") is False
+
+
+def test_partial_cfg_rollback_can_retry_after_capacity_is_released() -> None:
+ manager = _manager(5)
+ assert manager.reserve_request("running", (_request("running", 0),)) is not None
+ waiting = (_request("waiting", 0), _request("waiting", 1))
+
+ assert manager.reserve_request("waiting", waiting) is None
+ assert manager.has_request("waiting") is False
+
+ manager.free_request("running")
+ metadata = manager.reserve_request("waiting", waiting)
+
+ assert metadata is not None
+ assert [len(sequence.block_ids[0]) for sequence in metadata.sequences] == [2, 2]
+
+
+def test_reallocated_request_gets_a_new_allocation_generation() -> None:
+ manager = _manager(3)
+ request = (_request("public", 0),)
+
+ first = manager.reserve_request("public", request)
+ assert first is not None
+ manager.free_request("public")
+ second = manager.reserve_request("public", request)
+
+ assert second is not None
+ assert second.allocation_generation == first.allocation_generation + 1
+
+
+def test_cfg_allocation_exception_rolls_back(monkeypatch) -> None:
+ manager = _manager(8)
+ free_before = manager.native_manager.block_pool.get_num_free_blocks()
+ native_allocate = manager.native_manager.allocate_slots
+ num_calls = 0
+
+ def fail_second_allocation(*args, **kwargs):
+ nonlocal num_calls
+ num_calls += 1
+ if num_calls == 2:
+ raise RuntimeError("injected allocation failure")
+ return native_allocate(*args, **kwargs)
+
+ monkeypatch.setattr(manager.native_manager, "allocate_slots", fail_second_allocation)
+
+ with pytest.raises(RuntimeError, match="injected allocation failure"):
+ manager.reserve_request(
+ "public",
+ (_request("public", 0), _request("public", 1)),
+ )
+
+ assert manager.has_request("public") is False
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_rejects_independent_context_until_role_routing_is_implemented() -> None:
+ manager = _manager(8)
+ context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+
+ with pytest.raises(DiffusionKVAdmissionError, match="DiffusionKVContext"):
+ manager.reserve_request(
+ "public",
+ (_request("public", 0, kv_contexts=(context,)),),
+ )
+
+
+def test_rejects_sequence_longer_than_admission_bound() -> None:
+ manager = _manager(8, max_model_len=8)
+
+ with pytest.raises(DiffusionKVAdmissionError, match="exceeds max_model_len"):
+ manager.reserve_request("public", (_request("public", 0, seq_len=9),))
+
+
+@pytest.mark.parametrize(
+ "requests",
+ [
+ (_request("public", 1),),
+ (_request("public", 1), _request("public", 0)),
+ (_request("public", 0), _request("public", 2)),
+ ],
+)
+def test_rejects_ambiguous_sequence_order(requests) -> None:
+ manager = _manager(8)
+
+ with pytest.raises(ValueError, match="ordered by contiguous sequence_id"):
+ manager.reserve_request("public", requests)
+
+
+def test_close_releases_every_public_request() -> None:
+ manager = _manager(8)
+ free_before = manager.native_manager.block_pool.get_num_free_blocks()
+ assert manager.reserve_request("first", (_request("first", 0),)) is not None
+ assert manager.reserve_request("second", (_request("second", 0),)) is not None
+
+ manager.close()
+
+ assert manager.native_manager.block_pool.get_num_free_blocks() == free_before
+ assert manager.has_request("first") is False
+ assert manager.has_request("second") is False
diff --git a/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py
new file mode 100644
index 000000000..4848fd49a
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py
@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from importlib import import_module, util
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+
+def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None:
+ if not torch.cuda.is_available():
+ pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test")
+ module_name = "vllm.v1.worker.gpu.block_table"
+ try:
+ module_spec = util.find_spec(module_name)
+ except ModuleNotFoundError as exc:
+ if exc.name is None or not module_name.startswith(f"{exc.name}."):
+ raise
+ pytest.skip("installed vLLM does not provide GPU BlockTables")
+ if module_spec is None:
+ pytest.skip("installed vLLM does not provide GPU BlockTables")
+ block_table_module = import_module(module_name)
+ block_tables_cls = getattr(block_table_module, "BlockTables", None)
+ if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"):
+ pytest.skip("installed vLLM does not provide the required BlockTables API")
+
+ device = torch.device("cuda", torch.accelerator.current_device_index())
+ block_tables = block_tables_cls(
+ block_sizes=[4],
+ max_num_reqs=1,
+ max_num_batched_tokens=16,
+ max_num_blocks_per_group=[4],
+ device=device,
+ kernel_block_sizes=[4],
+ )
+ block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True)
+ block_tables.apply_staged_writes()
+ idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
+ query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device)
+ positions = torch.arange(9, dtype=torch.int64, device=device)
+
+ slot_mappings = block_tables.compute_slot_mappings(
+ idx_mapping,
+ query_start_loc,
+ positions,
+ num_tokens_padded=9,
+ )
+ torch.accelerator.synchronize(device)
+
+ assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]]
diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py
new file mode 100644
index 000000000..1f242420d
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py
@@ -0,0 +1,1406 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import math
+from contextlib import nullcontext
+from dataclasses import dataclass
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+ DiffusionPagedAttentionAdapter,
+ DiffusionPagedAttentionMetadata,
+ DiffusionPagedAttentionRow,
+ DiffusionPagedAttentionRowBinding,
+ DiffusionPagedAttentionRuntime,
+)
+from vllm_omni.diffusion.forward_context import (
+ get_forward_context,
+ override_paged_kv_adapter,
+ set_forward_context,
+ set_forward_context_denoise_step_idx,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeBlockTables:
+ def __init__(self) -> None:
+ self.max_num_reqs = 4
+ self.max_num_batched_tokens = 16
+ self.cp_size = 1
+ self.cp_rank = 0
+ self.cp_interleave = 1
+ self.block_sizes = [4]
+ self.blocks_per_kv_block = [1]
+ self.num_blocks = SimpleNamespace(np=np.full((1, self.max_num_reqs), 4, dtype=np.int32))
+ self.gather_calls: list[tuple[torch.Tensor, int]] = []
+ self.slot_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = []
+
+ def gather_block_tables(
+ self,
+ idx_mapping: torch.Tensor,
+ num_reqs_padded: int,
+ ) -> tuple[torch.Tensor, ...]:
+ self.gather_calls.append((idx_mapping.clone(), num_reqs_padded))
+ return (torch.tensor([[3, 4], [7, 8]], dtype=torch.int32),)
+
+ def compute_slot_mappings(
+ self,
+ idx_mapping: torch.Tensor,
+ query_start_loc: torch.Tensor,
+ positions: torch.Tensor,
+ num_tokens_padded: int,
+ ) -> torch.Tensor:
+ self.slot_calls.append(
+ (
+ idx_mapping.clone(),
+ query_start_loc.clone(),
+ positions.clone(),
+ num_tokens_padded,
+ )
+ )
+ return positions[:num_tokens_padded].to(torch.int64).unsqueeze(0)
+
+
+class _FakeSpec:
+ non_causal: bool
+ block_size = 4
+
+ def __init__(self, *, non_causal: bool) -> None:
+ self.non_causal = non_causal
+
+ def max_num_blocks_per_req(self, _vllm_config, max_len: int) -> int:
+ return math.ceil(max_len / self.block_size)
+
+
+class _FakeAttentionGroup:
+ def __init__(self, reorder_batch_threshold: int | None) -> None:
+ self.builder = SimpleNamespace(reorder_batch_threshold=reorder_batch_threshold)
+
+ def get_metadata_builder(self, _index: int):
+ return self.builder
+
+
+@dataclass(frozen=True, slots=True)
+class _FakeNativeMetadata:
+ build_id: int
+ causal: bool
+ seq_lens: torch.Tensor
+ query_start_loc_cpu: torch.Tensor
+ positions: torch.Tensor
+ slot_mappings: torch.Tensor
+
+
+class _FakeLayer:
+ def __init__(self, *, non_causal: bool, num_heads: int = 2, num_kv_heads: int = 2) -> None:
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads
+ self.head_size = 4
+ self.head_size_v = 4
+ self.spec = _FakeSpec(non_causal=non_causal)
+ self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+ self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = []
+ self.native_events: list[str] = []
+ self.layer_name = "layer-0"
+ self.kv_cache = object()
+ self.attn_backend = SimpleNamespace(forward_includes_kv_cache_update=False)
+ self.impl = _FakeNativeImpl(self)
+
+
+class _FakeNativeImpl:
+ def __init__(self, layer: _FakeLayer) -> None:
+ self.layer = layer
+
+ def do_kv_cache_update(
+ self,
+ _layer,
+ key: torch.Tensor | None,
+ value: torch.Tensor | None,
+ _kv_cache,
+ slot_mapping: torch.Tensor,
+ ) -> None:
+ self.layer.native_events.append("update")
+ self.layer.updates.append((key, value, slot_mapping))
+
+ def forward(
+ self,
+ _layer,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ _kv_cache,
+ metadata: object,
+ output: torch.Tensor,
+ ) -> torch.Tensor:
+ self.layer.native_events.append("forward")
+ self.layer.calls.append((query, key, value, metadata))
+ return output.copy_(query.reshape_as(output))
+
+
+def _make_adapter(
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ non_causal: bool = True,
+ reorder_batch_threshold: int | None = None,
+ num_heads: int = 2,
+ num_kv_heads: int = 2,
+ capacity: int = 16,
+) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
+ events: list[tuple] = []
+ monkeypatch.setattr(
+ adapter_module.current_omni_platform,
+ "requires_diffusion_paged_kv_prewrite",
+ lambda: False,
+ )
+ block_tables = _FakeBlockTables()
+ block_tables.max_num_batched_tokens = capacity
+ block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
+ layer = _FakeLayer(non_causal=non_causal, num_heads=num_heads, num_kv_heads=num_kv_heads)
+ kv_cache_config = SimpleNamespace(
+ kv_cache_groups=[SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=layer.spec)],
+ )
+ row_map = {
+ ("req-0", 0, None): DiffusionPagedAttentionRowBinding(2, capacity),
+ ("req-1", 1, None): DiffusionPagedAttentionRowBinding(3, capacity),
+ ("req-0", None, "text"): DiffusionPagedAttentionRowBinding(1, capacity),
+ }
+
+ monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+ def build_metadata(**kwargs):
+ metadata = _FakeNativeMetadata(
+ build_id=len(events),
+ causal=kwargs["causal"],
+ seq_lens=kwargs["seq_lens"].clone(),
+ query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
+ positions=kwargs["positions"].clone(),
+ slot_mappings=kwargs["slot_mappings"].clone(),
+ )
+ events.append(("build", kwargs, metadata))
+ return {"layer-0": metadata}
+
+ monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
+ monkeypatch.setattr(
+ adapter_module,
+ "build_slot_mappings_by_layer",
+ lambda slot_mappings, _config: {"layer-0": slot_mappings[0]},
+ )
+
+ config = SimpleNamespace(
+ name="vllm-config",
+ model_config=SimpleNamespace(dtype=torch.float32),
+ )
+ adapter = DiffusionPagedAttentionAdapter(
+ vllm_config=config,
+ device=torch.device("cpu"),
+ kv_cache_config=kv_cache_config,
+ block_tables=block_tables,
+ attn_groups=[[_FakeAttentionGroup(reorder_batch_threshold)]],
+ layers={"layer-0": layer},
+ resolve_row=lambda request_id, sequence_id, context_id: row_map[(request_id, sequence_id, context_id)],
+ )
+ return adapter, block_tables, layer, events
+
+
+def _run_omni_paged_backend(context):
+ backend = FlashAttentionImpl(
+ num_heads=context.layer.num_heads,
+ head_size=context.layer.head_size,
+ softmax_scale=0.5,
+ num_kv_heads=context.layer.num_kv_heads,
+ )
+ return backend.forward_paged(context)
+
+
+def test_prepare_batch_reuses_native_block_table_and_metadata_builders(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, block_tables, _, events = _make_adapter(monkeypatch)
+ rows = (
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=3,
+ seq_len=8,
+ kv_start_pos=4,
+ ),
+ DiffusionPagedAttentionRow(
+ request_id="req-1",
+ sequence_id=1,
+ query_len=2,
+ seq_len=2,
+ ),
+ )
+
+ batch = adapter.prepare_batch(rows)
+
+ assert batch.num_tokens == 5
+ assert batch.row_indices.tolist() == [2, 3]
+ assert batch.query_start_loc.tolist() == [0, 3, 5]
+ assert batch.seq_lens.tolist() == [8, 2]
+ assert batch.positions.tolist() == [4, 5, 6, 0, 1]
+ assert block_tables.gather_calls[0][1] == 2
+ assert block_tables.gather_calls[0][0].tolist() == [2, 3]
+ assert block_tables.slot_calls[0][3] == 5
+ build_kwargs = events[0][1]
+ assert build_kwargs["causal"] is False
+ assert build_kwargs["max_query_len"] == 3
+ assert build_kwargs["max_seq_len"] == 8
+ assert build_kwargs["positions"] is batch.positions
+ assert batch.attn_metadata["layer-0"] is events[0][2]
+ assert events[0][2].build_id == 0
+ assert batch.slot_mappings_by_layer["layer-0"].tolist() == [4, 5, 6, 0, 1]
+
+
+def test_omni_paged_backend_consumes_context_and_restores_diffusion_shape(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, layer, events = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=5,
+ seq_len=5,
+ )
+ ]
+ )
+ query = torch.randn(1, 5, 2, 4)
+ key = torch.randn_like(query)
+ value = torch.randn_like(query)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, key, value)
+ output = _run_omni_paged_backend(context)
+
+ assert output.shape == query.shape
+ assert torch.equal(output, query)
+ assert layer.calls[0][0].shape == (5, 2, 4)
+ assert layer.native_events == ["update", "forward"]
+ assert layer.calls[0][3] is events[0][2]
+
+
+def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch, num_heads=32, num_kv_heads=8)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
+ )
+ query = torch.randn(1, 3, 32, 4)
+ key = torch.randn(1, 3, 8, 4)
+ value = torch.randn_like(key)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+ assert layer.num_heads == 32
+ assert layer.num_kv_heads == 8
+ assert context.query.shape == (3, 32, 4)
+ assert context.key_write.shape == (3, 8, 4)
+ assert context.value_write.shape == (3, 8, 4)
+
+
+def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+ )
+ packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
+ query, key, value = packed_qkv.split(2 * 4, dim=-1)
+ query = query.reshape(1, 5, 2, 4)
+ key = key.reshape(1, 5, 2, 4)
+ value = value.reshape(1, 5, 2, 4)
+ assert not key.is_contiguous()
+ assert not value.is_contiguous()
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+ assert context.query.is_contiguous()
+ assert not context.key_write.is_contiguous()
+ assert not context.value_write.is_contiguous()
+ assert context.key_write.data_ptr() == key.data_ptr()
+ assert context.value_write.data_ptr() == value.data_ptr()
+
+
+def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+ )
+ query = torch.randn(1, 5, 2, 4)
+ key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+ value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
+ assert key.shape == (1, 5, 2, 4)
+ assert key.stride(-1) != 1
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+ assert context.key_write.is_contiguous()
+ assert context.value_write.is_contiguous()
+
+
+def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ layer.attn_backend.forward_includes_kv_cache_update = True
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
+ )
+ qkv = torch.randn(1, 3, 2, 4)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, qkv)
+ assert layer.native_events == ["forward"]
+ assert layer.updates == []
+
+
+def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, events = _make_adapter(monkeypatch)
+ layer.attn_backend.forward_includes_kv_cache_update = True
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, qkv)
+ assert layer.native_events == ["forward", "forward", "forward"]
+ assert layer.updates == []
+ assert [call[0].shape[0] for call in layer.calls] == [1, 2, 1]
+ assert [call[1].shape[0] for call in layer.calls] == [1, 2, 1]
+ assert [call[2].shape[0] for call in layer.calls] == [1, 2, 1]
+ assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
+ segment_metadata = [call[3] for call in layer.calls]
+ assert [metadata.causal for metadata in segment_metadata] == [True, False, True]
+ assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[1], [3], [4]]
+ assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
+ assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
+ assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
+ for event in events[1:]:
+ build_kwargs = event[1]
+ assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
+ assert (
+ build_kwargs["slot_mappings"].untyped_storage().data_ptr()
+ == batch.slot_mappings.untyped_storage().data_ptr()
+ )
+
+
+def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ layer.attn_backend.forward_includes_kv_cache_update = True
+ monkeypatch.setattr(
+ adapter_module.current_omni_platform,
+ "requires_diffusion_paged_kv_prewrite",
+ lambda: True,
+ )
+ adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
+ row_index=2,
+ max_seq_len=16,
+ block_ids=((3, 7, 9, 11),),
+ )
+ row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
+ batch = adapter.prepare_batch([row])
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, qkv)
+ assert layer.native_events == ["update", "forward", "forward", "forward"]
+ assert len(layer.updates) == 1
+ assert all(call[1] is None and call[2] is None for call in layer.calls)
+
+
+def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, events = _make_adapter(monkeypatch)
+ # Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
+ # output-buffer contract and is covered by the GPU adapter tests.
+ monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=6,
+ seq_len=9,
+ kv_start_pos=3,
+ ),
+ DiffusionPagedAttentionRow(
+ request_id="req-1",
+ sequence_id=1,
+ query_len=6,
+ seq_len=9,
+ kv_start_pos=3,
+ ),
+ ]
+ )
+ query = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4)
+ key = torch.arange(2 * 6 * 2 * 4, dtype=torch.float32).reshape(2, 6, 2, 4) + 100
+ value = key + 100
+ attn_mask = torch.ones(9, 9, dtype=torch.bool).tril().repeat(2, 1, 1)
+ attn_mask[:, 5:8, 5:8] = True
+ metadata = SimpleNamespace(
+ attn_mask=attn_mask.unsqueeze(1),
+ full_attn_spans=[[(5, 8)], [(5, 8)]],
+ query_ranges=None,
+ extra={},
+ )
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context(
+ "layer-0",
+ query,
+ key,
+ value,
+ omni_attn_metadata=metadata,
+ )
+ # Every model layer sees the same spans, so the native segment metadata
+ # is built once per active batch and then reused.
+ adapter.prepare_layer_context(
+ "layer-0",
+ query,
+ key,
+ value,
+ omni_attn_metadata=metadata,
+ )
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, query)
+ assert layer.native_events == ["update", "forward", "forward", "forward"]
+ assert len(layer.updates) == 1
+ assert layer.updates[0][0].shape[0] == 12
+ assert torch.equal(layer.updates[0][0], key.reshape(12, 2, 4))
+ assert torch.equal(layer.updates[0][1], value.reshape(12, 2, 4))
+ assert [call[0].shape[0] for call in layer.calls] == [4, 6, 2]
+ assert [call[1].shape[0] for call in layer.calls] == [4, 6, 2]
+ assert [call[2].shape[0] for call in layer.calls] == [4, 6, 2]
+ assert all(call[3] is event[2] for call, event in zip(layer.calls, events[1:], strict=True))
+ assert [metadata.build_id for metadata in context.piecewise_native_metadata] == [1, 2, 3]
+
+ # The first metadata build is the normal whole-query path. Piecewise calls
+ # then use causal [3, 5), full [5, 8), and causal [8, 9) segments.
+ assert len(events) == 4
+ segment_builds = [event[1] for event in events[1:]]
+ assert [build["causal"] for build in segment_builds] == [True, False, True]
+ assert [build["seq_lens"].tolist() for build in segment_builds] == [
+ [5, 5],
+ [8, 8],
+ [9, 9],
+ ]
+ assert [build["query_start_loc_cpu"].tolist() for build in segment_builds] == [
+ [0, 2, 4],
+ [0, 3, 6],
+ [0, 1, 2],
+ ]
+ segment_metadata = [call[3] for call in layer.calls]
+ assert [metadata.causal for metadata in segment_metadata] == [True, False, True]
+ assert [metadata.seq_lens.tolist() for metadata in segment_metadata] == [[5, 5], [8, 8], [9, 9]]
+ assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [
+ [0, 2, 4],
+ [0, 3, 6],
+ [0, 1, 2],
+ ]
+ assert [metadata.positions.tolist() for metadata in segment_metadata] == [
+ [3, 4, 3, 4],
+ [5, 6, 7, 5, 6, 7],
+ [8, 8],
+ ]
+ assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [
+ [[3, 4, 3, 4]],
+ [[5, 6, 7, 5, 6, 7]],
+ [[8, 8]],
+ ]
+
+
+def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
+ build_id = 0
+
+ def build_metadata(**kwargs):
+ nonlocal build_id
+ shared_scheduler_metadata.fill_(build_id)
+ metadata = SimpleNamespace(
+ build_id=build_id,
+ causal=kwargs["causal"],
+ seq_lens=kwargs["seq_lens"].clone(),
+ query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
+ positions=kwargs["positions"].clone(),
+ slot_mappings=kwargs["slot_mappings"].clone(),
+ scheduler_metadata=shared_scheduler_metadata,
+ prefix_scheduler_metadata=None,
+ )
+ build_id += 1
+ return {"layer-0": metadata}
+
+ monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
+ assert [item.item() for item in segment_schedules] == [1, 2, 3]
+ assert len({item.data_ptr() for item in segment_schedules}) == 3
+
+
+def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, events = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(
+ attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool).tril(),
+ full_attn_spans=[[]],
+ query_ranges=None,
+ extra={},
+ )
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, qkv)
+ assert layer.native_events == ["update", "forward"]
+ assert len(events) == 2
+ assert events[1][1]["causal"] is True
+ assert events[1][1]["seq_lens"].tolist() == [4]
+
+
+def test_omni_paged_backend_runs_unequal_hunyuan_prefixes_without_padding(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, layer, events = _make_adapter(monkeypatch, capacity=64)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=17,
+ seq_len=50,
+ kv_start_pos=33,
+ ),
+ DiffusionPagedAttentionRow(
+ request_id="req-1",
+ sequence_id=1,
+ query_len=17,
+ seq_len=54,
+ kv_start_pos=37,
+ ),
+ ]
+ )
+ query = torch.arange(2 * 17 * 2 * 4, dtype=torch.float32).reshape(2, 17, 2, 4)
+ key = query + 100
+ value = query + 200
+ metadata = SimpleNamespace(
+ attn_mask=torch.ones(2, 1, 17, 54, dtype=torch.bool),
+ full_attn_spans=[[(33, 50)], [(37, 54)]],
+ query_ranges=None,
+ extra={},
+ )
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context(
+ "layer-0",
+ query,
+ key,
+ value,
+ omni_attn_metadata=metadata,
+ )
+ output = _run_omni_paged_backend(context)
+
+ assert torch.equal(output, query)
+ assert batch.query_start_loc.tolist() == [0, 17, 34]
+ assert batch.seq_lens.tolist() == [50, 54]
+ assert batch.positions.tolist() == [*range(33, 50), *range(37, 54)]
+ assert torch.equal(context.key_write, key.reshape(34, 2, 4))
+ assert torch.equal(context.value_write, value.reshape(34, 2, 4))
+ assert layer.native_events == ["update", "forward"]
+ assert layer.updates[0][2].tolist() == batch.positions.tolist()
+ assert layer.calls[0][3].seq_lens.tolist() == [50, 54]
+ assert layer.calls[0][3].query_start_loc_cpu.tolist() == [0, 17, 34]
+ assert len(events) == 2
+
+
+def test_forward_rejects_unaligned_heterogeneous_piecewise_spans_before_cache_write(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4),
+ DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=4, seq_len=4),
+ ]
+ )
+ qkv = torch.randn(2, 4, 2, 4)
+ metadata = SimpleNamespace(
+ attn_mask=torch.ones(2, 1, 4, 4, dtype=torch.bool),
+ full_attn_spans=[[(1, 3)], [(2, 4)]],
+ query_ranges=None,
+ extra={},
+ )
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="produce aligned segments"):
+ adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ assert layer.native_events == []
+
+
+@pytest.mark.parametrize(
+ "spans",
+ [
+ [(1, 3), (2, 4)],
+ [(1, 5)],
+ ],
+)
+def test_forward_rejects_invalid_piecewise_spans_before_cache_write(
+ monkeypatch: pytest.MonkeyPatch,
+ spans: list[tuple[int, int]],
+) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(full_attn_spans=[spans], extra={})
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="sorted, non-overlapping, and within the sequence"):
+ adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ assert layer.native_events == []
+
+
+def test_forward_rejects_non_4d_mask_with_piecewise_spans_before_cache_write(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(
+ attn_mask=torch.ones(1, 4, 4, dtype=torch.bool),
+ full_attn_spans=[[(1, 3)]],
+ extra={},
+ )
+
+ with adapter.activate(batch), pytest.raises(NotImplementedError, match="requires a 4D tensor"):
+ adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ assert layer.native_events == []
+
+
+def test_forward_rejects_piecewise_query_ranges_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ qkv = torch.randn(1, 4, 2, 4)
+ metadata = SimpleNamespace(
+ attn_mask=torch.ones(1, 1, 4, 4, dtype=torch.bool),
+ full_attn_spans=[[(1, 3)]],
+ query_ranges=(object(),),
+ extra={},
+ )
+
+ with adapter.activate(batch), pytest.raises(NotImplementedError, match="query_ranges"):
+ adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ assert layer.native_events == []
+
+
+def test_forward_rejects_piecewise_non_suffix_query_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=2,
+ seq_len=6,
+ kv_start_pos=2,
+ )
+ ]
+ )
+ qkv = torch.randn(1, 2, 2, 4)
+ metadata = SimpleNamespace(full_attn_spans=[[(2, 4)]], extra={})
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"):
+ adapter.prepare_layer_context(
+ "layer-0",
+ qkv,
+ qkv,
+ qkv,
+ omni_attn_metadata=metadata,
+ )
+
+ assert layer.native_events == []
+
+
+def test_forward_accepts_non_aligned_current_write(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=2,
+ seq_len=5,
+ kv_start_pos=3,
+ )
+ ]
+ )
+ query = torch.randn(1, 2, 2, 4)
+ key = torch.arange(1 * 2 * 2 * 4, dtype=torch.float32).reshape(1, 2, 2, 4)
+ value = key + 100
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, key, value)
+
+ assert context.query.shape == (2, 2, 4)
+ assert torch.equal(context.key_write, key.reshape(2, 2, 4))
+ assert torch.equal(context.value_write, value.reshape(2, 2, 4))
+ assert context.slot_mapping.tolist() == [3, 4]
+ assert layer.native_events == []
+
+
+def test_forward_rejects_untranslated_omni_metadata_before_cache_write(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, layer, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+ qkv = torch.randn(1, 2, 4)
+ metadata = SimpleNamespace(attn_mask=torch.ones(1, 1), extra={})
+
+ with adapter.activate(batch), pytest.raises(NotImplementedError, match="cannot translate"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
+
+ assert layer.updates == []
+
+
+def test_forward_accepts_flattened_hidden_dimension(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
+ )
+ query = torch.randn(3, 8)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", query, query, query)
+
+ assert context.query.shape == (3, 2, 4)
+ assert context.restore_output(context.query).shape == query.shape
+
+
+def test_forward_accepts_one_flattened_token_for_single_head(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch, num_heads=1, num_kv_heads=1)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+ qkv = torch.randn(1, 4)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+ output = _run_omni_paged_backend(context)
+
+ assert context.query.shape == (1, 1, 4)
+ assert context.query_token_shape == (1,)
+ assert not context.query_has_head_dims
+ assert output.shape == qkv.shape
+
+
+def test_forward_requires_active_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ qkv = torch.randn(1, 2, 4)
+
+ with pytest.raises(RuntimeError, match="adapter.activate"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+
+def test_new_preparation_invalidates_native_buffer_views_in_older_batch(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ old_batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+ adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)])
+
+ with pytest.raises(ValueError, match="stale"):
+ with adapter.activate(old_batch):
+ pass
+
+
+def test_block_table_change_invalidates_prepared_native_buffer_views(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+
+ adapter.invalidate_prepared_batches()
+
+ with pytest.raises(ValueError, match="stale"):
+ with adapter.activate(batch):
+ pass
+
+
+def test_block_table_change_is_rejected_during_active_forward(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+
+ with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"):
+ adapter.invalidate_prepared_batches()
+
+
+def test_prepare_batch_rejects_active_native_buffers(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+
+ with adapter.activate(batch), pytest.raises(RuntimeError, match="during an active forward"):
+ adapter.prepare_batch([DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1)])
+
+
+def test_forward_rejects_tensor_count_different_from_prepared_batch(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=3, seq_len=3)]
+ )
+ qkv = torch.randn(2, 2, 4)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="token count"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+
+def test_forward_rejects_dtype_different_from_model_activation_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1)]
+ )
+ qkv = torch.randn(1, 2, 4, dtype=torch.float64)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="model activation dtype"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+
+def test_forward_rejects_invalid_key_token_count(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
+ )
+ query = torch.randn(4, 2, 4)
+ key = torch.randn(3, 2, 4)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="key token count"):
+ adapter.prepare_layer_context("layer-0", query, key, query)
+
+
+def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
+ DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
+ ]
+ )
+ qkv = torch.randn(2, 2, 2, 4)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="must match prepared rows"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+
+def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
+ DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
+ ]
+ )
+ qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
+
+ with adapter.activate(batch):
+ context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+ assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
+ torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
+ restored = context.restore_output(context.query)
+ assert restored.shape == qkv.shape
+ torch.testing.assert_close(restored[0, 0], qkv[0, 0])
+ assert torch.count_nonzero(restored[0, 1:]) == 0
+ torch.testing.assert_close(restored[1], qkv[1])
+
+
+def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=1,
+ seq_len=3,
+ kv_start_pos=2,
+ ),
+ DiffusionPagedAttentionRow(
+ request_id="req-1",
+ sequence_id=1,
+ query_len=1,
+ seq_len=5,
+ kv_start_pos=4,
+ ),
+ ]
+ )
+ query = torch.randn(2, 1, 2, 4)
+ full_kv = torch.randn(2, 4, 2, 4)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="batched key layout must match prepared rows"):
+ adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
+
+
+def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch)
+ metadata = DiffusionPagedAttentionMetadata(
+ prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
+ denoise_rows=(
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=2,
+ seq_len=5,
+ kv_start_pos=3,
+ ),
+ ),
+ )
+ runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
+
+ with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
+ set_forward_context_denoise_step_idx(0)
+ assert Attention._active_paged_kv_adapter() is adapter
+ assert adapter._active_batch.rows == metadata.prefill_rows
+
+ set_forward_context_denoise_step_idx(1)
+ assert Attention._active_paged_kv_adapter() is adapter
+ denoise_batch = adapter._active_batch
+ assert denoise_batch.rows == metadata.denoise_rows
+
+ set_forward_context_denoise_step_idx(2)
+ assert Attention._active_paged_kv_adapter() is adapter
+ assert adapter._active_batch is denoise_batch
+
+ assert adapter._active_batch is None
+
+
+def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
+ row = DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=2,
+ seq_len=8,
+ kv_start_pos=3,
+ )
+
+ batch = adapter.prepare_batch([row])
+ qkv = torch.randn(2, 2, 4)
+
+ with adapter.activate(batch), pytest.raises(ValueError, match="requires each query/write span to end"):
+ adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
+
+
+def test_prepare_batch_rejects_native_decode_after_prefill_order(monkeypatch: pytest.MonkeyPatch) -> None:
+ adapter, block_tables, _, _ = _make_adapter(
+ monkeypatch,
+ non_causal=False,
+ reorder_batch_threshold=1,
+ )
+ rows = [
+ DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=2, seq_len=2),
+ DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=1, seq_len=1),
+ ]
+
+ with pytest.raises(ValueError, match="decode/short-query rows before"):
+ adapter.prepare_batch(rows)
+
+ assert block_tables.gather_calls == []
+ assert block_tables.slot_calls == []
+
+
+def test_prepare_batch_rejects_seq_len_beyond_installed_row_capacity(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, block_tables, _, _ = _make_adapter(monkeypatch)
+ block_tables.num_blocks.np[0, 2] = 1
+ row = DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=5,
+ seq_len=5,
+ )
+
+ with pytest.raises(ValueError, match="installed Worker row contains only 1 blocks"):
+ adapter.prepare_batch([row])
+
+ assert block_tables.gather_calls == []
+ assert block_tables.slot_calls == []
+
+
+def test_prepare_batch_rejects_seq_len_beyond_installed_logical_length(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, block_tables, _, _ = _make_adapter(monkeypatch)
+ adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(2, 4)
+ row = DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=5,
+ seq_len=5,
+ )
+
+ with pytest.raises(ValueError, match="installed allocation has logical length 4"):
+ adapter.prepare_batch([row])
+
+ assert block_tables.gather_calls == []
+ assert block_tables.slot_calls == []
+
+
+def test_prepare_batch_passes_dcp_local_seq_lens_to_native_builder(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter, block_tables, _, events = _make_adapter(monkeypatch)
+ block_tables.cp_size = 2
+ block_tables.cp_rank = 1
+ block_tables.cp_interleave = 2
+ block_tables.num_blocks.np[0, 2] = 1
+
+ batch = adapter.prepare_batch(
+ [DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
+ )
+
+ assert batch.row_indices.tolist() == [2]
+ assert events[0][1]["dcp_local_seq_lens"].tolist() == [2]
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "message"),
+ [
+ ({"sequence_id": 0, "context_id": "text"}, "exactly one"),
+ ({}, "exactly one"),
+ ({"sequence_id": 0, "query_len": 0}, "query_len"),
+ ({"sequence_id": 0, "kv_start_pos": 4, "query_len": 2, "seq_len": 5}, "exceeds seq_len"),
+ ],
+)
+def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: str) -> None:
+ values = {
+ "request_id": "req-0",
+ "query_len": 1,
+ "seq_len": 1,
+ **kwargs,
+ }
+ with pytest.raises(ValueError, match=message):
+ DiffusionPagedAttentionRow(**values)
+
+
+def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ selected_backends = []
+ specialized_backends = []
+ impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
+ native_backend = SimpleNamespace(
+ get_name=lambda: "ASCEND",
+ indexes_kv_by_block_stride=lambda: True,
+ get_impl_cls=lambda: impl_cls,
+ forward_includes_kv_cache_update=True,
+ )
+
+ original_backend_per_kind = {"full": object()}
+ config = SimpleNamespace(
+ attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
+ parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
+ model_config=SimpleNamespace(dtype=torch.float16),
+ cache_config=SimpleNamespace(cache_dtype="auto"),
+ )
+
+ def select_backend(**_kwargs):
+ selected_backends.append(
+ (
+ config.attention_config.backend,
+ config.attention_config.backend_per_kind,
+ config.parallel_config.prefill_context_parallel_size,
+ )
+ )
+ return native_backend
+
+ monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
+ monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+ monkeypatch.setattr(
+ adapter_module.current_omni_platform,
+ "get_diffusion_paged_kv_attn_backend",
+ lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+ )
+ layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
+ spec = FullAttentionSpec(
+ block_size=16,
+ num_kv_heads=4,
+ head_size=8,
+ dtype=torch.float16,
+ non_causal=True,
+ )
+
+ native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(
+ layer_name="layer-0",
+ layer=layer,
+ spec=spec,
+ vllm_config=config,
+ device=torch.device("cpu"),
+ ulysses_degree=2,
+ )
+
+ assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
+ assert specialized_backends == [(native_backend, 2)]
+ assert config.attention_config.backend is None
+ assert config.attention_config.backend_per_kind is original_backend_per_kind
+ assert config.parallel_config.prefill_context_parallel_size == 2
+ assert native_layer.num_heads == 4
+ assert native_layer.num_kv_heads == 2
+ assert native_layer.spec.num_kv_heads == 2
+ assert native_layer.spec.indexes_kv_by_block_stride is True
+ assert native_layer._q_scale_float == 1.0
+ assert native_layer._k_scale_float == 1.0
+ assert native_layer._v_scale_float == 1.0
+ impl_kwargs = impl_cls.call_args.kwargs
+ assert impl_kwargs["num_heads"] == 4
+ assert impl_kwargs["num_kv_heads"] == 2
+
+
+def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
+ events: list[str] = []
+
+ class Strategy:
+ name = "ulysses"
+
+ def pre_attention(self, query, key, value, metadata):
+ events.append("pre")
+ return query + 1, key + 2, value + 3, metadata, object()
+
+ def post_attention(self, output, _ctx):
+ events.append("post")
+ return output + 4
+
+ class Adapter:
+ def prepare_layer_context(self, layer_name, query, key, value, *, omni_attn_metadata):
+ events.append("prepare")
+ assert layer_name == "layer-0"
+ assert omni_attn_metadata is None
+ assert torch.equal(key, original + 2)
+ assert torch.equal(value, original + 3)
+ return SimpleNamespace(query=query)
+
+ class Backend:
+ def forward_paged(self, context):
+ events.append("kernel")
+ return context.query * 2
+
+ layer = Attention.__new__(Attention)
+ nn.Module.__init__(layer)
+ layer.prefix = "layer-0"
+ layer.paged_kv_cache_role = "primary"
+ layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
+ layer.attention = Backend()
+ layer.use_ring = False
+ layer._no_parallel_strategy = object()
+ layer._get_active_parallel_strategy = lambda: Strategy()
+ layer._with_kv_cache_dtype = lambda metadata: metadata
+ original = torch.zeros(1, 2, 2, 4)
+
+ with set_forward_context(), override_paged_kv_adapter(Adapter()):
+ assert layer.is_paged_kv_active()
+ output = layer._forward_impl(original, original, original)
+
+ assert events == ["pre", "prepare", "kernel", "post"]
+ assert torch.equal(output, torch.full_like(original, 6))
+
+
+def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
+ events: list[tuple[str, tuple[int, ...]]] = []
+
+ class Strategy:
+ name = "ulysses"
+
+ def pre_attention(self, query, key, value, metadata):
+ return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
+
+ def post_attention(self, output, _ctx):
+ events.append(("post", tuple(output.shape)))
+ return output
+
+ class Adapter:
+ def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
+ assert omni_attn_metadata is None
+ events.append(("kernel_input", tuple(query.shape)))
+ assert query[:, -1].tolist() == [[[3.0]]]
+ assert torch.equal(query, key)
+ assert torch.equal(query, value)
+ return SimpleNamespace(query=query)
+
+ class Backend:
+ def forward_paged(self, context):
+ return context.query
+
+ layer = Attention.__new__(Attention)
+ nn.Module.__init__(layer)
+ layer.prefix = "layer-0"
+ layer.paged_kv_cache_role = "primary"
+ layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
+ layer.attention = Backend()
+ layer.use_ring = False
+ layer._no_parallel_strategy = object()
+ layer._get_active_parallel_strategy = lambda: Strategy()
+ tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
+
+ with set_forward_context(), override_paged_kv_adapter(Adapter()):
+ get_forward_context().sp_padding_size = 1
+ output = layer._forward_impl(tensor, tensor, tensor)
+
+ assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
+ assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
+
+
+def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
+ class Adapter:
+ def prepare_layer_context(self, *args, **kwargs):
+ raise AssertionError("an unsupported backend must fail before adapter preparation")
+
+ layer = Attention.__new__(Attention)
+ nn.Module.__init__(layer)
+ layer.prefix = "layer-0"
+ layer.paged_kv_cache_role = "primary"
+ layer.attn_backend = SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA")
+ layer.use_ring = False
+ layer._no_parallel_strategy = object()
+ layer._get_active_parallel_strategy = lambda: layer._no_parallel_strategy
+
+ with (
+ set_forward_context(),
+ override_paged_kv_adapter(Adapter()),
+ pytest.raises(NotImplementedError, match="backend with paged support"),
+ ):
+ layer._forward_impl(torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4), torch.zeros(1, 2, 2, 4))
+
+
+def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
+ events: list[str] = []
+
+ class Strategy:
+ name = "ulysses"
+
+ def pre_attention(self, query, key, value, metadata):
+ events.append("pre")
+ return query, key, value, metadata, object()
+
+ def post_attention(self, output, _ctx):
+ events.append("post")
+ return output
+
+ layer = Attention.__new__(Attention)
+ nn.Module.__init__(layer)
+ layer.paged_kv_cache_role = "primary"
+ layer.use_ring = False
+ layer._no_parallel_strategy = object()
+ layer._get_active_parallel_strategy = lambda: Strategy()
+ layer._with_kv_cache_dtype = lambda metadata: metadata
+ layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query
+ qkv = torch.zeros(1, 2, 2, 4)
+
+ assert not layer.is_paged_kv_active()
+ output = layer._forward_impl(qkv, qkv, qkv)
+
+ assert events == ["pre", "dense", "post"]
+ assert output is qkv
+
+
+def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
+ layer = Attention.__new__(Attention)
+ nn.Module.__init__(layer)
+ layer._scheduler_paged_kv = True
+ layer.paged_kv_cache_role = "primary"
+ layer._get_active_parallel_strategy = lambda: object()
+ qkv = torch.zeros(1, 2, 2, 4)
+
+ with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
+ layer._forward_impl(qkv, qkv, qkv)
diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py
new file mode 100644
index 000000000..ef5d75f4d
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F
+from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+ DiffusionPagedAttentionAdapter,
+ DiffusionPagedAttentionLayerAdapter,
+ DiffusionPagedAttentionRow,
+ DiffusionPagedAttentionRowBinding,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+_LAYER_NAME = "model.layers.0.attn"
+_NUM_HEADS = 2
+_HEAD_SIZE = 64
+_BLOCK_SIZE = 16
+
+
+class _SmokeDiffusionAttention(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.num_heads = _NUM_HEADS
+ self.num_kv_heads = _NUM_HEADS
+ self.head_size = _HEAD_SIZE
+ self.softmax_scale = _HEAD_SIZE**-0.5
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+ device = torch.device("cuda", torch.accelerator.current_device_index())
+ vllm_config = VllmConfig(
+ cache_config=CacheConfig(
+ block_size=_BLOCK_SIZE,
+ cache_dtype="float16",
+ enable_prefix_caching=False,
+ ),
+ attention_config=AttentionConfig(use_non_causal=True),
+ )
+ model_config = _DiffusionVllmModelConfig(
+ model=None,
+ dtype=torch.float16,
+ max_model_len=32,
+ original_max_model_len=32,
+ hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"),
+ )
+ vllm_config.model_config = model_config
+ model_config.set_attention_geometry(
+ num_heads=_NUM_HEADS,
+ num_kv_heads=_NUM_HEADS,
+ head_size=_HEAD_SIZE,
+ )
+ vllm_config.compilation_config.static_forward_context.clear()
+
+ diffusion_layer = _SmokeDiffusionAttention().to(device)
+ spec = FullAttentionSpec(
+ block_size=_BLOCK_SIZE,
+ num_kv_heads=_NUM_HEADS,
+ head_size=_HEAD_SIZE,
+ dtype=torch.float16,
+ non_causal=True,
+ )
+ with set_current_vllm_config(vllm_config):
+ native_layer = DiffusionPagedAttentionLayerAdapter(
+ layer_name=_LAYER_NAME,
+ layer=diffusion_layer,
+ spec=spec,
+ vllm_config=vllm_config,
+ device=device,
+ )
+ vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer
+
+ canonical_spec = native_layer.spec
+ kv_cache_config = KVCacheConfig(
+ num_blocks=4,
+ kv_cache_tensors=[
+ KVCacheTensor(
+ size=canonical_spec.page_size_bytes * 4,
+ shared_by=[_LAYER_NAME],
+ )
+ ],
+ kv_cache_groups=[
+ KVCacheGroupSpec(
+ layer_names=[_LAYER_NAME],
+ kv_cache_spec=canonical_spec,
+ )
+ ],
+ )
+ with set_current_vllm_config(vllm_config):
+ attn_groups, _, kernel_block_sizes = init_attn_backend(
+ kv_cache_config,
+ vllm_config,
+ device,
+ )
+ runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+ init_kv_cache(
+ runner_kv_caches,
+ vllm_config.compilation_config.static_forward_context,
+ kv_cache_config,
+ attn_groups,
+ device,
+ vllm_config.cache_config.cache_dtype,
+ kernel_block_sizes,
+ vllm_config,
+ )
+ block_tables = BlockTables(
+ block_sizes=[_BLOCK_SIZE],
+ max_num_reqs=1,
+ max_num_batched_tokens=32,
+ max_num_blocks_per_group=[2],
+ device=device,
+ kernel_block_sizes=kernel_block_sizes,
+ )
+
+ block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
+ block_tables.apply_staged_writes()
+ adapter = DiffusionPagedAttentionAdapter(
+ vllm_config=vllm_config,
+ device=device,
+ kv_cache_config=kv_cache_config,
+ block_tables=block_tables,
+ attn_groups=attn_groups,
+ layers={_LAYER_NAME: native_layer},
+ resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
+ row_index=0,
+ max_seq_len=19,
+ ),
+ )
+ prefix_batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=17,
+ seq_len=17,
+ )
+ ]
+ )
+
+ torch.manual_seed(1)
+ query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
+ packed_qkv = torch.randn(
+ 19,
+ 3 * _NUM_HEADS * _HEAD_SIZE,
+ dtype=torch.float16,
+ device=device,
+ )
+ _, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
+ key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+ value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
+ assert not key.is_contiguous()
+ assert not value.is_contiguous()
+ omni_backend = FlashAttentionImpl(
+ num_heads=_NUM_HEADS,
+ head_size=_HEAD_SIZE,
+ softmax_scale=diffusion_layer.softmax_scale,
+ num_kv_heads=_NUM_HEADS,
+ )
+ with adapter.activate(prefix_batch):
+ prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17])
+ prefix_output = omni_backend.forward_paged(prefix_context)
+ prefix_slot_mappings = prefix_batch.slot_mappings.clone()
+
+ suffix_batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=2,
+ seq_len=19,
+ kv_start_pos=17,
+ )
+ ]
+ )
+ with adapter.activate(suffix_batch):
+ suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:])
+ suffix_output = omni_backend.forward_paged(suffix_context)
+ suffix_slot_mappings = suffix_batch.slot_mappings.clone()
+
+ piecewise_batch = adapter.prepare_batch(
+ [
+ DiffusionPagedAttentionRow(
+ request_id="req-0",
+ sequence_id=0,
+ query_len=19,
+ seq_len=19,
+ )
+ ]
+ )
+ mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril()
+ mixed_mask[5:10, 5:10] = True
+ mixed_metadata = AttentionMetadata(
+ attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+ full_attn_spans=[[(5, 10)]],
+ )
+ with adapter.activate(piecewise_batch):
+ piecewise_context = adapter.prepare_layer_context(
+ _LAYER_NAME,
+ query,
+ key,
+ value,
+ omni_attn_metadata=mixed_metadata,
+ )
+ piecewise_output = omni_backend.forward_paged(piecewise_context)
+
+ prefix_reference = (
+ F.scaled_dot_product_attention(
+ query[:17].transpose(0, 1).unsqueeze(0),
+ key[:17].transpose(0, 1).unsqueeze(0),
+ value[:17].transpose(0, 1).unsqueeze(0),
+ is_causal=False,
+ scale=diffusion_layer.softmax_scale,
+ )
+ .squeeze(0)
+ .transpose(0, 1)
+ )
+ suffix_reference = (
+ F.scaled_dot_product_attention(
+ query[17:].transpose(0, 1).unsqueeze(0),
+ key.transpose(0, 1).unsqueeze(0),
+ value.transpose(0, 1).unsqueeze(0),
+ is_causal=False,
+ scale=diffusion_layer.softmax_scale,
+ )
+ .squeeze(0)
+ .transpose(0, 1)
+ )
+ piecewise_reference = (
+ F.scaled_dot_product_attention(
+ query.transpose(0, 1).unsqueeze(0),
+ key.transpose(0, 1).unsqueeze(0),
+ value.transpose(0, 1).unsqueeze(0),
+ attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+ scale=diffusion_layer.softmax_scale,
+ )
+ .squeeze(0)
+ .transpose(0, 1)
+ )
+ torch.accelerator.synchronize(device)
+
+ assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]]
+ assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]]
+ torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
+ torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
+ torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)
diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py
index 43734ca31..ae5609b2a 100644
--- a/tests/diffusion/diffusion_kv/test_request.py
+++ b/tests/diffusion/diffusion_kv/test_request.py
@@ -107,11 +107,8 @@ def test_request_rejects_non_context_values() -> None:
def _manager(*, enable_caching: bool) -> KVCacheManager:
# vLLM 0.26 lazily registers its built-in KV specs. Another test may have
# already registered an out-of-tree spec, making the registry non-empty
- # before FullAttentionSpec is registered. Initialize the native mappings
- # explicitly; vLLM 0.25 has no registration hook and needs no setup.
- register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
- if register_all is not None:
- register_all(None)
+ # before FullAttentionSpec is registered. Initialize the native mappings.
+ native_kv_managers.register_all_kvcache_specs(None)
spec = FullAttentionSpec(
block_size=BLOCK_SIZE,
diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py
new file mode 100644
index 000000000..31faf4d08
--- /dev/null
+++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py
@@ -0,0 +1,302 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _Backend:
+ @classmethod
+ def indexes_kv_by_block_stride(cls) -> bool:
+ return True
+
+
+class _RunnerKVBackend:
+ def __init__(self) -> None:
+ self.kv_cache_config = None
+
+ def register_kv_cache_layers(self, layers):
+ return {layer_name: spec for layer_name, (_layer, spec) in layers.items()}
+
+ def initialize_kv_cache(self, config) -> None:
+ configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names}
+ if configured_layers != {"image_attention"}:
+ raise ValueError(
+ "Rank-local Diffusion KVCacheConfig layer mismatch: "
+ f"expected=['image_attention'], configured={sorted(configured_layers)}"
+ )
+ self.kv_cache_config = config
+
+
+def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention:
+ attention = Attention.__new__(Attention)
+ nn.Module.__init__(attention)
+ attention.prefix = prefix
+ attention.paged_kv_cache_role = "primary" if enabled else None
+ attention.paged_kv_cache_dtype = torch.bfloat16
+ attention.num_kv_heads = 2
+ attention.head_size = 8
+ attention.causal = False
+ attention.attn_backend = _Backend
+ return attention
+
+
+def _runner(attention: Attention) -> DiffusionModelRunner:
+ runner = object.__new__(DiffusionModelRunner)
+ runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+ runner.vllm_config = SimpleNamespace(
+ cache_config=SimpleNamespace(block_size=4),
+ model_config=SimpleNamespace(dtype=torch.float16),
+ )
+ pipeline = nn.Module()
+ pipeline.image_attention = attention
+ runner.pipeline = pipeline
+ runner.diffusion_kv_backend = _RunnerKVBackend()
+ return runner
+
+
+def test_runner_discovers_native_spec_from_loaded_attention() -> None:
+ specs = _runner(_attention(enabled=True)).get_kv_cache_spec()
+
+ assert set(specs) == {"image_attention"}
+ spec = specs["image_attention"]
+ assert isinstance(spec, FullAttentionSpec)
+ assert spec.block_size == 4
+ assert spec.num_kv_heads == 2
+ assert spec.head_size == 8
+ assert spec.dtype is torch.bfloat16
+ assert spec.indexes_kv_by_block_stride is True
+ assert spec.non_causal is True
+
+
+def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None:
+ runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn"))
+
+ specs = runner.get_kv_cache_spec()
+
+ assert set(specs) == {"layers.0.self_attn.image_attn.attn"}
+
+
+def test_runner_rejects_duplicate_attention_prefixes() -> None:
+ runner = _runner(_attention(enabled=True, prefix="shared"))
+ second_attention = _attention(enabled=True, prefix="shared")
+ runner.pipeline.second_image_attention = second_attention
+
+ with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"):
+ runner.get_kv_cache_spec()
+
+
+def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
+ runner = _runner(_attention(enabled=False))
+
+ with pytest.raises(RuntimeError, match="no cache-enabled Attention"):
+ runner.get_kv_cache_spec()
+
+
+def test_runner_retains_matching_rank_local_config() -> None:
+ runner = _runner(_attention(enabled=True))
+ spec = runner.get_kv_cache_spec()["image_attention"]
+ config = KVCacheConfig(
+ num_blocks=8,
+ kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])],
+ kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)],
+ )
+
+ runner.set_kv_cache_config(config)
+
+ assert runner.kv_cache_config is config
+
+
+def test_runner_rejects_rank_local_config_for_different_layers() -> None:
+ runner = _runner(_attention(enabled=True))
+ spec = runner.get_kv_cache_spec()["image_attention"]
+ config = KVCacheConfig(
+ num_blocks=8,
+ kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])],
+ kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)],
+ )
+
+ with pytest.raises(ValueError, match="layer mismatch"):
+ runner.set_kv_cache_config(config)
+
+
+def test_worker_selects_its_rank_local_config() -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.rank = 1
+ worker.od_config = SimpleNamespace(num_gpus=2)
+ worker.vllm_config = SimpleNamespace(
+ model_config=SimpleNamespace(max_model_len=None),
+ cache_config=SimpleNamespace(num_gpu_blocks=None),
+ )
+ worker._maybe_get_memory_pool_context = lambda _tag: nullcontext()
+ worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
+ configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)]
+
+ worker.set_kv_cache_configs(configs, 64)
+
+ assert worker.installed is configs[1]
+ assert worker.vllm_config.model_config.max_model_len == 64
+ assert worker.vllm_config.cache_config.num_gpu_blocks == 8
+ with pytest.raises(ValueError, match="rank count mismatch"):
+ worker.set_kv_cache_configs(configs[:1], 64)
+
+
+def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.rank = 0
+ worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=4096))
+ worker.init_snapshot = SimpleNamespace(free_memory=16384)
+ worker.requested_memory = 8192
+ worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=1024)
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+ profile_request = object()
+
+ assert worker.determine_available_kv_memory(profile_request) == [4096]
+ worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_treats_zero_kv_memory_budget_as_profiled_auto_sizing(monkeypatch) -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.rank = 0
+ worker.init_snapshot = SimpleNamespace(free_memory=900)
+ worker.requested_memory = 750
+ worker.vllm_config = SimpleNamespace(
+ cache_config=SimpleNamespace(
+ kv_cache_memory_bytes=0,
+ gpu_memory_utilization=0.75,
+ )
+ )
+ worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+ profile_result = SimpleNamespace(
+ non_kv_cache_memory=400,
+ after_profile=SimpleNamespace(free_memory=500),
+ )
+ monkeypatch.setattr(
+ diffusion_worker_module,
+ "memory_profiling",
+ lambda snapshot, weights_memory: nullcontext(profile_result),
+ )
+ profile_request = object()
+
+ assert worker.determine_available_kv_memory(profile_request) == [350]
+ worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+@pytest.mark.parametrize("non_kv_cache_memory", [750, 800])
+def test_worker_rejects_non_positive_profiled_kv_memory(monkeypatch, non_kv_cache_memory: int) -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.rank = 0
+ worker.init_snapshot = SimpleNamespace(free_memory=900)
+ worker.requested_memory = 750
+ worker.vllm_config = SimpleNamespace(
+ cache_config=SimpleNamespace(
+ kv_cache_memory_bytes=None,
+ gpu_memory_utilization=0.75,
+ )
+ )
+ worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=100)
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+ profile_result = SimpleNamespace(
+ non_kv_cache_memory=non_kv_cache_memory,
+ after_profile=SimpleNamespace(free_memory=100),
+ )
+ monkeypatch.setattr(
+ diffusion_worker_module,
+ "memory_profiling",
+ lambda snapshot, weights_memory: nullcontext(profile_result),
+ )
+ profile_request = object()
+
+ with pytest.raises(
+ RuntimeError,
+ match=(
+ "rank 0: RuntimeError: No memory remains for Diffusion KV cache after profiling: "
+ f"requested_memory=750 bytes, non_kv_cache_memory={non_kv_cache_memory} bytes"
+ ),
+ ):
+ worker.determine_available_kv_memory(profile_request)
+
+ worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_profiles_activation_headroom_instead_of_current_residency(monkeypatch) -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.rank = 0
+ worker.init_snapshot = SimpleNamespace(free_memory=1200)
+ worker.requested_memory = 1000
+ worker.vllm_config = SimpleNamespace(
+ cache_config=SimpleNamespace(
+ kv_cache_memory_bytes=None,
+ gpu_memory_utilization=0.8,
+ )
+ )
+ worker.model_runner = SimpleNamespace(profile_run=Mock(), model_memory_usage=250)
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+ profile_result = SimpleNamespace(
+ non_kv_cache_memory=650,
+ after_profile=SimpleNamespace(free_memory=550),
+ )
+ profiling_calls = []
+
+ def profile(snapshot, *, weights_memory):
+ profiling_calls.append((snapshot, weights_memory))
+ return nullcontext(profile_result)
+
+ monkeypatch.setattr(
+ diffusion_worker_module,
+ "memory_profiling",
+ profile,
+ )
+ profile_request = object()
+
+ assert worker.determine_available_kv_memory(profile_request) == [350]
+ assert profiling_calls == [(worker.init_snapshot, 250)]
+ worker.model_runner.profile_run.assert_called_once_with(profile_request)
+
+
+def test_worker_rejects_profile_without_preload_snapshot(monkeypatch) -> None:
+ worker = object.__new__(DiffusionWorker)
+ worker.init_snapshot = None
+ worker.requested_memory = None
+ worker.vllm_config = SimpleNamespace(cache_config=SimpleNamespace(kv_cache_memory_bytes=None))
+ worker.model_runner = SimpleNamespace()
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", lambda value: [value])
+
+ with pytest.raises(RuntimeError, match="snapshot was not captured before model loading"):
+ worker.determine_available_kv_memory(object())
+
+
+def test_rank_probe_gathers_local_failure_before_raising(monkeypatch) -> None:
+ gathered = []
+
+ def gather(local_result):
+ gathered.append(local_result)
+ return [local_result, (True, "peer-result")]
+
+ def fail_local_probe():
+ raise ValueError("local failure")
+
+ monkeypatch.setattr(diffusion_worker_module, "_all_gather_rank_values", gather)
+
+ with pytest.raises(RuntimeError, match="rank 0: ValueError: local failure"):
+ diffusion_worker_module._run_and_gather_rank_values(
+ "test probe",
+ fail_local_probe,
+ )
+
+ assert gathered == [(False, "ValueError: local failure")]
diff --git a/tests/diffusion/diffusion_kv/test_worker_contract.py b/tests/diffusion/diffusion_kv/test_worker_contract.py
index 0c78abbaf..5f2caab06 100644
--- a/tests/diffusion/diffusion_kv/test_worker_contract.py
+++ b/tests/diffusion/diffusion_kv/test_worker_contract.py
@@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
runner = object.__new__(DiffusionModelRunner)
- runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+ runner.od_config = SimpleNamespace(
+ diffusion_kv_mode=mode,
+ parallel_config=SimpleNamespace(cfg_parallel_size=1),
+ )
return runner
@@ -85,19 +88,60 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
assert new_req.diffusion_kv_metadata is metadata
-def test_paged_request_without_metadata_reaches_request_forward() -> None:
+def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
- expected_output = DiffusionOutput(output=None)
- runner._execute_request_list = Mock(
- return_value=SimpleNamespace(
- runner_outputs=[SimpleNamespace(result=expected_output)],
- )
+ metadata = DiffusionKVMetadata(
+ request_id="req-0",
+ allocation_generation=1,
+ sequences=(
+ DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+ DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+ ),
)
+
+ attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+ assert [
+ (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+ for row in attn_metadata.prefill_rows
+ ] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)]
+ assert [
+ (row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
+ for row in attn_metadata.denoise_rows
+ ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
+
+
+def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
+ runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+ runner.od_config.parallel_config.cfg_parallel_size = 2
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank",
+ lambda: 1,
+ )
+ metadata = DiffusionKVMetadata(
+ request_id="req-0",
+ allocation_generation=1,
+ sequences=(
+ DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
+ DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
+ ),
+ )
+
+ attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+ assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1]
+ assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1]
+
+
+def test_paged_request_without_metadata_fails_before_request_forward() -> None:
+ runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+ runner._execute_request_list = Mock()
req = SimpleNamespace(request_id="req-0")
- assert runner.execute_model(req) is expected_output
+ with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+ runner.execute_model(req)
- runner._execute_request_list.assert_called_once()
+ runner._execute_request_list.assert_not_called()
def test_dense_request_rejects_metadata_before_request_forward() -> None:
@@ -141,18 +185,18 @@ def test_request_rpc_rejects_envelope_request_identity_mismatch() -> None:
assert calls == []
-def test_batch_path_allows_missing_metadata_before_forward() -> None:
+def test_batch_path_rejects_missing_metadata_before_forward() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
- expected_output = object()
- runner._execute_request_list = Mock(return_value=expected_output)
+ runner._execute_request_list = Mock()
new_req = NewRequestData(
request_id="req-0",
req=SimpleNamespace(request_id="req-0"),
)
- assert runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config) is expected_output
+ with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
+ runner.execute_model_batch(make_scheduler_output(new_req), runner.od_config)
- runner._execute_request_list.assert_called_once()
+ runner._execute_request_list.assert_not_called()
def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -170,7 +214,7 @@ def test_batch_path_rejects_envelope_request_identity_mismatch() -> None:
runner._execute_request_list.assert_not_called()
-def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
+def test_step_path_rejects_missing_metadata_before_step_support_check() -> None:
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
runner.pipeline = object()
runner._supports_step_mode = Mock(return_value=False)
@@ -179,10 +223,10 @@ def test_step_path_allows_missing_metadata_before_step_support_check() -> None:
req=SimpleNamespace(request_id="req-0"),
)
- with pytest.raises(ValueError, match="does not support step execution"):
+ with pytest.raises(ValueError, match="requires Diffusion KV metadata"):
runner.execute_stepwise(make_scheduler_output(new_req))
- runner._supports_step_mode.assert_called_once()
+ runner._supports_step_mode.assert_not_called()
def test_step_path_rejects_envelope_request_identity_mismatch() -> None:
@@ -288,3 +332,42 @@ def test_dense_worker_call_does_not_extend_model_runner_signature() -> None:
worker.execute_model(req, SimpleNamespace())
assert calls == [(req, None)]
+
+
+def test_dlo_worker_selects_request_and_metadata_from_same_envelope(monkeypatch) -> None:
+ calls: list[tuple] = []
+
+ class ModelRunner:
+ def execute_model(self, req, **kwargs):
+ calls.append((req, kwargs))
+ return DiffusionOutput(output=None)
+
+ worker = object.__new__(DiffusionWorker)
+ worker.model_runner = ModelRunner()
+ worker.lora_manager = None
+ worker._get_profiler = lambda: None
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.distributed.parallel_state.get_data_parallel_rank",
+ lambda: 1,
+ )
+ req_0 = SimpleNamespace(request_id="req-0")
+ req_1 = SimpleNamespace(request_id="req-1")
+ metadata_0 = make_metadata("req-0")
+ metadata_1 = make_metadata("req-1")
+ envelopes = [
+ NewRequestData(request_id="req-0", req=req_0, diffusion_kv_metadata=metadata_0),
+ NewRequestData(request_id="req-1", req=req_1, diffusion_kv_metadata=metadata_1),
+ ]
+
+ result = worker.execute_model(envelopes, SimpleNamespace())
+
+ assert calls == [
+ (
+ req_1,
+ {
+ "kv_prefetch_job": None,
+ "diffusion_kv_metadata": metadata_1,
+ },
+ )
+ ]
+ assert result["dp_rank"] == 1
diff --git a/tests/diffusion/distributed/test_a2a_permute.py b/tests/diffusion/distributed/test_a2a_permute.py
new file mode 100644
index 000000000..a74f98378
--- /dev/null
+++ b/tests/diffusion/distributed/test_a2a_permute.py
@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed import a2a_permute
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatch) -> None:
+ cu13_include = tmp_path / "nvidia" / "cu13" / "include"
+ nccl_include = tmp_path / "nvidia" / "nccl" / "include"
+ nccl_lib = tmp_path / "nvidia" / "nccl" / "lib" / "libnccl.so.2"
+ cu13_include.mkdir(parents=True)
+ nccl_include.mkdir(parents=True)
+ nccl_lib.parent.mkdir(parents=True)
+ (cu13_include / "cusparse.h").touch()
+ (nccl_include / "nccl.h").touch()
+ nccl_lib.touch()
+
+ load_kwargs = {}
+ monkeypatch.setattr(a2a_permute.sysconfig, "get_paths", lambda: {"purelib": str(tmp_path)})
+ monkeypatch.setattr(
+ torch.utils.cpp_extension,
+ "load",
+ lambda **kwargs: load_kwargs.update(kwargs),
+ )
+ monkeypatch.setattr(a2a_permute.symm_mem, "set_backend", lambda _backend: None)
+ monkeypatch.setattr(a2a_permute, "_BUILT", False)
+
+ a2a_permute.ensure_a2a_permute_available()
+
+ assert set(load_kwargs["extra_include_paths"]) == {str(cu13_include), str(nccl_include)}
+ assert load_kwargs["extra_ldflags"] == [str(nccl_lib)]
+
+
+@dataclass
+class _FakeAllocation:
+ """Byte workspace stand-in that slices and reinterprets like ``symm_mem.empty()``."""
+
+ size: int
+ device: torch.device
+
+ def __post_init__(self) -> None:
+ self.buffer = torch.empty(self.size, dtype=torch.uint8)
+
+ def __getitem__(self, item: slice) -> torch.Tensor:
+ return self.buffer[item]
+
+
+def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
+ allocations: list[_FakeAllocation] = []
+ all_reduce_calls: list[object] = []
+ synchronize_calls: list[torch.device] = []
+ stream_id = [1]
+ a2a_permute._SYMM_WORKSPACES.clear()
+
+ monkeypatch.setattr(a2a_permute.torch.cuda, "is_current_stream_capturing", lambda: False)
+ monkeypatch.setattr(
+ a2a_permute.torch.cuda,
+ "current_stream",
+ lambda _device: SimpleNamespace(cuda_stream=stream_id[0]),
+ )
+ monkeypatch.setattr(a2a_permute, "_resolve_process_group", lambda _name: object())
+ monkeypatch.setattr(a2a_permute.torch, "ones", lambda *_args, **_kwargs: object())
+ monkeypatch.setattr(
+ a2a_permute.dist,
+ "all_reduce",
+ lambda _tensor, group: all_reduce_calls.append(group),
+ )
+ monkeypatch.setattr(
+ a2a_permute.torch.accelerator,
+ "synchronize",
+ lambda device: synchronize_calls.append(device),
+ )
+
+ def empty(size, *, dtype, device):
+ assert dtype == torch.uint8
+ allocation = _FakeAllocation(size=size, device=torch.device(device))
+ allocations.append(allocation)
+ return allocation
+
+ monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
+ # The handle is only retained to keep the rendezvous registration alive.
+ monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda _allocation, _group: object())
+
+ device = torch.device("cuda:0")
+ first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
+ smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
+ larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
+
+ assert (first.shape, first.dtype) == ((2, 3), torch.float16)
+ assert (smaller.shape, smaller.dtype) == ((1, 4), torch.float16)
+ assert (larger.shape, larger.dtype) == ((4, 3), torch.float16)
+ # A sub-capacity request reuses the workspace; only growth allocates again.
+ assert smaller.untyped_storage().data_ptr() == first.untyped_storage().data_ptr()
+ assert larger.untyped_storage().data_ptr() != first.untyped_storage().data_ptr()
+ assert [allocation.size for allocation in allocations] == [12, 24]
+ assert len(a2a_permute._SYMM_WORKSPACES) == 1
+ assert len(all_reduce_calls) == 1
+ assert synchronize_calls == [device, device]
+
+ stream_id[0] = 2
+ with pytest.raises(RuntimeError, match="single CUDA stream"):
+ a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
+
+ a2a_permute.clear_a2a_permute_workspaces()
+ assert not a2a_permute._SYMM_WORKSPACES
+ assert synchronize_calls == [device, device, device]
diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py
index b2ee7c10d..dcaa8e01a 100644
--- a/tests/diffusion/distributed/test_distributed_vae_executor.py
+++ b/tests/diffusion/distributed/test_distributed_vae_executor.py
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
from types import SimpleNamespace
import pytest
@@ -58,6 +59,11 @@ class E2EOperator:
return torch.cat(tiles, dim=0)
+@dataclass
+class FakeWorldGroup:
+ device_group: object
+
+
class DummyMixin(DistributedVaeMixin):
def __init__(self):
self.use_tiling = True
@@ -75,11 +81,13 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
@pytest.fixture(autouse=True)
-def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
+def mock_world_group(monkeypatch: pytest.MonkeyPatch):
+ group = object()
monkeypatch.setattr(
- "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
- lambda: None,
+ "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_world_group",
+ lambda: FakeWorldGroup(device_group=group),
)
+ return group
@pytest.fixture(autouse=True)
@@ -93,6 +101,11 @@ def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
# ============================
+def test_uses_dedicated_world_device_group(mock_world_group):
+ executor = DistributedVaeExecutor()
+ assert executor.group is mock_world_group
+
+
def test_balance_tasks():
executor = DistributedVaeExecutor()
workloads = [2, 5, 13, 8, 2, 4] # each is 17
diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py
index 74f30f156..bfbbca588 100644
--- a/tests/diffusion/distributed/test_expert_parallel_layout.py
+++ b/tests/diffusion/distributed/test_expert_parallel_layout.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -92,9 +92,8 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
- monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
- for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+ for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -195,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
+@pytest.mark.cpu
+@pytest.mark.core_model
+@pytest.mark.parametrize(
+ ("fully_shard_degree", "error_message"),
+ [(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")],
+)
+def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str):
+ """Reject invalid HSDP shard sizes before allocating any process groups."""
+ world_size = 4
+ created_groups: list[object] = []
+
+ def fail_if_group_created(*args, **kwargs):
+ del args, kwargs
+ created_groups.append(object())
+ raise AssertionError("a process group was created before HSDP validation")
+
+ fake_world_group = SimpleNamespace(
+ rank_in_group=0,
+ local_rank=0,
+ device_group=object(),
+ )
+ monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+ monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+ monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+ monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created)
+
+ for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"):
+ monkeypatch.setattr(parallel_state, name, None)
+ for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+ monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+ with pytest.raises(ValueError, match=error_message):
+ parallel_state.initialize_model_parallel(
+ fully_shard_degree=fully_shard_degree,
+ use_hsdp=True,
+ backend="gloo",
+ )
+
+ assert created_groups == []
+
+
@pytest.mark.cpu
@pytest.mark.core_model
def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
@@ -229,9 +269,8 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
- monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
- for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+ for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
@@ -259,6 +298,33 @@ def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
]
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_destroy_model_parallel_clears_vllm_pipeline_group(monkeypatch):
+ """A destroyed diffusion PP group must not block the next initialization."""
+
+ class _DestroyableGroup:
+ def __init__(self) -> None:
+ self.destroy_calls = 0
+
+ def destroy(self) -> None:
+ self.destroy_calls += 1
+
+ pipeline_group = _DestroyableGroup()
+ for name in ("_DP", "_CFG", "_SP"):
+ monkeypatch.setattr(parallel_state, name, None)
+ monkeypatch.setattr(parallel_state, "_PP", pipeline_group)
+ monkeypatch.setattr(parallel_state, "_EXPERT_PARALLEL_GROUP_RANKS", None)
+ for name in ("_DP", "_PCP", "_TP", "_EP"):
+ monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+ monkeypatch.setattr(parallel_state.vllm_parallel_state, "_PP", pipeline_group, raising=False)
+
+ parallel_state.destroy_model_parallel()
+
+ assert pipeline_group.destroy_calls == 1
+ assert parallel_state.vllm_parallel_state._PP is None
+
+
@pytest.mark.cpu
@pytest.mark.core_model
def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
@@ -295,9 +361,8 @@ def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
- monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
- for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+ for name in ("_DP", "_CFG", "_SP", "_PP", "_EXPERT_PARALLEL_GROUP_RANKS"):
monkeypatch.setattr(parallel_state, name, None)
for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
diff --git a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py
index e861bd61a..8de3a42f8 100644
--- a/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py
+++ b/tests/diffusion/distributed/test_flux2_vae_patch_parallel.py
@@ -53,14 +53,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
- if not hasattr(first_output, "request_output") or not first_output.request_output:
- raise ValueError("No request_output found in OmniRequestOutput")
+ if not isinstance(first_output, OmniRequestOutput):
+ raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
- req_out = first_output.request_output
- if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
- raise ValueError("Invalid request_output structure or missing 'images' key")
-
- images = req_out.images
+ images = first_output.images
if images is None or len(images) != 1:
raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
return images[0]
diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py
index 91c6784fc..25762c356 100644
--- a/tests/diffusion/distributed/test_hsdp.py
+++ b/tests/diffusion/distributed/test_hsdp.py
@@ -1,10 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
import gc
-import os
-import socket
import pytest
import torch
@@ -12,6 +10,7 @@ import torch.distributed as dist
import torch.nn as nn
from torch.distributed.tensor import DeviceMesh, DTensor
+from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.distributed.hsdp import (
HSDPInferenceConfig,
@@ -37,35 +36,17 @@ class _PackedModel(nn.Module):
self.root_weight = nn.Parameter(torch.ones(2))
-def _find_free_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
- sock.bind(("127.0.0.1", 0))
- return int(sock.getsockname()[1])
-
-
@pytest.fixture(scope="module")
def cpu_process_group():
if dist.is_initialized():
yield
return
- master_port = _find_free_port()
- os.environ.update(
- {
- "RANK": "0",
- "LOCAL_RANK": "0",
- "WORLD_SIZE": "1",
- "MASTER_ADDR": "127.0.0.1",
- "MASTER_PORT": str(master_port),
- }
- )
- dist.init_process_group("gloo", rank=0, world_size=1)
+ dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
try:
yield
finally:
dist.destroy_process_group()
- for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
- os.environ.pop(key, None)
gc.collect()
@@ -234,7 +215,7 @@ class TestDiffusionParallelConfigHSDP:
def test_hsdp_cannot_use_with_tp(self):
"""Test that HSDP and Tensor Parallelism cannot be used together."""
- with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+ with pytest.raises(ValueError, match="not compatible with TP"):
DiffusionParallelConfig(
tensor_parallel_size=2,
use_hsdp=True,
@@ -243,7 +224,7 @@ class TestDiffusionParallelConfigHSDP:
def test_hsdp_cannot_use_with_dp(self):
"""Test that HSDP and Data Parallelism cannot be used together."""
- with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+ with pytest.raises(ValueError, match="not compatible with DP"):
DiffusionParallelConfig(
data_parallel_size=2,
use_hsdp=True,
@@ -264,158 +245,6 @@ class TestDiffusionParallelConfigHSDP:
assert config.hsdp_shard_size == 2 # auto: 4 // 2
-class TestStandaloneHSDPDetection:
- """Tests for standalone HSDP detection and dit_parallel_size calculation.
-
- These tests verify the logic used in initialize_model_parallel() to detect
- standalone HSDP mode and calculate effective parallel sizes.
-
- Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
- """
-
- @staticmethod
- def compute_standalone_hsdp_params(
- data_parallel_size: int = 1,
- cfg_parallel_size: int = 1,
- sequence_parallel_size: int = 1,
- pipeline_parallel_size: int = 1,
- tensor_parallel_size: int = 1,
- fully_shard_degree: int = 1,
- hsdp_replicate_size: int = 1,
- ) -> dict:
- """Compute standalone HSDP detection parameters.
-
- This mirrors the logic in initialize_model_parallel().
- """
- dit_parallel_size = (
- data_parallel_size
- * cfg_parallel_size
- * sequence_parallel_size
- * pipeline_parallel_size
- * tensor_parallel_size
- )
-
- # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
- is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
-
- # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
- if is_standalone_hsdp:
- effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
- else:
- effective_dit_parallel_size = dit_parallel_size
-
- effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
-
- return {
- "original_dit_parallel_size": dit_parallel_size,
- "is_standalone_hsdp": is_standalone_hsdp,
- "effective_dit_parallel_size": effective_dit_parallel_size,
- "effective_dp_size": effective_dp_size,
- }
-
- def test_standalone_hsdp_basic(self):
- """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
- result = self.compute_standalone_hsdp_params(
- fully_shard_degree=4,
- hsdp_replicate_size=1,
- )
- assert result["original_dit_parallel_size"] == 1
- assert result["is_standalone_hsdp"] is True
- assert result["effective_dit_parallel_size"] == 4
- assert result["effective_dp_size"] == 4
-
- def test_standalone_hsdp_with_replicate(self):
- """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
- result = self.compute_standalone_hsdp_params(
- fully_shard_degree=4,
- hsdp_replicate_size=2,
- )
- assert result["original_dit_parallel_size"] == 1
- assert result["is_standalone_hsdp"] is True
- assert result["effective_dit_parallel_size"] == 8 # 4 * 2
- assert result["effective_dp_size"] == 8
-
- def test_combined_hsdp_sp_not_standalone(self):
- """Test HSDP combined with SP is NOT detected as standalone.
-
- This is a regression test for the bug where the condition
- `dit_parallel_size == fully_shard_degree` incorrectly matched
- combined modes like SP=4 + HSDP=4.
- """
- result = self.compute_standalone_hsdp_params(
- sequence_parallel_size=4,
- fully_shard_degree=4,
- hsdp_replicate_size=1,
- )
- assert result["original_dit_parallel_size"] == 4
- assert result["is_standalone_hsdp"] is False
- # Should NOT override dp_size for combined mode
- assert result["effective_dp_size"] == 1 # original data_parallel_size
-
- def test_combined_hsdp_cfg_not_standalone(self):
- """Test HSDP combined with CFG is NOT detected as standalone."""
- result = self.compute_standalone_hsdp_params(
- cfg_parallel_size=2,
- fully_shard_degree=4,
- hsdp_replicate_size=1,
- )
- assert result["original_dit_parallel_size"] == 2
- assert result["is_standalone_hsdp"] is False
- assert result["effective_dp_size"] == 1
-
- def test_combined_hsdp_pp_not_standalone(self):
- """Test HSDP combined with PP is NOT detected as standalone."""
- result = self.compute_standalone_hsdp_params(
- pipeline_parallel_size=2,
- fully_shard_degree=4,
- hsdp_replicate_size=1,
- )
- assert result["original_dit_parallel_size"] == 2
- assert result["is_standalone_hsdp"] is False
- assert result["effective_dp_size"] == 1
-
- def test_no_hsdp_not_standalone(self):
- """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
- result = self.compute_standalone_hsdp_params(
- fully_shard_degree=1,
- )
- assert result["original_dit_parallel_size"] == 1
- assert result["is_standalone_hsdp"] is False
- assert result["effective_dp_size"] == 1
-
- def test_combined_multiple_parallelism_not_standalone(self):
- """Test HSDP combined with multiple parallelism is NOT standalone."""
- result = self.compute_standalone_hsdp_params(
- sequence_parallel_size=2,
- cfg_parallel_size=2,
- fully_shard_degree=4,
- hsdp_replicate_size=1,
- )
- assert result["original_dit_parallel_size"] == 4 # 2 * 2
- assert result["is_standalone_hsdp"] is False
- assert result["effective_dp_size"] == 1
-
- def test_standalone_hsdp_large_shard(self):
- """Test standalone HSDP with large shard size."""
- result = self.compute_standalone_hsdp_params(
- fully_shard_degree=8,
- hsdp_replicate_size=1,
- )
- assert result["is_standalone_hsdp"] is True
- assert result["effective_dit_parallel_size"] == 8
- assert result["effective_dp_size"] == 8
-
- def test_standalone_hsdp_large_replicate(self):
- """Test standalone HSDP with large replicate size."""
- result = self.compute_standalone_hsdp_params(
- fully_shard_degree=4,
- hsdp_replicate_size=4,
- )
- assert result["is_standalone_hsdp"] is True
- assert result["effective_dit_parallel_size"] == 16 # 4 * 4
- assert result["effective_dp_size"] == 16
-
-
class TestHSDPShardConditions:
"""Tests for _hsdp_shard_conditions matching logic."""
diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py
index 0b48ed266..cccf116c3 100644
--- a/tests/diffusion/distributed/test_pipeline_parallel.py
+++ b/tests/diffusion/distributed/test_pipeline_parallel.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for PipelineParallelMixin and related PP communication helpers.
CPU unit tests exercise mixin logic without a process group. Distributed CPU
@@ -9,8 +9,6 @@ real multi-GPU NCCL collectives.
from __future__ import annotations
-import os
-import socket
from typing import Literal
import pytest
@@ -21,6 +19,7 @@ from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
@@ -40,31 +39,18 @@ DeviceKind = Literal["cpu", "cuda"]
_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-def _find_free_port() -> str:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- return str(s.getsockname()[1])
-
-
-def update_environment_variables(envs_dict: dict[str, str]) -> None:
- for k, v in envs_dict.items():
- os.environ[k] = v
-
-
def _cleanup_distributed() -> None:
- """Destroy omni distributed state and clear test-process leftovers.
+ """Destroy omni distributed state left dangling after a test.
``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
aliased to the destroyed group. Parent-process baselines must clear that
- dangling reference (and RANK/MASTER_* env) so later unit tests that call
+ dangling reference so later unit tests that call
``vllm.distributed.initialize_model_parallel`` do not see stale state.
"""
destroy_distributed_env()
import vllm.distributed.parallel_state as vllm_parallel_state
vllm_parallel_state._PP = None
- for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
- os.environ.pop(key, None)
# ---------------------------------------------------------------------------
@@ -416,24 +402,20 @@ def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
def init_dist(
local_rank: int,
world_size: int,
- master_port: str,
+ init_method: str,
device_kind: DeviceKind,
) -> torch.device:
"""Initialise the distributed environment for a spawned worker."""
device = _worker_device(local_rank, device_kind)
if device_kind == "cuda":
current_omni_platform.set_device(device)
- update_environment_variables(
- {
- "RANK": str(local_rank),
- "LOCAL_RANK": str(local_rank),
- "WORLD_SIZE": str(world_size),
- "MASTER_ADDR": "localhost",
- "MASTER_PORT": master_port,
- }
- )
backend = "gloo" if device_kind == "cpu" else None
- init_distributed_environment(backend=backend)
+ init_distributed_environment(
+ world_size=world_size,
+ rank=local_rank,
+ distributed_init_method=init_method,
+ backend=backend,
+ )
return device
@@ -493,11 +475,11 @@ def make_pipeline_and_inputs(
def isend_irecv_worker(
local_rank: int,
world_size: int,
- master_port: str,
+ init_method: str,
device_kind: DeviceKind,
result_queue,
):
- device = init_dist(local_rank, world_size, master_port, device_kind)
+ device = init_dist(local_rank, world_size, init_method, device_kind)
initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
pp_group = get_pp_group()
@@ -517,13 +499,13 @@ def isend_irecv_worker(
_cleanup_distributed()
-def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
+def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
torch.multiprocessing.spawn(
isend_irecv_worker,
- args=(pp_size, master_port, device_kind, q),
+ args=(pp_size, init_method, device_kind, q),
nprocs=pp_size,
)
results = {label: tensor for label, tensor in [q.get(), q.get()]}
@@ -538,7 +520,7 @@ def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) ->
@pytest.mark.parametrize("pp_size", [2])
def test_isend_irecv_tensor_dict(pp_size: int):
"""isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
- _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
+ _run_isend_irecv(pp_size, device_kind="cpu", init_method=get_distributed_init_method())
@pytest.mark.full_model
@@ -547,7 +529,7 @@ def test_isend_irecv_tensor_dict(pp_size: int):
@pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
def test_isend_irecv_tensor_dict_parity(pp_size: int):
"""Nightly: isend/irecv on real multi-GPU NCCL."""
- _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
+ _run_isend_irecv(pp_size, device_kind="cuda", init_method=get_distributed_init_method())
# ---------------------------------------------------------------------------
@@ -578,7 +560,7 @@ def compute_single_gpu_baseline(
if key in _baseline_cache:
return _baseline_cache[key]
- device = init_dist(0, 1, _find_free_port(), device_kind)
+ device = init_dist(0, 1, get_distributed_init_method(), device_kind)
try:
initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
@@ -604,7 +586,7 @@ def compute_single_gpu_baseline(
def predict_noise_worker(
local_rank: int,
world_size: int,
- master_port: str,
+ init_method: str,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
@@ -614,7 +596,7 @@ def predict_noise_worker(
result_queue,
):
"""Generic predict-noise worker parameterized by PP and CFG topology."""
- device = init_dist(local_rank, world_size, master_port, device_kind)
+ device = init_dist(local_rank, world_size, init_method, device_kind)
initialize_model_parallel(
pipeline_parallel_size=pp_size,
cfg_parallel_size=cfg_size,
@@ -661,7 +643,7 @@ def _run_predict_noise(
rtol: float,
atol: float,
device_kind: DeviceKind,
- master_port: str,
+ init_method: str,
) -> None:
test_config = {
"num_layers": num_layers,
@@ -681,7 +663,7 @@ def _run_predict_noise(
world_size = pp_size * cfg_size
torch.multiprocessing.spawn(
predict_noise_worker,
- args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
+ args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
nprocs=world_size,
)
@@ -786,7 +768,7 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, r
rtol=rtol,
atol=atol,
device_kind="cpu",
- master_port=_find_free_port(),
+ init_method=get_distributed_init_method(),
)
@@ -809,7 +791,7 @@ def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers,
rtol=rtol,
atol=atol,
device_kind="cuda",
- master_port=_find_free_port(),
+ init_method=get_distributed_init_method(),
)
@@ -824,7 +806,7 @@ def compute_scheduler_step_baseline(
device_kind: DeviceKind,
) -> torch.Tensor:
"""Single-process reference: predict_noise + scheduler_step."""
- device = init_dist(0, 1, _find_free_port(), device_kind)
+ device = init_dist(0, 1, get_distributed_init_method(), device_kind)
try:
initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
@@ -853,7 +835,7 @@ def compute_scheduler_step_baseline(
def scheduler_step_worker(
local_rank: int,
world_size: int,
- master_port: str,
+ init_method: str,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
@@ -861,7 +843,7 @@ def scheduler_step_worker(
device_kind: DeviceKind,
result_queue,
):
- device = init_dist(local_rank, world_size, master_port, device_kind)
+ device = init_dist(local_rank, world_size, init_method, device_kind)
initialize_model_parallel(
pipeline_parallel_size=pp_size,
cfg_parallel_size=cfg_size,
@@ -905,7 +887,7 @@ def _run_scheduler_step(
do_true_cfg: bool,
input_seed: int,
device_kind: DeviceKind,
- master_port: str,
+ init_method: str,
) -> None:
test_config = {
"num_layers": 4,
@@ -925,7 +907,7 @@ def _run_scheduler_step(
world_size = pp_size * cfg_size
torch.multiprocessing.spawn(
scheduler_step_worker,
- args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
+ args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
nprocs=world_size,
)
@@ -957,7 +939,7 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
do_true_cfg=do_true_cfg,
input_seed=input_seed,
device_kind="cpu",
- master_port=_find_free_port(),
+ init_method=get_distributed_init_method(),
)
@@ -979,5 +961,5 @@ def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
do_true_cfg=do_true_cfg,
input_seed=input_seed,
device_kind="cuda",
- master_port=_find_free_port(),
+ init_method=get_distributed_init_method(),
)
diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py
index d43d18c49..f083e1d26 100644
--- a/tests/diffusion/distributed/test_sequence_parallel.py
+++ b/tests/diffusion/distributed/test_sequence_parallel.py
@@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str):
@pytest.mark.skipif(
- not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
- reason="Only tested on CUDA and XPU",
+ not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+ reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.skipif(
- not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
- reason="Only tested on CUDA and XPU",
+ not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+ reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No
@pytest.mark.skipif(
- not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
- reason="Only tested on CUDA and XPU",
+ not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
+ reason="Only tested on CUDA, ROCM and XPU",
)
@pytest.mark.diffusion
@pytest.mark.parallel
diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py
index f95dbfe2b..4ef564571 100644
--- a/tests/diffusion/distributed/test_sp_plan_hooks.py
+++ b/tests/diffusion/distributed/test_sp_plan_hooks.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the Sequence Parallelism (SP) framework.
These tests verify the SP plan mechanism and hooks work correctly without
@@ -17,6 +17,7 @@ import pytest
import torch
import torch.nn as nn
+from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelInput,
SequenceParallelOutput,
@@ -419,7 +420,6 @@ class TestHookRegistration:
# =============================================================================
-@pytest.mark.L4
class TestModelSpPlans:
"""Test that model _sp_plan definitions are valid.
@@ -427,6 +427,8 @@ class TestModelSpPlans:
May require GPU for model imports.
"""
+ pytestmark = hardware_marks(res={"cuda": "L4"})
+
def test_zimage_transformer_sp_plan(self):
"""Test ZImageTransformer2DModel _sp_plan structure.
diff --git a/tests/diffusion/distributed/test_tensor_parallel.py b/tests/diffusion/distributed/test_tensor_parallel.py
index c50c4336d..ae97c9900 100644
--- a/tests/diffusion/distributed/test_tensor_parallel.py
+++ b/tests/diffusion/distributed/test_tensor_parallel.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""DiT tensor-parallel parity and perf checks for diffusion models.
@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
from __future__ import annotations
+import gc
import os
import time
from pathlib import Path
@@ -18,8 +19,8 @@ import pytest
import torch
from PIL import Image
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from tests.helpers.runtime import OmniRunner
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -75,8 +76,13 @@ def _run_generate(
if num_requests < 2:
raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
+ gc.collect()
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
+ with current_omni_platform.device(device_index):
+ free_bytes, total_bytes = current_omni_platform.mem_get_info()
+ initial_used_mb = (total_bytes - free_bytes) / (1024**2)
+
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
try:
@@ -113,7 +119,10 @@ def _run_generate(
pass
median_time_s = float(np.median(per_request_times_s))
- peak_memory_mb = monitor.peak_used_mb
+ # DeviceMemoryMonitor samples absolute device usage. Normalize by
+ # this run's starting usage so compiler and allocator caches retained
+ # by the preceding TP=1 run are not charged to the TP=2 run.
+ peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
return _extract_single_image([last_output]), median_time_s, peak_memory_mb
finally:
monitor.stop()
diff --git a/tests/diffusion/distributed/test_ulysses_uaa_perf.py b/tests/diffusion/distributed/test_ulysses_uaa_perf.py
index 2a16a9ae5..707800cc0 100644
--- a/tests/diffusion/distributed/test_ulysses_uaa_perf.py
+++ b/tests/diffusion/distributed/test_ulysses_uaa_perf.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Perf smoke tests for Ulysses advanced_uaa communication overhead.
This test is intended for CI monitoring only:
@@ -9,8 +9,6 @@ This test is intended for CI monitoring only:
from __future__ import annotations
-import os
-import socket
from dataclasses import dataclass
import pytest
@@ -18,6 +16,7 @@ import torch
import torch.distributed as dist
from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.attention.parallel.ulysses import (
_all_gather_int,
_ulysses_all_to_all_any_o,
@@ -33,20 +32,6 @@ from vllm_omni.diffusion.distributed.parallel_state import (
from vllm_omni.platforms import current_omni_platform
-def _find_free_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
- os.environ["RANK"] = str(rank)
- os.environ["LOCAL_RANK"] = str(rank)
- os.environ["WORLD_SIZE"] = str(world_size)
- os.environ["MASTER_ADDR"] = "127.0.0.1"
- os.environ["MASTER_PORT"] = str(master_port)
-
-
def _max_all_reduce(pg: dist.ProcessGroup, value: float, *, device: torch.device) -> float:
t = torch.tensor([value], device=device, dtype=torch.float32)
dist.all_reduce(t, op=dist.ReduceOp.MAX, group=pg)
@@ -77,22 +62,20 @@ def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
if available_gpus < case.world_size:
pytest.skip(f"Requires {case.world_size} GPUs, got {available_gpus}")
- master_port = _find_free_port()
+ init_method = get_distributed_init_method()
torch.multiprocessing.spawn(
_perf_worker,
- args=(case.world_size, master_port, case.ulysses_degree, case.ring_degree),
+ args=(case.world_size, init_method, case.ulysses_degree, case.ring_degree),
nprocs=case.world_size,
)
-def _perf_worker(local_rank: int, world_size: int, master_port: int, ulysses_degree: int, ring_degree: int) -> None:
+def _perf_worker(local_rank: int, world_size: int, init_method: str, ulysses_degree: int, ring_degree: int) -> None:
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
- _set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port)
-
try:
- init_distributed_environment(world_size=world_size, rank=local_rank)
+ init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
initialize_model_parallel(
data_parallel_size=1,
cfg_parallel_size=1,
diff --git a/tests/diffusion/distributed/test_vae_decode_parallelism.py b/tests/diffusion/distributed/test_vae_decode_parallelism.py
index c3996fc23..28ded4639 100644
--- a/tests/diffusion/distributed/test_vae_decode_parallelism.py
+++ b/tests/diffusion/distributed/test_vae_decode_parallelism.py
@@ -72,8 +72,7 @@ def _to_float_array(data: Any) -> np.ndarray:
def _extract_output_array(output: Any) -> np.ndarray:
- req_out = output.request_output
- payload = req_out.images[0]
+ payload = output.images[0]
return _to_float_array(payload)
diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py
index babbc7abc..49622aecd 100644
--- a/tests/diffusion/kernels/mot/test_mot_linear.py
+++ b/tests/diffusion/kernels/mot/test_mot_linear.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
# ruff: noqa: N803, E741
"""Layer-level correctness & performance test for MoT parallel linear layers.
@@ -237,6 +240,53 @@ def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
return output
+def _assert_gen_expert_parallel_metadata(layer):
+ """The secondary expert must satisfy vLLM's parallel-linear contract."""
+ for attr in (
+ "tp_rank",
+ "tp_size",
+ "input_size",
+ "input_size_per_partition",
+ "output_size",
+ "output_size_per_partition",
+ ):
+ assert getattr(layer.gen_exp, attr) == getattr(layer, attr)
+
+
+def test_mot_gen_expert_parallel_metadata():
+ """Keep secondary experts compatible with vLLM online quantizers."""
+ from vllm.model_executor.layers.quantization.online.fp8 import (
+ _is_tp_sharded,
+ )
+
+ with set_current_vllm_config(VllmConfig()):
+ layers = (
+ MoTQKVParallelLinear(
+ hidden_size=16,
+ head_size=4,
+ total_num_heads=4,
+ total_num_kv_heads=2,
+ bias=False,
+ vae_bias=False,
+ params_dtype=torch.bfloat16,
+ disable_tp=True,
+ ),
+ MoTRowParallelLinear(
+ 16,
+ 8,
+ bias=False,
+ vae_bias=False,
+ input_is_parallel=True,
+ params_dtype=torch.bfloat16,
+ disable_tp=True,
+ ),
+ )
+
+ for layer in layers:
+ _assert_gen_expert_parallel_metadata(layer)
+ assert _is_tp_sharded(layer.gen_exp) is False
+
+
def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
"""Compare outputs, print metrics, and assert correctness.
@@ -335,6 +385,8 @@ def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool
disable_tp=True,
).cuda()
+ _assert_gen_expert_parallel_metadata(mot_linear)
+
assert text_linear.output_size_per_partition == N, (
f"Expected output_size_per_partition={N}, "
f"got {text_linear.output_size_per_partition}. "
@@ -423,6 +475,8 @@ def test_mot_o_proj(
disable_tp=True,
).cuda()
+ _assert_gen_expert_parallel_metadata(mot_linear)
+
_sync_weights(text_linear, vae_linear, mot_linear)
text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
@@ -501,6 +555,13 @@ def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda()
with torch.no_grad():
+ # vLLM linear layers leave weights uninitialized (torch.empty), so
+ # a freshly zeroed GPU page would make both outputs zero and the
+ # cosine-similarity gate degenerate. Seed deterministic values so
+ # und-mode is compared against a meaningful reference.
+ ref_linear.weight.normal_(mean=0.0, std=0.02)
+ if ref_linear.bias is not None:
+ ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias)
@@ -549,6 +610,13 @@ def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
).cuda()
with torch.no_grad():
+ # vLLM linear layers leave weights uninitialized (torch.empty), so
+ # a freshly zeroed GPU page would make both outputs zero and the
+ # cosine-similarity gate degenerate. Seed deterministic values so
+ # und-mode is compared against a meaningful reference.
+ ref_linear.weight.normal_(mean=0.0, std=0.02)
+ if ref_linear.bias is not None:
+ ref_linear.bias.normal_(mean=0.0, std=0.02)
mot_linear.weight.copy_(ref_linear.weight)
if bias and ref_linear.bias is not None:
mot_linear.bias.copy_(ref_linear.bias)
diff --git a/tests/diffusion/layers/test_activation.py b/tests/diffusion/layers/test_activation.py
new file mode 100644
index 000000000..c3205502f
--- /dev/null
+++ b/tests/diffusion/layers/test_activation.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.layers.activation import SiluAndMul
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.cpu
+def test_silu_and_mul_native_matches_packed_reference() -> None:
+ packed = torch.randn(257, 256, dtype=torch.float32)
+ gate, up = packed.chunk(2, dim=-1)
+
+ torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)
+
+
+@hardware_test(res={"npu": "A3"}, num_cards=1)
+def test_silu_and_mul_npu_matches_packed_reference() -> None:
+ packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16)
+ gate, up = packed.chunk(2, dim=-1)
+
+ torch.testing.assert_close(
+ SiluAndMul()(packed),
+ F.silu(gate) * up,
+ atol=2e-2,
+ rtol=2e-2,
+ )
diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py
new file mode 100644
index 000000000..40888556a
--- /dev/null
+++ b/tests/diffusion/layers/test_fused_qk_norm_rope.py
@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+_HEAD_DIM = 128
+_ROTARY_DIM = 96
+_EPS = 1e-5
+
+
+def _reference(q, k, q_weight, k_weight, rope_table):
+ q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
+ k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
+ half = _ROTARY_DIM // 2
+ cos = rope_table[..., :half].unsqueeze(1)
+ sin = rope_table[..., half:].unsqueeze(1)
+
+ def apply(x):
+ first = x[..., :half]
+ second = x[..., half:_ROTARY_DIM]
+ return torch.cat(
+ (
+ first * cos - second * sin,
+ second * cos + first * sin,
+ x[..., _ROTARY_DIM:],
+ ),
+ dim=-1,
+ )
+
+ return apply(q), apply(k)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("seq_len", [1, 257, 1024])
+def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
+ from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+ fused_qk_norm_rope,
+ )
+
+ torch.manual_seed(17)
+ heads = 14
+ qkv = torch.randn(
+ seq_len,
+ heads * _HEAD_DIM * 3,
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
+ k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
+ seq_len,
+ heads,
+ _HEAD_DIM,
+ )
+ q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+ k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+ freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
+ rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
+
+ expected_q, expected_k = _reference(
+ q,
+ k,
+ q_weight,
+ k_weight,
+ rope_table,
+ )
+ actual_q, actual_k = fused_qk_norm_rope(
+ q,
+ k,
+ q_weight,
+ k_weight,
+ rope_table,
+ _EPS,
+ )
+
+ torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+ torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
new file mode 100644
index 000000000..35d990afa
--- /dev/null
+++ b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+import types
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+_HEAD_DIM = 128
+_ROTARY_DIM = 96
+_EPS = 1e-5
+
+
+def _reference(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ q_weight: torch.Tensor,
+ k_weight: torch.Tensor,
+ rope_table: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ half = _ROTARY_DIM // 2
+
+ def apply(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+ x = F.rms_norm(x, (_HEAD_DIM,), weight, _EPS)
+ cos = rope_table[..., :half].unsqueeze(1)
+ sin = rope_table[..., half:].unsqueeze(1)
+ return torch.cat(
+ (
+ x[..., :half] * cos - x[..., half:_ROTARY_DIM] * sin,
+ x[..., half:_ROTARY_DIM] * cos + x[..., :half] * sin,
+ x[..., _ROTARY_DIM:],
+ ),
+ dim=-1,
+ )
+
+ return apply(q, q_weight), apply(k, k_weight)
+
+
+def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ torch.manual_seed(7)
+ tokens, heads = 11, 4
+ q = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
+ k = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16)
+ q_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
+ k_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16)
+ freqs = torch.randn(tokens, _ROTARY_DIM // 2)
+ rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
+ return q, k, q_weight, k_weight, rope_table
+
+
+@pytest.mark.cpu
+def test_npu_qk_norm_rope_uses_torch_npu_fused_primitives_without_mindiesd(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
+
+ calls: dict[str, int] = {"rms_norm": 0, "rotary_mul": 0}
+
+ def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
+ calls["rms_norm"] += 1
+ return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
+
+ def npu_rotary_mul(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+ calls["rotary_mul"] += 1
+ assert x.shape[-1] == _ROTARY_DIM
+ assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
+ assert sin.shape == cos.shape
+ half = _ROTARY_DIM // 2
+ return torch.cat(
+ (
+ x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
+ x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
+ ),
+ dim=-1,
+ )
+
+ monkeypatch.setitem(
+ sys.modules,
+ "torch_npu",
+ types.SimpleNamespace(npu_rms_norm=npu_rms_norm, npu_rotary_mul=npu_rotary_mul),
+ )
+ monkeypatch.setattr(fused, "find_spec", lambda _: None)
+ q, k, q_weight, k_weight, rope_table = _inputs()
+
+ actual_q, actual_k = fused._npu_qk_norm_rope(
+ q,
+ k,
+ q_weight,
+ k_weight,
+ rope_table,
+ _EPS,
+ _ROTARY_DIM,
+ )
+ expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
+
+ assert calls == {"rms_norm": 2, "rotary_mul": 2}
+ torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
+ torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
+
+
+@pytest.mark.cpu
+def test_npu_qk_norm_rope_uses_mindiesd_and_normalizes_packed_layout(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
+
+ calls: dict[str, int] = {"rms_norm": 0, "mindiesd_rope": 0}
+
+ def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
+ calls["rms_norm"] += 1
+ return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),)
+
+ def rotary_position_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs):
+ calls["mindiesd_rope"] += 1
+ assert x.dim() == 4
+ assert x.shape[0] == 1
+ assert x.shape[-1] == _ROTARY_DIM
+ assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM)
+ assert sin.shape == cos.shape
+ assert kwargs == {
+ "rotated_mode": "rotated_half",
+ "head_first": False,
+ "fused": True,
+ }
+ half = _ROTARY_DIM // 2
+ return torch.cat(
+ (
+ x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half],
+ x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:],
+ ),
+ dim=-1,
+ )
+
+ monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
+ monkeypatch.setitem(
+ sys.modules,
+ "mindiesd",
+ types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
+ )
+ monkeypatch.setattr(fused, "find_spec", lambda _: object())
+ q, k, q_weight, k_weight, rope_table = _inputs()
+
+ actual_q, actual_k = fused._npu_qk_norm_rope(
+ q,
+ k,
+ q_weight,
+ k_weight,
+ rope_table,
+ _EPS,
+ _ROTARY_DIM,
+ )
+ expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
+
+ assert calls == {"rms_norm": 2, "mindiesd_rope": 2}
+ torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0)
+ torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
+
+
+@hardware_test(res={"npu": "A3"}, num_cards=1)
+def test_fused_qk_norm_rope_npu_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Exercise the public Ascend dispatch with real torch_npu/MindIE operators."""
+ from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused
+
+ q, k, q_weight, k_weight, rope_table = (tensor.to("npu") for tensor in _inputs())
+ calls = 0
+ original_npu_impl = fused._npu_qk_norm_rope
+
+ def capture_npu_impl(*args, **kwargs):
+ nonlocal calls
+ calls += 1
+ return original_npu_impl(*args, **kwargs)
+
+ monkeypatch.setattr(fused, "_npu_qk_norm_rope", capture_npu_impl)
+ actual_q, actual_k = fused.fused_qk_norm_rope(
+ q,
+ k,
+ q_weight,
+ k_weight,
+ rope_table,
+ _EPS,
+ head_dim=_HEAD_DIM,
+ rotary_dim=_ROTARY_DIM,
+ )
+ expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table)
+
+ assert calls == 1
+ torch.testing.assert_close(actual_q, expected_q, atol=2e-2, rtol=2e-2)
+ torch.testing.assert_close(actual_k, expected_k, atol=2e-2, rtol=2e-2)
diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py
index 71125684b..4794ca1c2 100644
--- a/tests/diffusion/layers/test_norm.py
+++ b/tests/diffusion/layers/test_norm.py
@@ -340,6 +340,97 @@ def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.Mo
assert captured["epsilon"].item() == pytest.approx(1e-5)
+# ── RMSNorm residual tests ──
+
+
+def test_rmsnorm_matches_native_residual_contract():
+ """The non-NPU path preserves separate residual-add and RMSNorm math."""
+ from vllm_omni.diffusion.layers.norm import RMSNorm
+
+ eps = 1e-6
+ norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
+ norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
+ x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
+ residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)
+
+ expected_residual = residual + x
+ expected_fp32 = expected_residual.float()
+ expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
+ expected_output = (expected_output * norm.weight.float()).to(x.dtype)
+ output, updated_residual = norm.forward_native(x, residual)
+
+ torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
+ torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+ "forward_name",
+ ["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
+)
+def test_rmsnorm_non_npu_residual_delegates_to_native(
+ monkeypatch: pytest.MonkeyPatch,
+ forward_name: str,
+):
+ """Platform wrappers keep residual arithmetic in forward_native."""
+ from vllm_omni.diffusion.layers.norm import RMSNorm
+
+ norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
+ x = torch.randn(2, 4, dtype=torch.bfloat16)
+ residual = torch.randn_like(x)
+ expected_output = torch.randn_like(x)
+ expected_residual = torch.randn_like(x)
+ captured: dict[str, object] = {}
+
+ def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
+ captured.update(x=native_x, residual=native_residual)
+ return expected_output, expected_residual
+
+ monkeypatch.setattr(norm, "forward_native", forward_native)
+
+ output, updated_residual = getattr(norm, forward_name)(x, residual)
+
+ assert output is expected_output
+ assert updated_residual is expected_residual
+ assert captured["x"] is x
+ assert captured["residual"] is residual
+
+
+def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
+ """NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
+ from vllm_omni.diffusion.layers.norm import RMSNorm
+
+ captured: dict[str, object] = {}
+ fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
+ fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)
+
+ def npu_add_rms_norm(
+ x: torch.Tensor,
+ residual: torch.Tensor,
+ gamma: torch.Tensor,
+ epsilon: float,
+ ):
+ captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
+ return fused_output, torch.empty(0), fused_residual
+
+ monkeypatch.setitem(
+ sys.modules,
+ "torch_npu",
+ types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
+ )
+ norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
+ x = torch.randn(2, 4, dtype=torch.bfloat16)
+ residual = torch.randn(2, 4, dtype=torch.bfloat16)
+
+ output, updated_residual = norm.forward_npu(x, residual)
+
+ assert output is fused_output
+ assert updated_residual is fused_residual
+ assert captured["x"] is x
+ assert captured["residual"] is residual
+ assert captured["gamma"] is norm.weight
+ assert captured["epsilon"] == 1e-5
+
+
# ── RMSNorm compile-path regression tests ──
diff --git a/tests/diffusion/layers/test_norm_cuda.py b/tests/diffusion/layers/test_norm_cuda.py
new file mode 100644
index 000000000..b791eab71
--- /dev/null
+++ b/tests/diffusion/layers/test_norm_cuda.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+ pytest.mark.core_model,
+ pytest.mark.diffusion,
+ pytest.mark.cuda,
+ pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
+]
+
+
+def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
+ eps = 1e-6
+ device = torch.device(current_omni_platform.device_type)
+ torch.manual_seed(0)
+ norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
+ x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
+ residual = torch.randn_like(x)
+ x_before = x.clone()
+ residual_before = residual.clone()
+
+ expected_output, expected_residual = norm.forward_native(x_before, residual_before)
+ output, updated_residual = norm.forward_cuda(x, residual)
+
+ torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
+ torch.testing.assert_close(output, expected_output, atol=0, rtol=0)
+ torch.testing.assert_close(x, x_before, atol=0, rtol=0)
+ torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)
diff --git a/tests/diffusion/lora/test_base_linear.py b/tests/diffusion/lora/test_base_linear.py
index 9386d0909..b2f76be05 100644
--- a/tests/diffusion/lora/test_base_linear.py
+++ b/tests/diffusion/lora/test_base_linear.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -89,6 +89,29 @@ def test_diffusion_base_linear_apply_multi_slice():
assert torch.allclose(out, expected)
+def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
+ from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+ layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+ torch.nn.Module.__init__(layer)
+ layer.base_layer = torch.nn.Module()
+ layer.n_slices = 1
+ layer.lora_a_stacked = (torch.zeros(2),)
+ layer.lora_b_stacked = (torch.zeros(3),)
+
+ layer._set_diffusion_lora_buffer_device(torch.device("meta"))
+ assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+ def recreate_on_cpu(self, max_loras, lora_config, model_config):
+ self.lora_a_stacked = (torch.zeros(4),)
+ self.lora_b_stacked = (torch.zeros(5),)
+
+ monkeypatch.setattr(BaseLinearLayerWithLoRA, "create_lora_weights", recreate_on_cpu)
+ layer.create_lora_weights(max_loras=1, lora_config=object())
+
+ assert all(tensor.is_meta for tensor in (*layer.lora_a_stacked, *layer.lora_b_stacked))
+
+
def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
# Verify that after reset_lora(), apply() skips LoRA matmuls even if the
# LoRA tensors are still allocated and non-empty.
diff --git a/tests/diffusion/lora/test_images_generations_lora.py b/tests/diffusion/lora/test_images_generations_lora.py
index dbdf59647..a01bcdae3 100644
--- a/tests/diffusion/lora/test_images_generations_lora.py
+++ b/tests/diffusion/lora/test_images_generations_lora.py
@@ -24,11 +24,20 @@ from safetensors.torch import save_file
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OmniServerParams
+from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "Tongyi-MAI/Z-Image-Turbo"
+# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI
+# worker, JIT-building its MHA kernel can take longer than the server readiness
+# timeout. This test validates LoRA switching rather than attention-backend or
+# compiled-execution performance, so use the eager SDPA path on ROCm only.
+ROCM_SERVER_ARGS = (
+ ["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else []
+)
+
PROMPT = "a photo of a cat sitting on a laptop keyboard"
SIZE = "256x256"
@@ -42,7 +51,14 @@ server_params = [
server_args=[
"--num-gpus",
"1",
+ *ROCM_SERVER_ARGS,
],
+ # A cold-cache Z-Image load can take slightly more than the
+ # orchestrator's 900s default on CI (weights alone have been
+ # observed at ~690s and full orchestrator readiness at ~906s).
+ # Match the large-model offline runner's startup allowance without
+ # weakening timeout detection for every online-serving test.
+ init_timeout=1800,
),
id="zimage_turbo_lora",
),
diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py
index 723e4b13a..8edc6ce4b 100644
--- a/tests/diffusion/lora/test_lora_manager.py
+++ b/tests/diffusion/lora/test_lora_manager.py
@@ -1,8 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
+from types import SimpleNamespace
+
import pytest
import torch
from vllm.lora.lora_weights import LoRALayerWeights
@@ -20,9 +22,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _DummyLoRALayer:
- def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
+ def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
self.n_slices = n_slices
self.output_slices = output_slices
+ self.tp_size = tp_size
self.set_calls: list[
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
] = []
@@ -142,6 +145,37 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
assert replace_calls == ["foo"]
+def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
+ import vllm_omni.diffusion.lora.manager as manager_mod
+
+ class _DLOLoRALayer(DummyBaseLayerWithLoRA):
+ lora_buffer_device = None
+
+ def _set_diffusion_lora_buffer_device(self, device):
+ self.lora_buffer_device = device
+
+ monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
+ monkeypatch.setattr(
+ manager_mod,
+ "from_layer_diffusion",
+ lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
+ )
+
+ pipeline = torch.nn.Module()
+ pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
+ pipeline.transformer = torch.nn.Module()
+ pipeline.transformer.foo = _FakeLinearBase()
+ manager = DiffusionLoRAManager(
+ pipeline=pipeline,
+ device=torch.device("cpu"),
+ dtype=torch.bfloat16,
+ )
+
+ manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
+
+ assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
+
+
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
import vllm_omni.diffusion.lora.manager as manager_mod
@@ -236,6 +270,14 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
def test_lora_manager_activates_packed_lora_from_sublayers():
pipeline = torch.nn.Module()
+ bound_names = None
+
+ def validate_binding(*, lora_model, bound_lora_names):
+ nonlocal bound_names
+ assert lora_model.id == 1
+ bound_names = bound_lora_names
+
+ pipeline._validate_diffusion_lora_binding = validate_binding
pipeline.stacked_params_mapping = [
(".to_qkv", ".to_q", "q"),
(".to_qkv", ".to_k", "k"),
@@ -279,6 +321,70 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+ assert bound_names == frozenset(loras)
+
+
+def test_lora_manager_rolls_back_all_layers_when_activation_fails():
+ class _FailingLoRALayer(_DummyLoRALayer):
+ fail = False
+
+ def set_lora(self, index: int, lora_a, lora_b):
+ if self.fail:
+ raise RuntimeError("bind failed")
+ super().set_lora(index, lora_a, lora_b)
+
+ manager = DiffusionLoRAManager(
+ pipeline=torch.nn.Module(),
+ device=torch.device("cpu"),
+ dtype=torch.bfloat16,
+ max_cached_adapters=2,
+ )
+ first = _DummyLoRALayer(n_slices=1, output_slices=(2,))
+ second = _FailingLoRALayer(n_slices=1, output_slices=(2,))
+ manager._lora_modules = {
+ "transformer.first": first,
+ "transformer.second": second,
+ }
+
+ def adapter(adapter_id: int):
+ loras = {
+ name: LoRALayerWeights(
+ module_name=name,
+ rank=2,
+ lora_alpha=2,
+ lora_a=torch.full((2, 2), float(adapter_id)),
+ lora_b=torch.full((2, 2), float(adapter_id)),
+ )
+ for name in manager._lora_modules
+ }
+ return type(
+ "LM",
+ (),
+ {
+ "id": adapter_id,
+ "loras": loras,
+ "get_lora": lambda self, key: self.loras.get(key),
+ },
+ )()
+
+ manager._registered_adapters = {1: adapter(1), 2: adapter(2)}
+ manager._activate_adapter(1, scale=1.0)
+ first_calls_after_success = len(first.set_calls)
+
+ second.fail = True
+ with pytest.raises(RuntimeError, match="bind failed"):
+ manager._activate_adapter(2, scale=1.0)
+
+ assert manager._active_adapter_id is None
+ assert first.reset_calls == 1
+ assert second.reset_calls == 1
+
+ # The failed activation must not leave the old adapter eligible for the
+ # fast path: reactivating it binds every layer again.
+ second.fail = False
+ manager._activate_adapter(1, scale=1.0)
+ assert manager._active_adapter_id == 1
+ assert len(first.set_calls) == first_calls_after_success + 2
def _dummy_lora_request(adapter_id: int) -> LoRARequest:
diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py
index f370d5b49..bd59ee94b 100644
--- a/tests/diffusion/model_loader/test_diffusers_loader.py
+++ b/tests/diffusion/model_loader/test_diffusers_loader.py
@@ -1,22 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Tests for the DiffusersPipelineLoader.
"""
+import json
+from pathlib import Path
from types import SimpleNamespace
import pytest
import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
+from safetensors.torch import save_file
from vllm.config.load import LoadConfig
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+ HostWeightPlan,
+ HostWeightPlanResult,
+ TensorBinding,
+)
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
from vllm_omni.diffusion.models.helios import HeliosPipeline
+from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
from vllm_omni.diffusion.registry import initialize_model
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -67,6 +77,73 @@ class _DummyPipelineModel(nn.Module):
return loaded
+class _HWRTransformer(nn.Module):
+ host_weight_restore_contract = FinalLayoutModelContract(
+ implementation_id="test-hwr-transformer",
+ version="1",
+ )
+
+ def __init__(self):
+ super().__init__()
+ self.weight = nn.Parameter(torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))
+
+ def validate_restored_host_weights(self):
+ assert self.weight.dtype is torch.bfloat16
+
+
+class _HWRPipeline(nn.Module):
+ def __init__(self, source_root: Path):
+ super().__init__()
+ self.transformer = _HWRTransformer()
+ self.load_count = 0
+ self.weights_sources = [
+ DiffusersPipelineLoader.ComponentSource(
+ model_or_path=str(source_root),
+ subfolder=None,
+ revision=None,
+ prefix="transformer.",
+ fall_back_to_pt=False,
+ )
+ ]
+
+ def load_weights(self, weights):
+ loaded: set[str] = set()
+ params = dict(self.named_parameters())
+ for name, tensor in weights:
+ if name in params:
+ params[name].data.copy_(tensor.to(dtype=params[name].dtype))
+ loaded.add(name)
+ self.load_count += 1
+ return loaded
+
+
+def _hwr_config(model: str | Path, root: Path, *, mode: str = "preferred") -> SimpleNamespace:
+ return SimpleNamespace(
+ model=str(model),
+ dtype=torch.bfloat16,
+ host_weight_runtime_mode=mode,
+ host_weight_runtime_root=str(root),
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ lora_path=None,
+ quantization_config=None,
+ diffusion_attention_config=None,
+ parallel_config=SimpleNamespace(
+ use_hsdp=False,
+ data_parallel_size=1,
+ sequence_parallel_size=1,
+ tensor_parallel_size=1,
+ pipeline_parallel_size=1,
+ cfg_parallel_size=1,
+ enable_expert_parallel=False,
+ ulysses_degree=1,
+ ring_degree=1,
+ allgather_degree=1,
+ ulysses_mode="strict",
+ ),
+ )
+
+
def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
od_config = SimpleNamespace(
dtype=torch.float32,
@@ -86,6 +163,322 @@ def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoade
return loader
+def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading(
+ tmp_path: Path,
+ monkeypatch,
+):
+ canonical_root = tmp_path / "canonical"
+ canonical_root.mkdir()
+ save_file(
+ {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+ str(canonical_root / "model.safetensors"),
+ )
+ store_root = tmp_path / "hwr-store"
+ hash_calls = 0
+ original_sha256 = source_identity_module._sha256_file
+
+ def counted_sha256(path: Path, state: object) -> str:
+ nonlocal hash_calls
+ hash_calls += 1
+ return original_sha256(path, state) # type: ignore[arg-type]
+
+ monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
+
+ def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]:
+ loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root))
+ pipeline = _HWRPipeline(canonical_root)
+ monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+ return loader, pipeline
+
+ cold_loader, cold_model = make_loader()
+ cold = cold_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+ assert cold is cold_model
+ assert cold_model.load_count == 1
+ assert cold_loader._hwr_state is not None
+
+ warm_loader, warm_model = make_loader()
+ monkeypatch.setattr(
+ warm_loader,
+ "_process_weights_after_loading",
+ lambda *args, **kwargs: pytest.fail("warm HWR restore re-entered byte-changing finalization"),
+ )
+ warm = warm_loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+ assert warm is warm_model
+ assert warm_model.load_count == 0
+ assert torch.equal(warm_model.transformer.weight, cold_model.transformer.weight)
+ from vllm_omni.diffusion.offloader.startup import take_offload_startup_state
+
+ startup_state = take_offload_startup_state(warm)
+ assert startup_state is not None
+ warm_plan = startup_state.host_weight_plan
+ assert warm_plan is not None
+ assert warm_plan.lease_carrier is not None
+ warm_plan.lease_carrier.close()
+ assert hash_calls == 1
+ assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1
+
+
+def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch):
+ from vllm_omni.diffusion.model_loader import diffusers_loader as loader_module
+
+ loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(tmp_path, tmp_path / "store"))
+ models: list[_DummyPipelineModel] = []
+
+ def init_model(*args, **kwargs):
+ del args, kwargs
+ model = _DummyPipelineModel(source_prefix="transformer.")
+ models.append(model)
+ return model
+
+ def commit_error(*args, **kwargs):
+ del args, kwargs
+ raise loader_module._HWRCommitError("restore commit failed")
+
+ monkeypatch.setattr(loader, "_init_from_load_format", init_model)
+ monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: ())
+ monkeypatch.setattr(loader, "_resolve_hwr", commit_error)
+ monkeypatch.setattr(loader, "load_weights", lambda *args, **kwargs: None)
+ monkeypatch.setattr(loader, "_process_weights_after_loading", lambda *args, **kwargs: None)
+ monkeypatch.setattr(loader, "_apply_skip_softmax_calibration", lambda *args, **kwargs: None)
+
+ recovered = loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+ assert len(models) == 2
+ assert recovered is models[1]
+ assert loader.take_host_weight_plan() is None
+
+
+def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
+ tmp_path: Path,
+ monkeypatch,
+):
+ canonical_root = tmp_path / "canonical"
+ canonical_root.mkdir()
+ save_file(
+ {"weight": torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2)},
+ str(canonical_root / "model.safetensors"),
+ )
+ loader = DiffusersPipelineLoader(
+ LoadConfig(),
+ _hwr_config(canonical_root, tmp_path / "empty-store", mode="required"),
+ )
+ pipeline = _HWRPipeline(canonical_root)
+ monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: pipeline)
+
+ with pytest.raises(RuntimeError, match="Host Weight Runtime resolution failed"):
+ loader.load_model(load_device="cpu", device=torch.device("cpu"))
+
+ assert pipeline.load_count == 0
+ assert loader.take_host_weight_plan() is None
+
+
+def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
+ return OmniDiffusionConfig(
+ model="",
+ dtype=torch.float32,
+ quantization_config="fp8",
+ parallel_config=DiffusionParallelConfig(
+ data_parallel_size=dp_size,
+ sequence_parallel_size=1,
+ ),
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=True,
+ )
+
+
+@pytest.mark.parametrize(
+ ("dist_offload", "use_allgather", "mode"),
+ [
+ (False, False, "preferred"),
+ (True, True, "preferred"),
+ (True, False, "disabled"),
+ ],
+)
+def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
+ monkeypatch,
+ tmp_path,
+ dist_offload,
+ use_allgather,
+ mode,
+):
+ """Disabled and AllGather paths must never construct or probe HWR."""
+ from vllm_omni.host_weight_runtime import HostWeightRuntime
+
+ root = tmp_path / "must-not-be-touched"
+ loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config("dummy-model", root, mode=mode))
+ model = _DummyPipelineModel(source_prefix="transformer.")
+ modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
+
+ def unexpected_store_construction(*args, **kwargs):
+ raise AssertionError(f"HWR store interaction was not eligible: {args}, {kwargs}")
+
+ monkeypatch.setattr(HostWeightRuntime, "from_config", unexpected_store_construction)
+ assert (
+ loader._resolve_hwr(
+ model,
+ modules,
+ dist_offload=dist_offload,
+ use_allgather=use_allgather,
+ load_format="default",
+ sources=tuple(model.weights_sources),
+ )
+ is None
+ )
+ assert not root.exists()
+
+
+def test_required_hwr_rejects_a_model_without_a_restore_contract(tmp_path):
+ loader = DiffusersPipelineLoader(
+ LoadConfig(),
+ _hwr_config("dummy-model", tmp_path / "store", mode="required"),
+ )
+ model = _DummyPipelineModel(source_prefix="transformer.")
+ modules = SimpleNamespace(dit_names=("transformer",), dits=(model.transformer,))
+
+ with pytest.raises(ValueError, match="restore contract"):
+ loader._resolve_hwr(
+ model,
+ modules,
+ dist_offload=True,
+ use_allgather=False,
+ load_format="default",
+ sources=tuple(model.weights_sources),
+ )
+
+
+@pytest.mark.parametrize("offline", [False, True])
+def test_prepare_weights_honors_component_index_and_explicit_override(tmp_path, mocker, offline):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ snapshot = tmp_path / "snapshot"
+ transformer = snapshot / "transformer"
+ transformer.mkdir(parents=True)
+ indexed_files = [f"diffusion_pytorch_model-{index:05d}-of-00002.safetensors" for index in (1, 2)]
+ stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+ index_path = transformer / "diffusion_pytorch_model.safetensors.index.json"
+ index_path.write_text(
+ json.dumps({"weight_map": {f"weight.{index}": filename for index, filename in enumerate(indexed_files)}})
+ )
+ for filename in indexed_files + [stale_file]:
+ (transformer / filename).touch()
+
+ mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", offline)
+
+ def download_index(*, filename, **_kwargs):
+ if filename == "transformer/diffusion_pytorch_model.safetensors.index.json":
+ return str(index_path)
+ raise loader_mod.huggingface_hub.errors.EntryNotFoundError(filename)
+
+ hub_api = mocker.Mock()
+ hub_api.hf_hub_download.side_effect = download_index
+ mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+ indexed_download = mocker.patch.object(loader_mod, "download_weights_from_hf_specific", return_value=str(snapshot))
+ generic_download = mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+ loader = _make_loader_with_weights([])
+ cache_dir = str(tmp_path / "cache")
+ loader.load_config.download_dir = cache_dir
+ folder, files, use_safetensors = loader._prepare_weights(
+ "org/model",
+ subfolder="transformer",
+ revision="revision",
+ fall_back_to_pt=True,
+ allow_patterns_overrides=None,
+ )
+
+ assert folder == str(transformer)
+ assert [str(transformer / filename) for filename in indexed_files] == files
+ assert use_safetensors
+ assert hub_api.hf_hub_download.call_count == 2
+ hub_api.hf_hub_download.assert_any_call(
+ repo_id="org/model",
+ filename="transformer/diffusion_pytorch_model.safetensors.index.json",
+ cache_dir=cache_dir,
+ revision="revision",
+ local_files_only=offline,
+ )
+ indexed_download.assert_called_once_with(
+ model_name_or_path="org/model",
+ cache_dir=cache_dir,
+ allow_patterns=[f"transformer/{filename}" for filename in indexed_files],
+ revision="revision",
+ ignore_patterns=loader.load_config.ignore_patterns,
+ require_all=True,
+ )
+
+ _, override_files, _ = loader._prepare_weights(
+ "org/model",
+ subfolder="transformer",
+ revision="revision",
+ fall_back_to_pt=True,
+ allow_patterns_overrides=[stale_file],
+ )
+ assert override_files == [str(transformer / stale_file)]
+ generic_download.assert_called_once_with(
+ "org/model",
+ cache_dir,
+ [stale_file],
+ "revision",
+ subfolder="transformer",
+ ignore_patterns=loader.load_config.ignore_patterns,
+ )
+
+
+def test_prepare_local_weights_honors_component_index(tmp_path):
+ transformer = tmp_path / "transformer"
+ transformer.mkdir()
+ indexed_file = "diffusion_pytorch_model-00001-of-00001.safetensors"
+ stale_file = "diffusion_pytorch_model-00001-of-00008.safetensors"
+ (transformer / "diffusion_pytorch_model.safetensors.index.json").write_text(
+ json.dumps({"weight_map": {"weight": indexed_file}})
+ )
+ for filename in (indexed_file, stale_file):
+ (transformer / filename).touch()
+
+ folder, files, use_safetensors = _make_loader_with_weights([])._prepare_weights(
+ tmp_path,
+ subfolder="transformer",
+ revision=None,
+ fall_back_to_pt=True,
+ allow_patterns_overrides=None,
+ )
+
+ assert folder == str(transformer)
+ assert files == [str(transformer / indexed_file)]
+ assert use_safetensors
+
+
+def test_prepare_weights_rejects_polluted_offline_cache_without_index(tmp_path, mocker):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ snapshot = tmp_path / "snapshot"
+ transformer = snapshot / "transformer"
+ transformer.mkdir(parents=True)
+ for shard_count in (4, 8):
+ for shard in range(1, shard_count + 1):
+ (transformer / f"diffusion_pytorch_model-{shard:05d}-of-{shard_count:05d}.safetensors").touch()
+
+ mocker.patch.object(loader_mod.huggingface_hub.constants, "HF_HUB_OFFLINE", True)
+ hub_api = mocker.Mock()
+ hub_api.hf_hub_download.side_effect = loader_mod.huggingface_hub.errors.EntryNotFoundError("index is not cached")
+ mocker.patch.object(loader_mod, "hf_api", return_value=hub_api)
+ mocker.patch.object(loader_mod, "download_weights_from_hf", return_value=str(snapshot))
+
+ loader = _make_loader_with_weights([])
+ with pytest.raises(ValueError, match="conflicting shard totals"):
+ loader._prepare_weights(
+ "org/model",
+ subfolder="transformer",
+ revision="revision",
+ fall_back_to_pt=True,
+ allow_patterns_overrides=None,
+ )
+
+ assert hub_api.hf_hub_download.call_count == 2
+ assert all(call.kwargs["local_files_only"] for call in hub_api.hf_hub_download.call_args_list)
+
+
def test_strict_check_only_validates_source_prefix_parameters():
model = _DummyPipelineModel(source_prefix="transformer.")
loader = _make_loader_with_weights(["transformer.weight"])
@@ -110,6 +503,85 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
loader.load_weights(model)
+def test_stream_online_quant_weights_offloads_layers_after_processing():
+ from vllm.model_executor.model_loader.reload.layerwise import (
+ get_layerwise_info,
+ )
+
+ events: list[str] = []
+
+ class _OnlineQuantMethod:
+ uses_meta_device = True
+
+ class _TrackedLayer(nn.Linear):
+ def __init__(self, name: str):
+ super().__init__(2, 2, bias=False)
+ self.name = name
+ self.quant_method = _OnlineQuantMethod()
+ get_layerwise_info(self).load_numel_total = self.weight.numel()
+
+ def to(self, *args, **kwargs):
+ events.append(self.name)
+ return super().to(*args, **kwargs)
+
+ class _StreamingModel(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.first = _TrackedLayer("first")
+ self.second = _TrackedLayer("second")
+
+ model = _StreamingModel()
+ weights = iter(
+ [
+ ("first.weight", torch.zeros((2, 2))),
+ ("second.weight", torch.zeros((2, 2))),
+ ]
+ )
+ streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
+
+ assert next(streamed)[0] == "first.weight"
+ get_layerwise_info(model.first).reset()
+ assert next(streamed)[0] == "second.weight"
+ assert events == ["first"]
+
+ get_layerwise_info(model.second).reset()
+ with pytest.raises(StopIteration):
+ next(streamed)
+ assert events == ["first", "second"]
+
+
+def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
+ from unittest.mock import Mock
+
+ from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
+ from vllm.model_executor.model_loader.reload import layerwise
+
+ class _TrackedLayer(nn.Linear):
+ def __init__(self):
+ super().__init__(2, 2, bias=False)
+ self.to_calls: list[object] = []
+
+ def to(self, *args, **kwargs):
+ self.to_calls.append(args[0] if args else kwargs.get("device"))
+ return super().to(*args, **kwargs)
+
+ model = nn.Module()
+ model.layer = _TrackedLayer()
+ quant_method = Mock(spec=QuantizeMethodBase)
+ quant_method.uses_meta_device = True
+ model.layer.quant_method = quant_method
+ model.layer._already_called_process_weights_after_loading = True
+ finalize = Mock()
+ monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
+
+ loader = _make_loader_with_weights([])
+ loader._process_weights_after_loading(model, torch.device("cuda"))
+
+ finalize.assert_called_once_with(model, model_config=None)
+ quant_method.process_weights_after_loading.assert_not_called()
+ assert model.layer.to_calls == []
+
+
class _ConfigAwareModel(nn.Module):
def __init__(self, *, od_config):
super().__init__()
@@ -179,6 +651,340 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
assert get_current_diffusion_config_or_none() is None
+def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = SimpleNamespace(
+ dtype=torch.float32,
+ parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+ quantization_config=None,
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ model="unused",
+ )
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ plan = HostWeightPlan(
+ backing_kind="checkpoint_mmap",
+ bindings={},
+ )
+ loaded_ordinary_weights = False
+
+ def load_weights(_model):
+ nonlocal loaded_ordinary_weights
+ loaded_ordinary_weights = True
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ loader.load_weights = load_weights # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+ )
+
+ assert loader.load_model(load_device="cpu") is model
+ assert not loaded_ordinary_weights
+ assert loader.take_host_weight_plan() is plan
+ assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = SimpleNamespace(
+ dtype=torch.float32,
+ parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+ quantization_config=None,
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ model="unused",
+ )
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
+ class MixedSourceModel(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = nn.Linear(2, 2, bias=False)
+ self.text_encoder = nn.Linear(2, 2, bias=False)
+ self.weights_sources = (
+ DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
+ DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
+ )
+ self.loaded_weight_names: list[str] = []
+
+ def load_weights(self, weights):
+ self.loaded_weight_names = [name for name, _ in weights]
+ return set(self.loaded_weight_names)
+
+ model = MixedSourceModel()
+ plan = HostWeightPlan(
+ backing_kind="checkpoint_mmap",
+ bindings={
+ "transformer.weight": TensorBinding(
+ checkpoint_key="weight",
+ file_path="unused",
+ )
+ },
+ planned_source_prefixes=frozenset({"transformer."}),
+ )
+ requested_prefixes: list[str] = []
+
+ def get_weights(source, model=None):
+ del model
+ requested_prefixes.append(source.prefix)
+ yield source.prefix + "weight", torch.ones(2, 2)
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ loader._get_weights_iterator = get_weights # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(plan),
+ )
+
+ assert loader.load_model(load_device="cpu") is model
+ assert requested_prefixes == ["text_encoder."]
+ assert model.loaded_weight_names == ["text_encoder.weight"]
+ assert loader.take_host_weight_plan() is plan
+
+
+def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = SimpleNamespace(
+ dtype=torch.float32,
+ parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
+ quantization_config=None,
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ model="unused",
+ )
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ calls: list[str] = []
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign]
+ loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
+ )
+
+ assert loader.load_model(load_device="cpu") is model
+ assert calls == ["load", "process"]
+ assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
+ from vllm.model_executor.layers.quantization.online.fp8 import (
+ Fp8PerTensorOnlineLinearMethod,
+ )
+
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = _make_dlo_online_quant_config()
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ model.transformer.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+ model.transformer.quant_method.uses_meta_device = True
+ calls: list[object] = []
+ allowlist_models: list[nn.Module] = []
+
+ original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+ def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+ allowlist_models.append(candidate)
+ return original_allowlist_check(candidate)
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+ loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
+ loader.load_weights = ( # type: ignore[method-assign]
+ lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+ )
+ loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(
+ None,
+ "online quantization requires the ordinary loader",
+ ),
+ )
+
+ assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+ assert allowlist_models == [model]
+ assert calls == [("load", True), "process"]
+ assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+ from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+ od_config = _make_dlo_online_quant_config()
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
+ model.transformer.quant_method.uses_meta_device = True
+ calls: list[object] = []
+ allowlist_models: list[nn.Module] = []
+
+ original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+ def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+ allowlist_models.append(candidate)
+ return original_allowlist_check(candidate)
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+ loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
+ loader.load_weights = ( # type: ignore[method-assign]
+ lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+ )
+ loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(
+ None,
+ "online quantization requires the ordinary loader",
+ ),
+ )
+
+ assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+ assert allowlist_models == [model]
+ assert calls == [("load", True), "process"]
+ assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
+ """A DLO group of one runs no weight collective, so an otherwise
+ unvalidated online method must not be rejected by the AllGather gate."""
+
+ class UnsupportedOnlineMethod:
+ uses_meta_device = True
+
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = _make_dlo_online_quant_config(dp_size=1)
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ model.transformer.quant_method = UnsupportedOnlineMethod()
+ calls: list[object] = []
+ allowlist_models: list[nn.Module] = []
+
+ original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+ def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+ allowlist_models.append(candidate)
+ return original_allowlist_check(candidate)
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+ loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
+ loader.load_weights = ( # type: ignore[method-assign]
+ lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+ )
+ loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(
+ None,
+ "online quantization requires the ordinary loader",
+ ),
+ )
+
+ assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+ assert allowlist_models == []
+ assert calls == [("load", True), "process"]
+
+
+def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ class UnsupportedOnlineMethod:
+ uses_meta_device = True
+
+ od_config = _make_dlo_online_quant_config()
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ model.transformer.quant_method = UnsupportedOnlineMethod()
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(
+ None,
+ "online quantization requires the ordinary loader",
+ ),
+ )
+
+ with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
+ loader.load_model(load_device="cpu")
+
+
+def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
+ mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
+ NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
+
+ import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+ od_config = _make_dlo_online_quant_config()
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+ model = nn.Module()
+ model.transformer = nn.Linear(2, 2, bias=False)
+ model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
+ model.transformer.quant_method.uses_meta_device = True
+ calls: list[object] = []
+ allowlist_models: list[nn.Module] = []
+
+ original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+ def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+ allowlist_models.append(candidate)
+ return original_allowlist_check(candidate)
+
+ loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign]
+ monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+ loader._request_offload_after_quant = lambda _model: 1 # type: ignore[method-assign]
+ loader.load_weights = ( # type: ignore[method-assign]
+ lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+ )
+ loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign]
+ loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
+ monkeypatch.setattr(
+ loader_mod,
+ "build_checkpoint_mmap_plan",
+ lambda *_args, **_kwargs: HostWeightPlanResult(
+ None,
+ "online quantization requires the ordinary loader",
+ ),
+ )
+
+ assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+ assert allowlist_models == [model]
+ assert calls == [("load", True), "process"]
+ assert loader.take_host_weight_plan() is None
+
+
def test_hsdp_processes_quantized_weights_before_sharding(mocker):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
from vllm_omni.diffusion.offloader.module_collector import PipelineModules
diff --git a/tests/diffusion/model_loader/test_final_layout_host_weights.py b/tests/diffusion/model_loader/test_final_layout_host_weights.py
new file mode 100644
index 000000000..0b64286cc
--- /dev/null
+++ b/tests/diffusion/model_loader/test_final_layout_host_weights.py
@@ -0,0 +1,1428 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Contract tests for final-layout BF16 Host Weight Runtime artifacts."""
+
+from __future__ import annotations
+
+import dataclasses
+import gc
+import json
+import threading
+from collections.abc import Sequence
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.host_weights import (
+ FINAL_LAYOUT_BF16_POLICY,
+ FINAL_LAYOUT_BF16_SPEC,
+ FinalLayoutBF16Producer,
+ FinalLayoutIdentityContext,
+ FinalLayoutLoaderIdentity,
+ FinalLayoutParallelIdentity,
+ FinalLayoutRequest,
+ FinalLayoutTensorRestorer,
+ ImplementationIdentity,
+ NodeSourceDigestCache,
+ PreparedWeightSource,
+ WeightSourceKind,
+ build_final_layout_identity,
+)
+from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module
+from vllm_omni.diffusion.model_loader.host_weights.contracts import (
+ FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
+ FinalLayoutArtifactSpec,
+ FinalLayoutContractCode,
+ FinalLayoutContractError,
+ FinalLayoutTensorPolicy,
+ implementation_abi_fingerprint,
+)
+from vllm_omni.diffusion.model_loader.host_weights.tensor_layout import (
+ RuntimeTensorTarget,
+ collect_final_layout_targets,
+)
+from vllm_omni.diffusion.models.host_weight_contract import (
+ FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+ FinalLayoutModelContract,
+)
+from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3DiTModel
+from vllm_omni.host_weight_runtime import (
+ AdaptationIdentity,
+ CanonicalJson,
+ CoordinationScope,
+ FailureCode,
+ HostWeightRuntime,
+ HostWeightRuntimeConfig,
+ LookupPhase,
+ PostLoadPublicationOutcome,
+ ProducerIdentity,
+ ProductionPolicy,
+ ProductionSourceMode,
+ ResolutionOutcome,
+ ResolutionStage,
+ RuntimeMode,
+ StorageDomainPolicy,
+ StorageScope,
+ TensorKind,
+ WaitPolicy,
+ WeightRepresentation,
+)
+from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _TinyDiT(nn.Module):
+ host_weight_restore_contract = FinalLayoutModelContract(
+ implementation_id="test-tiny-dit",
+ version="1",
+ )
+
+ def __init__(self, *, device: torch.device | str = "cpu") -> None:
+ super().__init__()
+ target = torch.device(device)
+ self.restore_validations = 0
+ self.proj = nn.Linear(3, 2, dtype=torch.bfloat16, device=target)
+ self.fp32_gain = nn.Parameter(torch.empty(1, dtype=torch.float32, device=target))
+ self.register_buffer("scale", torch.empty(1, dtype=torch.float32, device=target))
+ self.register_buffer(
+ "derived",
+ torch.tensor([7.0], dtype=torch.float32),
+ persistent=False,
+ )
+
+ def validate_restored_host_weights(self) -> None:
+ self.restore_validations += 1
+ assert not self.proj.weight.is_meta
+ assert self.proj.weight.dtype is torch.bfloat16
+ assert self.fp32_gain.dtype is torch.float32
+ assert self.scale.dtype is torch.float32
+
+
+class _TinyPipeline(nn.Module):
+ def __init__(self, *, dit_device: torch.device | str = "cpu") -> None:
+ super().__init__()
+ self.transformer = _TinyDiT(device=dit_device)
+ self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
+ self.vae = nn.Module()
+ self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
+
+
+class _AlternateTinyPipeline(_TinyPipeline):
+ pass
+
+
+class _MultiTinyPipeline(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.transformer_a = _TinyDiT()
+ self.transformer_b = _TinyDiT()
+
+
+class _Float32DiT(nn.Module):
+ host_weight_restore_contract = FinalLayoutModelContract(
+ implementation_id="test-float32-dit",
+ version="1",
+ )
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj = nn.Linear(2, 2, dtype=torch.float32)
+
+ def validate_restored_host_weights(self) -> None:
+ assert self.proj.weight.dtype is torch.float32
+
+
+class _Float32Pipeline(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.transformer = _Float32DiT()
+
+
+_SYNTHETIC_FP32_ABI = CanonicalJson.from_value(
+ {
+ "artifact_identity": "diffusion-final-layout-identity-v1",
+ "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+ "producer": "test-fp32-producer-v1",
+ "representation_policy": "test-all-fp32-v1",
+ "restorer": "exact-final-layout-tensor-rebind-v1",
+ "source_identity": "prepared-diffusion-weight-source-v2",
+ "tensor_contract": "complete-strided-tensor-ownership-v1",
+ }
+)
+_SYNTHETIC_FP32_SPEC = FinalLayoutArtifactSpec(
+ representation=WeightRepresentation(
+ name="test-final-layout-fp32",
+ dtype=str(torch.float32),
+ metadata=CanonicalJson.from_value({"policy": "all-fp32"}),
+ ),
+ producer=ProducerIdentity(
+ producer_id="test.final-layout-fp32",
+ version="1",
+ implementation_fingerprint=implementation_abi_fingerprint(_SYNTHETIC_FP32_ABI),
+ manifest_schema="test-final-layout-fp32-manifest-v1",
+ restorer_schema=FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA,
+ ),
+ implementation_abi=_SYNTHETIC_FP32_ABI,
+)
+
+
+class _SyntheticFP32Policy:
+ @property
+ def spec(self) -> FinalLayoutArtifactSpec:
+ return _SYNTHETIC_FP32_SPEC
+
+ def validate_request(self, request: FinalLayoutRequest) -> None:
+ if request.load_format != "default":
+ raise ValueError("synthetic policy requires default loading")
+
+ def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str:
+ del name, tensor
+ return "weight" if kind is TensorKind.PARAMETER else "persistent_buffer"
+
+ def validate_target(self, target: RuntimeTensorTarget) -> None:
+ if target.tensor.dtype is not torch.float32:
+ raise ValueError("synthetic policy requires FP32 tensors")
+
+ def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None:
+ if not targets:
+ raise ValueError("synthetic policy requires tensors")
+
+ def build_format_metadata(
+ self,
+ *,
+ component_names: tuple[str, ...],
+ tensor_contract_digest: str,
+ tensor_count: int,
+ ) -> CanonicalJson:
+ return CanonicalJson.from_value(
+ {
+ "component_names": list(component_names),
+ "format": "test-final-layout-fp32",
+ "tensor_contract_sha256": tensor_contract_digest,
+ "tensor_count": tensor_count,
+ }
+ )
+
+ def validate_format_metadata(
+ self,
+ metadata: CanonicalJson,
+ *,
+ component_names: tuple[str, ...],
+ tensor_contract_digest: str,
+ tensor_count: int,
+ ) -> None:
+ assert metadata == self.build_format_metadata(
+ component_names=component_names,
+ tensor_contract_digest=tensor_contract_digest,
+ tensor_count=tensor_count,
+ )
+
+
+_SYNTHETIC_FP32_POLICY = _SyntheticFP32Policy()
+
+
+class _PolicyWithSpec:
+ def __init__(self, base: FinalLayoutTensorPolicy, spec: FinalLayoutArtifactSpec) -> None:
+ self._base = base
+ self._spec = spec
+
+ @property
+ def spec(self) -> FinalLayoutArtifactSpec:
+ return self._spec
+
+ def validate_request(self, request: FinalLayoutRequest) -> None:
+ self._base.validate_request(request)
+
+ def tensor_role(self, name: str, tensor: torch.Tensor, kind: TensorKind) -> str:
+ return self._base.tensor_role(name, tensor, kind)
+
+ def validate_target(self, target: RuntimeTensorTarget) -> None:
+ self._base.validate_target(target)
+
+ def validate_collection(self, targets: Sequence[RuntimeTensorTarget]) -> None:
+ self._base.validate_collection(targets)
+
+ def build_format_metadata(
+ self,
+ *,
+ component_names: tuple[str, ...],
+ tensor_contract_digest: str,
+ tensor_count: int,
+ ) -> CanonicalJson:
+ return self._base.build_format_metadata(
+ component_names=component_names,
+ tensor_contract_digest=tensor_contract_digest,
+ tensor_count=tensor_count,
+ )
+
+ def validate_format_metadata(
+ self,
+ metadata: CanonicalJson,
+ *,
+ component_names: tuple[str, ...],
+ tensor_contract_digest: str,
+ tensor_count: int,
+ ) -> None:
+ self._base.validate_format_metadata(
+ metadata,
+ component_names=component_names,
+ tensor_contract_digest=tensor_contract_digest,
+ tensor_count=tensor_count,
+ )
+
+
+class _FakeH3Linear(nn.Module):
+ def __init__(
+ self,
+ *args: object,
+ bias: bool = False,
+ params_dtype: torch.dtype | None = None,
+ total_num_heads: int | None = None,
+ total_num_kv_heads: int | None = None,
+ **kwargs: object,
+ ) -> None:
+ del args, kwargs
+ super().__init__()
+ dtype = params_dtype or torch.get_default_dtype()
+ self.weight = nn.Parameter(torch.empty(1, dtype=dtype))
+ if bias:
+ self.bias = nn.Parameter(torch.empty(1, dtype=dtype))
+ else:
+ self.register_parameter("bias", None)
+ self.num_heads = total_num_heads
+ self.num_kv_heads = total_num_kv_heads
+
+
+class _FakeH3Attention(nn.Module):
+ def __init__(self, **kwargs: object) -> None:
+ del kwargs
+ super().__init__()
+
+
+class _FakeFluxLinear(_FakeH3Linear):
+ def __init__(
+ self,
+ *args: object,
+ bias: bool = False,
+ params_dtype: torch.dtype | None = None,
+ total_num_heads: int | None = None,
+ total_num_kv_heads: int | None = None,
+ **kwargs: object,
+ ) -> None:
+ super().__init__(
+ *args,
+ bias=bias,
+ params_dtype=params_dtype,
+ total_num_heads=total_num_heads,
+ total_num_kv_heads=total_num_kv_heads,
+ **kwargs,
+ )
+ self.weight.weight_loader = self._load_weight
+
+ @staticmethod
+ def _load_weight(
+ parameter: nn.Parameter,
+ loaded_weight: torch.Tensor,
+ shard_id: str | None = None,
+ ) -> None:
+ del shard_id
+ parameter.data.copy_(loaded_weight.to(parameter).reshape_as(parameter))
+
+
+class _FakeFluxAttention(nn.Module):
+ def __init__(self, **kwargs: object) -> None:
+ del kwargs
+ super().__init__()
+ self.register_buffer("beta", torch.tensor([0.5]))
+ self.register_buffer("eps", torch.tensor([-1e-6]))
+
+
+class _Flux2KleinPipeline(nn.Module):
+ def __init__(self, transformer: nn.Module) -> None:
+ super().__init__()
+ self.transformer = transformer
+ self.text_encoder = nn.Linear(2, 2, dtype=torch.bfloat16)
+ self.vae = nn.Module()
+ self.vae.register_buffer("gain", torch.tensor([9.0], dtype=torch.float32))
+
+
+def _reduced_flux2_klein(
+ monkeypatch: pytest.MonkeyPatch,
+ *,
+ device: torch.device | str = "cpu",
+) -> nn.Module:
+ from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as flux2
+
+ monkeypatch.setattr(flux2, "ColumnParallelLinear", _FakeFluxLinear)
+ monkeypatch.setattr(flux2, "MergedColumnParallelLinear", _FakeFluxLinear)
+ monkeypatch.setattr(flux2, "QKVParallelLinear", _FakeFluxLinear)
+ monkeypatch.setattr(flux2, "RowParallelLinear", _FakeFluxLinear)
+ monkeypatch.setattr(flux2, "Attention", _FakeFluxAttention)
+
+ transformer = flux2.Flux2Transformer2DModel(
+ patch_size=1,
+ in_channels=8,
+ out_channels=8,
+ num_layers=1,
+ num_single_layers=1,
+ attention_head_dim=8,
+ num_attention_heads=1,
+ joint_attention_dim=8,
+ timestep_guidance_channels=8,
+ mlp_ratio=1.0,
+ axes_dims_rope=(2, 2, 2, 2),
+ guidance_embeds=False,
+ )
+ return transformer.to(device=device, dtype=torch.bfloat16)
+
+
+class _H3Pipeline(nn.Module):
+ def __init__(self, transformer: nn.Module) -> None:
+ super().__init__()
+ self.transformer = transformer
+
+
+def _small_h3_config() -> SimpleNamespace:
+ return SimpleNamespace(
+ tf_model_config={
+ "num_layers": 1,
+ "token_refiner_num_layers": 1,
+ "hidden_size": 8,
+ "num_attention_heads": 2,
+ "attention_head_dim": 4,
+ "ffn_hidden_size": 16,
+ "latents_dim": 2,
+ "audio_latents_dim": 2,
+ "patch_size": (1, 2, 2),
+ "text_dim": 6,
+ "timestep_input_dim": 4,
+ "time_embed_hidden_size": 8,
+ "time_embed_dim": 4,
+ "adaln_out_features": 18 * 8,
+ "final_adaln_out_features": 2 * 8,
+ "rope_inv_freq_len": 2,
+ },
+ parallel_config=SimpleNamespace(ulysses_degree=1),
+ )
+
+
+def _prepared_source(
+ tmp_path: Path,
+ *,
+ content: bytes = b"canonical-source-for-identity",
+ requested_revision: str | None = None,
+ resolved_revision: str | None = None,
+ directory: str = "canonical",
+ prefix: str = "transformer.",
+) -> PreparedWeightSource:
+ source_root = tmp_path / directory
+ if resolved_revision is not None:
+ source_root = source_root / "snapshots" / resolved_revision
+ source_root.mkdir(parents=True, exist_ok=True)
+ weight_file = source_root / "model.safetensors"
+ if not weight_file.exists():
+ weight_file.write_bytes(content)
+ elif weight_file.read_bytes() != content:
+ weight_file.write_bytes(content)
+ return PreparedWeightSource(
+ model_or_path="test-org/tiny-diffusion",
+ subfolder=None,
+ requested_revision=requested_revision,
+ prefix=prefix,
+ resolved_root=source_root,
+ weight_files=(weight_file,),
+ use_safetensors=True,
+ )
+
+
+def _request(**changes: object) -> FinalLayoutRequest:
+ request = FinalLayoutRequest(
+ model_id="test-org/tiny-diffusion",
+ loader=FinalLayoutLoaderIdentity(
+ implementation=ImplementationIdentity(
+ implementation_id="test-diffusion-loader",
+ version="1",
+ fingerprint="test-loader-implementation-v1",
+ ),
+ model_config_fingerprint="test-model-config-v1",
+ weight_transform_fingerprint="test-weight-transform-v1",
+ ),
+ )
+ return dataclasses.replace(request, **changes)
+
+
+def _identity(
+ model: nn.Module,
+ source: PreparedWeightSource,
+ *,
+ request: FinalLayoutRequest | None = None,
+ policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY,
+ source_digest_cache: NodeSourceDigestCache | None = None,
+) -> FinalLayoutIdentityContext:
+ transformer = model.get_submodule("transformer")
+ return build_final_layout_identity(
+ model,
+ dit_modules=(("transformer", transformer),),
+ prepared_sources=(source,),
+ request=request or _request(),
+ policy=policy,
+ source_digest_cache=source_digest_cache,
+ )
+
+
+def _multi_identity(
+ model: _MultiTinyPipeline,
+ sources: tuple[PreparedWeightSource, ...],
+) -> FinalLayoutIdentityContext:
+ return build_final_layout_identity(
+ model,
+ dit_modules=(
+ ("transformer_a", model.transformer_a),
+ ("transformer_b", model.transformer_b),
+ ),
+ prepared_sources=sources,
+ request=_request(),
+ policy=FINAL_LAYOUT_BF16_POLICY,
+ )
+
+
+def _runtime(root: Path) -> HostWeightRuntime:
+ return HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=StorageDomainPolicy(
+ root=root,
+ scope=StorageScope.NODE,
+ domain_id="node",
+ storage_class=detect_storage_class(root.parent),
+ ),
+ production=ProductionPolicy(
+ allow_local_build=False,
+ allow_post_load_publish=True,
+ ),
+ wait=WaitPolicy(coordination_timeout_seconds=5.0),
+ )
+ )
+
+
+def _dit_modules(model: _TinyPipeline) -> tuple[tuple[str, nn.Module], ...]:
+ return (("transformer", model.transformer),)
+
+
+def _fill_final_weights(model: _TinyPipeline) -> None:
+ with torch.no_grad():
+ model.transformer.proj.weight.copy_(torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3))
+ model.transformer.proj.bias.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16))
+ model.transformer.fp32_gain.copy_(torch.tensor([6.5]))
+ model.transformer.scale.copy_(torch.tensor([2.5]))
+ model.text_encoder.weight.fill_(11)
+ model.text_encoder.bias.fill_(12)
+
+
+def _pointer_snapshot(model: nn.Module) -> dict[str, tuple[int, int, str]]:
+ snapshot: dict[str, tuple[int, int, str]] = {}
+ for name, tensor in list(model.named_parameters()) + list(model.named_buffers()):
+ pointer = 0 if tensor.is_meta else tensor.untyped_storage().data_ptr()
+ snapshot[name] = (id(tensor), pointer, tensor.device.type)
+ return snapshot
+
+
+def test_preload_identity_is_stable_for_cpu_and_meta_skeletons(tmp_path: Path) -> None:
+ source = _prepared_source(tmp_path)
+ cpu_model = _TinyPipeline()
+ meta_model = _TinyPipeline(dit_device="meta")
+
+ cpu = _identity(cpu_model, source)
+ meta = _identity(meta_model, source)
+
+ assert cpu.identity == meta.identity
+ assert cpu.tensor_names == {
+ "transformer.fp32_gain",
+ "transformer.proj.bias",
+ "transformer.proj.weight",
+ "transformer.scale",
+ }
+ assert "transformer.derived" not in cpu.tensor_names
+ assert all("text_encoder" not in name and "vae" not in name for name in cpu.tensor_names)
+
+ request_fields = {field.name for field in dataclasses.fields(FinalLayoutRequest)}
+ assert not any("data_parallel" in name for name in request_fields)
+ assert not any("metadata" in name for name in request_fields)
+ identity_bytes = cpu.identity.canonical_bytes
+ assert b"data_parallel" not in identity_bytes
+ assert b"dlo_use_allgather" not in identity_bytes
+ assert b"registration" not in identity_bytes
+ with pytest.raises(TypeError, match="unexpected keyword argument"):
+ dataclasses.replace(_request(), loader_metadata=CanonicalJson.empty())
+
+
+def test_generic_identity_and_tensor_discovery_accept_a_second_policy(tmp_path: Path) -> None:
+ model = _Float32Pipeline()
+ source = _prepared_source(tmp_path, directory="fp32-source")
+
+ context = _identity(model, source, policy=_SYNTHETIC_FP32_POLICY)
+ targets = collect_final_layout_targets(
+ model,
+ (("transformer", model.transformer),),
+ policy=_SYNTHETIC_FP32_POLICY,
+ require_materialized=True,
+ )
+
+ assert context.spec is _SYNTHETIC_FP32_SPEC
+ assert context.identity.representation.name == "test-final-layout-fp32"
+ assert context.identity.representation.dtype == str(torch.float32)
+ assert all(target.tensor.dtype is torch.float32 for target in targets)
+ assert context.identity.producer.implementation_fingerprint != (
+ FINAL_LAYOUT_BF16_SPEC.producer.implementation_fingerprint
+ )
+
+
+def test_source_identity_requires_complete_multi_component_coverage(tmp_path: Path) -> None:
+ model = _MultiTinyPipeline()
+ source_a = _prepared_source(
+ tmp_path,
+ directory="source-a",
+ prefix="transformer_a.",
+ )
+
+ with pytest.raises(FinalLayoutContractError, match="transformer_b") as exc_info:
+ _multi_identity(model, (source_a,))
+ assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
+
+ source_b = _prepared_source(
+ tmp_path,
+ directory="source-b",
+ prefix="transformer_b.",
+ )
+ context = _multi_identity(model, (source_a, source_b))
+ metadata = context.identity.source.metadata.to_value()
+ assert isinstance(metadata, dict)
+ bindings = metadata["target_bindings"]
+ assert isinstance(bindings, list)
+ assert {binding["source_prefix"] for binding in bindings} == {
+ "transformer_a.",
+ "transformer_b.",
+ }
+ assert {binding["target_name"].split(".", 1)[0] for binding in bindings} == {
+ "transformer_a",
+ "transformer_b",
+ }
+
+
+def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path) -> None:
+ model = _MultiTinyPipeline()
+ root_source = _prepared_source(
+ tmp_path,
+ directory="root-source",
+ prefix="",
+ )
+ source_a = _prepared_source(
+ tmp_path,
+ directory="specific-a",
+ prefix="transformer_a.",
+ )
+
+ context = _multi_identity(model, (root_source, source_a))
+ reordered = _multi_identity(model, (source_a, root_source))
+ assert context.identity == reordered.identity
+ metadata = context.identity.source.metadata.to_value()
+ assert isinstance(metadata, dict)
+ bindings = metadata["target_bindings"]
+ assert isinstance(bindings, list)
+ by_target = {binding["target_name"]: binding["source_prefix"] for binding in bindings}
+ assert all(prefix == "transformer_a." for name, prefix in by_target.items() if name.startswith("transformer_a."))
+ assert all(prefix == "" for name, prefix in by_target.items() if name.startswith("transformer_b."))
+
+ conflicting_a = _prepared_source(
+ tmp_path,
+ directory="conflicting-a",
+ prefix="transformer_a.",
+ )
+ with pytest.raises(FinalLayoutContractError, match="equally specific") as exc_info:
+ _multi_identity(model, (root_source, source_a, conflicting_a))
+ assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID
+
+
+def test_node_source_digest_cache_reuses_recovers_and_invalidates(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model = _TinyPipeline()
+ source = _prepared_source(tmp_path)
+ cache_root = tmp_path / "store"
+ calls = 0
+ original_sha256 = source_identity_module._sha256_file
+
+ def counted_sha256(path: Path, state: object) -> str:
+ nonlocal calls
+ calls += 1
+ return original_sha256(path, state) # type: ignore[arg-type]
+
+ monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256)
+ first = _identity(
+ model,
+ source,
+ source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+ )
+ second = _identity(
+ model,
+ source,
+ source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+ )
+
+ assert second.identity == first.identity
+ assert calls == 1
+
+ cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json"))
+ cache_entry.write_text("not-json", encoding="utf-8")
+ recovered = _identity(
+ model,
+ source,
+ source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+ )
+ assert recovered.identity == first.identity
+ assert calls == 2
+
+ cache_document = json.loads(cache_entry.read_text(encoding="utf-8"))
+ cache_document["content_id"] = f"sha256:{'0' * 64}"
+ cache_entry.write_text(json.dumps(cache_document), encoding="utf-8")
+ checksum_recovered = _identity(
+ model,
+ source,
+ source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+ )
+ assert checksum_recovered.identity == first.identity
+ assert calls == 3
+
+ source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size")
+ changed = _identity(
+ model,
+ source,
+ source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0),
+ )
+ assert changed.identity != first.identity
+ assert calls == 4
+
+
+def test_source_identity_hashes_checkpoint_shards_in_parallel(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model = _TinyPipeline()
+ source = _prepared_source(tmp_path)
+ second_file = source.resolved_root / "model-00002.safetensors"
+ second_file.write_bytes(b"second-checkpoint-shard")
+ source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file))
+ barrier = threading.Barrier(2)
+ thread_ids: set[int] = set()
+ thread_ids_lock = threading.Lock()
+ original_sha256 = source_identity_module._sha256_file
+
+ def synchronized_sha256(path: Path, state: object) -> str:
+ with thread_ids_lock:
+ thread_ids.add(threading.get_ident())
+ barrier.wait(timeout=5)
+ return original_sha256(path, state) # type: ignore[arg-type]
+
+ monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256)
+ _identity(model, source)
+
+ assert len(thread_ids) == 2
+
+
+def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model = _TinyPipeline()
+ source = _prepared_source(tmp_path)
+ cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+
+ def unavailable_lock(*_args: object, **_kwargs: object) -> object:
+ raise source_identity_module.FileLockTimeoutError("busy")
+
+ monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock)
+ context = _identity(model, source, source_digest_cache=cache)
+
+ assert context.identity.source.revision.startswith("content-")
+ assert not tuple(cache.entries.glob("*.json"))
+
+
+def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model = _TinyPipeline()
+ source = _prepared_source(tmp_path)
+ original_mkdir = Path.mkdir
+
+ def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None:
+ if "source-digests-v1" in path.parts:
+ raise OSError("injected cache directory failure")
+ original_mkdir(path, *args, **kwargs) # type: ignore[arg-type]
+
+ monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir)
+ cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01)
+ context = _identity(model, source, source_digest_cache=cache)
+
+ assert context.identity.source.revision.startswith("content-")
+
+
+def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None:
+ model = _TinyPipeline()
+ source_root = tmp_path / "local-symlink-source"
+ blobs_root = tmp_path / "local-blobs"
+ source_root.mkdir()
+ blobs_root.mkdir()
+ blob_path = blobs_root / ("a" * 64)
+ blob_path.write_bytes(b"a" * 32)
+ weight_path = source_root / "model.safetensors"
+ weight_path.symlink_to(blob_path)
+ source = PreparedWeightSource(
+ model_or_path=str(source_root),
+ subfolder=None,
+ requested_revision=None,
+ prefix="transformer.",
+ resolved_root=source_root,
+ weight_files=(weight_path,),
+ use_safetensors=True,
+ )
+
+ first = _identity(model, source)
+ first_metadata = first.identity.source.metadata.to_value()
+ assert isinstance(first_metadata, dict)
+ first_sources = first_metadata["sources"]
+ assert isinstance(first_sources, list)
+ assert first_sources[0]["files"][0]["content_id"].startswith("sha256:")
+
+ blob_path.write_bytes(b"b" * 32)
+ second = _identity(model, source)
+
+ assert second.identity != first.identity
+
+
+def test_hf_blob_shortcut_requires_validated_snapshot_topology(tmp_path: Path) -> None:
+ model = _TinyPipeline()
+ model_id = "test-org/tiny-diffusion"
+ repo_root = tmp_path / "models--test-org--tiny-diffusion"
+ blobs_root = repo_root / "blobs"
+ revision = "1" * 40
+ snapshot_root = repo_root / "snapshots" / revision
+ blobs_root.mkdir(parents=True)
+ snapshot_root.mkdir(parents=True)
+ blob_name = "2" * 64
+ blob_path = blobs_root / blob_name
+ blob_path.write_bytes(b"trusted-hf-blob")
+ weight_path = snapshot_root / "model.safetensors"
+ weight_path.symlink_to(Path("..") / ".." / "blobs" / blob_name)
+ source = PreparedWeightSource(
+ model_or_path=model_id,
+ subfolder=None,
+ requested_revision="main",
+ prefix="transformer.",
+ resolved_root=snapshot_root,
+ weight_files=(weight_path,),
+ use_safetensors=True,
+ source_kind=WeightSourceKind.HUGGING_FACE_HUB,
+ )
+
+ context = _identity(model, source)
+ metadata = context.identity.source.metadata.to_value()
+ assert isinstance(metadata, dict)
+ sources = metadata["sources"]
+ assert isinstance(sources, list)
+ assert sources[0]["files"][0]["content_id"] == f"immutable-blob:{blob_name}"
+
+
+def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) -> None:
+ model = _TinyPipeline()
+ commit = "0123456789abcdef0123456789abcdef01234567"
+ source = _prepared_source(tmp_path, requested_revision="main", resolved_revision=commit)
+ by_alias = _identity(model, source)
+ by_commit = _identity(model, dataclasses.replace(source, requested_revision=commit))
+
+ assert by_alias.identity == by_commit.identity
+
+ base = _request()
+ variants = (
+ dataclasses.replace(
+ base,
+ loader=dataclasses.replace(base.loader, model_config_fingerprint="test-model-config-v2"),
+ ),
+ dataclasses.replace(
+ base,
+ loader=dataclasses.replace(
+ base.loader,
+ implementation=dataclasses.replace(base.loader.implementation, version="2"),
+ ),
+ ),
+ dataclasses.replace(
+ base,
+ parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0),
+ ),
+ dataclasses.replace(
+ base,
+ parallel=FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1),
+ ),
+ dataclasses.replace(
+ base,
+ parallel=FinalLayoutParallelIdentity(
+ sequence_parallel_size=2,
+ ulysses_degree=2,
+ ulysses_mode="strict",
+ ),
+ ),
+ )
+ base_identity = _identity(model, source, request=base).identity
+ assert all(_identity(model, source, request=variant).identity != base_identity for variant in variants)
+
+ model_v2 = _TinyPipeline()
+ model_v2.transformer.host_weight_restore_contract = dataclasses.replace(
+ model_v2.transformer.host_weight_restore_contract,
+ version="2",
+ )
+ assert _identity(model_v2, source).identity != base_identity
+
+ changed_source = _prepared_source(
+ tmp_path,
+ directory="changed",
+ content=b"different-canonical-source",
+ )
+ assert _identity(model, changed_source).identity != base_identity
+ source_with_adapter = dataclasses.replace(
+ source,
+ checkpoint_adapter=ImplementationIdentity(
+ implementation_id="test-checkpoint-adapter",
+ version="1",
+ fingerprint="test-checkpoint-adapter-v1",
+ ),
+ )
+ assert _identity(model, source_with_adapter).identity != base_identity
+
+ abi_v2 = CanonicalJson.from_value(
+ {
+ "artifact_identity": "diffusion-final-layout-identity-v1",
+ "model_contract": FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA,
+ "producer": "finalized-bf16-tensor-writer-v1",
+ "representation_policy": "bf16-with-preserved-fp32-v1",
+ "restorer": "exact-final-layout-tensor-rebind-v2",
+ "source_identity": "prepared-diffusion-weight-source-v2",
+ "tensor_contract": "complete-strided-tensor-ownership-v1",
+ }
+ )
+ with pytest.raises(ValueError, match="fingerprint differs"):
+ dataclasses.replace(FINAL_LAYOUT_BF16_SPEC, implementation_abi=abi_v2)
+ spec_v2 = dataclasses.replace(
+ FINAL_LAYOUT_BF16_SPEC,
+ producer=dataclasses.replace(
+ FINAL_LAYOUT_BF16_SPEC.producer,
+ version="2",
+ implementation_fingerprint=implementation_abi_fingerprint(abi_v2),
+ ),
+ implementation_abi=abi_v2,
+ )
+ identity_v2 = _identity(
+ model,
+ source,
+ policy=_PolicyWithSpec(FINAL_LAYOUT_BF16_POLICY, spec_v2),
+ ).identity
+ assert identity_v2.representation == base_identity.representation
+ assert identity_v2.key != base_identity.key
+
+
+@pytest.mark.parametrize(
+ "parallel",
+ [
+ FinalLayoutParallelIdentity(),
+ FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=0),
+ FinalLayoutParallelIdentity(tensor_parallel_size=2, tensor_parallel_rank=1),
+ FinalLayoutParallelIdentity(sequence_parallel_size=2, ulysses_degree=2),
+ FinalLayoutParallelIdentity(sequence_parallel_size=2, ring_degree=2),
+ FinalLayoutParallelIdentity(
+ tensor_parallel_size=2,
+ tensor_parallel_rank=1,
+ sequence_parallel_size=2,
+ ulysses_degree=2,
+ ),
+ ],
+)
+def test_parallel_identity_matrix_is_exact_for_tp_and_sp(tmp_path: Path, parallel: FinalLayoutParallelIdentity) -> None:
+ """TP rank/size and SP backend are identity coordinates for HWR artifacts."""
+ model = _TinyPipeline()
+ source = _prepared_source(tmp_path)
+ context = _identity(model, source, request=_request(parallel=parallel))
+ layout = context.identity.layout
+
+ assert layout.tensor_parallel_size == parallel.tensor_parallel_size
+ assert layout.tensor_parallel_rank == parallel.tensor_parallel_rank
+ assert layout.sequence_parallel_size == parallel.sequence_parallel_size
+ assert layout.sequence_parallel_backend == parallel.sequence_parallel_backend
+ assert "data_parallel" not in context.identity.layout.metadata.to_value().get("parallel", {})
+
+
+@pytest.mark.parametrize(
+ ("semantic_request", "message"),
+ [
+ (_request(load_format="diffusers"), "load_format='default'"),
+ (
+ _request(parallel=FinalLayoutParallelIdentity(pipeline_parallel_size=2)),
+ "pipeline-parallel",
+ ),
+ (
+ _request(parallel=FinalLayoutParallelIdentity(cfg_parallel_size=2)),
+ "CFG-parallel",
+ ),
+ (_request(parallel=FinalLayoutParallelIdentity(use_hsdp=True)), "HSDP"),
+ (
+ _request(parallel=FinalLayoutParallelIdentity(enable_expert_parallel=True)),
+ "expert-parallel",
+ ),
+ (
+ _request(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256")),
+ "LoRA",
+ ),
+ ],
+)
+def test_bf16_policy_fails_closed_for_unsupported_semantics(
+ tmp_path: Path,
+ semantic_request: FinalLayoutRequest,
+ message: str,
+) -> None:
+ model = _TinyPipeline()
+ with pytest.raises(ValueError, match=message):
+ _identity(model, _prepared_source(tmp_path), request=semantic_request)
+
+
+def test_tensor_ownership_is_complete_mixed_precision_and_alias_free(tmp_path: Path) -> None:
+ model = _TinyPipeline()
+ records = collect_final_layout_targets(
+ model,
+ _dit_modules(model),
+ policy=FINAL_LAYOUT_BF16_POLICY,
+ require_materialized=True,
+ )
+ by_name = {record.name: record for record in records}
+
+ assert set(by_name) == {
+ "transformer.fp32_gain",
+ "transformer.proj.bias",
+ "transformer.proj.weight",
+ "transformer.scale",
+ }
+ assert by_name["transformer.proj.weight"].tensor.dtype is torch.bfloat16
+ assert by_name["transformer.fp32_gain"].tensor.dtype is torch.float32
+ assert by_name["transformer.scale"].role == "persistent_buffer"
+
+ model.transformer.register_parameter("alias", model.transformer.proj.weight)
+ with pytest.raises(FinalLayoutContractError, match="aliases tensor object"):
+ _identity(model, _prepared_source(tmp_path))
+
+
+def test_identity_requires_explicit_model_restore_contract(tmp_path: Path) -> None:
+ model = _TinyPipeline()
+ model.transformer.host_weight_restore_contract = "unsupported-contract"
+
+ with pytest.raises(FinalLayoutContractError, match="does not declare"):
+ _identity(model, _prepared_source(tmp_path))
+
+
+def test_publication_is_warm_only_and_restore_is_transactional(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ source = _prepared_source(tmp_path)
+ cold_model = _TinyPipeline()
+ _fill_final_weights(cold_model)
+ context = _identity(cold_model, source)
+ runtime = _runtime(tmp_path / "store")
+ producer = FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model), max_shard_bytes=16)
+ assert producer.spec.lookup_phase is LookupPhase.POST_LOAD_ONLY
+ assert producer.spec.source_mode is ProductionSourceMode.FINALIZED_MODEL
+ assert producer.spec.coordination_scope is CoordinationScope.SINGLE_PROCESS
+
+ miss = runtime.resolve(context.identity)
+ assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+ cold_pointers = _pointer_snapshot(cold_model)
+ cold_values = {name: tensor.detach().clone() for name, tensor in cold_model.named_parameters()}
+
+ publication = runtime.publish_after_load(context.identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+ assert publication.failure is None
+ assert _pointer_snapshot(cold_model) == cold_pointers
+ assert all(torch.equal(dict(cold_model.named_parameters())[name], value) for name, value in cold_values.items())
+ assert cold_model.transformer.restore_validations == 1
+
+ def unexpected_produce(_writer: object) -> object:
+ raise AssertionError("already-present publication must not invoke the producer")
+
+ monkeypatch.setattr(producer, "produce", unexpected_produce)
+ already_present = runtime.publish_after_load(context.identity, producer=producer)
+ assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT
+
+ warm_model = _TinyPipeline(dit_device="meta")
+ warm_context = _identity(warm_model, source)
+ assert warm_context.identity == context.identity
+ hit = runtime.resolve(warm_context.identity)
+ assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
+ assert hit.lease is not None
+
+ before_plan = _pointer_snapshot(warm_model)
+ text_weight = warm_model.text_encoder.weight
+ vae_gain = warm_model.vae.gain
+ plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+ assert _pointer_snapshot(warm_model) == before_plan
+ assert warm_model.transformer.restore_validations == 0
+
+ plan.commit()
+ with pytest.raises(RuntimeError, match="already committed"):
+ plan.commit()
+
+ expected = {
+ "transformer.proj.weight": torch.arange(6, dtype=torch.float32).to(torch.bfloat16).reshape(2, 3),
+ "transformer.proj.bias": torch.tensor([3.0, 4.0], dtype=torch.bfloat16),
+ "transformer.fp32_gain": torch.tensor([6.5]),
+ "transformer.scale": torch.tensor([2.5]),
+ }
+ restored = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
+ assert all(torch.equal(restored[name], value) for name, value in expected.items())
+ assert all(not restored[name].is_meta for name in warm_context.tensor_names)
+ assert warm_model.transformer.restore_validations == 1
+ assert warm_model.text_encoder.weight is text_weight
+ assert warm_model.vae.gain is vae_gain
+ assert warm_model.transformer.derived.item() == 7.0
+ for name in warm_context.tensor_names:
+ assert restored[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
+
+ del restored, warm_model, plan
+ gc.collect()
+ hit.lease.close()
+ assert isinstance(runtime.store, FilesystemHostWeightStore)
+ assert runtime.store.cleanup(context.identity) is None
+
+
+def test_independent_publications_are_content_deterministic(tmp_path: Path) -> None:
+ source = _prepared_source(tmp_path)
+ model = _TinyPipeline()
+ _fill_final_weights(model)
+ context = _identity(model, source)
+ first_runtime = _runtime(tmp_path / "first-store")
+ second_runtime = _runtime(tmp_path / "second-store")
+
+ first = first_runtime.publish_after_load(
+ context.identity,
+ producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=16),
+ )
+ second = second_runtime.publish_after_load(
+ context.identity,
+ producer=FinalLayoutBF16Producer(context, model, _dit_modules(model), max_shard_bytes=1024),
+ )
+ assert first.outcome is PostLoadPublicationOutcome.PUBLISHED
+ assert second.outcome is PostLoadPublicationOutcome.PUBLISHED
+
+ first_hit = first_runtime.resolve(context.identity)
+ second_hit = second_runtime.resolve(context.identity)
+ assert first_hit.lease is not None and second_hit.lease is not None
+ assert first_hit.lease.manifest.artifact_content_sha256 == second_hit.lease.manifest.artifact_content_sha256
+ assert first_hit.lease.manifest.format_metadata == second_hit.lease.manifest.format_metadata
+ assert {entry.name: entry.sha256 for entry in first_hit.lease.manifest.tensors} == {
+ entry.name: entry.sha256 for entry in second_hit.lease.manifest.tensors
+ }
+ first_hit.lease.close()
+ second_hit.lease.close()
+
+
+def test_source_replacement_prevents_publication(tmp_path: Path) -> None:
+ source = _prepared_source(tmp_path)
+ model = _TinyPipeline()
+ _fill_final_weights(model)
+ context = _identity(model, source)
+ runtime = _runtime(tmp_path / "store")
+ source.weight_files[0].write_bytes(b"replacement-source")
+
+ publication = runtime.publish_after_load(
+ context.identity,
+ producer=FinalLayoutBF16Producer(context, model, _dit_modules(model)),
+ )
+
+ assert publication.outcome is PostLoadPublicationOutcome.FAILED
+ assert publication.failure is not None
+ assert publication.failure.stage is ResolutionStage.CANONICAL_LOADING
+ assert publication.failure.code is FailureCode.CANONICAL_SOURCE_FAILED
+ assert publication.failure.details.to_value() == {"final_layout_code": "source_changed"}
+ assert not context.sources_unchanged()
+ assert runtime.resolve(context.identity).report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+
+def test_source_replacement_between_restore_plan_and_commit_causes_no_mutation(tmp_path: Path) -> None:
+ source = _prepared_source(tmp_path)
+ cold_model = _TinyPipeline()
+ _fill_final_weights(cold_model)
+ context = _identity(cold_model, source)
+ runtime = _runtime(tmp_path / "store")
+ assert (
+ runtime.publish_after_load(
+ context.identity,
+ producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)),
+ ).outcome
+ is PostLoadPublicationOutcome.PUBLISHED
+ )
+ hit = runtime.resolve(context.identity)
+ assert hit.lease is not None
+
+ warm_model = _TinyPipeline()
+ warm_context = _identity(warm_model, source)
+ plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+ before = _pointer_snapshot(warm_model)
+ source.weight_files[0].write_bytes(b"source-changed-before-commit")
+
+ with pytest.raises(FinalLayoutContractError, match="canonical source changed"):
+ plan.commit()
+ assert _pointer_snapshot(warm_model) == before
+ del plan
+ gc.collect()
+ hit.lease.close()
+
+
+def test_restore_rejects_wrong_identity_and_coverage_without_mutation(tmp_path: Path) -> None:
+ source = _prepared_source(tmp_path)
+ cold_model = _TinyPipeline()
+ _fill_final_weights(cold_model)
+ context = _identity(cold_model, source)
+ runtime = _runtime(tmp_path / "store")
+ assert (
+ runtime.publish_after_load(
+ context.identity,
+ producer=FinalLayoutBF16Producer(context, cold_model, _dit_modules(cold_model)),
+ ).outcome
+ is PostLoadPublicationOutcome.PUBLISHED
+ )
+ hit = runtime.resolve(context.identity)
+ assert hit.lease is not None
+
+ alternate = _AlternateTinyPipeline()
+ _fill_final_weights(alternate)
+ with pytest.raises(ValueError, match="producer model implementation differs"):
+ FinalLayoutBF16Producer(context, alternate, _dit_modules(alternate))
+ with pytest.raises(ValueError, match="restore model implementation differs"):
+ FinalLayoutTensorRestorer(context).plan_restore(alternate, hit.lease)
+
+ model = _TinyPipeline()
+ wrong_context = _identity(
+ model,
+ source,
+ request=_request(
+ parallel=FinalLayoutParallelIdentity(
+ tensor_parallel_size=2,
+ tensor_parallel_rank=1,
+ )
+ ),
+ )
+ with pytest.raises(ValueError, match="semantic identity differs"):
+ FinalLayoutTensorRestorer(wrong_context).plan_restore(model, hit.lease)
+
+ exact_context = _identity(model, source)
+ before = _pointer_snapshot(model)
+ model.transformer.register_buffer("new_persistent_state", torch.tensor([1.0]))
+ with pytest.raises(FinalLayoutContractError, match="tensor contract differs"):
+ FinalLayoutTensorRestorer(exact_context).plan_restore(model, hit.lease)
+ after = _pointer_snapshot(model)
+ assert all(after[name] == value for name, value in before.items())
+ hit.lease.close()
+
+
+def test_reduced_minimax_h3_satisfies_real_tensor_ownership_contract(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeH3Linear)
+ monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeH3Linear)
+ monkeypatch.setattr(h3, "QKVParallelLinear", _FakeH3Linear)
+ monkeypatch.setattr(h3, "RowParallelLinear", _FakeH3Linear)
+ monkeypatch.setattr(h3, "Attention", _FakeH3Attention)
+ monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
+
+ transformer = h3.MiniMaxH3DiTModel(_small_h3_config(), quant_config=None)
+ pipeline = _H3Pipeline(transformer)
+ source = _prepared_source(tmp_path, directory="minimax-h3-source")
+ context = _identity(pipeline, source)
+ targets = collect_final_layout_targets(
+ pipeline,
+ (("transformer", transformer),),
+ policy=FINAL_LAYOUT_BF16_POLICY,
+ require_materialized=True,
+ )
+ by_name = {target.name: target for target in targets}
+
+ assert context.model_contracts == (MiniMaxH3DiTModel.host_weight_restore_contract,)
+ assert "transformer.video_patch_proj.weight" in by_name
+ assert "transformer.blocks.0.attn.qkv_proj.weight" in by_name
+ assert "transformer.rope.inv_freq" in by_name
+ assert by_name["transformer.video_patch_proj.weight"].tensor.dtype is torch.float32
+ assert by_name["transformer.blocks.0.attn.qkv_proj.weight"].tensor.dtype is torch.bfloat16
+ assert by_name["transformer.rope.inv_freq"].role == "persistent_buffer"
+ transformer.validate_restored_host_weights()
+
+
+def test_minimax_h3_declares_the_final_layout_restore_contract() -> None:
+ contract = MiniMaxH3DiTModel.host_weight_restore_contract
+ assert isinstance(contract, FinalLayoutModelContract)
+ assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+ assert contract.implementation_id == "minimax-h3-dit"
+ assert callable(MiniMaxH3DiTModel.validate_restored_host_weights)
+ assert "data_parallel" not in FinalLayoutRequest.__dataclass_fields__
+
+
+def test_reduced_flux2_klein_declares_complete_final_layout_contract(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
+ Flux2Transformer2DModel,
+ )
+
+ transformer = _reduced_flux2_klein(monkeypatch)
+ pipeline = _Flux2KleinPipeline(transformer)
+ source = _prepared_source(tmp_path, directory="flux2-klein-source")
+ context = _identity(pipeline, source)
+ targets = collect_final_layout_targets(
+ pipeline,
+ (("transformer", transformer),),
+ policy=FINAL_LAYOUT_BF16_POLICY,
+ require_materialized=True,
+ )
+ target_names = {target.name for target in targets}
+
+ contract = Flux2Transformer2DModel.host_weight_restore_contract
+ assert isinstance(contract, FinalLayoutModelContract)
+ assert contract.schema == FINAL_LAYOUT_TENSOR_MODEL_CONTRACT_SCHEMA
+ assert contract.implementation_id == "flux2-klein-dit"
+ assert contract.version == "1"
+ assert context.model_contracts == (contract,)
+ assert transformer.stacked_params_mapping == [
+ (".to_qkv.", ".to_q.", "q"),
+ (".to_qkv.", ".to_k.", "k"),
+ (".to_qkv.", ".to_v.", "v"),
+ (".add_kv_proj", ".add_q_proj", "q"),
+ (".add_kv_proj", ".add_k_proj", "k"),
+ (".add_kv_proj", ".add_v_proj", "v"),
+ ]
+ loaded = transformer.load_weights(
+ [
+ ("transformer_blocks.0.attn.to_q.weight", torch.tensor([1.0], dtype=torch.bfloat16)),
+ ("transformer_blocks.0.attn.add_q_proj.weight", torch.tensor([2.0], dtype=torch.bfloat16)),
+ (
+ "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+ torch.tensor([3.0], dtype=torch.bfloat16),
+ ),
+ ]
+ )
+ assert loaded == {
+ "transformer_blocks.0.attn.to_qkv.weight",
+ "transformer_blocks.0.attn.add_kv_proj.weight",
+ "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight",
+ }
+ assert any(name.startswith("transformer.transformer_blocks.0.") for name in target_names)
+ assert any(name.startswith("transformer.single_transformer_blocks.0.") for name in target_names)
+ assert "transformer.transformer_blocks.0.attn.to_qkv.weight" in target_names
+ assert "transformer.transformer_blocks.0.attn.add_kv_proj.weight" in target_names
+ assert "transformer.single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight" in target_names
+ assert "transformer.transformer_blocks.0.attn.attn.beta" in target_names
+ assert "transformer.transformer_blocks.0.attn.attn.eps" in target_names
+ assert all("text_encoder" not in name and "vae" not in name for name in target_names)
+ transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_restore_validator_rejects_non_ready_state(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer = _reduced_flux2_klein(monkeypatch)
+
+ mapping = transformer.stacked_params_mapping
+ transformer.stacked_params_mapping = None
+ with pytest.raises(ValueError, match="packed QKV mapping"):
+ transformer.validate_restored_host_weights()
+ transformer.stacked_params_mapping = mapping
+
+ weight = transformer.x_embedder.weight
+ weight.data = weight.data.float()
+ with pytest.raises(ValueError, match="must stay bf16"):
+ transformer.validate_restored_host_weights()
+ weight.data = weight.data.to(torch.bfloat16)
+
+ weight.data = weight.data.T
+ with pytest.raises(ValueError, match="contiguous strided layout"):
+ transformer.validate_restored_host_weights()
+ weight.data = weight.data.contiguous()
+
+ attention = transformer.transformer_blocks[0].attn.attn
+ attention._non_persistent_buffers_set.add("beta")
+ with pytest.raises(ValueError, match="must be persistent"):
+ transformer.validate_restored_host_weights()
+ attention._non_persistent_buffers_set.remove("beta")
+
+ attention.beta.data = attention.beta.data.float()
+ with pytest.raises(ValueError, match="scalar bf16"):
+ transformer.validate_restored_host_weights()
+
+
+def test_flux2_klein_publication_restore_fallback_and_teardown(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ source = _prepared_source(tmp_path, directory="flux2-klein-source")
+ cold_transformer = _reduced_flux2_klein(monkeypatch)
+ cold_model = _Flux2KleinPipeline(cold_transformer)
+ with torch.no_grad():
+ for index, parameter in enumerate(cold_transformer.parameters(), start=1):
+ parameter.fill_(index)
+ for name, buffer in cold_transformer.named_buffers():
+ if name.endswith((".beta", ".eps")):
+ buffer.fill_(0.5 if name.endswith(".beta") else -1e-6)
+
+ context = _identity(cold_model, source)
+ runtime = _runtime(tmp_path / "flux2-klein-store")
+ miss = runtime.resolve(context.identity)
+ assert miss.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+ producer = FinalLayoutBF16Producer(
+ context,
+ cold_model,
+ (("transformer", cold_transformer),),
+ max_shard_bytes=64,
+ )
+ publication = runtime.publish_after_load(context.identity, producer=producer)
+ assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+
+ warm_transformer = _reduced_flux2_klein(monkeypatch, device="meta")
+ warm_model = _Flux2KleinPipeline(warm_transformer)
+ warm_context = _identity(warm_model, source)
+ assert warm_context.identity == context.identity
+ hit = runtime.resolve(warm_context.identity)
+ assert hit.report.outcome is ResolutionOutcome.LOCAL_HIT
+ assert hit.lease is not None
+
+ plan = FinalLayoutTensorRestorer(warm_context).plan_restore(warm_model, hit.lease)
+ assert all(tensor.is_meta for name, tensor in warm_model.named_parameters() if name.startswith("transformer."))
+ plan.commit()
+ warm_transformer.validate_restored_host_weights()
+
+ cold_tensors = dict(cold_model.named_parameters()) | dict(cold_model.named_buffers())
+ warm_tensors = dict(warm_model.named_parameters()) | dict(warm_model.named_buffers())
+ for name in warm_context.tensor_names:
+ assert torch.equal(warm_tensors[name], cold_tensors[name])
+ assert warm_tensors[name].untyped_storage().data_ptr() == hit.lease.tensors[name].untyped_storage().data_ptr()
+ assert warm_transformer.stacked_params_mapping == cold_transformer.stacked_params_mapping
+
+ del warm_tensors, warm_model, warm_transformer, plan
+ gc.collect()
+ hit.lease.close()
+ assert hit.lease.closed
+ assert isinstance(runtime.store, FilesystemHostWeightStore)
+ assert runtime.store.cleanup(context.identity) is None
diff --git a/tests/diffusion/model_loader/test_hub_prefetch.py b/tests/diffusion/model_loader/test_hub_prefetch.py
new file mode 100644
index 000000000..39eb2b277
--- /dev/null
+++ b/tests/diffusion/model_loader/test_hub_prefetch.py
@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+
+import huggingface_hub
+import pytest
+
+from vllm_omni.diffusion.model_loader import hub_prefetch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_prefetch_subfolders_propagates_revision(monkeypatch):
+ calls = []
+
+ def fake_snapshot_download(**kwargs):
+ calls.append(kwargs)
+
+ monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+ monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
+
+ hub_prefetch.prefetch_subfolders(
+ "org/model",
+ ["transformer"],
+ revision="pinned-revision",
+ )
+
+ assert calls == [
+ {
+ "repo_id": "org/model",
+ "revision": "pinned-revision",
+ "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
+ }
+ ]
+
+
+def test_from_pretrained_retry_preserves_revision(monkeypatch):
+ factory_calls = []
+ prefetch_calls = []
+
+ def fake_factory(model, **kwargs):
+ factory_calls.append((model, kwargs))
+ if len(factory_calls) == 1:
+ raise OSError("partial cache")
+ return "loaded"
+
+ def fake_prefetch(model, subfolders, **kwargs):
+ prefetch_calls.append((model, tuple(subfolders), kwargs))
+
+ monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
+ monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
+
+ result = hub_prefetch.from_pretrained_with_prefetch(
+ fake_factory,
+ "org/model",
+ subfolder="transformer",
+ prefetch_list=["transformer", "vae"],
+ revision="pinned-revision",
+ max_attempts=2,
+ )
+
+ assert result == "loaded"
+ assert len(factory_calls) == 2
+ assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
+ assert prefetch_calls == [
+ (
+ "org/model",
+ ("transformer", "vae"),
+ {"local_files_only": False, "revision": "pinned-revision"},
+ )
+ ]
diff --git a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py
index d38be97fa..999fdbb41 100644
--- a/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py
+++ b/tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L1 unit tests for the native Boogu-Image pipeline.
@@ -133,6 +133,51 @@ def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
assert boogu_pipeline.mllm is mock_dependencies["inner_encoder"]
+def test_constructor_forwards_revision_to_all_component_loaders(mock_dependencies, mocker):
+ from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+ BooguImagePipeline,
+ )
+
+ revision = "hotfix-1k-20260708"
+ prefetch = mocker.patch(f"{_MODULE}.prefetch_subfolders")
+ scheduler_loader = mocker.patch(
+ f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
+ return_value=mock_dependencies["scheduler"],
+ )
+ mllm_loader = mocker.patch(
+ f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
+ return_value=mock_dependencies["mllm_wrapper"],
+ )
+ processor_loader = mocker.patch(
+ f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
+ return_value=mock_dependencies["processor"],
+ )
+ vae_loader = mocker.patch(
+ f"{_MODULE}.AutoencoderKL.from_pretrained",
+ return_value=mock_dependencies["vae"],
+ )
+ od_config = OmniDiffusionConfig(
+ model="dummy-boogu",
+ revision=revision,
+ tf_model_config=TransformerConfig(params={}),
+ dtype=torch.float32,
+ num_gpus=1,
+ )
+
+ pipeline = BooguImagePipeline(od_config=od_config)
+
+ (source,) = pipeline.weights_sources
+ assert source.revision == revision
+ prefetch.assert_called_once_with(
+ "dummy-boogu",
+ ["scheduler", "vae", "mllm", "processor"],
+ local_files_only=True,
+ revision=revision,
+ )
+ for loader in (scheduler_loader, mllm_loader, processor_loader, vae_loader):
+ assert loader.call_args.kwargs["revision"] == revision
+
+
def test_constructor_weights_sources(boogu_pipeline):
(source,) = boogu_pipeline.weights_sources
assert source.model_or_path == "dummy-boogu"
diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index 796de06f1..f2527bd99 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -1019,11 +1019,19 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest
cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True}))
assert cfg is not None
assert list(cfg.hints) == ["edge"]
+ assert cfg.hints["edge"].control_weight == 1.0
+ assert cfg.normalized_control_weights == [1.0]
assert cfg.guidance_scale == 3.0
assert cfg.control_guidance == 1.5
assert cfg.flow_shift == 10.0
assert cfg.num_video_frames_per_chunk == 93
assert cfg.share_vision_temporal_positions is True
+ assert cfg.emphasize_control_in_prompt is True
+ prompt_ablation_cfg = transfer.resolve_transfer_config(
+ make_sampling_params(extra_args={"edge": True, "emphasize_control_in_prompt": False})
+ )
+ assert prompt_ablation_cfg is not None
+ assert prompt_ablation_cfg.emphasize_control_in_prompt is False
# fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies.
defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True}))
assert defaulted_fps_cfg is not None
@@ -1087,11 +1095,46 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest
)
additional = preprocess(request).prompt["additional_information"]
assert (request.sampling_params.height, request.sampling_params.width) == (192, 320)
+ assert request.sampling_params.extra_args["_cosmos3_transfer_requested_size"] == {"width": 32, "height": 16}
assert tuple(additional["preprocessed_transfer_video"].shape) == (1, 3, 4, 192, 320)
assert additional["transfer_input_fps"] == 12.5
assert "preprocessed_video" not in additional
+def test_transfer_control_weight_validation_and_normalization() -> None:
+ from vllm_omni.diffusion.models.cosmos3 import transfer
+
+ single = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": 0.4}}))
+ assert single is not None
+ assert single.hints["edge"].control_weight == 0.4
+ assert single.normalized_control_weights == [1.0]
+
+ multiple = transfer.resolve_transfer_config(
+ make_sampling_params(
+ extra_args={
+ "edge": {"control_weight": 1.0},
+ "depth": {"control_weight": 3.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)},
+ }
+ )
+ )
+ assert multiple is not None
+ assert multiple.normalized_control_weights == [0.25, 0.75]
+
+ with pytest.raises(ValueError, match="Unsupported.*weight"):
+ transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"weight": 0.5}}))
+ with pytest.raises(ValueError, match="finite and non-negative"):
+ transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": -0.5}}))
+ with pytest.raises(ValueError, match="positive sum"):
+ transfer.resolve_transfer_config(
+ make_sampling_params(
+ extra_args={
+ "edge": {"control_weight": 0.0},
+ "depth": {"control_weight": 0.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)},
+ }
+ )
+ )
+
+
def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
"""When fps and frame_rate differ, transfer must pick frame_rate -- the same
precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real
@@ -1352,6 +1395,63 @@ def test_format_and_tokenize_prompts_applies_video_templates_and_system_override
]
+def test_format_and_tokenize_prompts_applies_transfer_prompt_contract(make_cosmos3_pipeline) -> None:
+ from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+ COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE,
+ COSMOS3_TRANSFER_SYSTEM_PROMPT,
+ )
+
+ pipeline = make_cosmos3_pipeline()
+ calls = _capture_tokenize_calls(pipeline)
+ directive = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names="edge")
+
+ pipeline._format_and_tokenize_prompts(
+ "A robot",
+ "bad",
+ num_frames=48,
+ frame_rate=24,
+ height=720,
+ width=1280,
+ max_sequence_length=32,
+ sp=SimpleNamespace(
+ extra_args={
+ "use_resolution_template": True,
+ "system_prompt": "request-level system prompt",
+ }
+ ),
+ use_system_prompt=True,
+ is_t2i=False,
+ system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT,
+ prompt_suffix=directive,
+ use_duration_template=True,
+ use_resolution_template=True,
+ negative_metadata_mode="same",
+ )
+
+ metadata = "The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."
+ assert calls[0]["text"] == f"A robot. {metadata} {directive}"
+ assert calls[1]["text"] == f"bad. {metadata}"
+ assert all(call["use_system_prompt"] is True for call in calls)
+ assert all(call["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT for call in calls)
+
+
+def test_format_and_tokenize_prompts_rejects_unknown_negative_metadata_mode(make_cosmos3_pipeline) -> None:
+ pipeline = make_cosmos3_pipeline()
+
+ with pytest.raises(ValueError, match="negative_metadata_mode"):
+ pipeline._format_and_tokenize_prompts(
+ "A robot",
+ "bad",
+ num_frames=48,
+ frame_rate=24,
+ height=720,
+ width=1280,
+ max_sequence_length=32,
+ sp=SimpleNamespace(extra_args={}),
+ negative_metadata_mode="unexpected",
+ )
+
+
def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT
@@ -1421,6 +1521,140 @@ def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_
)
+def test_format_and_tokenize_prompts_transfer_aspect_ratio_override(make_cosmos3_pipeline) -> None:
+ import json
+
+ pipeline = make_cosmos3_pipeline()
+ calls = _capture_tokenize_calls(pipeline)
+
+ pipeline._format_and_tokenize_prompts(
+ '{"caption": "A robot", "aspect_ratio": "4,3"}',
+ "",
+ num_frames=48,
+ frame_rate=24,
+ height=480,
+ width=832,
+ max_sequence_length=32,
+ sp=SimpleNamespace(extra_args={"aspect_ratio": "4:3"}),
+ is_t2i=False,
+ aspect_ratio_override="16,9",
+ )
+
+ assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9"
+
+
+@pytest.mark.parametrize("rank_zero", [True, False])
+def test_format_and_tokenize_prompts_corrects_conflicting_aspect_ratio_metadata_on_every_rank(
+ make_cosmos3_pipeline,
+ monkeypatch: pytest.MonkeyPatch,
+ mocker,
+ rank_zero: bool,
+) -> None:
+ import json
+
+ from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+ pipeline = make_cosmos3_pipeline()
+ calls = _capture_tokenize_calls(pipeline)
+ monkeypatch.setattr(pipeline_cosmos3, "_is_rank_zero", lambda: rank_zero)
+ warning = mocker.patch.object(pipeline_cosmos3.logger, "warning")
+
+ pipeline._format_and_tokenize_prompts(
+ '{"caption": "A robot", "aspect_ratio": "4,3"}',
+ "",
+ num_frames=48,
+ frame_rate=24,
+ height=720,
+ width=1280,
+ max_sequence_length=32,
+ sp=SimpleNamespace(extra_args={}),
+ is_t2i=False,
+ )
+
+ assert json.loads(calls[0]["text"])["aspect_ratio"] == "16,9"
+ if rank_zero:
+ rendered_warning = warning.call_args.args[0] % warning.call_args.args[1:]
+ assert "JSON prompt aspect_ratio='4,3' conflicts with the generated 1280x720 canvas" in rendered_warning
+ else:
+ warning.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("width", "height", "aspect_ratio"),
+ [
+ (1104, 816, "4,3"),
+ (832, 468, "16,9"),
+ ],
+)
+def test_format_and_tokenize_prompts_keeps_nearest_canonical_aspect_ratio(
+ make_cosmos3_pipeline,
+ mocker,
+ width: int,
+ height: int,
+ aspect_ratio: str,
+) -> None:
+ import json
+
+ from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+ pipeline = make_cosmos3_pipeline()
+ calls = _capture_tokenize_calls(pipeline)
+ warning = mocker.patch.object(pipeline_cosmos3.logger, "warning")
+
+ pipeline._format_and_tokenize_prompts(
+ json.dumps({"caption": "A robot", "aspect_ratio": aspect_ratio}),
+ "",
+ num_frames=48,
+ frame_rate=24,
+ height=height,
+ width=width,
+ max_sequence_length=32,
+ sp=SimpleNamespace(extra_args={}),
+ is_t2i=False,
+ )
+
+ assert json.loads(calls[0]["text"])["aspect_ratio"] == aspect_ratio
+ warning.assert_not_called()
+
+
+def test_transfer_bucket_selection_warns_on_size_and_aspect_ratio_conflicts(
+ make_cosmos3_pipeline,
+ mocker,
+) -> None:
+ from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+
+ pipeline = make_cosmos3_pipeline()
+ warning = mocker.patch.object(pipeline_cosmos3.logger, "warning")
+ sp = make_sampling_params(
+ width=832,
+ height=480,
+ extra_args={
+ "resolution": "480",
+ "aspect_ratio": "4:3",
+ "_cosmos3_transfer_requested_size": {"width": 1280, "height": 720},
+ },
+ )
+
+ height, width, aspect_ratio = pipeline._transfer_bucket_size(sp, (480, 832))
+ assert (height, width, aspect_ratio) == (480, 832, "16,9")
+
+ pipeline._warn_transfer_bucket_conflicts(
+ sp,
+ "A transfer prompt",
+ source_hw=(480, 832),
+ height=height,
+ width=width,
+ aspect_ratio=aspect_ratio,
+ )
+
+ rendered_warnings = [call.args[0] % call.args[1:] for call in warning.call_args_list]
+ assert any("ignores requested size=1280x720 (WxH)" in message for message in rendered_warnings)
+ assert any(
+ "requested aspect_ratio='4:3' conflicts with the control-selected 16,9 bucket" in message
+ for message in rendered_warnings
+ )
+
+
def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None:
import json
@@ -1740,6 +1974,7 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
control_guidance=1.5,
control_guidance_interval=None,
control_latents=[torch.zeros_like(latents)],
+ control_weights=[1.0],
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
velocity_mask=velocity_mask,
condition_latents=torch.zeros_like(latents),
@@ -1750,6 +1985,9 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
(2, False),
(1, True),
]
+ assert pipeline.transformer.calls[0]["kwargs"]["control_weights"] == [1.0]
+ assert "control_weights" not in pipeline.transformer.calls[1]["kwargs"]
+ assert pipeline.transformer.calls[2]["kwargs"]["control_weights"] == [1.0]
torch.testing.assert_close(result, torch.full_like(latents, 254.0))
@@ -1857,14 +2095,36 @@ def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline,
torch.testing.assert_close(result, torch.full_like(latents, 508.0))
-@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)])
-def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None:
+@pytest.mark.parametrize(
+ ("hint_key", "emphasize_control", "expected_fps", "negative_prompt"),
+ [
+ ("edge", None, 8.0, None),
+ ("wsm", None, 10.0, ""),
+ ("edge", False, 8.0, "custom negative"),
+ ],
+)
+def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_wsm(
+ make_cosmos3_pipeline,
+ hint_key: str,
+ emphasize_control: bool | None,
+ expected_fps: float,
+ negative_prompt: str | None,
+) -> None:
+ from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+ COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE,
+ COSMOS3_TRANSFER_SYSTEM_PROMPT,
+ )
+
pipeline = make_cosmos3_pipeline()
captured: dict[str, Any] = {}
def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
- del prompt, negative_prompt, num_frames, height, width, args, kwargs
+ del num_frames, height, width, args
+ captured["prompt"] = prompt
+ captured["negative_prompt"] = negative_prompt
captured["format_frame_rate"] = frame_rate
+ captured["format_kwargs"] = kwargs
+ captured["format_aspect_ratio"] = kwargs["aspect_ratio_override"]
return _ids(2), _mask(), _ids(1), _mask()
def fake_encode(video: torch.Tensor) -> torch.Tensor:
@@ -1887,7 +2147,7 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
captured["shared_kwargs"] = kwargs["shared_kwargs"]
return kwargs["latents"]
- pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
+ pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
pipeline._format_and_tokenize_prompts = fake_format
pipeline._encode_video_tensor = fake_encode
pipeline._prepare_transfer_latents = fake_prepare
@@ -1901,33 +2161,49 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8)
+ extra_args = {
+ hint_key: {"control": control},
+ "max_frames": 5,
+ "num_video_frames_per_chunk": 5,
+ "show_control_condition": True,
+ }
+ if emphasize_control is not None:
+ extra_args["emphasize_control_in_prompt"] = emphasize_control
+ prompt_data = {
+ "prompt": "transfer",
+ "modalities": ["video"],
+ "additional_information": {
+ "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16),
+ "transfer_input_fps": 8.0,
+ },
+ }
+ if negative_prompt is not None:
+ prompt_data["negative_prompt"] = negative_prompt
request = SimpleNamespace(
- prompts=[
- {
- "prompt": "transfer",
- "modalities": ["video"],
- "additional_information": {
- "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16),
- "transfer_input_fps": 8.0,
- },
- }
- ],
+ prompts=[prompt_data],
sampling_params=make_sampling_params(
height=16,
width=16,
# fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10).
- extra_args={
- hint_key: {"control": control},
- "max_frames": 5,
- "num_video_frames_per_chunk": 5,
- "show_control_condition": True,
- },
+ extra_args=extra_args,
),
)
output = pipeline.forward(request)
assert captured["format_frame_rate"] == expected_fps
+ assert captured["format_aspect_ratio"] == "1,1"
+ expected_negative_prompt = "" if negative_prompt is None else negative_prompt
+ assert captured["negative_prompt"] == expected_negative_prompt
+ assert captured["format_kwargs"]["use_system_prompt"] is True
+ assert captured["format_kwargs"]["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT
+ assert captured["format_kwargs"]["use_duration_template"] is True
+ assert captured["format_kwargs"]["use_resolution_template"] is True
+ assert captured["format_kwargs"]["negative_metadata_mode"] == "same"
+ expected_suffix = (
+ None if emphasize_control is False else COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_key)
+ )
+ assert captured["format_kwargs"]["prompt_suffix"] == expected_suffix
assert captured["shared_kwargs"]["fps"] == expected_fps
assert captured["flow_shifts"] == [10.0]
# Transfer applies the V2V flow shift when building its timestep schedule.
@@ -1943,7 +2219,7 @@ def test_forward_transfer_runs_multichunk_overlap_path(
pipeline = make_cosmos3_pipeline()
captured: dict[str, Any] = {"targets": [], "conditional_frames": []}
- pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16)
+ pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target)
diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
index 5dd8a875a..2bfa7b558 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
@@ -504,6 +504,35 @@ def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
assert generator_param.device == accelerator_device
+def test_multi_control_attention_weights_target_outputs() -> None:
+ from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
+
+ class ControlValueAttention:
+ def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+ del k
+ # Each per-control K/V layout is [text, control_i, target].
+ return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1)
+
+ attention = SimpleNamespace(multi_control_attn=ControlValueAttention())
+ q = torch.zeros(1, 3, 1, 1)
+ k = torch.zeros_like(q)
+ v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]])
+ text_kv = torch.zeros(1, 1, 1, 1)
+
+ output = Cosmos3CrossAttention._forward_multi_control(
+ attention,
+ q,
+ k,
+ v,
+ text_kv,
+ text_kv,
+ control_token_sizes=(1, 1),
+ control_weights=(0.25, 0.75),
+ )
+
+ torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]]))
+
+
def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
@@ -521,9 +550,21 @@ def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatc
video_shape=(1, 2, 2),
fps=24.0,
control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
+ control_weights=[1.0, 3.0],
)
assert tuple(output.shape) == tuple(hidden_states.shape)
+ with pytest.raises(ValueError, match="control_weights length"):
+ model(
+ hidden_states=hidden_states,
+ timestep=torch.tensor([1.0]),
+ text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+ text_mask=torch.ones(1, 2, dtype=torch.long),
+ video_shape=(1, 2, 2),
+ fps=24.0,
+ control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)],
+ control_weights=[1.0],
+ )
with pytest.raises(ValueError, match="control latent shape"):
model(
hidden_states=hidden_states,
diff --git a/tests/cosmos3/test_session_memory_equivalence.py b/tests/diffusion/models/cosmos3/test_session_memory_equivalence.py
similarity index 100%
rename from tests/cosmos3/test_session_memory_equivalence.py
rename to tests/diffusion/models/cosmos3/test_session_memory_equivalence.py
diff --git a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py b/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py
deleted file mode 100644
index c7e977ecf..000000000
--- a/tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import pytest
-import torch.nn as nn
-
-pytest.importorskip("dreamid_omni.modules.model")
-
-from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
- model = object.__new__(FusionModel)
- nn.Module.__init__(model)
- model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
- model.attn = nn.Identity()
-
- conditions = getattr(model, "_hsdp_shard_conditions", None)
-
- assert conditions is not None
- assert len(conditions) == 1
-
- matched = []
- for name, module in model.named_modules():
- if any(cond(name, module) for cond in conditions):
- matched.append(name)
-
- assert matched == ["fused_blocks.0", "fused_blocks.1"]
diff --git a/tests/diffusion/models/dreamzero/test_pipeline_state.py b/tests/diffusion/models/dreamzero/test_pipeline_state.py
index 366876ef7..ff349d6af 100644
--- a/tests/diffusion/models/dreamzero/test_pipeline_state.py
+++ b/tests/diffusion/models/dreamzero/test_pipeline_state.py
@@ -6,7 +6,10 @@ from types import SimpleNamespace
import pytest
-from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
+ DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION,
+ DreamZeroPipeline,
+)
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +35,37 @@ def test_dreamzero_pipeline_state_is_session_keyed() -> None:
assert session_b.call_count == 3
+def test_dreamzero_kv_spec_accounts_for_model_owned_cuda_state(monkeypatch) -> None:
+ pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+ pipeline.transformer = SimpleNamespace(
+ frame_seqlen=16,
+ blocks=[
+ SimpleNamespace(
+ self_attn=SimpleNamespace(
+ max_attention_size=96,
+ tp_num_heads=4,
+ )
+ )
+ ],
+ text_len=8,
+ model_type="t2v",
+ num_layers=2,
+ num_heads=4,
+ dim=256,
+ num_frame_per_block=4,
+ num_action_per_block=2,
+ num_state_per_block=1,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero.get_classifier_free_guidance_world_size",
+ lambda: 1,
+ )
+
+ spec = pipeline.ar_diffusion_kv_cache_spec()
+
+ assert spec.model_owned_state_bytes_per_session == DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION
+
+
def test_dreamzero_pipeline_state_follows_runner_lifecycle_notifications() -> None:
pipeline = _empty_pipeline()
diff --git a/tests/diffusion/models/flux2/test_flux2_quantization.py b/tests/diffusion/models/flux2/test_flux2_quantization.py
new file mode 100644
index 000000000..3c461acf5
--- /dev/null
+++ b/tests/diffusion/models/flux2/test_flux2_quantization.py
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for FLUX.2-dev quantization routing."""
+
+import pytest
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _RecordingBlock(nn.Module):
+ def __init__(self, *, quant_config=None, prefix: str = "", **kwargs):
+ super().__init__()
+ self.quant_config = quant_config
+ self.prefix = prefix
+
+
+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
+
+ monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
+ monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
+
+ quant_config = Fp8Config(
+ is_checkpoint_fp8_serialized=False,
+ activation_scheme="dynamic",
+ )
+ model = flux2_transformer.Flux2Transformer2DModel(
+ in_channels=4,
+ out_channels=4,
+ num_layers=2,
+ num_single_layers=3,
+ attention_head_dim=8,
+ num_attention_heads=2,
+ joint_attention_dim=16,
+ timestep_guidance_channels=8,
+ axes_dims_rope=(2, 2, 2, 2),
+ quant_config=quant_config,
+ )
+
+ assert [block.quant_config for block in model.transformer_blocks] == [quant_config, quant_config]
+ assert [block.prefix for block in model.transformer_blocks] == [
+ "transformer_blocks.0",
+ "transformer_blocks.1",
+ ]
+ assert [block.quant_config for block in model.single_transformer_blocks] == [
+ quant_config,
+ quant_config,
+ quant_config,
+ ]
+ assert [block.prefix for block in model.single_transformer_blocks] == [
+ "single_transformer_blocks.0",
+ "single_transformer_blocks.1",
+ "single_transformer_blocks.2",
+ ]
+
+ assert isinstance(model.x_embedder, nn.Linear)
+ assert isinstance(model.context_embedder, nn.Linear)
+ assert isinstance(model.proj_out, nn.Linear)
diff --git a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py
index 3900c0c94..ff9e0e110 100644
--- a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py
+++ b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py
@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
Flux2Pipeline,
_resolve_component_quant_config,
)
+from vllm_omni.quantization import ComponentQuantizationConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_resolve_component_quant_config_routes_components():
text_encoder_config = object()
- component_config = Mock()
- component_config.resolve.side_effect = {
- "text_encoder": text_encoder_config,
- "transformer": None,
- }.get
+ component_config = ComponentQuantizationConfig(
+ component_configs={"text_encoder": text_encoder_config},
+ default_config=None,
+ )
assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
assert _resolve_component_quant_config(component_config, "transformer") is None
diff --git a/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py b/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py
new file mode 100644
index 000000000..e5581cd77
--- /dev/null
+++ b/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py
@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+from cache_dit import ForwardPattern
+from torch import nn
+
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import (
+ HiDreamO1TextModel,
+)
+from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import (
+ HiDreamO1ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _empty_text_model() -> HiDreamO1TextModel:
+ model = HiDreamO1TextModel.__new__(HiDreamO1TextModel)
+ nn.Module.__init__(model)
+ model.layers = nn.ModuleList([nn.Identity()])
+ return model
+
+
+def test_cache_dit_targets_the_decoder_module_that_owns_layers():
+ text_model = _empty_text_model()
+ pipeline = SimpleNamespace(
+ _dit_modules=HiDreamO1ImagePipeline._dit_modules,
+ model=SimpleNamespace(
+ model=SimpleNamespace(language_model=text_model),
+ ),
+ )
+
+ with (
+ patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") as block_adapter_cls,
+ patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") as cache_dit,
+ ):
+ backend = CacheDiTBackend()
+ backend.enable(pipeline)
+ backend.refresh(pipeline, num_inference_steps=50)
+
+ adapter_config = HiDreamO1TextModel._cache_dit_adapter_config
+ assert isinstance(adapter_config, CacheDiTAdapterConfig)
+ assert adapter_config.block_forward_patterns == {"layers": ForwardPattern.Pattern_3}
+ assert adapter_config.has_separate_cfg is True
+
+ block_adapter_cls.assert_called_once_with(
+ transformer=text_model,
+ blocks=[text_model.layers],
+ forward_pattern=[ForwardPattern.Pattern_3],
+ has_separate_cfg=True,
+ check_forward_pattern=True,
+ )
+ cache_dit.enable_cache.assert_called_once()
+ enable_call = cache_dit.enable_cache.call_args
+ assert enable_call.args == (block_adapter_cls.return_value,)
+ assert enable_call.kwargs["cache_config"].Fn_compute_blocks == 1
+ assert enable_call.kwargs["calibrator_config"] is None
+ cache_dit.refresh_context.assert_called_once_with(
+ text_model,
+ num_inference_steps=50,
+ verbose=True,
+ )
diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
index dfb383e87..b62fe950e 100644
--- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
+++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend():
pipeline._ensure_grouped_attention_backend_supported(2)
+def test_scheduler_paged_step_execution_is_rejected():
+ pipeline = _pipeline()
+ pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER
+ state = _state("paged-step", 0)
+
+ with pytest.raises(ValueError, match="request-level execution only"):
+ pipeline.prepare_encode(state)
+ with pytest.raises(ValueError, match="request-level execution only"):
+ pipeline.denoise_step(InputBatch.make_batch([state]))
+
+
def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
pipeline = _pipeline()
pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})
diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
index 543b93a25..65fc0e67b 100644
--- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
+++ b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for ImageKVCacheManager.
Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
@@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer"
-NUM_HEADS = 4
-NUM_KV_HEADS = 2
+NUM_HEADS = 32
+NUM_KV_HEADS = 8
HEAD_DIM = 16
IMAGE_TOKEN_LEN = 8
SUFFIX_TOKEN_LEN = 3
@@ -35,8 +35,23 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM)
class MockAttention(nn.Module):
def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs):
super().__init__()
+ self.paged_kv_active = False
+ self.calls = []
+ self.paged_calls = []
+ self.paged_metadata = []
+
+ def is_paged_kv_active(self):
+ return self.paged_kv_active
def forward(self, query, key, value, attn_metadata=None, **kwargs):
+ self.calls.append((query, key, value))
+ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+ if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None:
+ self.paged_calls.append((query, key, value))
+ self.paged_metadata.append(attn_metadata)
+ if attn_metadata is not None and attn_metadata.joint_query is not None:
+ return torch.cat((attn_metadata.joint_query, query), dim=1)
return query
@@ -99,6 +114,7 @@ def _call_mgr(
shard_image_size=None,
gen_timestep_scatter_index=None,
position_ids=None,
+ full_attn_spans=None,
):
query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -115,9 +131,127 @@ def _call_mgr(
shard_image_size=shard_image_size,
gen_timestep_scatter_index=gen_timestep_scatter_index,
position_ids=position_ids,
+ full_attn_spans=full_attn_spans,
)
+def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
+ mgr = _make_cache_mgr()
+
+ assert isinstance(mgr, nn.Module)
+ assert dict(mgr.named_modules())["attn"] is mgr.attn
+ assert mgr.state_dict() == {}
+
+
+def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
+ from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+ mgr = _make_cache_mgr()
+ mgr.attn.paged_kv_active = True
+ bs = 2
+ prefix_lens = torch.tensor([[2], [4]], dtype=torch.long)
+ first_q_len = 12
+ first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0)
+ full_attn_spans = [[(2, 12)], [(2, 12)]]
+
+ with set_forward_context(), override_paged_kv_adapter(object()):
+ _call_mgr(
+ mgr,
+ bs,
+ q_len=first_q_len,
+ seq_len=first_q_len,
+ key_flat=first_k,
+ value_flat=first_v,
+ first_step=True,
+ gen_timestep_scatter_index=prefix_lens,
+ full_attn_spans=full_attn_spans,
+ )
+
+ assert mgr.image_kv_cache_map is None
+ assert mgr.image_kv_cache_lens is None
+ assert len(mgr.attn.paged_calls) == 1
+ assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans
+
+ current_q_len = IMAGE_TOKEN_LEN
+ current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0)
+ _call_mgr(
+ mgr,
+ bs,
+ q_len=current_q_len,
+ seq_len=prefix_lens[-1].item() + current_q_len,
+ key_flat=current_k,
+ value_flat=current_v,
+ position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1),
+ full_attn_spans=full_attn_spans,
+ )
+
+ assert len(mgr.attn.paged_calls) == 2
+ assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
+ assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans
+
+
+def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None:
+ from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+ mgr = _make_cache_mgr()
+ mgr.attn.paged_kv_active = True
+ prefix_len = 4
+ prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0)
+ mgr._injected_ar_kv = [(prefix_key, prefix_value)]
+ current_key, current_value = _make_known_kv(3, base=1.0)
+
+ with set_forward_context(), override_paged_kv_adapter(object()):
+ with pytest.raises(NotImplementedError, match="does not support imported AR KV"):
+ _call_mgr(
+ mgr,
+ bs=1,
+ q_len=3,
+ seq_len=7,
+ key_flat=current_key,
+ value_flat=current_value,
+ first_step=True,
+ gen_timestep_scatter_index=torch.tensor([[0]]),
+ full_attn_spans=[[(4, 7)]],
+ )
+
+
+@pytest.mark.parametrize(
+ ("paged_kv_active", "expected_kv_heads"),
+ [(False, NUM_HEADS), (True, NUM_KV_HEADS)],
+)
+def test_gqa_kv_stays_compressed_for_paged_attention(
+ paged_kv_active: bool,
+ expected_kv_heads: int,
+) -> None:
+ mgr = _make_cache_mgr()
+ mgr.attn.paged_kv_active = paged_kv_active
+ q_len = 3
+ key, value = _make_known_kv(q_len)
+
+ _call_mgr(
+ mgr,
+ bs=1,
+ q_len=q_len,
+ seq_len=q_len,
+ key_flat=key,
+ value_flat=value,
+ first_step=True,
+ num_image_tokens=0,
+ gen_timestep_scatter_index=_gen_timestep_index(1, 0),
+ full_attn_spans=[[(0, q_len)]] if paged_kv_active else None,
+ )
+
+ _, key_input, value_input = mgr.attn.calls[-1]
+ expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM)
+ assert key_input.shape == expected_shape
+ assert value_input.shape == expected_shape
+
+
# ============================================================
# Test 1: No AR KV — basic cache → reuse
# ============================================================
diff --git a/vllm_omni/diffusion/models/magi_human/__init__.py b/tests/diffusion/models/lingbot_world/__init__.py
similarity index 100%
rename from vllm_omni/diffusion/models/magi_human/__init__.py
rename to tests/diffusion/models/lingbot_world/__init__.py
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture
new file mode 100644
index 000000000..015a5aa91
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_model_index.json.fixture
@@ -0,0 +1,13 @@
+{
+ "_class_name": "LingBotWorldCausalDMDPipeline",
+ "_diffusers_version": "0.35.0.dev0",
+ "boundary_ratio": 0.9,
+ "image_encoder": [null, null],
+ "image_processor": [null, null],
+ "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+ "text_encoder": ["transformers", "UMT5EncoderModel"],
+ "tokenizer": ["transformers", "T5TokenizerFast"],
+ "transformer": ["diffusers", "CausalLingBotWorldTransformer3DModel"],
+ "transformer_2": [null, null],
+ "vae": ["diffusers", "AutoencoderKLWan"]
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture
new file mode 100644
index 000000000..d07408a04
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_official_shapes.json.fixture
@@ -0,0 +1,42 @@
+{
+ "sources": {
+ "config": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/config.json",
+ "shard_1_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00001-of-00008.safetensors",
+ "shard_8_header": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model-00008-of-00008.safetensors"
+ },
+ "dtype": "F32",
+ "shard_representatives": {
+ "diffusion_pytorch_model-00001-of-00008.safetensors": "blocks.0.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00002-of-00008.safetensors": "blocks.4.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00003-of-00008.safetensors": "blocks.10.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00004-of-00008.safetensors": "blocks.14.cam_scale_layer.bias",
+ "diffusion_pytorch_model-00005-of-00008.safetensors": "blocks.20.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00006-of-00008.safetensors": "blocks.25.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00007-of-00008.safetensors": "blocks.30.cam_injector_layer1.bias",
+ "diffusion_pytorch_model-00008-of-00008.safetensors": "blocks.35.cam_injector_layer1.bias"
+ },
+ "parameter_shapes": {
+ "blocks.0.cam_injector_layer1.bias": [5120],
+ "blocks.4.cam_injector_layer1.bias": [5120],
+ "blocks.10.cam_injector_layer1.bias": [5120],
+ "blocks.14.cam_scale_layer.bias": [5120],
+ "blocks.20.cam_injector_layer1.bias": [5120],
+ "blocks.25.cam_injector_layer1.bias": [5120],
+ "blocks.30.cam_injector_layer1.bias": [5120],
+ "blocks.35.cam_injector_layer1.bias": [5120],
+ "patch_embedding.weight": [5120, 36, 1, 2, 2],
+ "patch_embedding_wancamctrl.weight": [5120, 1536],
+ "c2ws_hidden_states_layer1.weight": [5120, 5120],
+ "c2ws_hidden_states_layer2.weight": [5120, 5120],
+ "blocks.0.modulation": [1, 6, 5120],
+ "blocks.0.self_attn.q.weight": [5120, 5120],
+ "blocks.0.ffn.0.weight": [13824, 5120],
+ "blocks.0.ffn.2.weight": [5120, 13824],
+ "blocks.0.cam_injector_layer1.weight": [5120, 5120],
+ "blocks.0.cam_injector_layer2.weight": [5120, 5120],
+ "blocks.0.cam_scale_layer.weight": [5120, 5120],
+ "blocks.0.cam_shift_layer.weight": [5120, 5120],
+ "head.modulation": [1, 2, 5120],
+ "head.head.weight": [64, 5120]
+ }
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture
new file mode 100644
index 000000000..950d26fae
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_scheduler_config.json.fixture
@@ -0,0 +1,30 @@
+{
+ "_class_name": "UniPCMultistepScheduler",
+ "_diffusers_version": "0.35.0.dev0",
+ "beta_end": 0.02,
+ "beta_schedule": "linear",
+ "beta_start": 0.0001,
+ "disable_corrector": [],
+ "dynamic_thresholding_ratio": 0.995,
+ "final_sigmas_type": "zero",
+ "flow_shift": 5.0,
+ "lower_order_final": true,
+ "num_train_timesteps": 1000,
+ "predict_x0": true,
+ "prediction_type": "flow_prediction",
+ "rescale_betas_zero_snr": false,
+ "sample_max_value": 1.0,
+ "solver_order": 2,
+ "solver_p": null,
+ "solver_type": "bh2",
+ "steps_offset": 0,
+ "thresholding": false,
+ "time_shift_type": "exponential",
+ "timestep_spacing": "linspace",
+ "trained_betas": null,
+ "use_beta_sigmas": false,
+ "use_dynamic_shifting": false,
+ "use_exponential_sigmas": false,
+ "use_flow_sigmas": true,
+ "use_karras_sigmas": false
+}
diff --git a/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture
new file mode 100644
index 000000000..2a2c93efd
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/fixtures/lingbot_world_weight_index_names.json.fixture
@@ -0,0 +1,112 @@
+{
+ "source": "https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers/resolve/main/transformer/diffusion_pytorch_model.safetensors.index.json",
+ "index_parameter_count": 1421,
+ "canonical_name_set_sha256": "5d31b7f64fd5e5e271b789b2dbe50fb7bbf146bd8b750646bdc9abb4e5b8139d",
+ "canonical_name_to_shard_sha256": "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992",
+ "canonical_serialization": "sorted lines formatted as parameter_name=shard_filename with LF separators and no trailing LF",
+ "shards": {
+ "diffusion_pytorch_model-00001-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "81a3d4778de604e75609aae7dac3d8dcbc081811e7748267fe953bbeb375e79d"
+ },
+ "diffusion_pytorch_model-00002-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "4e36e2596a6468f29179fa5795396f2e9a7f7d8debc912b34a6f6c2c7c7c610d"
+ },
+ "diffusion_pytorch_model-00003-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "6f74a7b2cf7444c2daa2b3fc32bd3787589b35543c230c5c13b89f0f0f2a0998"
+ },
+ "diffusion_pytorch_model-00004-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "ec5c967ed5271ba7834337c4ee93a73f8ed0eead4492d4920624251c65c058ac"
+ },
+ "diffusion_pytorch_model-00005-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "ad1fdab5830dd9ce43ea007f112b4be223327b7c8197c766d9cf9219e74d619b"
+ },
+ "diffusion_pytorch_model-00006-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "fd58f0d86be42d8afcf55d3e07d8a22f7f6ef5064da931a600c890b71b6dfbae"
+ },
+ "diffusion_pytorch_model-00007-of-00008.safetensors": {
+ "parameter_count": 180,
+ "parameter_names_sha256": "305a1714b5f4fb54d5877c593b451273ab1490d90fd0896a299d96e140295047"
+ },
+ "diffusion_pytorch_model-00008-of-00008.safetensors": {
+ "parameter_count": 161,
+ "parameter_names_sha256": "4c18d94d96168682eb212f93e3cdf7313928776d8d9e6e4de0255fe899decb25"
+ }
+ },
+ "top_level_modules": [
+ "blocks",
+ "c2ws_hidden_states_layer1",
+ "c2ws_hidden_states_layer2",
+ "head",
+ "patch_embedding",
+ "patch_embedding_wancamctrl",
+ "text_embedding",
+ "time_embedding",
+ "time_projection"
+ ],
+ "non_block_parameters": [
+ "c2ws_hidden_states_layer1.bias",
+ "c2ws_hidden_states_layer1.weight",
+ "c2ws_hidden_states_layer2.bias",
+ "c2ws_hidden_states_layer2.weight",
+ "head.head.bias",
+ "head.head.weight",
+ "head.modulation",
+ "patch_embedding.bias",
+ "patch_embedding.weight",
+ "patch_embedding_wancamctrl.bias",
+ "patch_embedding_wancamctrl.weight",
+ "text_embedding.0.bias",
+ "text_embedding.0.weight",
+ "text_embedding.2.bias",
+ "text_embedding.2.weight",
+ "time_embedding.0.bias",
+ "time_embedding.0.weight",
+ "time_embedding.2.bias",
+ "time_embedding.2.weight",
+ "time_projection.1.bias",
+ "time_projection.1.weight"
+ ],
+ "block_parameter_suffixes": [
+ "cam_injector_layer1.bias",
+ "cam_injector_layer1.weight",
+ "cam_injector_layer2.bias",
+ "cam_injector_layer2.weight",
+ "cam_scale_layer.bias",
+ "cam_scale_layer.weight",
+ "cam_shift_layer.bias",
+ "cam_shift_layer.weight",
+ "cross_attn.k.bias",
+ "cross_attn.k.weight",
+ "cross_attn.norm_k.weight",
+ "cross_attn.norm_q.weight",
+ "cross_attn.o.bias",
+ "cross_attn.o.weight",
+ "cross_attn.q.bias",
+ "cross_attn.q.weight",
+ "cross_attn.v.bias",
+ "cross_attn.v.weight",
+ "ffn.0.bias",
+ "ffn.0.weight",
+ "ffn.2.bias",
+ "ffn.2.weight",
+ "modulation",
+ "norm3.bias",
+ "norm3.weight",
+ "self_attn.k.bias",
+ "self_attn.k.weight",
+ "self_attn.norm_k.weight",
+ "self_attn.norm_q.weight",
+ "self_attn.o.bias",
+ "self_attn.o.weight",
+ "self_attn.q.bias",
+ "self_attn.q.weight",
+ "self_attn.v.bias",
+ "self_attn.v.weight"
+ ]
+}
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py
new file mode 100644
index 000000000..df012777f
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+ LINGBOT_CAMERA_ACTION_SCHEMA,
+ LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+ LingBotCameraControlReducer,
+ integrate_lingbot_camera_actions,
+ parse_lingbot_camera_action_frames,
+)
+from vllm_omni.experimental.ar_diffusion.session import (
+ ARDiffusionSession,
+ ARDiffusionSessionEvent,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionChunkMetadata,
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _action_control(mode: str, **data) -> ARDiffusionControlInput:
+ return ARDiffusionControlInput(
+ track="camera",
+ schema=LINGBOT_CAMERA_ACTION_SCHEMA,
+ data={"mode": mode, **data},
+ )
+
+
+def _prepared_frames(prepared) -> tuple[tuple[str, ...], ...]:
+ assert len(prepared.controls) == 1
+ return prepared.controls[0].data["frames"]
+
+
+def test_state_controls_hold_release_and_preserve_short_pulse() -> None:
+ reducer = LingBotCameraControlReducer()
+ pressed = reducer.prepare(
+ current_controls={},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": ["W"], "client_ts_ms": 10}],
+ ),
+ ),
+ ),
+ ),
+ chunk_index=0,
+ )
+ assert _prepared_frames(pressed) == (("w",), ("w",), ("w",))
+ assert pressed.controls[0].to_dict()["data"]["frames"] == [["w"], ["w"], ["w"]]
+ reducer.commit(pressed)
+
+ held = reducer.prepare(
+ current_controls={"camera": pressed.controls[0]},
+ events=(),
+ chunk_index=1,
+ )
+ assert _prepared_frames(held) == (("w",), ("w",), ("w",))
+ reducer.commit(held)
+
+ pulse = reducer.prepare(
+ current_controls={"camera": held.controls[0]},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=2,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[
+ {"actions": ["j"], "client_ts_ms": 20},
+ {"actions": [], "client_ts_ms": 21},
+ ],
+ ),
+ ),
+ ),
+ ),
+ chunk_index=2,
+ )
+ assert _prepared_frames(pulse) == (("j",), (), ())
+ reducer.commit(pulse)
+
+ released = reducer.prepare(
+ current_controls={"camera": pulse.controls[0]},
+ events=(),
+ chunk_index=3,
+ )
+ assert _prepared_frames(released) == ((), (), ())
+
+
+def test_script_is_transactional_consumed_once_and_neutral_padded() -> None:
+ reducer = LingBotCameraControlReducer()
+ event = ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(
+ _action_control(
+ "script",
+ frames=[["w"], ["d"], ["l"], ["i"]],
+ ),
+ ),
+ )
+ first_attempt = reducer.prepare(
+ current_controls={},
+ events=(event,),
+ chunk_index=0,
+ )
+ assert _prepared_frames(first_attempt) == (("w",), ("d",), ("l",))
+
+ # An uncommitted failed attempt must not consume the reducer's script.
+ second_attempt = reducer.prepare(
+ current_controls={},
+ events=(event,),
+ chunk_index=0,
+ )
+ assert _prepared_frames(second_attempt) == (("w",), ("d",), ("l",))
+ reducer.commit(second_attempt)
+
+ tail = reducer.prepare(
+ current_controls={"camera": second_attempt.controls[0]},
+ events=(),
+ chunk_index=1,
+ )
+ assert _prepared_frames(tail) == (("i",), (), ())
+ reducer.commit(tail)
+ exhausted = reducer.prepare(
+ current_controls={"camera": tail.controls[0]},
+ events=(),
+ chunk_index=2,
+ )
+ assert _prepared_frames(exhausted) == ((), (), ())
+
+
+def test_trajectory_control_remains_an_expert_override() -> None:
+ reducer = LingBotCameraControlReducer()
+ trajectory = ARDiffusionControlInput(
+ track="camera",
+ schema=LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
+ data={"poses": [[[1.0]]], "intrinsics": [[1.0]]},
+ )
+ prepared = reducer.prepare(
+ current_controls={},
+ events=(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ controls=(trajectory,),
+ ),
+ ),
+ chunk_index=0,
+ )
+ assert prepared.controls == (trajectory,)
+
+
+def test_action_integrator_matches_lingbot_motion_and_calibration() -> None:
+ trajectory, pitch = integrate_lingbot_camera_actions(
+ [["w"], ["d"], ["l"]],
+ width=832,
+ height=480,
+ )
+
+ assert pitch == 0.0
+ torch.testing.assert_close(
+ trajectory.poses[0, :3, 3],
+ torch.tensor([0.0, 0.0, 0.05], dtype=torch.float64),
+ atol=1e-6,
+ rtol=0,
+ )
+ torch.testing.assert_close(
+ trajectory.poses[1, :3, 3],
+ torch.tensor([0.05, 0.0, 0.05], dtype=torch.float64),
+ atol=1e-6,
+ rtol=0,
+ )
+ expected_yaw = torch.tensor(
+ [
+ [0.9945219, 0.0, 0.10452846],
+ [0.0, 1.0, 0.0],
+ [-0.10452846, 0.0, 0.9945219],
+ ],
+ dtype=torch.float64,
+ )
+ torch.testing.assert_close(
+ trajectory.poses[2, :3, :3],
+ expected_yaw,
+ atol=1e-6,
+ rtol=0,
+ )
+ torch.testing.assert_close(
+ trajectory.intrinsics[0],
+ torch.tensor([500.0, 500.0, 416.0, 240.0]),
+ )
+
+ resized, _ = integrate_lingbot_camera_actions(
+ [[]],
+ width=416,
+ height=240,
+ )
+ scaled = resized.intrinsics[0] * torch.tensor([416 / 832, 240 / 480, 416 / 832, 240 / 480])
+ torch.testing.assert_close(
+ scaled,
+ torch.tensor([500.0, 500.0, 208.0, 120.0]),
+ )
+
+
+@pytest.mark.parametrize(
+ ("data", "message"),
+ [
+ ({"mode": "frames", "frames": [["x"], [], []]}, "W/A/S/D/I/J/K/L"),
+ ({"mode": "frames", "frames": [["w"]]}, "exactly one action list"),
+ ({"mode": "state", "frames": [[], [], []]}, "mode='frames'"),
+ ],
+)
+def test_model_tick_action_payload_validation(data, message: str) -> None:
+ with pytest.raises(ValueError, match=message):
+ parse_lingbot_camera_action_frames(data, expected_frames=3)
+
+
+class _TickConsumer:
+ def __init__(self) -> None:
+ self.ticks: list[ARDiffusionTickRequest] = []
+
+ async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
+ self.ticks.append(tick)
+ return ARDiffusionChunkMetadata.from_tick(tick)
+
+ def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
+ return output
+
+
+class _Lifecycle:
+ async def reset_session(self, session_id: str) -> None:
+ del session_id
+
+ async def close_session(self, session_id: str) -> None:
+ del session_id
+
+
+def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
+ async def exercise() -> None:
+ consumer = _TickConsumer()
+ reducer = LingBotCameraControlReducer()
+ session = ARDiffusionSession(
+ "world-actions",
+ tick_consumer=consumer,
+ lifecycle=_Lifecycle(),
+ max_pending_events=8,
+ request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
+ control_reducer=reducer,
+ )
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=1,
+ prompt="turn toward the window",
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": ["d"]}],
+ ),
+ ),
+ )
+ )
+ await session.next_chunk()
+
+ tick = consumer.ticks[0]
+ assert tick.prompt == "turn toward the window"
+ assert tick.applied_event_ids == (1,)
+ assert tick.controls[0].data["frames"] == (("d",), ("d",), ("d",))
+
+ await session.reset()
+ await session.accept_event(
+ ARDiffusionSessionEvent(
+ event_id=2,
+ controls=(
+ _action_control(
+ "state",
+ transitions=[{"actions": []}],
+ ),
+ ),
+ )
+ )
+ await session.next_chunk()
+ assert consumer.ticks[1].chunk_index == 0
+ assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
+
+ asyncio.run(exercise())
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py
new file mode 100644
index 000000000..f5567a706
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py
@@ -0,0 +1,543 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/transformer.py"
+_MODULE_NAME = "_lingbot_world_attention_under_test"
+_MISSING_MODULE = object()
+_STUBBED_MODULE_NAMES = (
+ "vllm",
+ "vllm.distributed",
+ "vllm.model_executor",
+ "vllm.model_executor.layers",
+ "vllm.model_executor.layers.conv",
+ "vllm.model_executor.layers.linear",
+ "vllm.model_executor.model_loader",
+ "vllm.model_executor.model_loader.weight_utils",
+ "vllm.model_executor.utils",
+ "vllm_omni",
+ "vllm_omni.diffusion",
+ "vllm_omni.diffusion.attention",
+ "vllm_omni.diffusion.attention.layer",
+ "vllm_omni.diffusion.layers",
+ "vllm_omni.diffusion.layers.norm",
+ "vllm_omni.diffusion.layers.rope",
+ "vllm_omni.experimental",
+ "vllm_omni.experimental.ar_diffusion",
+ "vllm_omni.experimental.ar_diffusion.kv_cache",
+ "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention",
+)
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+ with torch.inference_mode():
+ yield
+
+
+def _install_vllm_stubs() -> None:
+ for name in _STUBBED_MODULE_NAMES:
+ sys.modules[name] = ModuleType(name)
+
+ distributed = sys.modules["vllm.distributed"]
+ distributed.get_tensor_model_parallel_rank = lambda: 0
+ distributed.get_tensor_model_parallel_world_size = lambda: 1
+ distributed.tensor_model_parallel_all_reduce = lambda value: value
+
+ def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None:
+ for name, value in attrs.items():
+ setattr(weight, name, value)
+
+ sys.modules["vllm.model_executor.utils"].set_weight_attrs = set_weight_attrs
+ sys.modules["vllm.model_executor.model_loader.weight_utils"].default_weight_loader = lambda param, loaded_weight: (
+ param.data.copy_(loaded_weight)
+ )
+
+ class _Linear(nn.Module):
+ def __init__(
+ self,
+ input_size: int,
+ output_size: int,
+ *,
+ bias: bool = True,
+ return_bias: bool = False,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ del kwargs
+ self.return_bias = return_bias
+ self.weight = nn.Parameter(torch.empty(output_size, input_size))
+ self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if self.bias is not None:
+ nn.init.zeros_(self.bias)
+ self.calls = 0
+
+ def forward(self, value: torch.Tensor):
+ self.calls += 1
+ output = F.linear(value, self.weight, self.bias)
+ return (output, self.bias) if self.return_bias else output
+
+ linear = sys.modules["vllm.model_executor.layers.linear"]
+ linear.ColumnParallelLinear = _Linear
+ linear.RowParallelLinear = _Linear
+
+ class _QKVLinear(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ head_size: int,
+ total_num_heads: int,
+ total_num_kv_heads: int | None = None,
+ *,
+ bias: bool = True,
+ **kwargs,
+ ) -> None:
+ super().__init__()
+ del kwargs
+ self.num_heads = total_num_heads
+ self.num_kv_heads = total_num_kv_heads or total_num_heads
+ output_size = (self.num_heads + 2 * self.num_kv_heads) * head_size
+ self.weight = nn.Parameter(torch.empty(output_size, hidden_size))
+ self.bias = nn.Parameter(torch.empty(output_size)) if bias else None
+ nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+ if self.bias is not None:
+ nn.init.zeros_(self.bias)
+ self.calls = 0
+
+ def load_shard(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+ shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+ shard_size = param.shape[0] // 3
+ param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+ self.weight.weight_loader = load_shard
+ if self.bias is not None:
+ self.bias.weight_loader = load_shard
+
+ def forward(self, value: torch.Tensor):
+ self.calls += 1
+ return F.linear(value, self.weight, self.bias), None
+
+ linear.QKVParallelLinear = _QKVLinear
+
+ class _Conv3dLayer(nn.Conv3d):
+ def __init__(self, in_channels: int, out_channels: int, kernel_size, stride, **kwargs) -> None:
+ super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride, **kwargs)
+
+ sys.modules["vllm.model_executor.layers.conv"].Conv3dLayer = _Conv3dLayer
+
+ class _LayerNorm(nn.LayerNorm):
+ def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True) -> None:
+ super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
+
+ sys.modules["vllm_omni.diffusion.layers.norm"].LayerNorm = _LayerNorm
+
+ class _Attention(nn.Module):
+ def __init__(self, *args, softmax_scale: float, **kwargs) -> None:
+ super().__init__()
+ del args
+ self.softmax_scale = softmax_scale
+ self.causal = kwargs["causal"]
+ self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata=None,
+ ) -> torch.Tensor:
+ del attn_metadata
+ self.calls.append((query.detach().clone(), key.detach().clone(), value.detach().clone()))
+ scores = torch.einsum("bqhd,bkhd->bhqk", query, key) * self.softmax_scale
+ weights = scores.softmax(dim=-1)
+ return torch.einsum("bhqk,bkhd->bqhd", weights, value)
+
+ sys.modules["vllm_omni.diffusion.attention.layer"].Attention = _Attention
+
+ class _RotaryEmbeddingWan(nn.Module):
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__()
+ del args, kwargs
+ self.calls: list[tuple[torch.Tensor, torch.Tensor]] = []
+
+ def forward(self, value: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ self.calls.append((cos.detach().clone(), sin.detach().clone()))
+ return value
+
+ sys.modules["vllm_omni.diffusion.layers.rope"].RotaryEmbeddingWan = _RotaryEmbeddingWan
+
+ paged_attention = sys.modules["vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"]
+
+ class _PagedLayerContext:
+ pass
+
+ class _PagedLayerInputs:
+ pass
+
+ paged_attention.ARDiffusionPagedLayerContext = _PagedLayerContext
+ paged_attention.ARDiffusionPagedLayerInputs = _PagedLayerInputs
+ paged_attention.ar_diffusion_paged_attention = lambda *args, **kwargs: (_ for _ in ()).throw(
+ AssertionError("CUDA paged attention was not expected in this CPU test")
+ )
+ paged_attention.paged_write_attn = lambda *args, **kwargs: (_ for _ in ()).throw(
+ AssertionError("paged path was not expected in this direct-cache test")
+ )
+
+
+def _load_module():
+ assert _MODULE_PATH.exists(), "LingBot attention module has not been implemented"
+ module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+ previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+ try:
+ _install_vllm_stubs()
+ spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = module
+ spec.loader.exec_module(module)
+ return module
+ finally:
+ for name, original in previous.items():
+ if original is _MISSING_MODULE:
+ sys.modules.pop(name, None)
+ else:
+ sys.modules[name] = original
+
+
+def test_load_module_restores_process_module_state() -> None:
+ module_names = (*_STUBBED_MODULE_NAMES, _MODULE_NAME)
+ previous = {name: sys.modules.get(name, _MISSING_MODULE) for name in module_names}
+
+ _load_module()
+
+ for name, original in previous.items():
+ assert sys.modules.get(name, _MISSING_MODULE) is original
+
+
+def test_allocate_lingbot_cache_creates_request_local_layer_storage() -> None:
+ module = _load_module()
+
+ cache = module.allocate_lingbot_cache(
+ batch_size=2,
+ num_layers=3,
+ max_tokens=7,
+ num_local_heads=2,
+ head_dim=4,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ assert len(cache.self_attention) == 3
+ assert len(cache.cross_attention) == 3
+ assert cache.cross_attention == [None, None, None]
+ assert cache.self_attention[0].key.shape == (2, 7, 2, 4)
+ assert cache.self_attention[0].value.shape == (2, 7, 2, 4)
+ assert cache.self_attention[0].key.dtype == torch.float32
+ assert cache.self_attention[0].end == 0
+ assert cache.self_attention[0].key.data_ptr() != cache.self_attention[1].key.data_ptr()
+
+
+def _allocate_single_layer(module, *, max_tokens: int):
+ return module.allocate_lingbot_cache(
+ batch_size=1,
+ num_layers=1,
+ max_tokens=max_tokens,
+ num_local_heads=1,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+
+def _set_identity_attention(attention: nn.Module) -> None:
+ qkv = getattr(attention, "qkv", None)
+ if qkv is not None:
+ projections = (qkv, attention.o)
+ else:
+ projections = (attention.q, attention.k, attention.v, attention.o)
+ for projection in projections:
+ with torch.no_grad():
+ identity = torch.eye(2)
+ projection.weight.copy_(torch.cat((identity, identity, identity)) if projection is qkv else identity)
+ if projection.bias is not None:
+ projection.bias.zero_()
+ attention.norm_q = nn.Identity()
+ attention.norm_k = nn.Identity()
+
+
+def _tokens(*values: float) -> torch.Tensor:
+ return torch.tensor([[[value, 0.0] for value in values]])
+
+
+def _record_inputs(module: nn.Module) -> list[tuple]:
+ records: list[tuple] = []
+ module.register_forward_pre_hook(lambda _module, args: records.append(args))
+ return records
+
+
+def _record_outputs(module: nn.Module) -> list[torch.Tensor]:
+ records: list[torch.Tensor] = []
+
+ def record(_module, _args, output) -> None:
+ records.append(output[0] if isinstance(output, tuple) else output)
+
+ module.register_forward_hook(record)
+ return records
+
+
+def _cache_snapshot(cache) -> tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]:
+ return (
+ cache.key.detach().view(torch.uint8).clone(),
+ cache.value.detach().view(torch.uint8).clone(),
+ (cache.end, cache.absolute_end, cache.last_start, cache.sink_end),
+ )
+
+
+def _assert_cache_unchanged(cache, snapshot) -> None:
+ key, value, metadata = snapshot
+ assert torch.equal(cache.key.detach().view(torch.uint8), key)
+ assert torch.equal(cache.value.detach().view(torch.uint8), value)
+ assert (cache.end, cache.absolute_end, cache.last_start, cache.sink_end) == metadata
+
+
+def test_self_attention_repeated_offset_overwrites_then_later_offset_appends() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=6).self_attention[0]
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+ attention(_tokens(10, 20), cache=cache, current_start=0, sink_tokens=0)
+
+ assert cache.end == 2
+ assert cache.absolute_end == 2
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0]))
+
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+
+ assert cache.end == 4
+ assert cache.absolute_end == 4
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([10.0, 20.0, 3.0, 4.0]))
+
+
+def test_self_attention_uses_one_fused_qkv_projection() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+
+ assert attention.qkv.calls == 1
+
+
+def test_self_attention_routes_paged_inputs_through_engine_owned_cache() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ paged_inputs = module.ARDiffusionPagedLayerInputs()
+ calls = []
+
+ def paged_write(inputs, query, key, value, k_act, v_act, scale):
+ calls.append((inputs, query.clone(), key.clone(), value.clone(), k_act, v_act, scale))
+ return query
+
+ module.paged_write_attn = paged_write
+ output = attention(
+ _tokens(1, 2),
+ cache=paged_inputs,
+ current_start=0,
+ sink_tokens=0,
+ update_cache=False,
+ )
+
+ assert output.shape == (1, 2, 2)
+ assert len(calls) == 1
+ assert calls[0][0] is paged_inputs
+ assert calls[0][4:6] == (None, None)
+ assert attention.attn.calls == []
+ assert not any(hasattr(attention, name) for name in ("q", "k", "v"))
+
+
+def test_self_attention_is_chunk_causal_without_masking_inside_current_chunk() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=8).self_attention[0]
+ attention_inputs = _record_inputs(attention.attn)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=0)
+ first_keys = attention_inputs[-1][1].detach().clone()
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=0)
+ second_keys = attention_inputs[-1][1].detach().clone()
+
+ assert attention.attn.causal is False
+ torch.testing.assert_close(first_keys[0, :, 0, 0], torch.tensor([1.0, 2.0]))
+ torch.testing.assert_close(second_keys[0, :, 0, 0], torch.tensor([1.0, 2.0, 3.0, 4.0]))
+
+
+def test_self_attention_retains_sink_and_latest_local_history_after_eviction() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ cache = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+ attention_inputs = _record_inputs(attention.attn)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, sink_tokens=1)
+ attention(_tokens(3, 4), cache=cache, current_start=2, sink_tokens=1)
+ attention(_tokens(5, 6), cache=cache, current_start=4, sink_tokens=1)
+
+ visible_keys = attention_inputs[-1][1].detach().clone()
+ assert cache.end == 4
+ assert cache.absolute_end == 6
+ torch.testing.assert_close(visible_keys[0, :, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+ torch.testing.assert_close(cache.key[0, : cache.end, 0, 0], torch.tensor([1.0, 4.0, 5.0, 6.0]))
+
+
+def test_cross_attention_projects_encoder_kv_once_per_request() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ key_outputs = _record_outputs(attention.k)
+ value_outputs = _record_outputs(attention.v)
+ attention_inputs = _record_inputs(attention.attn)
+
+ output, cache = attention(_tokens(1), _tokens(2, 3), cache=None)
+ second_output, reused_cache = attention(_tokens(4), _tokens(20, 30), cache=cache)
+
+ assert output.shape == second_output.shape == (1, 1, 2)
+ assert reused_cache is cache
+ assert len(key_outputs) == 1
+ assert len(value_outputs) == 1
+ assert cache.end == 2
+ torch.testing.assert_close(cache.key[0, :, 0, 0], torch.tensor([2.0, 3.0]))
+ torch.testing.assert_close(attention_inputs[-1][1], cache.key)
+
+
+def test_self_attention_cache_is_isolated_between_requests() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ _set_identity_attention(attention)
+ first = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+ second = _allocate_single_layer(module, max_tokens=4).self_attention[0]
+
+ attention(_tokens(1, 2), cache=first, current_start=0, sink_tokens=0)
+
+ assert first.end == 2
+ assert second.end == 0
+ assert torch.count_nonzero(second.key) == 0
+ attention(_tokens(8), cache=second, current_start=0, sink_tokens=0)
+ torch.testing.assert_close(first.key[0, : first.end, 0, 0], torch.tensor([1.0, 2.0]))
+ torch.testing.assert_close(second.key[0, : second.end, 0, 0], torch.tensor([8.0]))
+
+
+def test_cross_attention_keeps_checkpoint_qkvo_and_norm_parameter_names() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+
+ names = set(attention.state_dict())
+
+ assert {
+ "q.weight",
+ "k.weight",
+ "v.weight",
+ "o.weight",
+ "norm_q.weight",
+ "norm_k.weight",
+ } <= names
+ assert callable(attention.norm_q.weight.weight_loader)
+
+
+def test_tp_world_size_one_attention_output_shapes() -> None:
+ module = _load_module()
+ self_attention = module.LingBotSelfAttention(dim=4, num_heads=2)
+ cross_attention = module.LingBotCrossAttention(dim=4, num_heads=2)
+ cache = module.allocate_lingbot_cache(
+ batch_size=2,
+ num_layers=1,
+ max_tokens=5,
+ num_local_heads=2,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+ hidden_states = torch.randn(2, 3, 4)
+ encoder_hidden_states = torch.randn(2, 5, 4)
+
+ self_output = self_attention(hidden_states, cache=cache.self_attention[0], current_start=0, sink_tokens=0)
+ cross_output, cache.cross_attention[0] = cross_attention(
+ hidden_states,
+ encoder_hidden_states,
+ cache=cache.cross_attention[0],
+ )
+
+ assert self_output.shape == hidden_states.shape
+ assert cross_output.shape == hidden_states.shape
+ assert cache.cross_attention[0] is not None
+
+
+def test_self_attention_applies_rotary_embedding_to_current_query_and_key() -> None:
+ module = _load_module()
+ attention = module.LingBotSelfAttention(dim=2, num_heads=1)
+ cache = _allocate_single_layer(module, max_tokens=2).self_attention[0]
+ cos = torch.ones(2, 1)
+ sin = torch.zeros(2, 1)
+ rotary_inputs = _record_inputs(attention.rotary_embedding)
+
+ attention(_tokens(1, 2), cache=cache, current_start=0, rotary_emb=(cos, sin), sink_tokens=0)
+
+ assert len(rotary_inputs) == 2
+ torch.testing.assert_close(rotary_inputs[0][1], cos)
+ torch.testing.assert_close(rotary_inputs[1][2], sin)
+
+
+def test_attention_block_owns_checkpoint_parent_namespaces() -> None:
+ module = _load_module()
+ block = module.LingBotAttentionBlock(dim=4, num_heads=2, prefix="blocks.0")
+
+ names = set(block.state_dict())
+ expected = {
+ *(f"self_attn.{child}.weight" for child in ("qkv", "o", "norm_q", "norm_k")),
+ *(f"cross_attn.{child}.weight" for child in ("q", "k", "v", "o", "norm_q", "norm_k")),
+ }
+
+ assert expected <= names
+
+
+def test_cross_attention_caches_are_isolated_between_requests() -> None:
+ module = _load_module()
+ attention = module.LingBotCrossAttention(dim=2, num_heads=1)
+
+ _, first = attention(_tokens(1), _tokens(2, 3), cache=None)
+ _, second = attention(_tokens(1), _tokens(8, 9), cache=None)
+
+ assert first.key.untyped_storage().data_ptr() != second.key.untyped_storage().data_ptr()
+ assert first.value.untyped_storage().data_ptr() != second.value.untyped_storage().data_ptr()
+ second_key = second.key.clone()
+ first.key.zero_()
+ torch.testing.assert_close(second.key, second_key)
+
+
+def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyPatch) -> None:
+ module = _load_module()
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+ monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: 1)
+ norm = module._LingBotRMSNorm(2, eps=1e-6)
+
+ norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0]))
+
+ torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0]))
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py
new file mode 100644
index 000000000..121364c07
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_camera.py
@@ -0,0 +1,525 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+_MODULE_PATH = Path(__file__).parents[4] / "vllm_omni/diffusion/models/lingbot_world/camera.py"
+_SPEC = importlib.util.spec_from_file_location("_lingbot_world_camera_under_test", _MODULE_PATH)
+assert _SPEC is not None and _SPEC.loader is not None
+_CAMERA_MODULE = importlib.util.module_from_spec(_SPEC)
+sys.modules[_SPEC.name] = _CAMERA_MODULE
+_SPEC.loader.exec_module(_CAMERA_MODULE)
+
+CameraTrajectory = _CAMERA_MODULE.CameraTrajectory
+build_plucker_embedding = _CAMERA_MODULE.build_plucker_embedding
+interpolate_camera_trajectory = _CAMERA_MODULE.interpolate_camera_trajectory
+load_camera_trajectory = _CAMERA_MODULE.load_camera_trajectory
+resolve_trusted_action_directory = _CAMERA_MODULE.resolve_trusted_action_directory
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _identity_poses(num_frames: int) -> np.ndarray:
+ return np.repeat(np.eye(4, dtype=np.float64)[None], num_frames, axis=0)
+
+
+def _write_trajectory(
+ directory: Path,
+ poses: np.ndarray,
+ intrinsics: np.ndarray,
+) -> None:
+ np.save(directory / "poses.npy", poses)
+ np.save(directory / "intrinsics.npy", intrinsics)
+
+
+def _trusted(directory: Path):
+ return resolve_trusted_action_directory(".", directory)
+
+
+def _trajectory(
+ poses: torch.Tensor,
+ intrinsics: torch.Tensor | None = None,
+) -> CameraTrajectory:
+ if intrinsics is None:
+ intrinsics = torch.tensor([[832.0 / 3.0, 160.0, 416.0, 240.0]]).repeat(poses.shape[0], 1)
+ return CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+
+@pytest.mark.parametrize("present_file", [None, "poses.npy"])
+def test_load_camera_trajectory_rejects_missing_files(tmp_path: Path, present_file: str | None) -> None:
+ if present_file == "poses.npy":
+ np.save(tmp_path / present_file, _identity_poses(1))
+
+ missing_file = "poses.npy" if present_file is None else "intrinsics.npy"
+ with pytest.raises(FileNotFoundError, match=missing_file):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize(
+ ("poses", "intrinsics", "message"),
+ [
+ (np.zeros((2, 3, 4)), np.zeros((2, 4)), "poses.npy"),
+ (_identity_poses(2), np.zeros((2, 3)), "intrinsics.npy"),
+ (_identity_poses(0), np.zeros((0, 4)), "1 and 4096 source frames"),
+ (_identity_poses(2), np.zeros((1, 4)), "same number of frames"),
+ ],
+)
+def test_load_camera_trajectory_rejects_invalid_shapes(
+ tmp_path: Path,
+ poses: np.ndarray,
+ intrinsics: np.ndarray,
+ message: str,
+) -> None:
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match=message):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+@pytest.mark.parametrize("array_name", ["poses", "intrinsics"])
+def test_load_camera_trajectory_rejects_non_finite_values(tmp_path: Path, array_name: str) -> None:
+ poses = _identity_poses(1)
+ intrinsics = np.ones((1, 4), dtype=np.float64)
+ if array_name == "poses":
+ poses[0, 0, 0] = np.nan
+ else:
+ intrinsics[0, 0] = np.inf
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match="finite"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_preserves_raw_c2w_poses(tmp_path: Path) -> None:
+ first_c2w = np.array(
+ [
+ [0.0, -1.0, 0.0, 2.0],
+ [1.0, 0.0, 0.0, 3.0],
+ [0.0, 0.0, 1.0, 4.0],
+ [0.0, 0.0, 0.0, 1.0],
+ ]
+ )
+ second_c2w = np.array(
+ [
+ [0.0, 0.0, 1.0, 5.0],
+ [0.0, 1.0, 0.0, 4.0],
+ [-1.0, 0.0, 0.0, 6.0],
+ [0.0, 0.0, 0.0, 1.0],
+ ]
+ )
+ poses = np.stack((first_c2w, second_c2w))
+ intrinsics = np.array([[100.0, 110.0, 120.0, 130.0]] * 2)
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert trajectory.poses.dtype == torch.float32
+ assert trajectory.intrinsics.dtype == torch.float32
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+ torch.testing.assert_close(trajectory.intrinsics, torch.from_numpy(intrinsics).float())
+
+
+def test_load_camera_trajectory_rejects_object_dtype_before_materializing(tmp_path: Path) -> None:
+ poses = _identity_poses(1).astype(object)
+ intrinsics = np.ones((1, 4), dtype=np.float32)
+ _write_trajectory(tmp_path, poses, intrinsics)
+
+ with pytest.raises(ValueError, match="poses.npy.*real numeric"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_accepts_official_length_and_materializes_request_prefix(tmp_path: Path) -> None:
+ _write_trajectory(
+ tmp_path,
+ _identity_poses(269),
+ np.arange(269 * 4, dtype=np.float32).reshape(269, 4),
+ )
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert trajectory.poses.shape == (117, 4, 4)
+ assert trajectory.intrinsics.shape == (117, 4)
+ np.testing.assert_array_equal(
+ trajectory.intrinsics.numpy(),
+ np.arange(269 * 4, dtype=np.float32).reshape(269, 4)[:117],
+ )
+
+
+def test_load_camera_trajectory_rejects_unbounded_source_frame_count(tmp_path: Path) -> None:
+ _write_trajectory(
+ tmp_path,
+ _identity_poses(4097),
+ np.ones((4097, 4), dtype=np.float32),
+ )
+
+ with pytest.raises(ValueError, match="poses.npy.*bounded camera-array limit|4096 source frames"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_truncated_payload_before_loading(tmp_path: Path) -> None:
+ _write_trajectory(tmp_path, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ poses_path = tmp_path / "poses.npy"
+ poses_path.write_bytes(poses_path.read_bytes()[:-1])
+
+ with pytest.raises(ValueError, match="poses.npy byte size"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_oversized_npy_header(tmp_path: Path) -> None:
+ header_size = 20_000
+ (tmp_path / "poses.npy").write_bytes(
+ np.lib.format.magic(2, 0) + header_size.to_bytes(4, "little") + b" " * (header_size - 1) + b"\n"
+ )
+ np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+
+ with pytest.raises(ValueError, match="poses.npy.*oversized NPY header"):
+ load_camera_trajectory(_trusted(tmp_path))
+
+
+def test_load_camera_trajectory_rejects_symlink_swap_after_resolution(tmp_path: Path) -> None:
+ trusted_root = tmp_path / "trusted"
+ action_directory = trusted_root / "action"
+ outside = tmp_path / "outside"
+ action_directory.mkdir(parents=True)
+ outside.mkdir()
+ _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ _write_trajectory(outside, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ resolved = resolve_trusted_action_directory("action", trusted_root)
+ action_directory.rename(trusted_root / "original-action")
+ action_directory.symlink_to(outside, target_is_directory=True)
+
+ with pytest.raises(ValueError, match="securely open.*action directory"):
+ load_camera_trajectory(resolved)
+
+
+def test_load_camera_trajectory_rejects_symlinked_camera_file(tmp_path: Path) -> None:
+ trusted_root = tmp_path / "trusted"
+ action_directory = trusted_root / "action"
+ outside = tmp_path / "outside"
+ action_directory.mkdir(parents=True)
+ outside.mkdir()
+ _write_trajectory(action_directory, _identity_poses(1), np.ones((1, 4), dtype=np.float32))
+ np.save(outside / "poses.npy", _identity_poses(1))
+ resolved = resolve_trusted_action_directory("action", trusted_root)
+ (action_directory / "poses.npy").unlink()
+ (action_directory / "poses.npy").symlink_to(outside / "poses.npy")
+
+ with pytest.raises(ValueError, match="poses.npy.*opened securely"):
+ load_camera_trajectory(resolved)
+
+
+@pytest.mark.skipif(not hasattr(__import__("os"), "mkfifo"), reason="requires POSIX FIFO support")
+def test_load_camera_trajectory_rejects_fifo_without_blocking(tmp_path: Path) -> None:
+ import os
+
+ os.mkfifo(tmp_path / "poses.npy")
+ np.save(tmp_path / "intrinsics.npy", np.ones((1, 4), dtype=np.float32))
+ script = """
+import importlib.util
+import sys
+
+spec = importlib.util.spec_from_file_location("_lingbot_camera_fifo_test", sys.argv[1])
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)
+action = module.resolve_trusted_action_directory(".", sys.argv[2])
+try:
+ module.load_camera_trajectory(action)
+except ValueError as error:
+ print(error)
+else:
+ raise AssertionError("FIFO camera input was accepted")
+"""
+
+ completed = subprocess.run(
+ [sys.executable, "-c", script, str(_MODULE_PATH), str(tmp_path)],
+ capture_output=True,
+ text=True,
+ timeout=30,
+ check=False,
+ )
+
+ assert completed.returncode == 0, completed.stderr
+ assert "poses.npy must be a regular NPY file" in completed.stdout
+
+
+def test_load_camera_trajectory_decodes_the_validated_file_snapshot(tmp_path: Path, monkeypatch) -> None:
+ original_poses = _identity_poses(1)
+ replacement_poses = original_poses.copy()
+ replacement_poses[0, 0, 3] = 123.0
+ _write_trajectory(tmp_path, original_poses, np.ones((1, 4), dtype=np.float32))
+ poses_path = tmp_path / "poses.npy"
+ original_load = _CAMERA_MODULE.np.load
+ mutation_count = 0
+
+ def mutate_source_then_decode(source, *args, **kwargs):
+ nonlocal mutation_count
+ if mutation_count == 0:
+ mutation_count += 1
+ np.save(poses_path, replacement_poses)
+ return original_load(source, *args, **kwargs)
+
+ monkeypatch.setattr(_CAMERA_MODULE.np, "load", mutate_source_then_decode)
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ assert mutation_count == 1
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(original_poses).float())
+
+
+def test_load_camera_trajectory_ignores_non_camera_event_files(tmp_path: Path) -> None:
+ poses = _identity_poses(1)
+ intrinsics = np.ones((1, 4), dtype=np.float32)
+ _write_trajectory(tmp_path, poses, intrinsics)
+ np.save(tmp_path / "wasd_events.npy", np.array([{"key": "w"}], dtype=object))
+
+ trajectory = load_camera_trajectory(_trusted(tmp_path))
+
+ torch.testing.assert_close(trajectory.poses, torch.from_numpy(poses).float())
+
+
+def test_interpolate_camera_trajectory_uses_linear_and_spherical_interpolation() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[1, :3, :3] = torch.tensor(
+ [
+ [0.0, -1.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[1, :3, 3] = torch.tensor([2.0, 4.0, 6.0])
+ intrinsics = torch.tensor([[10.0, 20.0, 30.0, 40.0], [20.0, 40.0, 60.0, 80.0]])
+ trajectory = CameraTrajectory(poses=poses, intrinsics=intrinsics)
+
+ result = interpolate_camera_trajectory(trajectory, 3)
+
+ torch.testing.assert_close(result.poses[0], poses[0])
+ torch.testing.assert_close(result.poses[-1], poses[-1])
+ torch.testing.assert_close(result.poses[1, :3, 3], torch.tensor([1.0, 2.0, 3.0]))
+ expected_mid_rotation = torch.tensor(
+ [
+ [math.sqrt(0.5), -math.sqrt(0.5), 0.0],
+ [math.sqrt(0.5), math.sqrt(0.5), 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ torch.testing.assert_close(result.poses[1, :3, :3], expected_mid_rotation)
+ torch.testing.assert_close(result.intrinsics[1], torch.tensor([15.0, 30.0, 45.0, 60.0]))
+
+
+def test_build_plucker_embedding_has_forward_center_ray() -> None:
+ trajectory = _trajectory(torch.eye(4).unsqueeze(0))
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=3,
+ width=3,
+ target_height=3,
+ target_width=3,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ torch.testing.assert_close(result[0, :3, 1, 1], torch.zeros(3))
+ torch.testing.assert_close(result[0, 3:, 1, 1], torch.tensor([0.0, 0.0, 1.0]))
+
+
+def test_build_plucker_embedding_scales_reference_intrinsics() -> None:
+ trajectory = _trajectory(
+ torch.eye(4).unsqueeze(0),
+ torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=4,
+ width=8,
+ target_height=4,
+ target_width=8,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_corner_direction = torch.tensor([-0.875, -0.75, 1.0])
+ expected_corner_direction /= torch.linalg.vector_norm(expected_corner_direction)
+ torch.testing.assert_close(result[0, 3:, 0, 0], expected_corner_direction)
+
+
+def test_build_plucker_embedding_samples_target_pixel_centers() -> None:
+ trajectory = _trajectory(
+ torch.eye(4).unsqueeze(0),
+ torch.tensor([[416.0, 240.0, 416.0, 240.0]]),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=6,
+ width=10,
+ target_height=2,
+ target_width=5,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ # Target centers map through the source scale: (0.5 * 2, 0.5 * 3).
+ expected_direction = torch.tensor([-0.8, -0.5, 1.0])
+ expected_direction /= torch.linalg.vector_norm(expected_direction)
+ torch.testing.assert_close(result[0, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_uses_framewise_deltas_for_rotations_and_translations() -> None:
+ poses = torch.eye(4).repeat(3, 1, 1)
+ poses[0, :3, 3] = torch.tensor([5.0, 7.0, 11.0])
+ poses[1, :3, :3] = torch.tensor(
+ [
+ [0.0, -1.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[1, :3, 3] = torch.tensor([5.0, 9.0, 11.0])
+ poses[2, :3, :3] = torch.tensor(
+ [
+ [-1.0, 0.0, 0.0],
+ [0.0, -1.0, 0.0],
+ [0.0, 0.0, 1.0],
+ ]
+ )
+ poses[2, :3, 3] = torch.tensor([3.0, 9.0, 11.0])
+ trajectory = _trajectory(
+ poses,
+ torch.tensor([[416.0, 480.0, 208.0, 240.0]]).repeat(3, 1),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=1,
+ width=2,
+ target_height=1,
+ target_width=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_direction = torch.tensor([0.0, 1.0, 1.0]) / math.sqrt(2.0)
+ torch.testing.assert_close(result[0, :3, 0, 1], torch.zeros(3))
+ torch.testing.assert_close(result[1, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+ torch.testing.assert_close(result[2, :3, 0, 1], torch.tensor([0.0, 1.0, 0.0]))
+ torch.testing.assert_close(result[1, 3:, 0, 1], expected_direction)
+ torch.testing.assert_close(result[2, 3:, 0, 1], expected_direction)
+
+
+def test_build_plucker_embedding_normalizes_framewise_translation_by_max_norm() -> None:
+ poses = torch.eye(4).repeat(3, 1, 1)
+ poses[:, :3, 3] = torch.tensor(
+ [
+ [10.0, -2.0, 5.0],
+ [13.0, 2.0, 5.0],
+ [19.0, 10.0, 5.0],
+ ]
+ )
+
+ result = build_plucker_embedding(
+ _trajectory(poses),
+ height=1,
+ width=1,
+ target_height=1,
+ target_width=1,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_origins = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 0.0], [0.6, 0.8, 0.0]])
+ torch.testing.assert_close(result[:, :3, 0, 0], expected_origins)
+
+
+def test_build_plucker_embedding_interpolates_raw_poses_before_framewise_conversion() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[0, :3, 3] = torch.tensor([20.0, 4.0, -2.0])
+ poses[1, :3, 3] = torch.tensor([26.0, 12.0, -2.0])
+ raw = _trajectory(poses)
+
+ interpolated = interpolate_camera_trajectory(raw, 3)
+ result = build_plucker_embedding(
+ interpolated,
+ height=1,
+ width=1,
+ target_height=1,
+ target_width=1,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_delta = torch.tensor([0.6, 0.8, 0.0])
+ torch.testing.assert_close(result[0, :3, 0, 0], torch.zeros(3))
+ torch.testing.assert_close(result[1, :3, 0, 0], expected_delta)
+ torch.testing.assert_close(result[2, :3, 0, 0], expected_delta)
+
+
+def test_build_plucker_embedding_uses_half_pixel_centers_and_origin_direction_channel_order() -> None:
+ poses = torch.eye(4).repeat(2, 1, 1)
+ poses[1, :3, 3] = torch.tensor([2.0, 0.0, 0.0])
+ trajectory = _trajectory(
+ poses,
+ torch.tensor([[832.0, 480.0, 0.0, 0.0]]).repeat(2, 1),
+ )
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=1,
+ width=2,
+ target_height=1,
+ target_width=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ expected_direction = torch.tensor([0.25, 0.5, 1.0])
+ expected_direction /= torch.linalg.vector_norm(expected_direction)
+ torch.testing.assert_close(result[1, :3, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
+ torch.testing.assert_close(result[1, 3:, 0, 0], expected_direction)
+
+
+def test_build_plucker_embedding_honors_shape_dtype_and_device() -> None:
+ trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+
+ result = build_plucker_embedding(
+ trajectory,
+ height=12,
+ width=20,
+ target_height=3,
+ target_width=5,
+ device=torch.device("cpu"),
+ dtype=torch.float64,
+ )
+
+ assert result.shape == (2, 6, 3, 5)
+ assert result.dtype == torch.float64
+ assert result.device == torch.device("cpu")
+
+
+def test_build_plucker_embedding_is_deterministic() -> None:
+ trajectory = _trajectory(torch.eye(4).repeat(2, 1, 1))
+ kwargs = {
+ "height": 12,
+ "width": 20,
+ "target_height": 3,
+ "target_width": 5,
+ "device": torch.device("cpu"),
+ "dtype": torch.float32,
+ }
+
+ first = build_plucker_embedding(trajectory, **kwargs)
+ second = build_plucker_embedding(trajectory, **kwargs)
+
+ torch.testing.assert_close(first, second, rtol=0.0, atol=0.0)
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py
new file mode 100644
index 000000000..e2e2a08b6
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py
@@ -0,0 +1,685 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import json
+import math
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture"
+_SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture"
+
+
+@pytest.fixture(autouse=True)
+def _inference_context():
+ with torch.inference_mode():
+ yield
+
+
+def _tiny_model(
+ module,
+ *,
+ num_layers: int = 2,
+ num_frames_per_block: int = 1,
+ sliding_window_num_frames: int = 3,
+):
+ return module.CausalLingBotWorldTransformer3DModel(
+ patch_size=(1, 2, 2),
+ num_attention_heads=2,
+ attention_head_dim=2,
+ in_channels=36,
+ out_channels=2,
+ text_dim=6,
+ freq_dim=4,
+ ffn_dim=8,
+ num_layers=num_layers,
+ cross_attn_norm=True,
+ eps=1e-6,
+ rope_max_seq_len=16,
+ sink_size=1,
+ num_frames_per_block=num_frames_per_block,
+ sliding_window_num_frames=sliding_window_num_frames,
+ local_attn_size=-1,
+ )
+
+
+def _cache(module, model, *, max_tokens: int | None = None):
+ if max_tokens is None:
+ window_frames = (
+ model.config.local_attn_size
+ if model.config.local_attn_size != -1
+ else model.config.sliding_window_num_frames
+ )
+ max_tokens = window_frames * 2 * 2
+ return module.allocate_lingbot_cache(
+ batch_size=1,
+ num_layers=len(model.blocks),
+ max_tokens=max_tokens,
+ num_local_heads=2,
+ head_dim=2,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+
+def _checkpoint_parameter_specs(model) -> dict[str, tuple[torch.Size, torch.dtype]]:
+ specs: dict[str, tuple[torch.Size, torch.dtype]] = {}
+ for name, parameter in model.named_parameters():
+ marker = ".self_attn.qkv."
+ if marker not in name:
+ specs[name] = (parameter.shape, parameter.dtype)
+ continue
+ shard_shape = torch.Size((parameter.shape[0] // 3, *parameter.shape[1:]))
+ for projection in ("q", "k", "v"):
+ specs[name.replace(marker, f".self_attn.{projection}.")] = (shard_shape, parameter.dtype)
+ return specs
+
+
+def _checkpoint_weights_for_model(model) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+ checkpoint_weights: dict[str, torch.Tensor] = {}
+ expected_parameters: dict[str, torch.Tensor] = {}
+ next_value = 1
+ for name, parameter in model.named_parameters():
+ marker = ".self_attn.qkv."
+ if marker not in name:
+ value = torch.full_like(parameter, next_value)
+ checkpoint_weights[name] = value
+ expected_parameters[name] = value
+ next_value += 1
+ continue
+ shards = []
+ shard_shape = (parameter.shape[0] // 3, *parameter.shape[1:])
+ for projection in ("q", "k", "v"):
+ value = torch.full(shard_shape, next_value, dtype=parameter.dtype, device=parameter.device)
+ checkpoint_weights[name.replace(marker, f".self_attn.{projection}.")] = value
+ shards.append(value)
+ next_value += 1
+ expected_parameters[name] = torch.cat(shards)
+ return checkpoint_weights, expected_parameters
+
+
+def _self_cache_snapshot(cache) -> list[tuple[torch.Tensor, torch.Tensor, tuple[int, int, int | None, int]]]:
+ return [attention_tests._cache_snapshot(layer_cache) for layer_cache in cache.self_attention]
+
+
+def _assert_self_cache_unchanged(cache, snapshot) -> None:
+ for layer_cache, layer_snapshot in zip(cache.self_attention, snapshot, strict=True):
+ attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot)
+
+
+def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None:
+ torch.manual_seed(7)
+ module = attention_tests._load_module()
+ model = _tiny_model(module).eval()
+ cache = _cache(module, model)
+ encoder_hidden_states = torch.randn(1, 3, 6)
+ camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+ cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+ text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+ outputs = []
+
+ assert model.patch_embedding.in_channels == 36
+ assert model.patch_embedding_wancamctrl.in_features == 6 * 8 * 8 * 1 * 2 * 2
+
+ for start_frame in range(4):
+ hidden_states = torch.randn(1, 36, 1, 4, 4)
+ timestep = torch.tensor([float(start_frame + 1)])
+ snapshot = _self_cache_snapshot(cache)
+
+ transient_output = model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=start_frame,
+ update_cache=False,
+ )
+
+ assert transient_output.shape == (1, 2, 1, 4, 4)
+ _assert_self_cache_unchanged(cache, snapshot)
+ assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+ committed_output = model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=start_frame,
+ update_cache=True,
+ )
+ torch.testing.assert_close(committed_output, transient_output)
+ outputs.append(committed_output)
+
+ expected_token_end = (start_frame + 1) * 2 * 2
+ assert all(layer_cache.absolute_end == expected_token_end for layer_cache in cache.self_attention)
+
+ assert len(outputs) == 4
+ assert all(len(outputs) == 1 for outputs in cross_key_outputs)
+ assert len(text_projection_outputs) == 1
+
+ alternate_cache = _cache(module, model)
+ alternate_output = model(
+ torch.zeros(1, 36, 1, 4, 4),
+ torch.tensor([1.0]),
+ encoder_hidden_states,
+ torch.zeros_like(camera_hidden_states),
+ cache=alternate_cache,
+ start_frame=0,
+ update_cache=False,
+ )
+ camera_output = model(
+ torch.zeros(1, 36, 1, 4, 4),
+ torch.tensor([1.0]),
+ encoder_hidden_states,
+ torch.ones_like(camera_hidden_states),
+ cache=_cache(module, model),
+ start_frame=0,
+ update_cache=False,
+ )
+ assert not torch.equal(camera_output, alternate_output)
+
+
+def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=2).eval()
+ cache = _cache(module, model)
+ hidden_states = torch.randn(1, 36, 1, 4, 4)
+ timestep = torch.tensor([1.0])
+ encoder_hidden_states = torch.randn(1, 3, 6)
+ camera_hidden_states = torch.randn(1, 6 * 8 * 8, 1, 4, 4)
+
+ model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=0,
+ update_cache=False,
+ )
+ cache.cross_attention[1] = None
+ text_projection_outputs = attention_tests._record_outputs(model.text_embedding)
+ cross_key_outputs = [attention_tests._record_outputs(block.cross_attn.k) for block in model.blocks]
+
+ model(
+ hidden_states,
+ timestep,
+ encoder_hidden_states,
+ camera_hidden_states,
+ cache=cache,
+ start_frame=0,
+ update_cache=False,
+ )
+
+ assert len(text_projection_outputs) == 1
+ assert [len(outputs) for outputs in cross_key_outputs] == [0, 1]
+ assert all(layer_cache is not None for layer_cache in cache.cross_attention)
+
+
+@pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks"))
+def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=1,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+ cache = _cache(module, model)
+ snapshot = _self_cache_snapshot(cache)
+ patch_inputs = attention_tests._record_inputs(model.patch_embedding)
+
+ with pytest.raises(ValueError, match="exactly 3 post-patch frames"):
+ model(
+ torch.randn(1, 36, frames, 4, 4),
+ torch.tensor([1.0]),
+ torch.randn(1, 3, 6),
+ torch.randn(1, 6 * 8 * 8, frames, 4, 4),
+ cache=cache,
+ start_frame=0,
+ update_cache=True,
+ )
+
+ assert patch_inputs == []
+ _assert_self_cache_unchanged(cache, snapshot)
+ assert cache.cross_attention == [None]
+
+
+def test_forward_accepts_exactly_one_configured_frame_block() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=1,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+
+ output = model(
+ torch.randn(1, 36, 3, 4, 4),
+ torch.tensor([1.0]),
+ torch.randn(1, 3, 6),
+ torch.randn(1, 6 * 8 * 8, 3, 4, 4),
+ cache=_cache(module, model),
+ start_frame=0,
+ update_cache=False,
+ )
+
+ assert output.shape == (1, 2, 3, 4, 4)
+
+
+def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(
+ module,
+ num_layers=2,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ )
+
+ cache = model.allocate_cache(
+ batch_size=2,
+ latent_height=8,
+ latent_width=12,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ assert len(cache.self_attention) == 2
+ assert len(cache.cross_attention) == 2
+ assert cache.cross_attention == [None, None]
+ assert all(layer.key.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+ assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention)
+
+
+def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ hidden_states = torch.zeros(1, 36, 2, 2, 4)
+ hidden_states[0, 0] = torch.tensor(
+ [
+ [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]],
+ [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]],
+ ]
+ )
+ with torch.no_grad():
+ model.patch_embedding.weight.zero_()
+ model.patch_embedding.bias.zero_()
+ model.patch_embedding.weight[0, 0, 0, 0, 0] = 1
+
+ tokens = model.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
+
+ torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0]))
+
+
+def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ hidden_states = (
+ torch.stack(
+ (
+ torch.arange(0, 8),
+ torch.arange(10, 18),
+ torch.arange(20, 28),
+ torch.arange(30, 38),
+ )
+ )
+ .unsqueeze(0)
+ .float()
+ )
+ expected = torch.tensor(
+ [
+ [
+ [[0, 2, 10, 12], [4, 6, 14, 16]],
+ [[20, 22, 30, 32], [24, 26, 34, 36]],
+ ],
+ [
+ [[1, 3, 11, 13], [5, 7, 15, 17]],
+ [[21, 23, 31, 33], [25, 27, 35, 37]],
+ ],
+ ],
+ dtype=torch.float32,
+ ).unsqueeze(0)
+
+ output = model._unpatchify(
+ hidden_states,
+ batch_size=1,
+ frames=2,
+ height=1,
+ width=2,
+ )
+
+ torch.testing.assert_close(output, expected)
+
+
+def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1, num_frames_per_block=2)
+ model.head.norm = torch.nn.Identity()
+ with torch.no_grad():
+ model.head.modulation.zero_()
+ model.head.head.weight.zero_()
+ model.head.head.bias.zero_()
+ model.head.head.weight[:, 0] = 1
+ hidden_states = torch.ones(1, 4, model.dim)
+ timestep_embedding = torch.tensor([[[10.0, 0.0, 0.0, 0.0], [20.0, 0.0, 0.0, 0.0]]])
+ expected = torch.tensor([21.0, 21.0, 41.0, 41.0]).view(1, 4, 1).expand(1, 4, 8)
+
+ output = model.head(hidden_states, timestep_embedding)
+
+ torch.testing.assert_close(output, expected)
+
+
+def test_constructor_defaults_match_official_checkpoint_config() -> None:
+ module = attention_tests._load_module()
+ parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters
+ expected = {
+ "in_channels": 36,
+ "out_channels": 16,
+ "num_layers": 40,
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "patch_size": (1, 2, 2),
+ "text_dim": 4096,
+ "ffn_dim": 13824,
+ "freq_dim": 256,
+ "rope_max_seq_len": 1024,
+ "eps": 1e-6,
+ "cross_attn_norm": True,
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+
+ assert {name: parameters[name].default for name in expected} == expected
+
+
+def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module)
+
+ assert model.dtype == next(model.parameters()).dtype
+
+
+def test_transformer_declares_regional_compile_block() -> None:
+ module = attention_tests._load_module()
+
+ assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"]
+
+
+def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None:
+ module = attention_tests._load_module()
+ reduced_values: list[torch.Tensor] = []
+
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 2)
+
+ def all_reduce(value: torch.Tensor) -> torch.Tensor:
+ reduced_values.append(value.detach().clone())
+ return value * 2
+
+ monkeypatch.setattr(module, "tensor_model_parallel_all_reduce", all_reduce)
+ norm = module._LingBotRMSNorm(hidden_size=2, eps=0.0)
+ value = torch.tensor([[3.0, 4.0]])
+
+ output = norm(value)
+
+ torch.testing.assert_close(output, value / math.sqrt(12.5))
+ assert len(reduced_values) == 1
+ torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]]))
+
+
+def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None:
+ module = attention_tests._load_module()
+ monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3)
+
+ with pytest.raises(ValueError, match="num_heads=2.*tp_size=3"):
+ module.LingBotSelfAttention(dim=4, num_heads=2)
+
+
+def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(ValueError, match="qk_norm.*rms_norm_across_heads"):
+ module.CausalLingBotWorldTransformer3DModel(qk_norm="layer_norm")
+
+
+@pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"])
+def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(ValueError, match=field):
+ module.CausalLingBotWorldTransformer3DModel(**{field: 4})
+
+
+def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None:
+ module = attention_tests._load_module()
+
+ with pytest.raises(RuntimeError, match="quant_config.*not supported"):
+ module.CausalLingBotWorldTransformer3DModel(quant_config=object())
+
+
+def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None:
+ module = attention_tests._load_module()
+
+ class ConfigProbe(module.CausalLingBotWorldTransformer3DModel):
+ def __init__(self, **kwargs) -> None:
+ self.received_kwargs = kwargs
+
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "_diffusers_version": "0.35.0.dev0",
+ "patch_size": [1, 2, 2],
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 4096,
+ "freq_dim": 256,
+ "ffn_dim": 13824,
+ "num_layers": 40,
+ "cross_attn_norm": True,
+ "eps": 1e-6,
+ "image_dim": None,
+ "added_kv_proj_dim": None,
+ "rope_max_seq_len": 1024,
+ "pos_embed_seq_len": None,
+ "qk_norm": "rms_norm_across_heads",
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+
+ model = ConfigProbe.from_config(config, prefix="transformer")
+
+ assert model.received_kwargs["patch_size"] == (1, 2, 2)
+ assert model.received_kwargs["num_layers"] == 40
+ assert model.received_kwargs["prefix"] == "transformer"
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("in_channels", 35),
+ ("num_layers", 39),
+ ("num_attention_heads", 32),
+ ("attention_head_dim", 64),
+ ("ffn_dim", 10240),
+ ("num_frames_per_block", 4),
+ ("sliding_window_num_frames", 24),
+ ("sink_size", 0),
+ ],
+)
+def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None:
+ module = attention_tests._load_module()
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "_diffusers_version": "0.35.0.dev0",
+ "patch_size": [1, 2, 2],
+ "num_attention_heads": 40,
+ "attention_head_dim": 128,
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 4096,
+ "freq_dim": 256,
+ "ffn_dim": 13824,
+ "num_layers": 40,
+ "cross_attn_norm": True,
+ "eps": 1e-6,
+ "image_dim": None,
+ "added_kv_proj_dim": None,
+ "rope_max_seq_len": 1024,
+ "pos_embed_seq_len": None,
+ "qk_norm": "rms_norm_across_heads",
+ "sink_size": 9,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 18,
+ "local_attn_size": -1,
+ }
+ config[field] = value
+
+ with pytest.raises(ValueError, match=field):
+ module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+
+def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None:
+ module = attention_tests._load_module()
+ config = {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "architectures": ["CausalLingBotWorldTransformer3DModel"],
+ "future_diffusers_metadata": {"version": 1},
+ }
+
+ with torch.device("meta"):
+ model = module.CausalLingBotWorldTransformer3DModel.from_config(config)
+
+ assert model.config.in_channels == 36
+ assert model.config.num_layers == 40
+
+
+def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1)
+ checkpoint_weights, expected_parameters = _checkpoint_weights_for_model(model)
+ loader_calls: list[str] = []
+ qkv_weight = model.blocks[0].self_attn.qkv.weight
+
+ def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: str) -> None:
+ loader_calls.append(shard_id)
+ shard_index = {"q": 0, "k": 1, "v": 2}[shard_id]
+ shard_size = param.shape[0] // 3
+ param.data.narrow(0, shard_index * shard_size, shard_size).copy_(loaded_weight)
+
+ qkv_weight.weight_loader = record_loader
+ loaded = model.load_weights(iter(checkpoint_weights.items()))
+
+ assert set(checkpoint_weights) <= loaded
+ assert set(dict(model.named_parameters())) <= loaded
+ assert loader_calls == ["q", "k", "v"]
+ for name, param in model.named_parameters():
+ torch.testing.assert_close(param, expected_parameters[name])
+
+ first_name = next(iter(checkpoint_weights))
+ partial = _tiny_model(module, num_layers=1)
+ partial_loaded = partial.load_weights([(first_name, checkpoint_weights[first_name])])
+ expected_partial = {first_name}
+ for projection in ("q", "k", "v"):
+ marker = f".self_attn.{projection}."
+ if marker in first_name:
+ expected_partial.add(first_name.replace(marker, ".self_attn.qkv."))
+ break
+ assert partial_loaded == expected_partial
+
+ with pytest.raises(KeyError, match="unexpected_model.weight"):
+ model.load_weights([("unexpected_model.weight", torch.ones(1))])
+
+
+def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None:
+ module = attention_tests._load_module()
+ model = _tiny_model(module, num_layers=1)
+ first_name, first_param = next(iter(model.named_parameters()))
+ loaded_weight = torch.full_like(first_param, 17)
+
+ def weights():
+ yield first_name, loaded_weight
+ torch.testing.assert_close(first_param, loaded_weight)
+
+ assert model.load_weights(weights()) == {first_name}
+
+
+def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None:
+ utils = pytest.importorskip("vllm.model_executor.models.utils")
+ module = attention_tests._load_module()
+ transformer = _tiny_model(module, num_layers=1)
+ checkpoint_weights, _ = _checkpoint_weights_for_model(transformer)
+
+ class Pipeline(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.transformer = transformer
+
+ pipeline = Pipeline()
+ loaded = utils.AutoWeightsLoader(pipeline).load_weights(
+ (f"transformer.{name}", value) for name, value in checkpoint_weights.items()
+ )
+
+ assert set(dict(pipeline.named_parameters())) <= loaded
+
+
+def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None:
+ module = attention_tests._load_module()
+ fixture = json.loads(_FIXTURE_PATH.read_text())
+ model = _tiny_model(module, num_layers=1)
+ parameter_names = set(_checkpoint_parameter_specs(model))
+
+ assert {name.split(".", 1)[0] for name in parameter_names} == set(fixture["top_level_modules"])
+ assert {name for name in parameter_names if not name.startswith("blocks.")} == set(fixture["non_block_parameters"])
+ assert {name.removeprefix("blocks.0.") for name in parameter_names if name.startswith("blocks.0.")} == set(
+ fixture["block_parameter_suffixes"]
+ )
+ expanded_names = set(fixture["non_block_parameters"])
+ expanded_names.update(
+ f"blocks.{layer}.{suffix}" for layer in range(40) for suffix in fixture["block_parameter_suffixes"]
+ )
+ canonical_names = "\n".join(sorted(expanded_names))
+
+ assert len(expanded_names) == fixture["index_parameter_count"] == 1421
+ assert hashlib.sha256(canonical_names.encode()).hexdigest() == fixture["canonical_name_set_sha256"]
+ assert fixture["canonical_name_to_shard_sha256"] == (
+ "e1addc27d2b1fad6f10226a23a265ee3e3c61e74dce4d5859727bdf180a67992"
+ )
+ assert list(fixture["shards"]) == [
+ f"diffusion_pytorch_model-{index:05}-of-00008.safetensors" for index in range(1, 9)
+ ]
+ assert sum(shard["parameter_count"] for shard in fixture["shards"].values()) == 1421
+ assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values())
+
+
+def test_official_default_shapes_match_public_safetensors_header_fixture() -> None:
+ module = attention_tests._load_module()
+ fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text())
+
+ with torch.device("meta"):
+ model = module.CausalLingBotWorldTransformer3DModel()
+ parameter_specs = _checkpoint_parameter_specs(model)
+ expected_dtype = {"F32": torch.float32}[fixture["dtype"]]
+
+ assert len(fixture["shard_representatives"]) == 8
+ assert set(fixture["shard_representatives"].values()) <= set(fixture["parameter_shapes"])
+
+ for name, expected_shape in fixture["parameter_shapes"].items():
+ shape, dtype = parameter_specs[name]
+ assert tuple(shape) == tuple(expected_shape), name
+ assert dtype == expected_dtype
diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py
new file mode 100644
index 000000000..56c25fee3
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_typing.py
@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_ROOT = Path(__file__).parents[4]
+_PRODUCTION_FILES = (
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/camera.py",
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/transformer.py",
+ _ROOT / "vllm_omni/diffusion/models/lingbot_world/pipeline.py",
+)
+
+
+def test_all_lingbot_production_functions_have_complete_annotations() -> None:
+ missing: list[str] = []
+ for path in _PRODUCTION_FILES:
+ tree = ast.parse(path.read_text())
+ for node in ast.walk(tree):
+ if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+ continue
+ arguments = (*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs)
+ for argument in arguments:
+ if argument.arg not in {"self", "cls"} and argument.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:{argument.arg}")
+ if node.args.vararg is not None and node.args.vararg.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:*{node.args.vararg.arg}")
+ if node.args.kwarg is not None and node.args.kwarg.annotation is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:**{node.args.kwarg.arg}")
+ if node.returns is None:
+ missing.append(f"{path.name}:{node.lineno}:{node.name}:return")
+
+ assert missing == []
diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py
new file mode 100644
index 000000000..1a08e8270
--- /dev/null
+++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py
@@ -0,0 +1,1765 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import gc
+import weakref
+from pathlib import Path
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+import torch
+from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
+from PIL import Image
+from torch import nn
+
+import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
+from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
+from vllm_omni.diffusion.models.lingbot_world.actions import (
+ integrate_lingbot_camera_actions,
+)
+from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as _CameraTrajectory
+from vllm_omni.diffusion.models.lingbot_world.camera import (
+ build_plucker_embedding as _real_build_plucker_embedding,
+)
+from vllm_omni.experimental.ar_diffusion.tick_protocol import (
+ ARDiffusionControlInput,
+ ARDiffusionTickRequest,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_ROOT = Path(__file__).parents[4]
+_LINGBOT_INIT_PATH = _ROOT / "vllm_omni/diffusion/models/lingbot_world/__init__.py"
+_MODEL_INDEX_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_model_index.json.fixture"
+_SCHEDULER_FIXTURE = Path(__file__).with_name("fixtures") / "lingbot_world_scheduler_config.json.fixture"
+
+
+def _scheduler_config(**overrides):
+ values = {
+ "_class_name": "UniPCMultistepScheduler",
+ "num_train_timesteps": 1000,
+ "flow_shift": 5.0,
+ "prediction_type": "flow_prediction",
+ "predict_x0": True,
+ "use_flow_sigmas": True,
+ "use_dynamic_shifting": False,
+ "use_beta_sigmas": False,
+ "use_exponential_sigmas": False,
+ "use_karras_sigmas": False,
+ "final_sigmas_type": "zero",
+ "timestep_spacing": "linspace",
+ "solver_order": 2,
+ "solver_type": "bh2",
+ "lower_order_final": True,
+ "disable_corrector": [],
+ "time_shift_type": "exponential",
+ }
+ values.update(overrides)
+ return values
+
+
+class _AutoWeightsLoader:
+ def __init__(self, module):
+ self.module = module
+
+ def load_weights(self, weights):
+ loaded = self.module.transformer.load_weights(
+ (name.removeprefix("transformer."), value) for name, value in weights if name.startswith("transformer.")
+ )
+ return {f"transformer.{name}" for name in loaded}
+
+
+class _FakePretrained:
+ @classmethod
+ def from_pretrained(cls, *args, **kwargs):
+ del args, kwargs
+ return cls()
+
+ def to(self, *args, **kwargs):
+ self.to_calls = getattr(self, "to_calls", [])
+ self.to_calls.append((args, kwargs))
+ return self
+
+
+class _FakeScheduler:
+ def __init__(self, **kwargs):
+ self.config = SimpleNamespace(**kwargs)
+
+
+class _FakeTransformerFactory:
+ @classmethod
+ def from_config(cls, config, *, quant_config=None, prefix=""):
+ cls.last_call = (config, quant_config, prefix)
+ return _RecordingTransformer()
+
+
+class _Cache:
+ def __init__(self, *, num_layers: int, max_tokens: int, num_local_heads: int, head_dim: int):
+ shape = (1, max_tokens, num_local_heads, head_dim)
+ self.self_attention = [
+ SimpleNamespace(key=torch.zeros(shape), value=torch.zeros(shape)) for _ in range(num_layers)
+ ]
+ self.cross_attention = [None] * num_layers
+
+
+class _RecordingTransformer(nn.Module):
+ def __init__(self, *, raise_on_call: int | None = None, dtype: torch.dtype = torch.float32):
+ super().__init__()
+ self.config = SimpleNamespace(
+ patch_size=(1, 2, 2),
+ in_channels=36,
+ out_channels=16,
+ text_dim=8,
+ num_layers=2,
+ num_attention_heads=2,
+ attention_head_dim=4,
+ num_frames_per_block=3,
+ sliding_window_num_frames=6,
+ local_attn_size=-1,
+ sink_size=3,
+ )
+ self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()])
+ for block in self.blocks:
+ block.self_attn = SimpleNamespace(num_local_heads=2, head_dim=4)
+ self.calls: list[dict] = []
+ self.cache_allocations: list[dict] = []
+ self.raise_on_call = raise_on_call
+ self._dtype = dtype
+ self.loaded_weights: list[tuple[str, torch.Tensor]] = []
+
+ @property
+ def dtype(self) -> torch.dtype:
+ return self._dtype
+
+ def forward(self, **kwargs):
+ call = {
+ "hidden_states": kwargs["hidden_states"].detach().clone(),
+ "timestep": kwargs["timestep"].detach().clone(),
+ "encoder_hidden_states": kwargs["encoder_hidden_states"].detach().clone(),
+ "camera_hidden_states": kwargs["camera_hidden_states"].detach().clone(),
+ "cache_id": id(kwargs["cache"]),
+ "start_frame": kwargs["start_frame"],
+ "update_cache": kwargs["update_cache"],
+ }
+ self.calls.append(call)
+ if self.raise_on_call == len(self.calls):
+ raise RuntimeError("forced transformer failure")
+ return torch.ones_like(kwargs["hidden_states"][:, :16])
+
+ def allocate_cache(self, **kwargs):
+ self.cache_allocations.append(dict(kwargs))
+ patch_height, patch_width = self.config.patch_size[1:]
+ window_frames = (
+ self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames
+ )
+ max_tokens = window_frames * (kwargs["latent_height"] // patch_height) * (kwargs["latent_width"] // patch_width)
+ return _Cache(
+ num_layers=self.config.num_layers,
+ max_tokens=max_tokens,
+ num_local_heads=2,
+ head_dim=4,
+ )
+
+ def load_weights(self, weights):
+ self.loaded_weights = list(weights)
+ return {name for name, _ in self.loaded_weights}
+
+
+class _StubVAE(_FakePretrained):
+ dtype = torch.float32
+
+ def __init__(self):
+ self.config = SimpleNamespace(
+ z_dim=16,
+ scale_factor_temporal=4,
+ scale_factor_spatial=8,
+ latents_mean=[float(index) - 3.0 for index in range(16)],
+ latents_std=[1.0 + index / 4.0 for index in range(16)],
+ )
+ self.encode_inputs: list[torch.Tensor] = []
+ self.decode_inputs: list[torch.Tensor] = []
+ self.on_decode = None
+
+ def encode(self, video: torch.Tensor):
+ self.encode_inputs.append(video.detach().clone())
+ latent_frames = (video.shape[2] - 1) // 4 + 1
+ latents = torch.zeros(
+ video.shape[0],
+ 16,
+ latent_frames,
+ video.shape[-2] // 8,
+ video.shape[-1] // 8,
+ dtype=video.dtype,
+ device=video.device,
+ )
+ latents[:, :, 0] = 2.0
+ return SimpleNamespace(latents=latents)
+
+ def decode(self, latents: torch.Tensor, return_dict: bool = False):
+ del return_dict
+ if self.on_decode is not None:
+ self.on_decode()
+ self.decode_inputs.append(latents.detach().clone())
+ pixel_frames = (latents.shape[2] - 1) * 4 + 1
+ decoded = torch.zeros(
+ latents.shape[0],
+ 3,
+ pixel_frames,
+ latents.shape[-2] * 8,
+ latents.shape[-1] * 8,
+ dtype=latents.dtype,
+ device=latents.device,
+ )
+ return (decoded,)
+
+
+class _SamplingParams:
+ def __init__(
+ self,
+ *,
+ height: int | None = 16,
+ width: int | None = 16,
+ num_frames: int = 9,
+ num_inference_steps: int | None = 4,
+ num_outputs_per_prompt: int = 1,
+ seed: int | None = 17,
+ generator: torch.Generator | None = None,
+ output_type: str | None = "latent",
+ max_sequence_length: int | None = 512,
+ extra_args: dict | None = None,
+ include_action: bool = True,
+ ) -> None:
+ self.height = height
+ self.width = width
+ self.num_frames = num_frames
+ self.num_inference_steps = num_inference_steps
+ self.num_outputs_per_prompt = num_outputs_per_prompt
+ self.seed = seed
+ self.generator = (
+ generator
+ if generator is not None
+ else (torch.Generator(device="cpu").manual_seed(seed) if seed is not None else None)
+ )
+ self.output_type = output_type
+ self.max_sequence_length = max_sequence_length
+ self.extra_args = {"action_path": "."} if include_action else {}
+ self.extra_args["_lingbot_camera_trajectory"] = _CameraTrajectory(
+ poses=torch.eye(4).repeat(32, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1),
+ )
+ if extra_args is not None:
+ self.extra_args.update(extra_args)
+ self.latents = None
+
+
+class _RequestBatch:
+ def __init__(self, prompt, sampling_params: _SamplingParams, *, num_reqs: int = 1):
+ self._prompts = [prompt] * num_reqs
+ self._sampling = sampling_params
+ self.num_reqs = num_reqs
+
+ @property
+ def prompts(self):
+ return self._prompts
+
+ @property
+ def sampling_params(self):
+ return self._sampling
+
+ @property
+ def request_id(self):
+ return "legacy-lingbot-request"
+
+
+def _load_pipeline_module():
+ return lingbot_pipeline
+
+
+@pytest.fixture(autouse=True)
+def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Replace only heavyweight component construction, not the import graph."""
+
+ loader_state = SimpleNamespace(prefetch_calls=[])
+
+ def prefetch_subfolders(model, subfolders, *, local_files_only):
+ loader_state.prefetch_calls.append((model, tuple(subfolders), local_files_only))
+
+ def from_pretrained_with_prefetch(callable_, model, **kwargs):
+ kwargs.pop("prefetch_list", None)
+ return callable_(model, **kwargs)
+
+ def load_transformer_config(model, subfolder, local_files_only):
+ del model, subfolder, local_files_only
+ return {
+ "_class_name": "CausalLingBotWorldTransformer3DModel",
+ "patch_size": [1, 2, 2],
+ "in_channels": 36,
+ "out_channels": 16,
+ "text_dim": 8,
+ "num_layers": 2,
+ "num_attention_heads": 2,
+ "attention_head_dim": 4,
+ "num_frames_per_block": 3,
+ "sliding_window_num_frames": 6,
+ "local_attn_size": -1,
+ }
+
+ def retrieve_latents(value, sample_mode="argmax"):
+ assert sample_mode == "argmax"
+ return value.latents
+
+ def load_json(model, filename, local_files_only):
+ del model, local_files_only
+ assert filename == "scheduler/scheduler_config.json"
+ return _scheduler_config()
+
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(32, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(32, 1),
+ )
+
+ def load_camera_trajectory(action_path):
+ assert action_path
+ return trajectory
+
+ def interpolate_camera_trajectory(value, num_frames):
+ return SimpleNamespace(
+ poses=value.poses[:num_frames],
+ intrinsics=value.intrinsics[:num_frames],
+ )
+
+ def build_plucker_embedding(value, *, height, width, target_height, target_width, device, dtype):
+ del target_height, target_width
+ frames = value.poses.shape[0]
+ data = torch.arange(frames * 6 * height * width, device=device, dtype=torch.float32)
+ return data.reshape(frames, 6, height, width).to(dtype=dtype)
+
+ replacements = {
+ "AutoTokenizer": _FakePretrained,
+ "UMT5EncoderModel": _FakePretrained,
+ "AutoWeightsLoader": _AutoWeightsLoader,
+ "DistributedAutoencoderKLWan": _StubVAE,
+ "FlowUniPCMultistepScheduler": _FakeScheduler,
+ "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
+ "get_local_device": lambda: torch.device("cpu"),
+ "set_forward_context_denoise_step_idx": lambda index: None,
+ "prefetch_subfolders": prefetch_subfolders,
+ "from_pretrained_with_prefetch": from_pretrained_with_prefetch,
+ "load_transformer_config": load_transformer_config,
+ "retrieve_latents": retrieve_latents,
+ "_load_json": load_json,
+ "load_camera_trajectory": load_camera_trajectory,
+ "interpolate_camera_trajectory": interpolate_camera_trajectory,
+ "build_plucker_embedding": build_plucker_embedding,
+ "randn_tensor": _diffusers_randn_tensor,
+ }
+ for name, value in replacements.items():
+ monkeypatch.setattr(lingbot_pipeline, name, value)
+ monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
+
+
+def _od_config(**overrides):
+ values = {
+ "model": "checkpoint",
+ "dtype": torch.float32,
+ "flow_shift": None,
+ "quantization_config": None,
+ "enable_diffusion_pipeline_profiler": False,
+ "model_config": {
+ "lingbot_action_root": str(_ROOT),
+ "ar_diffusion_height": 16,
+ "ar_diffusion_width": 16,
+ },
+ "enable_cpu_offload": False,
+ "enable_layerwise_offload": False,
+ "enforce_eager": True,
+ "parallel_config": SimpleNamespace(
+ pipeline_parallel_size=1,
+ sequence_parallel_size=1,
+ cfg_parallel_size=1,
+ vae_patch_parallel_size=1,
+ use_hsdp=False,
+ enable_expert_parallel=False,
+ ),
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
+def _prompt(*, action_path: str | None = None, images=None):
+ if images is None:
+ images = Image.new("RGB", (16, 16), color=(255, 128, 0))
+ prompt = {
+ "prompt": "move through the room",
+ "multi_modal_data": {"image": images},
+ "additional_information": {},
+ }
+ if action_path is not None:
+ prompt["additional_information"]["action_path"] = action_path
+ return prompt
+
+
+def _pipeline(module, *, transformer=None, od_config=None):
+ pipeline = module.LingBotWorldCausalDMDPipeline(od_config=od_config or _od_config())
+ if transformer is not None:
+ pipeline.transformer = transformer
+ pipeline.encode_prompt = lambda *args, **kwargs: torch.ones(1, 512, 8)
+ pipeline.progress_bar = noop_progress_bar
+ return pipeline
+
+
+def _preprocess_request(module, *, prompt=None, sampling=None, od_config=None):
+ request = SimpleNamespace(
+ prompt=_prompt() if prompt is None else prompt,
+ sampling_params=_SamplingParams() if sampling is None else sampling,
+ )
+ preprocess = module.get_lingbot_world_pre_process_func(od_config or _od_config())
+ return preprocess(request)
+
+
+@pytest.mark.parametrize("offload_field", ["enable_cpu_offload", "enable_layerwise_offload"])
+def test_pipeline_respects_loader_managed_component_placement(offload_field: str) -> None:
+ module = _load_pipeline_module()
+
+ pipeline = _pipeline(module, od_config=_od_config(**{offload_field: True}))
+
+ assert getattr(pipeline.text_encoder, "to_calls", []) == []
+ assert getattr(pipeline.vae, "to_calls", []) == []
+
+
+def test_ar_diffusion_capability_uses_fixed_tp_local_lingbot_geometry() -> None:
+ module = _load_pipeline_module()
+ module.get_tensor_model_parallel_world_size = lambda: 1
+ pipeline = _pipeline(module)
+
+ spec = pipeline.ar_diffusion_kv_cache_spec()
+
+ assert spec.num_layers == 2
+ assert spec.num_kv_heads == 2
+ assert spec.head_size == 4
+ assert spec.tokens_per_frame == 1
+ assert spec.frames_per_block == 3
+ assert spec.window_frames == 3
+ assert spec.sink_frames == 3
+ assert [(branch.name, branch.local_index) for branch in spec.kv_branches] == [("main", 0)]
+ assert spec.cross_attention_lengths == {"text": 512}
+ assert spec.model_owned_state_bytes_per_session == 9_600
+
+
+def test_preprocess_materializes_external_inputs_before_worker_execution(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ action_root = tmp_path / "trusted-actions"
+ action_dir = action_root / "forward"
+ action_dir.mkdir(parents=True)
+ image_path = tmp_path / "first-frame.png"
+ Image.new("RGBA", (16, 16), color=(1, 2, 3, 4)).save(image_path)
+ sampling = _SamplingParams(extra_args={"action_path": "forward"})
+ sampling.extra_args.pop("_lingbot_camera_trajectory")
+ request = SimpleNamespace(
+ prompt={"prompt": "move", "multi_modal_data": {"image": str(image_path)}},
+ sampling_params=sampling,
+ )
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(9, 1, 1),
+ intrinsics=torch.ones(9, 4),
+ )
+ load_calls = []
+ module.load_camera_trajectory = lambda action: load_calls.append(action) or trajectory
+
+ preprocess = module.get_lingbot_world_pre_process_func(
+ _od_config(model_config={"lingbot_action_root": str(action_root)})
+ )
+ result = preprocess(request)
+
+ assert result is request
+ assert isinstance(request.prompt["multi_modal_data"]["image"], Image.Image)
+ assert request.prompt["multi_modal_data"]["image"].mode == "RGB"
+ assert sampling.extra_args["action_path"] == "forward"
+ assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory
+ assert len(load_calls) == 1
+
+
+def test_preprocess_materializes_camera_from_typed_tick_without_action_path() -> None:
+ module = _load_pipeline_module()
+ poses = torch.eye(4).repeat(9, 1, 1)
+ intrinsics = torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(9, 1)
+ tick = ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="tick-request-0",
+ chunk_index=0,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={
+ "poses": poses.tolist(),
+ "intrinsics": intrinsics.tolist(),
+ },
+ ),
+ ),
+ )
+ sampling = _SamplingParams(include_action=False)
+ sampling.extra_args.update(tick.to_extra_args())
+ request = SimpleNamespace(
+ request_id="tick-request-0",
+ prompt=_prompt(),
+ sampling_params=sampling,
+ )
+
+ result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+ trajectory = result.sampling_params.extra_args["_lingbot_camera_trajectory"]
+ torch.testing.assert_close(trajectory.poses, poses)
+ torch.testing.assert_close(trajectory.intrinsics, intrinsics)
+
+
+def test_preprocess_materializes_chunk_sized_actions_from_typed_tick() -> None:
+ module = _load_pipeline_module()
+ tick = ARDiffusionTickRequest(
+ session_id="world-actions",
+ request_id="tick-request-0",
+ chunk_index=0,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_actions.v1",
+ data={
+ "mode": "frames",
+ "frames": [["w"], ["w", "j"], []],
+ },
+ ),
+ ),
+ )
+ sampling = _SamplingParams(include_action=False)
+ sampling.extra_args.update(tick.to_extra_args())
+ request = SimpleNamespace(
+ request_id="tick-request-0",
+ prompt=_prompt(),
+ sampling_params=sampling,
+ )
+
+ result = module.get_lingbot_world_pre_process_func(_od_config())(request)
+
+ assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
+ assert result.sampling_params.extra_args["_lingbot_camera_actions"] == (
+ ("w",),
+ ("w", "j"),
+ (),
+ )
+
+
+def _request(*, sampling=None, prompt=None, num_reqs: int = 1):
+ return _RequestBatch(
+ _prompt() if prompt is None else prompt,
+ _SamplingParams() if sampling is None else sampling,
+ num_reqs=num_reqs,
+ )
+
+
+def _resolve_pipeline_through_real_registry(pipeline_module):
+ from vllm_omni.diffusion import registry as registry_module
+
+ resolved = registry_module.DiffusionModelRegistry._try_load_model_cls("LingBotWorldCausalDMDPipeline")
+ entry = registry_module._DIFFUSION_MODELS["LingBotWorldCausalDMDPipeline"]
+ cache_acceleration_disabled = "LingBotWorldCausalDMDPipeline" in registry_module._NO_CACHE_ACCELERATION
+ preprocess_name = registry_module._DIFFUSION_PRE_PROCESS_FUNCS["LingBotWorldCausalDMDPipeline"]
+ preprocess = registry_module.get_diffusion_pre_process_func(
+ SimpleNamespace(
+ model_class_name="LingBotWorldCausalDMDPipeline",
+ model_config={"lingbot_action_root": str(_ROOT)},
+ )
+ )
+ return resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess
+
+
+def test_component_discovery_uses_official_checkpoint_contract() -> None:
+ module = _load_pipeline_module()
+ pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config())
+
+ assert pipeline._dit_modules == ["transformer"]
+ assert pipeline._encoder_modules == ["text_encoder"]
+ assert pipeline._vae_modules == ["vae"]
+ assert pipeline.dummy_run_num_frames == 0
+ assert pipeline.weights_sources == [
+ module.DiffusersPipelineLoader.ComponentSource("checkpoint", "transformer", None, "transformer.", True)
+ ]
+ assert _FakeTransformerFactory.last_call[1:] == (None, "transformer")
+ assert pipeline.scheduler.config.shift == 5.0
+ assert pipeline.scheduler.config.num_train_timesteps == 1000
+ assert module._loader_state.prefetch_calls == [("checkpoint", ("tokenizer", "text_encoder", "vae"), False)]
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "feature"),
+ [
+ ("pipeline_parallel_size", 2, "pipeline parallelism"),
+ ("sequence_parallel_size", 2, "sequence parallelism"),
+ ("cfg_parallel_size", 2, "CFG parallelism"),
+ ("vae_patch_parallel_size", 2, "VAE parallelism"),
+ ("use_hsdp", True, "HSDP"),
+ ("enable_expert_parallel", True, "expert parallelism"),
+ ],
+)
+def test_unsupported_parallel_modes_fail_before_component_loading(field: str, value: object, feature: str) -> None:
+ module = _load_pipeline_module()
+ parallel_config = _od_config().parallel_config
+ setattr(parallel_config, field, value)
+
+ with pytest.raises(NotImplementedError, match=feature):
+ module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=parallel_config))
+
+ assert module._loader_state.prefetch_calls == []
+
+
+def test_unsupported_quantization_fails_before_component_loading() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(NotImplementedError, match="quantization"):
+ module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object()))
+
+ assert module._loader_state.prefetch_calls == []
+
+
+def test_official_scheduler_config_matches_fixed_dmd_contract() -> None:
+ import json
+
+ module = _load_pipeline_module()
+ scheduler_config = json.loads(_SCHEDULER_FIXTURE.read_text())
+ module._load_json = lambda *args, **kwargs: scheduler_config.copy()
+
+ pipeline = _pipeline(module)
+
+ assert pipeline.scheduler.config.num_train_timesteps == 1000
+ assert pipeline.scheduler.config.shift == 5.0
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [
+ ("_class_name", "FlowMatchEulerDiscreteScheduler"),
+ ("num_train_timesteps", 999),
+ ("prediction_type", "epsilon"),
+ ("predict_x0", False),
+ ("use_flow_sigmas", False),
+ ("use_dynamic_shifting", True),
+ ("final_sigmas_type", "sigma_min"),
+ ],
+)
+def test_scheduler_config_rejects_semantic_drift(field: str, value: object) -> None:
+ module = _load_pipeline_module()
+ module._load_json = lambda *args, **kwargs: _scheduler_config(**{field: value})
+
+ with pytest.raises(ValueError, match=field):
+ _pipeline(module)
+
+
+def test_official_model_index_discovers_only_declared_components() -> None:
+ import json
+
+ module = _load_pipeline_module()
+ model_index = json.loads(_MODEL_INDEX_FIXTURE.read_text())
+ resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
+ module
+ )
+
+ assert model_index["_class_name"] == "LingBotWorldCausalDMDPipeline"
+ assert resolved is module.LingBotWorldCausalDMDPipeline
+ assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline")
+ assert cache_acceleration_disabled
+ assert preprocess_name == "get_lingbot_world_pre_process_func"
+ assert callable(preprocess)
+ assert model_index["tokenizer"] == ["transformers", "T5TokenizerFast"]
+ assert model_index["text_encoder"] == ["transformers", "UMT5EncoderModel"]
+ assert model_index["vae"] == ["diffusers", "AutoencoderKLWan"]
+ assert model_index["scheduler"] == ["diffusers", "UniPCMultistepScheduler"]
+ assert model_index["transformer"] == ["diffusers", "CausalLingBotWorldTransformer3DModel"]
+ assert model_index["image_encoder"] == [None, None]
+ assert model_index["image_processor"] == [None, None]
+ assert model_index["transformer_2"] == [None, None]
+
+
+def test_postprocess_returns_latents_without_video_conversion() -> None:
+ module = _load_pipeline_module()
+ postprocess = module.get_lingbot_world_post_process_func(_od_config())
+ latents = torch.randn(1, 16, 3, 2, 2)
+
+ output = postprocess(latents, sampling_params=SimpleNamespace(output_type="latent"))
+
+ assert output is latents
+
+
+def test_postprocess_uses_the_standard_diffusion_output_envelope(monkeypatch) -> None:
+ import diffusers.video_processor as video_processor_module
+
+ module = _load_pipeline_module()
+ processed_video = object()
+ calls = []
+
+ class VideoProcessor:
+ def __init__(self, *, vae_scale_factor):
+ assert vae_scale_factor == 8
+
+ def postprocess_video(self, video, *, output_type):
+ calls.append((video, output_type))
+ return processed_video
+
+ monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
+ postprocess = module.get_lingbot_world_post_process_func(_od_config())
+ video = torch.randn(1, 3, 9, 8, 8)
+
+ output = postprocess(video, output_type="np")
+
+ assert output == {"payload": {"video": processed_video}, "metadata": {}}
+ assert calls == [(video, "np")]
+
+
+def test_forward_propagates_pipeline_profiler_stage_durations() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._profiler_lock = Lock()
+ pipeline._stage_durations = {"_generate_block": 0.25, "vae.decode": 0.5}
+
+ output = pipeline(_request())
+
+ assert output.stage_durations == pipeline.stage_durations
+
+
+def test_pipeline_weight_loader_preserves_component_parameter_prefixes() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ weight = torch.randn(2, 2)
+
+ loaded = pipeline.load_weights([("transformer.blocks.0.weight", weight)])
+
+ assert loaded == {"transformer.blocks.0.weight"}
+ assert pipeline.transformer.loaded_weights == [("blocks.0.weight", weight)]
+
+
+def test_request_requires_runner_provided_generator() -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(seed=None, generator=None)
+
+ with pytest.raises(ValueError, match="runner-provided torch.Generator"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+@pytest.mark.parametrize("unsupported", ["generator-list", "caller-latents"])
+def test_request_rejects_unsupported_sampling_state(unsupported: str) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams()
+ if unsupported == "generator-list":
+ sampling.generator = [sampling.generator]
+ message = "generator list"
+ else:
+ sampling.latents = torch.zeros(1, 16, 3, 2, 2)
+ message = "caller-provided latents"
+
+ with pytest.raises(ValueError, match=message):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer(dtype=torch.bfloat16)
+ pipeline = _pipeline(module, transformer=transformer)
+ requested_noise_dtypes: list[torch.dtype] = []
+
+ def randn(shape, *, generator, device, dtype):
+ del generator
+ assert device == pipeline.device
+ requested_noise_dtypes.append(dtype)
+ return torch.zeros(shape, device=device, dtype=dtype)
+
+ module.randn_tensor = randn
+ result = pipeline(_request())
+
+ assert requested_noise_dtypes == [torch.float32] * 4
+ assert result.output.dtype == torch.float32
+ assert all(call["hidden_states"].dtype == torch.bfloat16 for call in transformer.calls)
+
+
+def test_path_image_rejects_oversized_source_before_decode_or_convert(monkeypatch, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ source_path = tmp_path / "oversized-compressed.png"
+ Image.new("1", (4097, 4097)).save(source_path)
+ decode_calls: list[str] = []
+
+ def forbidden_convert(*args, **kwargs):
+ del args, kwargs
+ decode_calls.append("convert")
+ raise AssertionError("oversized source reached convert")
+
+ def forbidden_load(*args, **kwargs):
+ del args, kwargs
+ decode_calls.append("load")
+ raise AssertionError("oversized source reached load")
+
+ monkeypatch.setattr(module.PIL.Image.Image, "convert", forbidden_convert)
+ monkeypatch.setattr(module.PIL.Image.Image, "load", forbidden_load)
+
+ with pytest.raises(ValueError, match="source image.*4096.*4096"):
+ _preprocess_request(module, prompt=_prompt(images=str(source_path)))
+
+ assert decode_calls == []
+
+
+def test_normal_path_image_is_decoded_after_source_size_validation(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ source_path = tmp_path / "normal.png"
+ Image.new("RGBA", (32, 24), color=(1, 2, 3, 128)).save(source_path)
+
+ request = _preprocess_request(module, prompt=_prompt(images=str(source_path)))
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+
+ assert module._MAX_SOURCE_IMAGE_PIXELS == 4096 * 4096
+ assert isinstance(parsed.image, Image.Image)
+ assert parsed.image.mode == "RGB"
+ assert parsed.image.size == (32, 24)
+
+
+@pytest.mark.parametrize(
+ ("source_size", "expected_size"),
+ [
+ ((832, 480), (464, 832)),
+ ((480, 832), (832, 480)),
+ ((512, 512), (624, 624)),
+ ],
+)
+def test_default_resolution_matches_official_480p_aspect_derivation(source_size, expected_size) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(height=None, width=None)
+
+ parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(images=Image.new("RGB", source_size)), sampling))
+
+ assert (parsed.height, parsed.width) == expected_size
+ assert parsed.height % 16 == 0
+ assert parsed.width % 16 == 0
+ assert parsed.height * parsed.width <= 480 * 832
+
+
+def test_pil_and_tensor_images_share_official_bicubic_preprocess() -> None:
+ module = _load_pipeline_module()
+ height, width = 7, 11
+ source = torch.arange(3 * height * width, dtype=torch.uint8).reshape(3, height, width)
+ pil_image = Image.fromarray(source.permute(1, 2, 0).numpy(), mode="RGB")
+ sampling = _SamplingParams(height=16, width=32)
+
+ pil_pipeline = _pipeline(module)
+ tensor_pipeline = _pipeline(module)
+ pil_inputs = pil_pipeline._parse_request(_RequestBatch(_prompt(images=pil_image), sampling))
+ tensor_inputs = tensor_pipeline._parse_request(
+ _RequestBatch(_prompt(images=source.clone()), _SamplingParams(height=16, width=32))
+ )
+ pil_condition = pil_pipeline._prepare_image_tensor(pil_inputs.image, height=16, width=32)
+ tensor_condition = tensor_pipeline._prepare_image_tensor(tensor_inputs.image, height=16, width=32)
+ expected = torch.nn.functional.interpolate(
+ source.unsqueeze(0).float() / 255.0,
+ size=(16, 32),
+ mode="bicubic",
+ align_corners=False,
+ )
+ expected = expected * 2.0 - 1.0
+
+ torch.testing.assert_close(pil_condition, expected)
+ torch.testing.assert_close(tensor_condition, expected)
+
+
+def test_path_image_decode_error_is_sanitized(monkeypatch) -> None:
+ module = _load_pipeline_module()
+ unsafe_path = "/private/source/customer-secret.png"
+ monkeypatch.setattr(
+ module.PIL.Image,
+ "open",
+ lambda path: (_ for _ in ()).throw(module.PIL.Image.DecompressionBombError(f"unsafe {path}")),
+ )
+
+ with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info:
+ _preprocess_request(module, prompt=_prompt(images=unsafe_path))
+
+ assert unsafe_path not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+@pytest.mark.parametrize("source_size", [(4096, 4096), (4097, 4097)])
+def test_supplied_pil_image_obeys_documented_source_pixel_ceiling(source_size) -> None:
+ module = _load_pipeline_module()
+ source_image = Image.new("1", source_size)
+ close_calls: list[bool] = []
+ source_image.close = lambda: close_calls.append(True)
+
+ if source_size[0] * source_size[1] <= 4096 * 4096:
+ request = _preprocess_request(module, prompt=_prompt(images=source_image))
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+ assert parsed.image is not source_image
+ assert parsed.image.mode == "RGB"
+ assert close_calls == []
+ else:
+ with pytest.raises(ValueError, match="source image.*4096.*4096"):
+ _preprocess_request(module, prompt=_prompt(images=source_image))
+ assert close_calls == []
+
+
+def test_supplied_pil_decode_error_is_sanitized_without_closing_caller(monkeypatch) -> None:
+ module = _load_pipeline_module()
+ source_image = Image.new("RGB", (16, 16))
+ unsafe_detail = "/private/source/customer-secret.png"
+ close_calls: list[bool] = []
+ monkeypatch.setattr(source_image, "convert", lambda mode: (_ for _ in ()).throw(OSError(unsafe_detail)))
+ monkeypatch.setattr(source_image, "close", lambda: close_calls.append(True))
+
+ with pytest.raises(ValueError, match="Unable to load multi_modal_data.image") as exc_info:
+ _preprocess_request(module, prompt=_prompt(images=source_image))
+
+ assert unsafe_detail not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+ assert close_calls == []
+
+
+def test_action_path_resolves_from_sampling_extra_args(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ action_root = tmp_path / "trusted-actions"
+ contained_action = action_root / "actions"
+ contained_action.mkdir(parents=True)
+ sampling = _SamplingParams(extra_args={"action_path": "actions"})
+ prompt = _prompt()
+ original_prompt = prompt.copy()
+ trajectory = _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4))
+ resolved_actions = []
+ module.load_camera_trajectory = lambda action: resolved_actions.append(action) or trajectory
+
+ request = _preprocess_request(
+ module,
+ prompt=prompt,
+ sampling=sampling,
+ od_config=_od_config(model_config={"lingbot_action_root": str(action_root)}),
+ )
+ parsed = _pipeline(module)._parse_request(_RequestBatch(request.prompt, request.sampling_params))
+
+ assert resolved_actions[0].root == action_root.resolve()
+ assert resolved_actions[0].relative == Path("actions")
+ assert parsed.camera_trajectory is trajectory
+ assert sampling.extra_args["action_path"] == "actions"
+ assert sampling.extra_args["_lingbot_camera_trajectory"] is trajectory
+ assert prompt == original_prompt
+
+
+@pytest.mark.parametrize("with_extra", [False, True], ids=["additional-only", "extra-and-additional"])
+def test_additional_information_action_path_is_not_an_input_source(with_extra: bool) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(
+ extra_args={"action_path": "."} if with_extra else None,
+ include_action=with_extra,
+ )
+
+ with pytest.raises(ValueError, match="additional_information.*action_path.*not supported"):
+ _preprocess_request(module, prompt=_prompt(action_path="legacy-actions"), sampling=sampling)
+
+
+def test_action_path_is_required_in_sampling_extra_args() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match="sampling_params.extra_args.action_path"):
+ _preprocess_request(module, sampling=_SamplingParams(include_action=False))
+
+
+@pytest.mark.parametrize("escape_kind", ["traversal", "absolute", "symlink"])
+def test_online_action_path_rejects_escape_from_trusted_root(escape_kind: str, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ root.mkdir()
+ outside = tmp_path / "outside"
+ outside.mkdir()
+ if escape_kind == "traversal":
+ action_path = "../outside"
+ elif escape_kind == "absolute":
+ action_path = str(outside)
+ else:
+ (root / "escape-link").symlink_to(outside, target_is_directory=True)
+ action_path = "escape-link"
+ with pytest.raises(ValueError, match="trusted.*root|contained"):
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": action_path}),
+ od_config=_od_config(model_config={"lingbot_action_root": str(root)}),
+ )
+
+
+def test_online_action_path_uses_environment_root_fallback(monkeypatch, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ action_dir = root / "forward"
+ action_dir.mkdir(parents=True)
+ monkeypatch.setenv("VLLM_OMNI_LINGBOT_ACTION_ROOT", str(root))
+ resolved_actions = []
+ module.load_camera_trajectory = lambda action: (
+ resolved_actions.append(action) or _CameraTrajectory(torch.eye(4).repeat(9, 1, 1), torch.ones(9, 4))
+ )
+
+ request = _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": "forward"}),
+ od_config=_od_config(model_config={}),
+ )
+
+ assert request.sampling_params.extra_args["_lingbot_camera_trajectory"] is not None
+ assert resolved_actions[0].root == root.resolve()
+ assert resolved_actions[0].relative == Path("forward")
+
+
+def test_online_action_path_error_suppresses_path_bearing_filesystem_cause(tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ root = tmp_path / "trusted"
+ root.mkdir()
+
+ with pytest.raises(ValueError, match="trusted action root") as exc_info:
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": "does-not-exist"}),
+ od_config=_od_config(model_config={"lingbot_action_root": str(root)}),
+ )
+
+ assert str(root) not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+@pytest.mark.parametrize("source", ["root", "candidate"])
+def test_online_action_path_unknown_user_is_sanitized(source: str, tmp_path: Path) -> None:
+ module = _load_pipeline_module()
+ unknown_user_path = "~__vllm_omni_user_that_does_not_exist__/actions"
+ root = tmp_path / "trusted"
+ root.mkdir()
+ configured_root = unknown_user_path if source == "root" else str(root)
+ action_path = "actions" if source == "root" else unknown_user_path
+
+ with pytest.raises(ValueError, match="trusted action root") as exc_info:
+ _preprocess_request(
+ module,
+ sampling=_SamplingParams(extra_args={"action_path": action_path}),
+ od_config=_od_config(model_config={"lingbot_action_root": configured_root}),
+ )
+
+ assert "__vllm_omni_user_that_does_not_exist__" not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+def test_action_path_requires_a_trusted_root_for_every_request() -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match="lingbot_action_root|VLLM_OMNI_LINGBOT_ACTION_ROOT"):
+ _preprocess_request(module, od_config=_od_config(model_config={}))
+
+
+@pytest.mark.parametrize(
+ ("request_batch", "message"),
+ [
+ (_request(num_reqs=2), "single prompt"),
+ (_request(prompt=_prompt(images=[Image.new("RGB", (16, 16))])), "image.*list"),
+ (_request(prompt=_prompt(images=[])), "image.*list"),
+ (
+ _request(prompt=_prompt(images=[Image.new("RGB", (16, 16)), Image.new("RGB", (16, 16))])),
+ "image.*list",
+ ),
+ (_request(sampling=_SamplingParams(num_outputs_per_prompt=2)), "num_outputs_per_prompt"),
+ (_request(sampling=_SamplingParams(num_inference_steps=5)), "num_inference_steps"),
+ (_request(sampling=_SamplingParams(height=15)), "height.*divisible"),
+ (_request(sampling=_SamplingParams(width=15)), "width.*divisible"),
+ (_request(sampling=_SamplingParams(height=None, width=16)), "height and width.*both"),
+ (_request(sampling=_SamplingParams(num_frames=13)), "num_frames.*three-frame"),
+ ],
+)
+def test_request_validation_rejects_unsupported_contracts(request_batch, message: str) -> None:
+ module = _load_pipeline_module()
+
+ with pytest.raises(ValueError, match=message):
+ _pipeline(module)._parse_request(request_batch)
+
+
+def test_resource_limits_accept_exact_documented_boundaries() -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(height=480, width=832, num_frames=117, max_sequence_length=512)
+
+ parsed = _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+ assert (parsed.height, parsed.width, parsed.num_frames, parsed.max_sequence_length) == (480, 832, 117, 512)
+
+
+@pytest.mark.parametrize(
+ ("sampling", "message"),
+ [
+ (_SamplingParams(height=480, width=848), "pixel area|480.*832"),
+ (_SamplingParams(num_frames=129), "num_frames.*117"),
+ (_SamplingParams(max_sequence_length=511), "max_sequence_length.*512"),
+ (_SamplingParams(max_sequence_length=513), "max_sequence_length.*512"),
+ (_SamplingParams(max_sequence_length=512.0), "max_sequence_length.*512"),
+ ],
+)
+def test_resource_limits_reject_oversize_before_any_component_call(sampling, message: str) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ calls: list[str] = []
+ pipeline.encode_prompt = lambda *args, **kwargs: calls.append("text")
+ pipeline._prepare_condition = lambda *args, **kwargs: calls.append("vae")
+ pipeline._prepare_camera = lambda *args, **kwargs: calls.append("camera")
+ pipeline.transformer.allocate_cache = lambda *args, **kwargs: calls.append("cache")
+
+ with pytest.raises(ValueError, match=message):
+ pipeline(_RequestBatch(_prompt(), sampling))
+
+ assert calls == []
+
+
+def test_request_validation_rejects_insufficient_camera_frames() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ short_trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(8, 1, 1),
+ intrinsics=torch.ones(8, 4),
+ )
+ sampling = _SamplingParams(extra_args={"_lingbot_camera_trajectory": short_trajectory})
+
+ with pytest.raises(ValueError, match="camera.*frames.*num_frames"):
+ pipeline(_request(sampling=sampling))
+
+
+def test_camera_load_error_is_actionable_without_echoing_path_contents() -> None:
+ module = _load_pipeline_module()
+ module.load_camera_trajectory = lambda path: (_ for _ in ()).throw(FileNotFoundError(f"missing {path}/poses.npy"))
+
+ with pytest.raises(ValueError, match="camera trajectory.*action_path") as exc_info:
+ _preprocess_request(module)
+
+ assert str(_ROOT) not in str(exc_info.value)
+ assert exc_info.value.__cause__ is None
+
+
+def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ module.randn_tensor = lambda shape, **kwargs: torch.full(
+ shape,
+ -99.0,
+ device=kwargs["device"],
+ dtype=kwargs["dtype"],
+ )
+
+ result = pipeline(_request())
+
+ assert result.output.shape == (1, 16, 3, 2, 2)
+ first_input = transformer.calls[0]["hidden_states"]
+ assert first_input.shape == (1, 36, 3, 2, 2)
+ mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1)
+ std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1)
+ expected_first = (torch.full((1, 16, 2, 2), 2.0) - mean) / std
+ expected_future = (torch.zeros(1, 16, 2, 2, 2) - mean.unsqueeze(2)) / std.unsqueeze(2)
+ torch.testing.assert_close(first_input[:, :16], torch.full((1, 16, 3, 2, 2), -99.0))
+ torch.testing.assert_close(first_input[:, 16:20, 0], torch.ones(1, 4, 2, 2))
+ torch.testing.assert_close(first_input[:, 16:20, 1:], torch.zeros(1, 4, 2, 2, 2))
+ torch.testing.assert_close(first_input[:, 20:36, 0], expected_first)
+ torch.testing.assert_close(first_input[:, 20:36, 1:], expected_future)
+
+ raw_camera = module.build_plucker_embedding(
+ SimpleNamespace(poses=torch.empty(3, 4, 4)),
+ height=16,
+ width=16,
+ target_height=16,
+ target_width=16,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+ expected_camera = module._fold_camera_embedding(raw_camera)
+ reference_camera = torch.nn.functional.pixel_unshuffle(raw_camera, 8).permute(1, 0, 2, 3).unsqueeze(0)
+ torch.testing.assert_close(expected_camera, reference_camera)
+ torch.testing.assert_close(transformer.calls[0]["camera_hidden_states"], expected_camera)
+ assert expected_camera.shape == (1, 384, 3, 2, 2)
+
+
+def test_fixed_dmd_transition_and_cache_commit_trace() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ result = pipeline(_request())
+
+ generator = torch.Generator(device="cpu").manual_seed(17)
+ current = torch.randn((1, 16, 3, 2, 2), generator=generator)
+ warped_schedule = (
+ (1000.0, 1.0),
+ (937.5, 0.9375),
+ (2500.0 / 3.0, 5.0 / 6.0),
+ (625.0, 0.625),
+ )
+ for index, (_, sigma) in enumerate(warped_schedule):
+ x0 = current - sigma
+ if index + 1 < len(module.LINGBOT_DMD_TIMESTEPS):
+ next_sigma = warped_schedule[index + 1][1]
+ noise = torch.randn(current.shape, generator=generator)
+ current = (1.0 - next_sigma) * x0 + next_sigma * noise
+ else:
+ current = x0
+ torch.testing.assert_close(result.output, current)
+
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls]),
+ torch.tensor([*(timestep for timestep, _ in warped_schedule), 0.0]),
+ )
+ assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True]
+ assert [call["start_frame"] for call in transformer.calls] == [0, 0, 0, 0, 0]
+ assert len({call["cache_id"] for call in transformer.calls}) == 1
+ torch.testing.assert_close(transformer.calls[-1]["hidden_states"][:, :16], result.output)
+
+
+@pytest.mark.parametrize(
+ ("flow_shift", "expected_timesteps"),
+ [
+ (5.0, (1000.0, 937.5, 2500.0 / 3.0, 625.0)),
+ (2.5, (1000.0, 15000.0 / 17.0, 5000.0 / 7.0, 5000.0 / 11.0)),
+ ],
+)
+def test_request_flow_shift_warps_transformer_timesteps(flow_shift: float, expected_timesteps) -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ pipeline(
+ _request(
+ sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": flow_shift}),
+ )
+ )
+
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls[:4]]),
+ torch.tensor(expected_timesteps),
+ )
+
+
+def test_request_tiny_positive_flow_shift_keeps_schedule_finite() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+
+ pipeline(
+ _request(
+ sampling=_SamplingParams(extra_args={"action_path": ".", "flow_shift": 1e-20}),
+ )
+ )
+
+ timesteps = torch.cat([call["timestep"] for call in transformer.calls[:4]])
+ assert torch.isfinite(timesteps).all()
+ assert timesteps[0].item() == 1000.0
+
+
+def test_request_flow_shift_override_has_precedence_without_mutating_scheduler() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module, od_config=_od_config(flow_shift=3.0))
+
+ default_inputs = pipeline._parse_request(_request())
+ sampling = _SamplingParams(extra_args={"flow_shift": 2.0})
+ override_inputs = pipeline._parse_request(_RequestBatch(_prompt(), sampling))
+
+ assert default_inputs.flow_shift == 3.0
+ assert override_inputs.flow_shift == 2.0
+ assert sampling.extra_args["action_path"] == "."
+ assert sampling.extra_args["flow_shift"] == 2.0
+ assert pipeline.scheduler.config.shift == 3.0
+
+
+def test_checkpoint_flow_shift_default_and_engine_request_override_precedence() -> None:
+ module = _load_pipeline_module()
+ checkpoint_scheduler_config = _scheduler_config(flow_shift=6.25, shift=7.5)
+ module._load_json = lambda *args, **kwargs: checkpoint_scheduler_config.copy()
+
+ checkpoint_default = _pipeline(module)
+ engine_override = _pipeline(module, od_config=_od_config(flow_shift=3.5))
+ request_override = checkpoint_default._parse_request(
+ _RequestBatch(
+ _prompt(),
+ _SamplingParams(extra_args={"flow_shift": 2.25}),
+ )
+ )
+
+ assert checkpoint_default.scheduler.config.shift == 6.25
+ assert checkpoint_default._parse_request(_request()).flow_shift == 6.25
+ assert engine_override.scheduler.config.shift == 3.5
+ assert engine_override._parse_request(_request()).flow_shift == 3.5
+ assert request_override.flow_shift == 2.25
+
+
+@pytest.mark.parametrize(
+ ("scheduler_config", "expected"),
+ [
+ ({"flow_shift": 6.0, "shift": 7.0}, 6.0),
+ ({"shift": 7.0}, 7.0),
+ ({}, 5.0),
+ ],
+)
+def test_checkpoint_scheduler_shift_fallback_order(scheduler_config, expected: float) -> None:
+ module = _load_pipeline_module()
+ checkpoint_config = _scheduler_config()
+ checkpoint_config.pop("flow_shift")
+ checkpoint_config.update(scheduler_config)
+ module._load_json = lambda *args, **kwargs: checkpoint_config.copy()
+
+ pipeline = _pipeline(module)
+
+ assert pipeline.scheduler.config.shift == expected
+
+
+@pytest.mark.parametrize("flow_shift", [0, -1, float("nan"), float("inf"), "invalid"])
+def test_request_flow_shift_must_be_positive_and_finite(flow_shift) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(extra_args={"flow_shift": flow_shift})
+
+ with pytest.raises(ValueError, match="flow_shift.*positive.*finite"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+@pytest.mark.parametrize(
+ "flow_shift",
+ [True, False, 10**10000],
+ ids=["true", "false", "overflowing-integer"],
+)
+def test_request_flow_shift_rejects_booleans_and_normalizes_integer_overflow(flow_shift) -> None:
+ module = _load_pipeline_module()
+ sampling = _SamplingParams(extra_args={"flow_shift": flow_shift})
+
+ with pytest.raises(ValueError, match="flow_shift.*positive.*finite"):
+ _pipeline(module)._parse_request(_RequestBatch(_prompt(), sampling))
+
+
+def test_flow_shift_is_request_local_without_mutating_scheduler() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+
+ def generate(flow_shift=None):
+ extra_args = {} if flow_shift is None else {"flow_shift": flow_shift}
+ return pipeline(
+ _RequestBatch(
+ _prompt(),
+ _SamplingParams(num_frames=21, extra_args=extra_args),
+ )
+ ).output
+
+ default_before = generate()
+ shifted_two = generate(2.0)
+ shifted_seven = generate(7.0)
+ default_after = generate()
+
+ torch.testing.assert_close(default_before, default_after)
+ assert not torch.equal(default_before, shifted_two)
+ assert not torch.equal(shifted_two, shifted_seven)
+ assert pipeline.scheduler.config.shift == 5.0
+
+
+def test_encode_prompt_zeroes_padded_umt5_states_to_exactly_512_tokens() -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ attention_mask = torch.zeros(1, 512, dtype=torch.long)
+ attention_mask[:, :3] = 1
+ pipeline.tokenizer = lambda *args, **kwargs: SimpleNamespace(
+ input_ids=torch.arange(512).view(1, 512),
+ attention_mask=attention_mask,
+ )
+ raw_states = torch.arange(512 * 8, dtype=torch.float32).view(1, 512, 8) + 1.0
+ pipeline.text_encoder = lambda input_ids, mask: SimpleNamespace(last_hidden_state=raw_states.clone())
+
+ encoded = module.LingBotWorldCausalDMDPipeline.encode_prompt(
+ pipeline,
+ "move",
+ max_sequence_length=512,
+ dtype=torch.float32,
+ )
+
+ assert encoded.shape == (1, 512, 8)
+ torch.testing.assert_close(encoded[:, :3], raw_states[:, :3])
+ torch.testing.assert_close(encoded[:, 3:], torch.zeros_like(encoded[:, 3:]))
+
+
+def test_multi_chunk_generation_uses_one_request_local_cache_and_decodes_accumulated_latents() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ allocations = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ allocations.append((kwargs, weakref.ref(cache)))
+ return cache
+
+ transformer.allocate_cache = allocate
+ sampling = _SamplingParams(num_frames=21, output_type="np")
+
+ result = pipeline(_request(sampling=sampling))
+
+ assert result.output.shape == (1, 3, 21, 16, 16)
+ assert len(transformer.calls) == 10
+ assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
+ expected_timesteps = torch.tensor([1000.0, 937.5, 2500.0 / 3.0, 625.0, 0.0] * 2)
+ torch.testing.assert_close(
+ torch.cat([call["timestep"] for call in transformer.calls]),
+ expected_timesteps,
+ )
+ assert len(allocations) == 1
+ kwargs, cache_ref = allocations[0]
+ assert kwargs == {
+ "batch_size": 1,
+ "latent_height": 2,
+ "latent_width": 2,
+ "device": torch.device("cpu"),
+ "dtype": torch.float32,
+ }
+ assert not hasattr(pipeline, "cache")
+ assert not hasattr(pipeline, "transformer_cache")
+ assert pipeline.vae.decode_inputs[0].shape == (1, 16, 6, 2, 2)
+ normalized_latents = torch.cat(
+ (transformer.calls[4]["hidden_states"][:, :16], transformer.calls[9]["hidden_states"][:, :16]),
+ dim=2,
+ )
+ mean = torch.tensor(pipeline.vae.config.latents_mean).view(1, 16, 1, 1, 1)
+ std = torch.tensor(pipeline.vae.config.latents_std).view(1, 16, 1, 1, 1)
+ torch.testing.assert_close(pipeline.vae.decode_inputs[0], normalized_latents * std + mean)
+ transformer.calls.clear()
+ gc.collect()
+ assert cache_ref() is None
+
+
+def _tick_extra_args(*, chunk_index: int, prompt: str = "move through the room"):
+ return ARDiffusionTickRequest(
+ session_id="world-1",
+ request_id="legacy-lingbot-request",
+ chunk_index=chunk_index,
+ applied_event_ids=(chunk_index,),
+ prompt=prompt,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_trajectory.v1",
+ data={"poses": [], "intrinsics": []},
+ ),
+ ),
+ ).to_extra_args()
+
+
+def test_typed_ticks_generate_one_global_block_and_return_standard_metadata(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ cross_calls = []
+ generated = []
+
+ def ar_text_caches(prompt_embeds, *, invalidate):
+ del prompt_embeds
+ cross_calls.append(invalidate)
+ return [SimpleNamespace()]
+
+ def generate_block(**kwargs):
+ block = torch.randn(
+ (1, 16, 3, 2, 2),
+ generator=kwargs["generator"],
+ )
+ generated.append(
+ {
+ "start_frame": kwargs["start_frame"],
+ "condition": kwargs["condition"].clone(),
+ "block": block.clone(),
+ }
+ )
+ return block
+
+ monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches)
+ monkeypatch.setattr(pipeline, "_generate_block", generate_block)
+
+ first_sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=0))
+ first = pipeline(_request(sampling=first_sampling))
+ switched_prompt = _prompt()
+ switched_prompt["prompt"] = "enter the snowy valley"
+ second_sampling = _SamplingParams(
+ extra_args=_tick_extra_args(
+ chunk_index=1,
+ prompt=switched_prompt["prompt"],
+ )
+ )
+ second = pipeline(_request(sampling=second_sampling, prompt=switched_prompt))
+
+ assert generated[0]["start_frame"] == 0
+ assert generated[1]["start_frame"] == 3
+ assert torch.count_nonzero(generated[0]["condition"]) > 0
+ torch.testing.assert_close(
+ generated[1]["condition"][:, :4],
+ torch.zeros_like(generated[1]["condition"][:, :4]),
+ )
+ assert torch.count_nonzero(generated[1]["condition"][:, 4:]) > 0
+ assert cross_calls == [False, True]
+ assert first.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
+ assert second.output["metadata"]["ar_diffusion"] == {
+ "session_id": "world-1",
+ "request_id": "legacy-lingbot-request",
+ "chunk_index": 1,
+ "applied_event_ids": [1],
+ }
+ assert pipeline._ar_sessions["world-1"].next_chunk_index == 2
+ assert pipeline._ar_sessions["world-1"].prompt == "enter the snowy valley"
+ assert pipeline._ar_sessions["world-1"].camera_tail is not None
+ assert pipeline._ar_sessions["world-1"].camera_tail.poses.shape == (1, 4, 4)
+ expected_generator = torch.Generator(device="cpu").manual_seed(17)
+ expected_first = torch.randn((1, 16, 3, 2, 2), generator=expected_generator)
+ expected_second = torch.randn((1, 16, 3, 2, 2), generator=expected_generator)
+ torch.testing.assert_close(generated[0]["block"], expected_first)
+ torch.testing.assert_close(generated[1]["block"], expected_second)
+
+
+def test_typed_action_ticks_integrate_camera_across_chunks(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ monkeypatch.setattr(
+ pipeline,
+ "_ar_text_caches",
+ lambda *args, **kwargs: [SimpleNamespace()],
+ )
+ monkeypatch.setattr(
+ pipeline,
+ "_generate_block",
+ lambda **kwargs: torch.zeros_like(kwargs["condition"][:, :16]),
+ )
+
+ def action_sampling(chunk_index: int, frames: list[list[str]]) -> _SamplingParams:
+ tick = ARDiffusionTickRequest(
+ session_id="world-actions",
+ request_id="legacy-lingbot-request",
+ chunk_index=chunk_index,
+ controls=(
+ ARDiffusionControlInput(
+ track="camera",
+ schema="lingbot.camera_actions.v1",
+ data={"mode": "frames", "frames": frames},
+ ),
+ ),
+ )
+ sampling = _SamplingParams(extra_args=tick.to_extra_args())
+ sampling.extra_args["_lingbot_camera_trajectory"] = None
+ sampling.extra_args["_lingbot_camera_actions"] = tuple(tuple(frame) for frame in frames)
+ return sampling
+
+ pipeline(_request(sampling=action_sampling(0, [["w"], ["w"], ["w"]])))
+ first_tail = pipeline._ar_sessions["world-actions"].camera_tail.poses.clone()
+ pipeline(_request(sampling=action_sampling(1, [["d"], ["d"], ["d"]])))
+ state = pipeline._ar_sessions["world-actions"]
+
+ assert state.next_chunk_index == 2
+ assert state.camera_tail is not None
+ torch.testing.assert_close(state.camera_tail.poses[0, 2, 3], first_tail[0, 2, 3])
+ assert state.camera_tail.poses[0, 0, 3] > first_tail[0, 0, 3]
+
+
+def test_first_typed_yaw_action_uses_pre_action_identity_anchor(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ monkeypatch.setattr(
+ module,
+ "build_plucker_embedding",
+ _real_build_plucker_embedding,
+ )
+ action_trajectory, _ = integrate_lingbot_camera_actions(
+ [["j"], [], []],
+ width=16,
+ height=16,
+ )
+ neutral_trajectory, _ = integrate_lingbot_camera_actions(
+ [[], [], []],
+ width=16,
+ height=16,
+ )
+
+ def inputs(trajectory, actions):
+ return SimpleNamespace(
+ camera_trajectory=trajectory,
+ camera_actions=actions,
+ num_latent_frames=3,
+ num_frames=9,
+ height=16,
+ width=16,
+ )
+
+ action_embedding, tail = pipeline._prepare_camera(
+ inputs(action_trajectory, (("j",), (), ())),
+ dtype=torch.float32,
+ )
+ neutral_embedding, _ = pipeline._prepare_camera(
+ inputs(neutral_trajectory, ((), (), ())),
+ dtype=torch.float32,
+ )
+
+ identity_anchor = torch.eye(4, dtype=action_trajectory.poses.dtype).unsqueeze(0)
+ explicit_trajectory = _CameraTrajectory(
+ poses=torch.cat((identity_anchor, action_trajectory.poses), dim=0),
+ intrinsics=torch.cat(
+ (action_trajectory.intrinsics[:1], action_trajectory.intrinsics),
+ dim=0,
+ ),
+ )
+ explicit = _real_build_plucker_embedding(
+ explicit_trajectory,
+ height=16,
+ width=16,
+ target_height=16,
+ target_width=16,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )[1:]
+
+ assert not torch.equal(action_embedding, neutral_embedding)
+ torch.testing.assert_close(action_embedding, module._fold_camera_embedding(explicit))
+ torch.testing.assert_close(tail.poses, action_trajectory.poses[-1:])
+
+
+def test_typed_tick_rejects_non_contiguous_chunk_index(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ monkeypatch.setattr(
+ pipeline,
+ "_ar_text_caches",
+ lambda *args, **kwargs: [SimpleNamespace()],
+ )
+
+ sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=2))
+ with pytest.raises(ValueError, match="must be contiguous"):
+ pipeline(_request(sampling=sampling))
+
+
+def test_typed_tick_rejects_chunk_beyond_realtime_condition_horizon(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_pipeline_module()
+ pipeline = _pipeline(module)
+ pipeline._ar_height = 16
+ pipeline._ar_width = 16
+ pipeline._ar_diffusion_kv_state = object()
+ pipeline._ar_sessions["world-1"] = module._LingBotARSessionState(
+ next_chunk_index=10,
+ )
+ monkeypatch.setattr(
+ pipeline,
+ "encode_prompt",
+ lambda *args, **kwargs: pytest.fail("the realtime horizon must be validated before prompt encoding"),
+ )
+
+ sampling = _SamplingParams(extra_args=_tick_extra_args(chunk_index=10))
+ with pytest.raises(
+ ValueError,
+ match=r"at most 10 ticks.*chunk_index 0 through 9.*117 pixel frames",
+ ):
+ pipeline(_request(sampling=sampling))
+
+
+def test_request_cache_is_released_before_vae_decode() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ cache_refs: list[weakref.ReferenceType] = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ cache_refs.append(weakref.ref(cache))
+ return cache
+
+ transformer.allocate_cache = allocate
+
+ def assert_cache_released() -> None:
+ gc.collect()
+ assert len(cache_refs) == 1
+ assert cache_refs[0]() is None
+
+ pipeline.vae.on_decode = assert_cache_released
+
+ result = pipeline(_request(sampling=_SamplingParams(num_frames=21, output_type="np")))
+
+ assert result.output.shape[2] == 21
+
+
+def test_117_frame_request_generates_ten_complete_latent_blocks() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer()
+ pipeline = _pipeline(module, transformer=transformer)
+ trajectory = _CameraTrajectory(
+ poses=torch.eye(4).repeat(117, 1, 1),
+ intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(117, 1),
+ )
+ sampling = _SamplingParams(
+ num_frames=117,
+ output_type="latent",
+ extra_args={"_lingbot_camera_trajectory": trajectory},
+ )
+
+ result = pipeline(_request(sampling=sampling))
+
+ assert result.output.shape == (1, 16, 30, 2, 2)
+ assert len(transformer.calls) == 50
+ assert [call["start_frame"] for call in transformer.calls[::5]] == list(range(0, 30, 3))
+
+
+def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
+ module = _load_pipeline_module()
+ transformer = _RecordingTransformer(raise_on_call=2)
+ pipeline = _pipeline(module, transformer=transformer)
+ cache_refs = []
+ original_allocate = transformer.allocate_cache
+
+ def allocate(**kwargs):
+ cache = original_allocate(**kwargs)
+ cache_refs.append(weakref.ref(cache))
+ return cache
+
+ transformer.allocate_cache = allocate
+
+ with pytest.raises(RuntimeError, match="forced transformer failure") as exc_info:
+ pipeline(_request())
+
+ exc_info.value.__traceback__ = None
+ del exc_info
+ transformer.calls.clear()
+ gc.collect()
+ assert len(cache_refs) == 1
+ assert cache_refs[0]() is None
+ assert not hasattr(pipeline, "cache")
+ assert not hasattr(pipeline, "transformer_cache")
+
+
+def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
+ module = _load_pipeline_module()
+ resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
+ module
+ )
+
+ assert entry == ("lingbot_world", "pipeline", "LingBotWorldCausalDMDPipeline")
+ assert resolved is module.LingBotWorldCausalDMDPipeline
+ assert cache_acceleration_disabled
+ assert preprocess_name == "get_lingbot_world_pre_process_func"
+ request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
+ assert preprocess(request) is request
+ assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
+ assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"
+ lingbot_init = _LINGBOT_INIT_PATH.read_text()
+ assert "from .pipeline import" in lingbot_init
+ assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init
+ assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init
diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py
new file mode 100644
index 000000000..4b1fef7d5
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py
@@ -0,0 +1,1041 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Focused LTX-2.5 pipeline correctness tests."""
+
+import json
+from dataclasses import replace
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
+from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+ LTX25_DISTILLED_COMPONENT_PROFILE,
+ LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+ LTX25_FULL_COMPONENT_PROFILE,
+ LTX25_TWO_STAGE_COMPONENT_PROFILE,
+ detect_ltx_model_version,
+ resolve_ltx_checkpoint_kind,
+ resolve_ltx_component_profile,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
+ LTXPhaseResult,
+ _first_frame_keyframes_mask,
+ _official_ltx_sigmas,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+ LTX_GUIDANCE_EXECUTOR,
+ LTXGuidanceSpec,
+ LTXModalityGuidance,
+ velocity_from_x0,
+ x0_from_velocity,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+ LTX2_DISTILLED_TWO_STAGE_RECIPE,
+ LTX2_ONE_STAGE_RECIPE,
+ LTX23_ONE_STAGE_RECIPE,
+ LTX25_DISTILLED_ONE_STAGE_RECIPE,
+ LTX25_DISTILLED_TWO_STAGE_RECIPE,
+ LTX25_FULL_RECIPE,
+ LTX25_TWO_STAGE_RECIPE,
+ LTX_DISTILLED_SIGMAS,
+ LTX_POSITIVE_ONLY_RECIPE,
+ resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+ LTX2DistilledOneStagePipeline,
+ LTX2Pipeline,
+)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+ LTX2DistilledTwoStagePipeline,
+ LTX2TwoStagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
+def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
+ pipe = object.__new__(LTX2Pipeline)
+ pipe.model_version = model_version
+ pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+ ts = torch.tensor([750.0, 500.0])
+
+ kwargs = LTXRuntime._denoise_timestep_kwargs(
+ pipe,
+ ts,
+ SimpleNamespace(),
+ SimpleNamespace(),
+ video_token_count=4,
+ audio_token_count=3,
+ )
+
+ torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
+ torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
+ assert kwargs["sigma"] is ts
+ assert kwargs["audio_sigma"] is ts
+
+
+def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
+ state = LTXAVState(
+ video=torch.zeros(1, 2, 2),
+ audio=torch.zeros(1, 4, 2),
+ )
+ forward_ctx = SimpleNamespace(
+ original_audio_num_frames=3,
+ sampler="euler_ancestral",
+ )
+ denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
+ pipe = SimpleNamespace(
+ _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
+ )
+
+ def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
+ assert actual_forward_ctx.sampler == "euler_ancestral"
+ updated_audio = torch.ones_like(state.audio)
+ updated_audio[:, 3:] = 99.0
+ return torch.ones_like(state.video), updated_audio
+
+ monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
+
+ actual = LTXRuntime._denoise_step(
+ pipe,
+ 0,
+ torch.tensor(1.0),
+ state,
+ forward_ctx,
+ denoise_ctx,
+ )
+
+ torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
+ torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
+
+
+def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
+ reference = torch.zeros(2, 12, 8)
+
+ mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
+
+ assert mask.shape == (2, 12, 1)
+ torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
+ torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
+
+
+def _make_ltx_request_pipe(cls):
+ pipe = object.__new__(cls)
+ torch.nn.Module.__init__(pipe)
+ pipe.device = torch.device("cpu")
+ pipe.tokenizer_max_length = 99
+ pipe.vae_spatial_compression_ratio = 32
+ pipe.vae_temporal_compression_ratio = 8
+ return pipe
+
+
+def _resolve_request_inputs_for_test(
+ pipe,
+ req,
+ *,
+ guidance_scale=4.0,
+ negative_prompt=None,
+):
+ return pipe._resolve_request_inputs(
+ req,
+ prompt=None,
+ negative_prompt=negative_prompt,
+ height=None,
+ width=None,
+ num_frames=None,
+ frame_rate=None,
+ num_inference_steps=None,
+ guidance_scale=guidance_scale,
+ num_videos_per_prompt=1,
+ generator=None,
+ latents=None,
+ audio_latents=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_attention_mask=None,
+ negative_prompt_attention_mask=None,
+ decode_timestep=0.0,
+ decode_noise_scale=None,
+ output_type="np",
+ max_sequence_length=None,
+ )
+
+
+def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
+ pipe = SimpleNamespace(
+ component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
+ use_diffusion_decoder=True,
+ )
+ od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
+
+ monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+ monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
+ monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
+
+ with pytest.raises(ImportError) as exc_info:
+ ltx2_components.initialize_pipeline_components(pipe, od_config)
+
+ assert str(exc_info.value) == (
+ "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
+ )
+
+
+def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
+ revision = "pinned-revision"
+ calls: dict[str, Any] = {"components": []}
+ natten_processor = object()
+
+ class FakeDiffusionDecoder:
+ def set_attn_processor(self, processor):
+ calls["diffusion_decoder_processor"] = processor
+
+ def enable_tiling(self):
+ calls["diffusion_decoder_tiling"] = True
+
+ def set_parallel_size(self, parallel_size, mode):
+ calls["diffusion_decoder_parallel"] = (parallel_size, mode)
+
+ profile = replace(
+ LTX25_DISTILLED_COMPONENT_PROFILE,
+ text_encoder_cls=object,
+ video_vae_cls=object,
+ vocoder_cls=object,
+ vocoder_fallback_cls=None,
+ scheduler_use_dynamic_shifting=False,
+ )
+ pipeline = SimpleNamespace(
+ component_profile=profile,
+ pipeline_kind="distilled_two_stage",
+ use_diffusion_decoder=True,
+ )
+ od_config = SimpleNamespace(
+ model="org/converted-ltx25",
+ revision=revision,
+ dtype=torch.bfloat16,
+ quantization_config=None,
+ extras={},
+ vae_use_tiling=False,
+ parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
+ )
+
+ monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+ monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
+ monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
+
+ def fake_prefetch(model, subfolders, **kwargs):
+ calls["prefetch"] = (model, tuple(subfolders), kwargs)
+
+ def fake_tokenizer(*args, **kwargs):
+ calls["tokenizer"] = (args, kwargs)
+ return SimpleNamespace(model_max_length=1024)
+
+ def fake_component(component_cls, model, subfolder, **kwargs):
+ calls["components"].append((component_cls, model, subfolder, kwargs))
+ if subfolder == "vae":
+ return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
+ if subfolder == "audio_vae":
+ return SimpleNamespace(
+ mel_compression_ratio=4,
+ temporal_compression_ratio=4,
+ config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
+ )
+ return object()
+
+ def fake_native_diffusion_decoder(model, **kwargs):
+ calls["native_diffusion_decoder"] = (model, kwargs)
+ return FakeDiffusionDecoder()
+
+ def fake_transformer_config(model, subfolder, local_files_only, *, revision):
+ calls["transformer_config"] = (model, subfolder, local_files_only, revision)
+ return {"component": "transformer"}
+
+ monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
+ monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
+ monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+ monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
+ monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+ monkeypatch.setattr(
+ ltx2_components,
+ "LTX2VideoVaeNeighborhoodNattenProcessor",
+ lambda: natten_processor,
+ )
+ monkeypatch.setattr(
+ ltx2_components,
+ "create_transformer_from_config",
+ lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
+ )
+
+ def fake_scheduler(*args, **kwargs):
+ calls["scheduler"] = (args, kwargs)
+ return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
+
+ monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
+
+ ltx2_components.initialize_pipeline_components(pipeline, od_config)
+
+ assert pipeline.weights_sources[0].revision == revision
+ assert calls["prefetch"][2]["revision"] == revision
+ assert calls["tokenizer"][1]["revision"] == revision
+ assert {call[2] for call in calls["components"]} == {
+ "text_encoder",
+ "connectors",
+ "vae",
+ "audio_vae",
+ "vocoder",
+ "latent_upsampler",
+ }
+ assert all(call[3]["revision"] == revision for call in calls["components"])
+ assert "diffusion_decoder" not in calls["prefetch"][1]
+ assert calls["native_diffusion_decoder"] == (
+ od_config.model,
+ {"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
+ )
+ assert calls["diffusion_decoder_processor"] is natten_processor
+ assert calls["diffusion_decoder_tiling"] is True
+ assert calls["diffusion_decoder_parallel"] == (2, "tile")
+ assert calls["transformer_config"] == (
+ od_config.model,
+ profile.transformer_subfolder,
+ False,
+ revision,
+ )
+ assert calls["scheduler"][1]["revision"] == revision
+
+
+def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
+ revision = "diffusers-revision"
+ config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
+ calls: dict[str, Any] = {}
+
+ class FakeDecoder:
+ def init_distributed(self):
+ calls["init_distributed"] = True
+
+ expected = FakeDecoder()
+
+ def fake_load_config(model, **kwargs):
+ calls["config"] = (model, kwargs)
+ return config
+
+ def fake_resolve(model, repo_id, filename, **kwargs):
+ calls["artifact"] = (model, repo_id, filename, kwargs)
+ return "/models/native-diffvae.safetensors"
+
+ def fake_load(cls, path, passed_config, dtype):
+ calls["load"] = (path, passed_config, dtype)
+ return expected
+
+ monkeypatch.setattr(
+ ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+ "load_config",
+ staticmethod(fake_load_config),
+ )
+ monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
+ monkeypatch.setattr(
+ ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+ "from_ltx25_native_checkpoint",
+ classmethod(fake_load),
+ )
+
+ actual = ltx2_components._load_ltx25_native_diffusion_decoder(
+ "Lightricks/LTX-2.5-Diffusers",
+ local_files_only=False,
+ dtype=torch.bfloat16,
+ revision=revision,
+ )
+
+ assert actual is expected
+ assert calls["config"] == (
+ "Lightricks/LTX-2.5-Diffusers",
+ {
+ "subfolder": "diffusion_decoder",
+ "local_files_only": False,
+ "revision": revision,
+ },
+ )
+ assert calls["artifact"] == (
+ "Lightricks/LTX-2.5-Diffusers",
+ "Lightricks/LTX-2.5",
+ "vae/ltx-2.5-video-vae-bf16.safetensors",
+ {
+ "model_revision": revision,
+ "artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
+ },
+ )
+ assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
+ assert calls["init_distributed"] is True
+
+
+def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
+ class MissingNattenProcessor:
+ def __init__(self):
+ raise ImportError("kernels is unavailable")
+
+ monkeypatch.setattr(
+ ltx2_components,
+ "LTX2VideoVaeNeighborhoodNattenProcessor",
+ MissingNattenProcessor,
+ )
+
+ with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
+ ltx2_components._create_ltx25_natten_processor()
+
+
+def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
+ (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
+ (tmp_path / "transformer").mkdir()
+ (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
+
+ assert detect_ltx_model_version(str(tmp_path)) == "2.3"
+
+
+def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(
+ json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+ )
+
+ def stub_components(pipe, od_config):
+ pipe.od_config = od_config
+ pipe.vae_spatial_compression_ratio = 32
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+ monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+ pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+
+ assert pipe.model_version == "2.5"
+ assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
+ assert pipe.preserve_sp_padded_audio_duration
+ assert pipe.reports_stage_durations
+
+
+def test_ltx25_full_recipe_matches_official_sft_defaults():
+ (phase,) = LTX25_FULL_RECIPE.phases
+
+ assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
+ assert LTX25_FULL_RECIPE.num_frames == 121
+ assert LTX25_FULL_RECIPE.frame_rate == 24.0
+ assert LTX25_FULL_RECIPE.num_inference_steps == 30
+ assert phase.guidance.video == LTXModalityGuidance(
+ cfg_scale=3.0,
+ stg_scale=1.0,
+ modality_scale=3.0,
+ rescale_scale=0.7,
+ stg_blocks=(28,),
+ )
+ assert phase.guidance.audio == LTXModalityGuidance(
+ cfg_scale=7.0,
+ stg_scale=1.0,
+ modality_scale=3.0,
+ rescale_scale=0.7,
+ stg_blocks=(28,),
+ )
+ assert phase.sigmas is None
+ assert phase.noise_scale == 1.0
+ assert phase.use_official_sigma_schedule
+ assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+ assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
+ assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
+
+
+def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
+ scheduler = SimpleNamespace(
+ config={
+ "base_image_seq_len": 1024,
+ "max_image_seq_len": 4096,
+ "base_shift": 0.95,
+ "max_shift": 2.05,
+ "shift_terminal": 0.1,
+ }
+ )
+
+ sigmas = _official_ltx_sigmas(
+ scheduler,
+ steps=30,
+ device=torch.device("cpu"),
+ )
+
+ torch.testing.assert_close(
+ sigmas[[0, 1, -2, -1]],
+ torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
+ rtol=1e-5,
+ atol=1e-6,
+ )
+
+
+def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(
+ json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+ )
+
+ def stub_components(pipe, od_config):
+ pipe.od_config = od_config
+ pipe.vae_spatial_compression_ratio = 32
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+ monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+ pipe = LTX2DistilledTwoStagePipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ enable_diffusion_pipeline_profiler=False,
+ )
+ )
+
+ assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+
+def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
+ for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
+ path = tmp_path / subfolder
+ path.mkdir()
+ (path / "config.json").write_text(json.dumps({"marker": marker}))
+
+ config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
+ config.enrich_config()
+
+ assert config.tf_model_config.get("marker") == "full"
+
+
+def test_ltx25_four_public_pipeline_semantics_are_disjoint():
+ assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
+ assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
+ assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
+ assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+ assert LTX2Pipeline.pipeline_kind == "one_stage"
+ assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
+ assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
+ assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
+
+
+@pytest.mark.parametrize(
+ "pipeline_cls",
+ [
+ LTX2Pipeline,
+ LTX2TwoStagePipeline,
+ LTX2DistilledOneStagePipeline,
+ LTX2DistilledTwoStagePipeline,
+ ],
+)
+@pytest.mark.parametrize(
+ ("extras", "use_diffusion_decoder"),
+ [({}, True), ({"ltx2_use_conv_vae": True}, False)],
+)
+def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
+ tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
+):
+ (tmp_path / "model_index.json").write_text(
+ json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+ )
+
+ def stub_components(pipe, od_config):
+ pipe.od_config = od_config
+ pipe.vae_spatial_compression_ratio = 32
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+ monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
+ monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+ pipe = pipeline_cls(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ extras=extras,
+ enable_diffusion_pipeline_profiler=False,
+ )
+ )
+
+ assert pipe.use_diffusion_decoder is use_diffusion_decoder
+ assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
+
+
+@pytest.mark.parametrize(
+ ("pipeline_kind", "expected"),
+ [
+ ("one_stage", "regular"),
+ ("two_stage", "regular"),
+ ("distilled_one_stage", "distilled"),
+ ("distilled_two_stage", "distilled"),
+ ("dmd2", None),
+ ],
+)
+def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
+ assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
+
+
+def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
+ with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
+ resolve_ltx_checkpoint_kind("unknown")
+
+
+def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
+ stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
+
+ assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
+ assert stage1.spatial_downscale == 2
+ assert stage1.sigmas == LTX_DISTILLED_SIGMAS
+ assert stage1.sampler == "euler_ancestral"
+ assert stage2.sampler == "euler"
+ assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+ assert stage2.input_transform == "spatial_upsample"
+
+ assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
+ assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
+
+
+def test_ltx25_distilled_one_stage_uses_native_low_resolution():
+ (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
+
+ assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
+ assert phase.spatial_downscale == 1
+ assert phase.sigmas == LTX_DISTILLED_SIGMAS
+ assert phase.sampler == "euler_ancestral"
+ assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
+
+
+def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
+ stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
+
+ assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
+ assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
+ assert stage1.spatial_downscale == 2
+ assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
+ assert stage2.adapter_slot == "ltx_distilled"
+ assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
+ assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+ assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
+
+
+@pytest.mark.parametrize(
+ "recipe",
+ [
+ LTX2_ONE_STAGE_RECIPE,
+ LTX23_ONE_STAGE_RECIPE,
+ LTX_POSITIVE_ONLY_RECIPE,
+ ],
+)
+def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
+ assert recipe.supports_cache_dit
+
+
+@pytest.mark.parametrize(
+ "recipe",
+ [
+ LTX25_FULL_RECIPE,
+ LTX25_DISTILLED_ONE_STAGE_RECIPE,
+ LTX2_DISTILLED_TWO_STAGE_RECIPE,
+ LTX25_DISTILLED_TWO_STAGE_RECIPE,
+ ],
+)
+def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
+ assert not recipe.supports_cache_dit
+
+
+def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
+ sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+ velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+ sigma = torch.tensor(0.5)
+
+ actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+ sample,
+ {"cond": velocity},
+ sigma,
+ LTXModalityGuidance(),
+ )
+ expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
+
+ assert actual.dtype == torch.float32
+ torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+
+
+def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
+ item_calls = []
+
+ class TrackingTensor(torch.Tensor):
+ @staticmethod
+ def __new__(cls):
+ return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
+
+ def item(self):
+ item_calls.append(True)
+ return super().item()
+
+ sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
+ x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
+
+ velocity_from_x0(sample, x0, TrackingTensor())
+
+ assert item_calls == [True]
+
+
+def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
+ sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+ velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+
+ actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+ sample,
+ {"cond": velocity},
+ torch.tensor(0.5),
+ LTXModalityGuidance(),
+ preserve_positive_velocity=True,
+ )
+
+ assert actual is velocity
+
+
+@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
+def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
+ expected_dtype = torch.bfloat16
+ captured = {}
+
+ def prepare_latents(**kwargs):
+ captured["video_dtype"] = kwargs["dtype"]
+ return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
+
+ def prepare_audio_latents(*_args, **kwargs):
+ captured["audio_dtype"] = kwargs["dtype"]
+ return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
+
+ pipeline = SimpleNamespace(
+ model_version=model_version,
+ transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
+ prepare_latents=prepare_latents,
+ prepare_audio_latents=prepare_audio_latents,
+ audio_sampling_rate=48_000,
+ audio_hop_length=160,
+ audio_vae_temporal_compression_ratio=4,
+ audio_vae_mel_compression_ratio=4,
+ audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
+ _resolve_audio_latent_length=lambda length, _latents: length,
+ )
+ request = SimpleNamespace(
+ num_videos_per_prompt=1,
+ height=64,
+ width=64,
+ num_frames=1,
+ frame_rate=24.0,
+ generator=None,
+ latents=None,
+ audio_latents=None,
+ audio_latents_normalized=False,
+ )
+ prompt = SimpleNamespace(
+ batch_size=1,
+ positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+ positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+ )
+
+ LTXRuntime._prepare_video_latents_stage(
+ pipeline,
+ request,
+ prompt,
+ device=torch.device("cpu"),
+ noise_scale=1.0,
+ )
+ LTXRuntime._prepare_audio_latents_stage(
+ pipeline,
+ request,
+ prompt,
+ device=torch.device("cpu"),
+ noise_scale=1.0,
+ )
+
+ assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
+
+
+def test_ltx_seeded_latents_match_official_packed_rng_layout():
+ pipeline = SimpleNamespace(
+ model_version="2.5",
+ vae_spatial_compression_ratio=2,
+ vae_temporal_compression_ratio=2,
+ transformer_spatial_patch_size=1,
+ transformer_temporal_patch_size=1,
+ audio_vae_mel_compression_ratio=2,
+ od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
+ )
+ device = torch.device("cpu")
+ actual_generator = torch.Generator(device=device).manual_seed(123)
+ expected_generator = torch.Generator(device=device).manual_seed(123)
+
+ actual_video = ltx2_latents.prepare_video_latents(
+ pipeline,
+ batch_size=1,
+ num_channels_latents=2,
+ height=4,
+ width=6,
+ num_frames=3,
+ dtype=torch.float32,
+ device=device,
+ generator=actual_generator,
+ )
+ actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
+ pipeline,
+ batch_size=1,
+ num_channels_latents=2,
+ audio_latent_length=3,
+ num_mel_bins=4,
+ dtype=torch.float32,
+ device=device,
+ generator=actual_generator,
+ )
+
+ expected_video = torch.randn((1, 12, 2), generator=expected_generator)
+ expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
+
+ torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
+ torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
+ assert (original_length, padded_length) == (3, 3)
+
+
+def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
+ request_inputs = LTXRequestInputs(
+ prompt="prompt",
+ negative_prompt="",
+ height=64,
+ width=64,
+ num_frames=1,
+ frame_rate=24.0,
+ num_inference_steps=8,
+ guidance=LTXGuidanceSpec.positive_only(),
+ num_videos_per_prompt=1,
+ generator=None,
+ latents=None,
+ audio_latents=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_attention_mask=None,
+ negative_prompt_attention_mask=None,
+ decode_timestep=0.25,
+ decode_noise_scale=0.5,
+ output_type="np",
+ max_sequence_length=16,
+ )
+ prompt_context = object()
+ phase_calls: list[Any] = []
+
+ def resolve_request_inputs(req, **kwargs):
+ return request_inputs
+
+ def run_phase(req, inputs, *, prompt_context=None, **kwargs):
+ phase_recipe = kwargs["phase_recipe"]
+ assert kwargs["image"] is source_image
+ phase_calls.append((phase_recipe, inputs, prompt_context))
+ if len(phase_calls) == 1:
+ assert prompt_context is None
+ assert inputs.decode_timestep == 0.25
+ assert inputs.decode_noise_scale == 0.5
+ context = prompt_context_sentinel
+ assert phase_recipe.name == "generate_lowres"
+ assert (inputs.height, inputs.width) == (32, 32)
+ assert inputs.num_inference_steps == 2
+ assert kwargs["sigmas"] == stage_1_sigmas
+ assert kwargs["noise_scale"] == 1.0
+ video = torch.ones(1, 128, 1, 1, 1)
+ audio = torch.full((1, 8, 1, 2), 2.0)
+ audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
+ else:
+ assert prompt_context is prompt_context_sentinel
+ assert phase_recipe.name == "refine"
+ assert (inputs.height, inputs.width) == (64, 64)
+ assert inputs.latents.shape == (1, 128, 1, 2, 2)
+ torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
+ assert inputs.audio_latents_normalized
+ assert inputs.guidance_scale == 1.0
+ assert inputs.num_inference_steps == 2
+ assert kwargs["sigmas"] == stage_2_sigmas
+ assert kwargs["noise_scale"] == stage_2_sigmas[0]
+ assert inputs.decode_timestep == 0.0
+ assert inputs.decode_noise_scale is None
+ context = prompt_context
+ video = torch.full((1, 128, 1, 2, 2), 3.0)
+ audio = torch.full((1, 8, 1, 2), 4.0)
+ audio_for_next_phase = None
+ return LTXPhaseResult(
+ forward_context=SimpleNamespace(prompt_context=context),
+ video=video,
+ audio=audio,
+ audio_for_next_phase=audio_for_next_phase,
+ )
+
+ def decode_phase(phase):
+ return DiffusionOutput(output=(phase.video, phase.audio))
+
+ class FakeUpsampler(torch.nn.Module):
+ dtype = torch.float32
+
+ def forward(self, latents):
+ return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+
+ prompt_context_sentinel = prompt_context
+ stage_1_sigmas = [1.0, 0.5, 0.0]
+ stage_2_sigmas = [0.8, 0.2, 0.0]
+ pipeline = object.__new__(LTX2TwoStagePipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
+ pipeline.device = torch.device("cpu")
+ pipeline.vae_spatial_compression_ratio = 32
+ pipeline.vae_temporal_compression_ratio = 8
+ pipeline.latent_upsampler = FakeUpsampler()
+ object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
+ object.__setattr__(pipeline, "run_phase", run_phase)
+ object.__setattr__(pipeline, "decode_phase", decode_phase)
+
+ source_image = object()
+ output = pipeline.forward(
+ SimpleNamespace(sampling_params_list=[], prompts=[]),
+ image=source_image,
+ stage_1_sigmas=stage_1_sigmas,
+ stage_2_sigmas=stage_2_sigmas,
+ )
+
+ assert len(phase_calls) == 2
+ assert phase_calls[1][2] is prompt_context_sentinel
+ torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
+ torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
+
+
+class TestLTXRequestParsing:
+ def test_request_resolves_custom_sigmas_from_extra_args(self):
+ from vllm_omni.diffusion.request import OmniDiffusionRequest
+ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ sigmas = [1.0, 0.5, 0.0]
+ req = DiffusionRequestBatch(
+ [
+ OmniDiffusionRequest(
+ prompt="prompt",
+ sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
+ request_id="ltx-custom-sigmas",
+ )
+ ]
+ )
+
+ assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
+
+ def test_request_resolves_per_stage_sigmas_from_extra_args(self):
+ from vllm_omni.diffusion.request import OmniDiffusionRequest
+ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ stage_1_sigmas = [1.0, 0.4, 0.0]
+ req = DiffusionRequestBatch(
+ [
+ OmniDiffusionRequest(
+ prompt="prompt",
+ sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
+ request_id="ltx-custom-phase-sigmas",
+ )
+ ]
+ )
+
+ resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
+ req, [1.0, 0.0], [0.8, 0.0]
+ )
+
+ assert resolved == (stage_1_sigmas, [0.8, 0.0])
+
+ @pytest.mark.parametrize("value", [-1, 52, 1.5, True])
+ def test_request_rejects_invalid_image_crf(self, value):
+ from vllm_omni.diffusion.request import OmniDiffusionRequest
+ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+ req = DiffusionRequestBatch(
+ [
+ OmniDiffusionRequest(
+ prompt="prompt",
+ sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
+ request_id="ltx-invalid-image-crf",
+ )
+ ]
+ )
+
+ with pytest.raises(ValueError, match="image_crf"):
+ _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
+
+
+class TestLTXForwardStages:
+ def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
+ pipe = _make_ltx_request_pipe(LTX2Pipeline)
+
+ class FakeTokenizer:
+ bos_token_id = 2
+ eos_token_id = 1
+ eos_token = ""
+ pad_token_id = 0
+ pad_token = ""
+ padding_side = "right"
+
+ def __init__(self):
+ self.calls = []
+
+ def __call__(self, text, **kwargs):
+ self.calls.append((text, kwargs))
+ ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
+ return SimpleNamespace(input_ids=ids)
+
+ def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
+ assert padding == "max_length"
+ assert return_tensors == "pt"
+ assert return_attention_mask is True
+ rows = []
+ masks = []
+ for ids in encoded["input_ids"]:
+ pad = max_length - len(ids)
+ rows.append([self.pad_token_id] * pad + ids)
+ masks.append([0] * pad + [1] * len(ids))
+ return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
+
+ class FakeTextEncoder(torch.nn.Module):
+ dtype = torch.float32
+
+ def __init__(self):
+ super().__init__()
+ self.seen_input_ids = None
+ self.seen_attention_mask = None
+
+ def forward(self, *, input_ids, attention_mask, output_hidden_states):
+ assert output_hidden_states is True
+ self.seen_input_ids = input_ids.clone()
+ self.seen_attention_mask = attention_mask.clone()
+ return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
+
+ tokenizer = FakeTokenizer()
+ text_encoder = FakeTextEncoder()
+ object.__setattr__(pipe, "tokenizer", tokenizer)
+ object.__setattr__(pipe, "text_encoder", text_encoder)
+
+ pipe._get_gemma_prompt_embeds(
+ [" long ", "short"],
+ max_sequence_length=4,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ )
+
+ torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
+ torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
+ assert tokenizer.padding_side == "left"
+ assert [call[0] for call in tokenizer.calls] == ["long", "short"]
diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py
index e309dbf47..aa5654c74 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py
@@ -1,8 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LTX image-to-video input and conditioning behavior."""
+import sys
from types import SimpleNamespace
import pytest
@@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls):
class TestLTXImageToVideoForwardStages:
- def test_i2v_pil_preprocessing_matches_official_center_crop(self):
+ def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages:
torch.testing.assert_close(actual, expected)
+ def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+
+ pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
+ pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
+ image = Image.fromarray(pixels.numpy())
+ monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
+
+ actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
+ expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
+
+ torch.testing.assert_close(actual, expected)
+
+ def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ captured_crfs = []
+ captured_prepare_kwargs = []
+
+ def fake_preprocess(image, *, height, width, crf, device, dtype):
+ del image, device
+ captured_crfs.append(crf)
+ return torch.zeros(1, 3, height, width, dtype=dtype)
+
+ def fake_prepare_latents(**kwargs):
+ captured_prepare_kwargs.append(kwargs)
+ return kwargs["image"], None
+
+ monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.model_version = "2.5"
+ object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+ object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+ request_inputs = SimpleNamespace(
+ height=4,
+ width=4,
+ num_frames=3,
+ num_videos_per_prompt=1,
+ generator=None,
+ latents=None,
+ )
+ prompt_context = SimpleNamespace(
+ batch_size=1,
+ positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+ )
+
+ pipe._prepare_video_latents_stage(
+ request_inputs,
+ prompt_context,
+ device=torch.device("cpu"),
+ noise_scale=0.0,
+ image=Image.new("RGB", (8, 4)),
+ )
+
+ assert captured_crfs == [18]
+ assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
+ assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
+
+ def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.model_version = "2.5"
+
+ with pytest.raises(ValueError, match="requires PIL images"):
+ pipe._prepare_video_latents_stage(
+ SimpleNamespace(),
+ SimpleNamespace(),
+ device=torch.device("cpu"),
+ noise_scale=0.0,
+ image=torch.zeros(3, 4, 4),
+ )
+
+ def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ captured_prepare_kwargs = []
+
+ def fake_prepare_latents(**kwargs):
+ captured_prepare_kwargs.append(kwargs)
+ return kwargs["image"], None
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.model_version = "2.5"
+ object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+ object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+ request_inputs = SimpleNamespace(
+ height=4,
+ width=4,
+ num_frames=3,
+ num_videos_per_prompt=1,
+ generator=None,
+ latents=None,
+ image_crf=0,
+ )
+ prompt_context = SimpleNamespace(
+ batch_size=1,
+ positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+ )
+
+ pipe._prepare_video_latents_stage(
+ request_inputs,
+ prompt_context,
+ device=torch.device("cpu"),
+ noise_scale=0.0,
+ image=torch.zeros(3, 4, 4),
+ )
+
+ assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
+
+ def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+ image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
+ monkeypatch.setitem(sys.modules, "av", None)
+
+ assert _apply_image_conditioning_crf(image, 18) is image
+
+ def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+ monkeypatch.setitem(sys.modules, "av", None)
+
+ with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
+ _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
+
def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages:
with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
pipe._resolve_request_image(mixed_req, None, request_inputs)
- def test_denoise_timestep_kwargs_masks_video_only(self):
+ def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
@@ -127,6 +258,39 @@ class TestLTXImageToVideoConditioning:
assert LTX2Pipeline.support_image_input is True
+ def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.use_diffusion_decoder = True
+ pipe.vae = SimpleNamespace(
+ encode=lambda image: image,
+ latents_mean=torch.full((2,), 100.0),
+ latents_std=torch.full((2,), 10.0),
+ config=SimpleNamespace(scaling_factor=1.0),
+ )
+ pipe.diffusion_decoder = SimpleNamespace(
+ latents_mean=torch.tensor([1.0, 3.0]),
+ latents_std=torch.tensor([2.0, 4.0]),
+ config=SimpleNamespace(scaling_factor=2.0),
+ )
+ monkeypatch.setattr(
+ ltx2_conditioning,
+ "retrieve_latents",
+ lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+ )
+
+ actual = pipe._encode_i2v_image_latents(
+ torch.zeros(1, 3, 1, 1),
+ batch_size=1,
+ generator=None,
+ dtype=torch.float32,
+ )
+
+ torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
+
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -189,13 +353,18 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
- def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+ @pytest.mark.parametrize(
+ ("model_version", "sampled_shape"),
+ [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+ )
+ def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
+ pipe.model_version = model_version
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
@@ -237,14 +406,20 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
- assert sampled_shapes == [(1, 3, 2)]
+ assert sampled_shapes == [sampled_shape]
- def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
+ @pytest.mark.parametrize(
+ ("model_version", "sampled_shape"),
+ [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+ )
+ def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+ import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
+ pipe.model_version = model_version
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
pipe.transformer_spatial_patch_size = 1
@@ -266,6 +441,7 @@ class TestLTXImageToVideoConditioning:
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+ monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
out, conditioning_mask = pipe.prepare_latents(
image=torch.zeros(1, 3, 1, 1),
@@ -279,10 +455,54 @@ class TestLTXImageToVideoConditioning:
device=torch.device("cpu"),
)
- assert sampled_shapes == [(1, 3, 2)]
+ assert sampled_shapes == [sampled_shape]
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
+ def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.model_version = "2.5"
+ pipe.vae_spatial_compression_ratio = 1
+ pipe.vae_temporal_compression_ratio = 1
+ pipe.transformer_spatial_patch_size = 1
+ pipe.transformer_temporal_patch_size = 1
+ pipe.vae = SimpleNamespace(
+ latents_mean=torch.zeros(2),
+ latents_std=torch.ones(2),
+ config=SimpleNamespace(scaling_factor=1.0),
+ )
+ object.__setattr__(
+ pipe,
+ "_encode_i2v_image_latents",
+ lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
+ )
+ captured = {}
+
+ def fake_create_noised_state(latents, noise_scale, generator=None):
+ del generator
+ captured["noise_scale"] = noise_scale
+ return latents
+
+ monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
+ pipe.prepare_latents(
+ image=torch.zeros(1, 3, 1, 1),
+ batch_size=1,
+ num_channels_latents=2,
+ height=1,
+ width=1,
+ num_frames=3,
+ noise_scale=0.909375,
+ dtype=torch.bfloat16,
+ device=torch.device("cpu"),
+ latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
+ )
+
+ assert captured["noise_scale"] == 0.909375
+
def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -385,6 +605,99 @@ class TestLTXImageToVideoConditioning:
torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
assert pipeline.step_indices == [0, 1]
+ def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+ device = torch.device("cpu")
+ sigmas = torch.tensor([1.0, 0.5, 0.0])
+ pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
+ mask = torch.tensor([[1.0, 0.0]])
+ adapter = LTXVideoAudioStepAdapter(
+ pipeline,
+ SimpleNamespace(sigmas=sigmas.clone()),
+ latent_num_frames=1,
+ latent_height=1,
+ latent_width=1,
+ image_conditioned=True,
+ sampler="euler_ancestral",
+ generator=torch.Generator(device=device).manual_seed(42),
+ conditioning_mask=mask,
+ )
+ video = torch.tensor([[[3.0], [2.0]]])
+ audio = torch.tensor([[[4.0]]])
+ video_velocity = torch.tensor([[[7.0], [0.25]]])
+ audio_velocity = torch.tensor([[[0.5]]])
+
+ reference_generator = torch.Generator(device=device).manual_seed(10042)
+ video_noise = torch.randn(video.shape, generator=reference_generator)
+ audio_noise = torch.randn(audio.shape, generator=reference_generator)
+
+ def official_step(sample, velocity, noise):
+ sigma = sigmas[0]
+ sigma_next = sigmas[1]
+ denoised = sample.float() - velocity.float() * sigma
+ sigma_down = sigma_next * (sigma_next / sigma)
+ ratio = sigma_down / sigma
+ deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
+ alpha_next = 1.0 - sigma_next
+ alpha_down = 1.0 - sigma_down
+ coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+ return (alpha_next / alpha_down) * deterministic + noise * coeff
+
+ expected_video = official_step(video, video_velocity, video_noise)
+ expected_video[:, :1] = video[:, :1]
+ expected_audio = official_step(audio, audio_velocity, audio_noise)
+
+ ((actual_video, actual_audio),) = adapter.step(
+ (video_velocity, audio_velocity),
+ (sigmas[0], sigmas[0]),
+ (video, audio),
+ )
+ torch.testing.assert_close(actual_video, expected_video)
+ torch.testing.assert_close(actual_audio, expected_audio)
+
+ terminal_video_velocity = torch.full_like(actual_video, 0.125)
+ terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
+ ((terminal_video, terminal_audio),) = adapter.step(
+ (terminal_video_velocity, terminal_audio_velocity),
+ (sigmas[1], sigmas[1]),
+ (actual_video, actual_audio),
+ )
+ expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
+ expected_terminal_video[:, :1] = actual_video[:, :1]
+ expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
+ torch.testing.assert_close(terminal_video, expected_terminal_video)
+ torch.testing.assert_close(terminal_audio, expected_terminal_audio)
+
+ def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
+
+ sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
+ velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
+ sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
+ generator = torch.Generator(device="cpu").manual_seed(10042)
+ reference_generator = torch.Generator(device="cpu").manual_seed(10042)
+
+ sigma, sigma_next = sigmas[:2]
+ denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
+ sigma_down = sigma_next * (sigma_next / sigma)
+ sigma_down_ratio = sigma_down / sigma
+ expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
+ alpha_next = 1.0 - sigma_next
+ alpha_down = 1.0 - sigma_down
+ coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+ noise = torch.randn(
+ sample.shape,
+ generator=reference_generator,
+ dtype=sample.dtype,
+ device=sample.device,
+ )
+ expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
+
+ actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py
index 36f7e9687..fbac0bfe5 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_latents.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from types import SimpleNamespace
import pytest
@@ -23,7 +26,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
-def test_prepare_video_latents_samples_directly_in_packed_token_space():
+def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +47,40 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ assert actual.stride()[1:] == (1, actual.shape[1])
+
+
+def test_prepare_video_latents_uses_diffusion_decoder_statistics():
+ pipeline = _make_pipeline(LTX2Pipeline)
+ pipeline.use_diffusion_decoder = True
+ pipeline.vae_spatial_compression_ratio = 1
+ pipeline.vae_temporal_compression_ratio = 1
+ pipeline.transformer_spatial_patch_size = 1
+ pipeline.transformer_temporal_patch_size = 1
+ pipeline.vae = SimpleNamespace(
+ latents_mean=torch.full((2,), 100.0),
+ latents_std=torch.full((2,), 10.0),
+ config=SimpleNamespace(scaling_factor=1.0),
+ )
+ pipeline.diffusion_decoder = SimpleNamespace(
+ latents_mean=torch.tensor([1.0, 3.0]),
+ latents_std=torch.tensor([2.0, 4.0]),
+ config=SimpleNamespace(scaling_factor=2.0),
+ )
+
+ actual = pipeline.prepare_latents(
+ batch_size=1,
+ num_channels_latents=2,
+ height=1,
+ width=1,
+ num_frames=1,
+ noise_scale=0.0,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+ )
+
+ torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +118,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
+ torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
+
+
+def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
+ pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
+ pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
+ generator_sp1 = torch.Generator().manual_seed(42)
+ generator_sp4 = torch.Generator().manual_seed(42)
+
+ latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
+ batch_size=1,
+ num_channels_latents=2,
+ num_mel_bins=8,
+ audio_latent_length=3,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ generator=generator_sp1,
+ )
+ latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
+ batch_size=1,
+ num_channels_latents=2,
+ num_mel_bins=8,
+ audio_latent_length=3,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ generator=generator_sp4,
+ )
+
+ torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
+ torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
+ torch.testing.assert_close(
+ torch.randn(8, generator=generator_sp4),
+ torch.randn(8, generator=generator_sp1),
+ rtol=0,
+ atol=0,
+ )
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +210,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
- torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
+ packed = latent_ops.pack_audio_latents(latents)
+ torch.testing.assert_close(padded[:, :10], packed[:, :10])
+ torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
@@ -163,3 +238,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
device=torch.device("cpu"),
latents=latents,
)
+
+
+def test_create_noised_state_matches_official_fp32_lerp():
+ latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
+ expected_generator = torch.Generator().manual_seed(42)
+ noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+ expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
+ latents.dtype
+ )
+
+ actual = latent_ops.create_noised_state(
+ latents,
+ torch.tensor(0.15, dtype=torch.bfloat16),
+ torch.Generator().manual_seed(42),
+ )
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_create_conditioned_noised_state_matches_official_two_lerps():
+ latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
+ clean_latents = latents.clone()
+ clean_latents[:, 0] = 3.0
+ denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
+ expected_generator = torch.Generator().manual_seed(42)
+ noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+ noised = torch.lerp(latents.float(), noise.float(), 0.909375)
+ expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
+
+ actual = latent_ops.create_conditioned_noised_state(
+ latents,
+ clean_latents,
+ denoise_mask,
+ 0.909375,
+ torch.Generator().manual_seed(42),
+ )
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_prepare_supplied_video_latents_uses_official_token_major_layout():
+ pipeline = _make_pipeline(LTX2Pipeline)
+ supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
+ actual_video = pipeline.prepare_latents(
+ batch_size=1,
+ num_channels_latents=4,
+ height=64,
+ width=64,
+ num_frames=9,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ latents=supplied,
+ )
+
+ torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
+ assert actual_video.stride()[1:] == (1, actual_video.shape[1])
+
+
+def test_clear_audio_padding_keeps_padding_outside_sampler_state():
+ updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
+
+ actual = latent_ops.clear_audio_padding(updated, 2)
+
+ torch.testing.assert_close(actual[:, :2], updated[:, :2])
+ torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))
diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py
index 8d8c22e79..0573dfb37 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py
@@ -12,9 +12,31 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
+ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
+
+ sigma = mocker.MagicMock()
+ sigma.to.return_value.item.return_value = 0.5
+
+ actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
+
+ sigma.to.assert_called_once_with(torch.float32)
+ sigma.to.return_value.item.assert_called_once_with()
+ torch.testing.assert_close(actual, torch.tensor([2.0]))
+
+
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
+ def test_single_guidance_pass_preserves_official_token_layout(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
+
+ tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
+ repeated = _repeat_batch(tokens, 1)
+
+ assert repeated is tokens
+ assert repeated.stride() == tokens.stride()
+
def test_combine_cfg_noise_matches_x0_space_formula(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -199,6 +221,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
+ original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
diff --git a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
new file mode 100644
index 000000000..f9b3bedce
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
@@ -0,0 +1,487 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for the fixed LTX refinement-phase adapter."""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
+from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
+ AdapterTarget,
+ parse_ltx_adapter,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
+from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+ LTX2_TWO_STAGE_RECIPE,
+ LTX_DISTILLED_ADAPTER_SLOT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _target(
+ lora_a,
+ lora_b,
+ *,
+ source_module="proj",
+ module="proj",
+ target_slice=None,
+):
+ return AdapterTarget(
+ source_module=source_module,
+ module=module,
+ slice=target_slice,
+ a_key="a",
+ b_key="b",
+ a_shape=tuple(lora_a.shape),
+ b_shape=tuple(lora_b.shape),
+ )
+
+
+def _wrap(
+ layer,
+ lora_a,
+ lora_b,
+ *,
+ source_module="proj",
+ module="proj",
+ target_slice=None,
+ layer_fused=False,
+):
+ target = _target(
+ lora_a,
+ lora_b,
+ source_module=source_module,
+ module=module,
+ target_slice=target_slice,
+ )
+ wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
+ layer,
+ [target],
+ adapter_dtype=lora_a.dtype,
+ layer_fused=layer_fused,
+ )
+ wrapper.load_target(target, lora_a, lora_b)
+ wrapper.set_enabled(True)
+ return wrapper
+
+
+class _LoRAAttention(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.to_qkv = torch.nn.Identity()
+
+
+class _LoRATransformer(torch.nn.Module):
+ stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj_in = torch.nn.Linear(4, 6, bias=False)
+ block = torch.nn.Module()
+ block.attn1 = _LoRAAttention()
+ self.transformer_blocks = torch.nn.ModuleList([block])
+
+
+def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {
+ "diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
+ "diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
+ },
+ adapter_path,
+ )
+
+ manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
+ assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
+ ("proj_in", "proj_in", None),
+ ("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
+ }
+
+
+def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
+ adapter_path,
+ )
+
+ with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
+ parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
+
+
+def _pipeline(*, dtype, quantization_config):
+ return SimpleNamespace(
+ pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
+ component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
+ transformer=object(),
+ od_config=SimpleNamespace(
+ model="unused",
+ lora_path=None,
+ dtype=dtype,
+ quantization_config=quantization_config,
+ ),
+ )
+
+
+@pytest.mark.parametrize(
+ ("dtype", "quantization_config", "layer_fused"),
+ [
+ (torch.float32, object(), False),
+ (torch.bfloat16, None, True),
+ ],
+)
+def test_ltx_phase_adapter_selects_execution_from_quantization(
+ monkeypatch,
+ dtype,
+ quantization_config,
+ layer_fused,
+):
+ pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
+ manifest = object()
+ captured: dict[str, Any] = {}
+
+ def resolve_artifact(*_args, **kwargs):
+ assert kwargs == {
+ "model_revision": None,
+ "artifact_revision": None,
+ }
+ return "adapter"
+
+ def parse(transformer, path):
+ assert transformer is pipeline.transformer
+ assert path == "adapter"
+ return manifest
+
+ class Runtime:
+ def __init__(self, transformer, parsed_manifest, **kwargs):
+ captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
+
+ monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
+ monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
+ monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
+
+ runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
+
+ assert isinstance(runtime, Runtime)
+ assert captured == {
+ "transformer": pipeline.transformer,
+ "manifest": manifest,
+ "dtype": dtype,
+ "layer_fused": layer_fused,
+ }
+
+
+def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
+ pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
+ monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
+
+ with pytest.raises(ValueError, match="bfloat16"):
+ ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
+
+
+def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
+ class TinyTransformer(torch.nn.Module):
+ stacked_params_mapping = ()
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj = torch.nn.Linear(2, 2, bias=False)
+
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {
+ "diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
+ "diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
+ },
+ adapter_path,
+ )
+ transformer = TinyTransformer()
+ runtime = LTXPhaseAdapterRuntime(
+ transformer,
+ parse_ltx_adapter(transformer, str(adapter_path)),
+ dtype=torch.float32,
+ )
+
+ assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
+ assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
+ with torch.no_grad():
+ transformer.proj.base_layer.weight.copy_(torch.eye(2))
+ runtime.finalize()
+
+ x = torch.tensor([[1.0, 1.0]])
+ runtime.activate(None)
+ torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
+ runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
+ torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
+
+
+def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
+ layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
+ base_weight = torch.tensor(
+ [
+ [0.5, -0.25, 0.125, 0.75],
+ [-0.375, 0.625, 0.25, -0.5],
+ [0.875, -0.125, -0.75, 0.375],
+ ],
+ dtype=torch.bfloat16,
+ )
+ lora_a = torch.tensor(
+ [
+ [0.1015625, -0.203125, 0.3046875, -0.40625],
+ [0.5078125, -0.609375, 0.7109375, -0.8125],
+ ],
+ dtype=torch.bfloat16,
+ )
+ lora_b = torch.tensor(
+ [[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
+ dtype=torch.bfloat16,
+ )
+ with torch.no_grad():
+ layer.weight.copy_(base_weight)
+ layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b, lora_a)
+ expected_weight.add_(base_weight)
+ x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
+ expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ assert torch.equal(wrapper(x), expected)
+ assert torch.equal(layer.weight, base_weight)
+
+
+def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeRow(torch.nn.Module):
+ input_size = 4
+ output_size = 2
+ input_size_per_partition = 2
+ input_is_parallel = True
+ tp_rank = 0
+ tp_size = 2
+ reduce_results = True
+ skip_bias_add = False
+ return_bias = False
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
+ self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
+ reductions = []
+
+ def all_reduce(value):
+ reductions.append(value.clone())
+ return value + 1
+
+ monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
+ layer = FakeRow()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b, lora_a[:, :2])
+ expected_weight.add_(layer.weight)
+ x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
+ expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ actual = wrapper(x)
+ assert len(reductions) == 1
+ assert torch.equal(reductions[0], expected_local)
+ assert torch.equal(actual, expected_local + 1)
+
+
+def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeColumn(torch.nn.Module):
+ input_size = 3
+ output_size = 4
+ output_size_per_partition = 2
+ tp_rank = 1
+ tp_size = 2
+ gather_output = False
+ skip_bias_add = False
+ return_bias = False
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(
+ torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
+ )
+ self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
+ layer = FakeColumn()
+ base_weight = layer.weight.detach().clone()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b[2:], lora_a)
+ expected_weight.add_(base_weight)
+ x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
+ expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ assert torch.equal(wrapper(x), expected)
+ assert torch.equal(layer.weight, base_weight)
+
+
+def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeColumn(torch.nn.Module):
+ pass
+
+ class FakeQKV(FakeColumn):
+ input_size = 4
+ head_size = 2
+ v_head_size = 2
+ total_num_heads = 2
+ total_num_kv_heads = 2
+ num_heads = 1
+ num_kv_heads = 1
+ num_kv_head_replicas = 1
+ tp_rank = 1
+ tp_size = 2
+ gather_output = False
+ skip_bias_add = True
+ return_bias = True
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(
+ torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
+ )
+ self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
+ monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
+ layer = FakeQKV()
+ base_weight = layer.weight.detach().clone()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
+ wrapper = _wrap(
+ layer,
+ lora_a,
+ lora_b,
+ source_module="to_k",
+ module="to_qkv",
+ target_slice="k",
+ layer_fused=True,
+ )
+
+ expected_weight = base_weight.clone()
+ expected_k = torch.matmul(lora_b[2:], lora_a)
+ expected_k.add_(base_weight[2:4])
+ expected_weight[2:4].copy_(expected_k)
+ x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
+ expected_output = torch.nn.functional.linear(x, expected_weight)
+
+ actual_output, returned_bias = wrapper(x)
+ assert torch.equal(actual_output, expected_output)
+ assert returned_bias is layer.bias
+ assert torch.equal(layer.weight, base_weight)
+
+
+@pytest.mark.parametrize(
+ ("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
+ [
+ (
+ lambda: type(
+ "FakeColumn",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "output_size": 4,
+ "output_size_per_partition": 2,
+ "tp_rank": 1,
+ "gather_output": False,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
+ },
+ )(),
+ None,
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
+ torch.tensor([[12963.0, 17284.0]]),
+ ),
+ (
+ lambda: type(
+ "FakeRow",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "output_size": 2,
+ "input_size_per_partition": 2,
+ "tp_rank": 1,
+ "reduce_results": False,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
+ },
+ )(),
+ None,
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[2.0], [3.0]]),
+ torch.tensor([[8600.0, 12900.0]]),
+ ),
+ (
+ lambda: type(
+ "FakeQKV",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "head_size": 2,
+ "v_head_size": 2,
+ "total_num_heads": 2,
+ "total_num_kv_heads": 2,
+ "num_heads": 1,
+ "num_kv_heads": 1,
+ "num_kv_head_replicas": 1,
+ "tp_rank": 1,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
+ },
+ )(),
+ "k",
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
+ torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
+ ),
+ ],
+)
+def test_ltx_phase_adapter_uses_local_tp_lora_slices(
+ monkeypatch,
+ layer_factory,
+ target_slice,
+ expected_a,
+ expected_b,
+ expected_output,
+):
+ layer = layer_factory()
+ layer.register_buffer("device_anchor", torch.empty(0))
+ layer_type = type(layer)
+ if target_slice == "k":
+ monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
+ elif "Column" in layer_type.__name__:
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
+ else:
+ monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
+
+ wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
+ torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
+ expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
+ assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
index d70820f41..9cafa22b6 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the shared LTX pipeline runtime and public contracts."""
@@ -8,18 +8,28 @@ import os
import tempfile
from dataclasses import replace
from types import SimpleNamespace
+from typing import Any
import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.models.ltx2 import ltx2_components
+from vllm_omni.diffusion.models.ltx2 import (
+ ltx2_components,
+ ltx2_transformer,
+)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
+ LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+ LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
+ LTX23_DISTILLED_COMPONENT_PROFILE,
+ LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
+ resolve_ltx_artifact,
+ resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,25 +43,40 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
+ _repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+ LTX2_DISTILLED_ONE_STAGE_RECIPE,
LTX2_DISTILLED_TWO_STAGE_RECIPE,
LTX2_ONE_STAGE_RECIPE,
+ LTX2_TWO_STAGE_RECIPE,
+ LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
+ LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
+ resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import (
+ LTXRequestInputs,
+ validate_ltx_checkpoint,
+ validate_pipeline_request,
)
-from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+ LTX2DistilledOneStagePipeline,
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+ LTX2DistilledPipeline,
+ LTX2DistilledTwoStagePipeline,
+ LTX2TwoStagePipeline,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +129,28 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
+ assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
+ assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
+ assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
+ assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
+ assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
+ assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
+ assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+ assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
+ for pipeline_cls, profile in (
+ (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
+ (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
+ (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
+ (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
+ ):
+ assert pipeline_cls._dit_modules == list(profile.dit_modules)
+ assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
+ assert pipeline_cls._vae_modules == list(profile.vae_modules)
+ assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +166,131 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
+@pytest.mark.parametrize(
+ ("model", "repo_id", "expected_revision"),
+ [
+ ("Lightricks/test", "Lightricks/test", "model-revision"),
+ ("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
+ ],
+)
+def test_ltx_artifact_uses_source_revision_and_hub_fallback(
+ monkeypatch,
+ model,
+ repo_id,
+ expected_revision,
+):
+ filename = "ltx-sidecar.safetensors"
+ calls = []
+
+ def fake_download(**kwargs):
+ calls.append(kwargs)
+ return "/cache/ltx-sidecar.safetensors"
+
+ monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+
+ assert (
+ resolve_ltx_artifact(
+ model,
+ repo_id,
+ filename,
+ model_revision="model-revision",
+ artifact_revision="artifact-revision",
+ )
+ == "/cache/ltx-sidecar.safetensors"
+ )
+ assert calls == [
+ {
+ "repo_id": repo_id,
+ "filename": filename,
+ "revision": expected_revision,
+ }
+ ]
+
+
+def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
+ filename = "ltx-sidecar.safetensors"
+ expected = tmp_path / filename
+ expected.write_bytes(b"sidecar")
+ monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
+
+ assert resolve_ltx_artifact(
+ str(tmp_path),
+ "Lightricks/test",
+ filename,
+ model_revision="unused-model-revision",
+ artifact_revision="unused-artifact-revision",
+ ) == str(expected)
+
+
+def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
+ calls = []
+
+ def fake_download(**kwargs):
+ calls.append(kwargs)
+ return "/cache/ltx-sidecar.safetensors"
+
+ monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+
+ assert (
+ resolve_ltx_artifact(
+ str(tmp_path),
+ "Lightricks/test",
+ "ltx-sidecar.safetensors",
+ model_revision="local-model-revision",
+ artifact_revision="artifact-revision",
+ )
+ == "/cache/ltx-sidecar.safetensors"
+ )
+ assert calls == [
+ {
+ "repo_id": "Lightricks/test",
+ "filename": "ltx-sidecar.safetensors",
+ "revision": "artifact-revision",
+ }
+ ]
+
+
+@pytest.mark.parametrize(
+ ("expected_kind", "scheduler_config"),
+ [
+ ("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
+ ("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
+ ],
+)
+def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
+ validate_ltx_checkpoint(
+ scheduler_config,
+ expected_kind=expected_kind,
+ pipeline_name="TestLTXPipeline",
+ )
+
+
+@pytest.mark.parametrize(
+ ("expected_kind", "scheduler_config", "error"),
+ [
+ ("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
+ ("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
+ ],
+)
+def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
+ with pytest.raises(ValueError, match=error):
+ validate_ltx_checkpoint(
+ scheduler_config,
+ expected_kind=expected_kind,
+ pipeline_name="TestLTXPipeline",
+ )
+
+
+def test_ltx_rejects_advanced_uaa_before_component_initialization():
+ od_config = SimpleNamespace(
+ model="unused",
+ parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
+ )
+
+ with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
+ LTX2Pipeline(od_config=od_config)
+
+
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +312,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
-def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
+def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
+ initialized = False
+
+ def stub_components(*_args, **_kwargs):
+ nonlocal initialized
+ initialized = True
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+
+ with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
+ LTX2TwoStagePipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ cache_backend="cache_dit",
+ )
+ )
+
+ assert not initialized
+
+
+def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
+
+ def stub_components(pipe, od_config):
+ pipe.od_config = od_config
+ pipe.vae_spatial_compression_ratio = 32
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+ monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+ monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
+ pipe = LTX2TwoStagePipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ enable_diffusion_pipeline_profiler=False,
+ lora_path=None,
+ model_config={},
+ )
+ )
+
+ assert pipe.model_version == "2.3"
+ assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
+
+
+def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
+
+ def stub_components(pipe, od_config):
+ pipe.od_config = od_config
+ pipe.vae_spatial_compression_ratio = 32
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+ monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+ pipe = LTX2DistilledPipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ enable_diffusion_pipeline_profiler=False,
+ )
+ )
+
+ assert pipe.model_version == "2.3"
+ assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
+ assert pipe._phase_adapter is None
+ assert pipe.preserve_sp_padded_audio_duration
+ assert pipe.reports_stage_durations
+
+
+def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
+ assert not LTX2TwoStagePipeline.supports_request_batch
+ assert LTX2TwoStagePipeline.support_image_input
+ assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
+def test_ltx_two_stage_resolves_request_embedded_images():
+ pipe = object.__new__(LTX2TwoStagePipeline)
+ request_inputs = SimpleNamespace(latents=None)
+ image = torch.zeros(3, 8, 8)
+ req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
+
+ assert pipe._resolve_request_image(req, None, request_inputs) is image
+
+
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +586,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
+def test_ltx_positive_only_guidance_preserves_token_major_layout():
+ tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
+
+ repeated = _repeat_batch(tensor, 1)
+
+ assert repeated is tensor
+ assert repeated.stride() == tensor.stride()
+
+
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +645,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
+def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
+ logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
+ logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
+ guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
+ expected = combine_guided_x0(
+ cond=logical_cond,
+ uncond_text=logical_uncond,
+ uncond_perturbed=0.0,
+ uncond_modality=0.0,
+ guidance=guidance,
+ )
+ padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
+ padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
+
+ actual = combine_guided_x0(
+ cond=padded_cond,
+ uncond_text=padded_uncond,
+ uncond_perturbed=0.0,
+ uncond_modality=0.0,
+ guidance=guidance,
+ rescale_token_count=2,
+ )
+
+ torch.testing.assert_close(actual[:, :2], expected)
+
+
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -488,15 +779,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):
assert len(installed) == 2
assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
- assert all(processor.has_learned_registers for processor in installed)
assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
assert all(kernel["role_category"] == "self" for kernel in kernels)
assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
assert all(kernel["disable_kv_quant"] for kernel in kernels)
+ assert all(processor.has_learned_registers for processor in installed)
+ assert not any(processor.preserve_learned_register_mask for processor in installed)
-@pytest.mark.parametrize("has_learned_registers", [False, True])
-def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
+def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
+ installed = []
+
+ class OmniAttention:
+ def __init__(self, **_kwargs):
+ pass
+
+ class Attention:
+ heads = 32
+ head_dim = 128
+ inner_kv_dim = 4096
+
+ def set_processor(self, processor):
+ installed.append(processor)
+
+ monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
+ connectors = SimpleNamespace(
+ video_connector=SimpleNamespace(
+ learnable_registers=object(),
+ transformer_blocks=[SimpleNamespace(attn1=Attention())],
+ ),
+ audio_connector=SimpleNamespace(
+ learnable_registers=object(),
+ transformer_blocks=[SimpleNamespace(attn1=Attention())],
+ ),
+ )
+
+ _install_connector_attention(connectors, preserve_learned_register_mask=True)
+
+ assert all(processor.has_learned_registers for processor in installed)
+ assert all(processor.preserve_learned_register_mask for processor in installed)
+
+
+def test_ltx_connector_attention_dispatches_through_omni_kernel():
calls = []
class Backend:
@@ -511,12 +835,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
calls.append((query, key, value, metadata))
return query
+ class UpcastingNorm(torch.nn.Module):
+ def forward(self, tensor):
+ return tensor.float()
+
attention = SimpleNamespace(
heads=2,
head_dim=4,
inner_kv_dim=8,
- norm_q=torch.nn.Identity(),
- norm_k=torch.nn.Identity(),
+ norm_q=UpcastingNorm(),
+ norm_k=UpcastingNorm(),
to_q=torch.nn.Identity(),
to_k=torch.nn.Identity(),
to_v=torch.nn.Identity(),
@@ -525,33 +853,46 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
rope_type="split",
omni_attention=OmniAttention(),
)
- hidden_states = torch.randn(2, 3, 8)
+ hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
additive_mask = torch.zeros(2, 1, 3, 3)
+ rotary = (
+ torch.randn(2, 2, 3, 2, dtype=torch.float32),
+ torch.randn(2, 2, 3, 2, dtype=torch.float32),
+ )
- output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
+ output = ltx2_components._LTXConnectorAttnProcessor()(
attention,
hidden_states,
attention_mask=additive_mask,
+ query_rotary_emb=rotary,
)
query, key, value, metadata = calls[0]
assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
- if has_learned_registers:
- assert metadata is None
- else:
- assert metadata.attn_mask.shape == (2, 3)
- assert metadata.attn_mask.dtype == torch.bool
- assert metadata.attn_mask.all()
- torch.testing.assert_close(output, hidden_states)
+ assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+ assert metadata.attn_mask.shape == (2, 3)
+ assert metadata.attn_mask.dtype == torch.bool
+ assert metadata.attn_mask.all()
+ expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
+ expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
+ torch.testing.assert_close(output, expected, rtol=0, atol=0)
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
- assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
+ assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
+def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
+ assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
+ assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
+ assert "load_weights" not in LTX2TwoStagePipeline.__dict__
+ assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
+ assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
+
+
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +935,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
-def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
+@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -618,7 +960,7 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
max_sequence_length=16,
)
prompt_context = object()
- phase_calls = []
+ phase_calls: list[Any] = []
def resolve_request_inputs(req, **kwargs):
return request_inputs
@@ -671,12 +1013,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
- pipeline = object.__new__(LTX2DistilledPipeline)
+ pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
+ activated_slots: list[Any] = []
+ if pipeline_cls is LTX2TwoStagePipeline:
+ pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +1031,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
+ if pipeline_cls is LTX2TwoStagePipeline:
+ assert activated_slots == [None, "ltx_distilled"]
torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
- torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
+ expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
+ torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
-def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
+@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
- pipeline = object.__new__(LTX2DistilledPipeline)
+ pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +1099,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
+def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
+ assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
+ assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
+ assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
+ "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
+ )
+
+
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -832,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
)
+def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
+ request = LTXRequestInputs(
+ prompt="prompt",
+ negative_prompt="",
+ height=128,
+ width=128,
+ num_frames=9,
+ frame_rate=24.0,
+ num_inference_steps=30,
+ guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
+ num_videos_per_prompt=1,
+ generator=None,
+ latents=None,
+ audio_latents=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_attention_mask=None,
+ negative_prompt_attention_mask=None,
+ decode_timestep=0.0,
+ decode_noise_scale=None,
+ output_type="np",
+ max_sequence_length=16,
+ )
+
+ with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
+ validate_pipeline_request(
+ request,
+ pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+ vae_spatial_compression_ratio=32,
+ vae_temporal_compression_ratio=8,
+ pipeline_name="LTX2Pipeline",
+ min_video_latent_spatial_size=(7, 7),
+ )
+
+ # ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
+ validate_pipeline_request(
+ request,
+ pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+ vae_spatial_compression_ratio=32,
+ vae_temporal_compression_ratio=8,
+ pipeline_name="LTX2Pipeline",
+ )
+
+ validate_pipeline_request(
+ replace(request, height=256, width=256),
+ pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+ vae_spatial_compression_ratio=32,
+ vae_temporal_compression_ratio=8,
+ pipeline_name="LTX2Pipeline",
+ min_video_latent_spatial_size=(7, 7),
+ )
+
+
@pytest.mark.parametrize(
("direct_kwargs", "sampling_kwargs", "error"),
[
@@ -958,6 +1372,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
+@pytest.mark.parametrize(
+ ("recipe", "steps", "stg_block"),
+ [
+ (LTX2_TWO_STAGE_RECIPE, 40, 29),
+ (LTX23_TWO_STAGE_RECIPE, 30, 28),
+ ],
+)
+def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
+ stage1, stage2 = recipe.phases
+ assert (recipe.width, recipe.height) == (1536, 1024)
+ assert recipe.num_inference_steps == steps
+ assert stage1.spatial_downscale == 2
+ assert stage1.adapter_slot is None
+ assert stage1.guidance.video.stg_blocks == (stg_block,)
+ assert stage1.guidance.audio.stg_blocks == (stg_block,)
+ assert stage2.input_transform == "spatial_upsample"
+ assert stage2.adapter_slot == "ltx_distilled"
+ assert stage2.guidance == LTXGuidanceSpec.positive_only()
+ assert stage2.num_inference_steps == 3
+ assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+ # Official Stage 2 refines video only; its audio result is discarded.
+ assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
+
+
+def test_ltx_two_stage_entries_select_the_official_adapter_slots():
+ phase_switches: list[Any] = []
+ ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
+ ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
+ for phase in LTX2_TWO_STAGE_RECIPE.phases:
+ ordinary_pipeline._enter_phase(phase)
+ assert phase_switches == [None, "ltx_distilled"]
+
+ distilled_pipeline = object.__new__(LTX2DistilledPipeline)
+ for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
+ distilled_pipeline._enter_phase(phase)
+
+
+def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
+ class FakeColumn(torch.nn.Module):
+ def __init__(
+ self,
+ input_size,
+ output_size,
+ *,
+ bias,
+ gather_output,
+ return_bias,
+ quant_config,
+ prefix,
+ ):
+ super().__init__()
+ self.input_size = input_size
+ self.output_size = output_size
+ self.bias_enabled = bias
+ self.gather_output = gather_output
+ self.return_bias = return_bias
+ self.quant_config = quant_config
+ self.prefix = prefix
+
+ class FakeRow(torch.nn.Module):
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ class FakeAttention(torch.nn.Module):
+ def __init__(self, **kwargs):
+ super().__init__()
+
+ monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
+ monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
+ monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
+ monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+ attention = ltx2_transformer.LTX2Attention(
+ query_dim=16,
+ heads=4,
+ dim_head=4,
+ qk_norm="rms_norm",
+ apply_gated_attention=True,
+ pack_qkv=False,
+ prefix="transformer_blocks.0.attn1",
+ )
+
+ gate = attention.to_gate_logits
+ assert isinstance(gate, FakeColumn)
+ assert gate.input_size == 16
+ assert gate.output_size == 4
+ assert gate.gather_output is False
+ assert gate.return_bias is False
+ assert gate.quant_config is None
+ assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
+
+
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1366,6 +1872,7 @@ class TestLTXForwardStages:
assert seen["request_inputs"].output_type == "latent"
assert seen["kwargs"] == {
"request_sigmas": request_sigmas,
+ "request_phase_sigmas": None,
"image": None,
}
@@ -1432,6 +1939,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1985,21 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
+ assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
+ "ltx2",
+ "pipeline_ltx2_two_stage",
+ "LTX2TwoStagePipeline",
+ )
+ assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
+ "ltx2",
+ "pipeline_ltx2",
+ "LTX2DistilledOneStagePipeline",
+ )
+ assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
+ "ltx2",
+ "pipeline_ltx2_two_stage",
+ "LTX2DistilledTwoStagePipeline",
+ )
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +2011,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
+ "LTX2ImageToVideoTwoStagePipeline",
+ "LTX23TwoStagePipeline",
+ "LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +2023,9 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledOneStagePipeline",
+ "LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +2038,9 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledOneStagePipeline",
+ "LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
index 73e10463e..e9bba1f1f 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
@@ -1,17 +1,60 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
import pytest
+import torch
from torch import nn
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
-from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
+ LTX2FeedForward,
+ LTX2VideoTransformer3DModel,
+ LTX2VideoTransformerBlock,
+ _make_rms_norm,
+ apply_interleaved_rotary_emb,
+ apply_keyframes_absolute_embedding,
+ apply_split_rotary_emb,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
+ torch.manual_seed(0)
+ x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+ cos = torch.randn(2, 3, 8, dtype=x.dtype)
+ sin = torch.randn(2, 3, 8, dtype=x.dtype)
+ x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
+ rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+ expected = x * cos + rotated * sin
+
+ actual = apply_interleaved_rotary_emb(x, (cos, sin))
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
+ torch.manual_seed(1)
+ x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+ cos = torch.randn(2, 3, 4, dtype=x.dtype)
+ sin = torch.randn(2, 3, 4, dtype=x.dtype)
+ split_x = x.reshape(2, 3, 2, 4)
+ first_x, second_x = split_x.chunk(2, dim=-2)
+ cos_u = cos.unsqueeze(-2)
+ sin_u = sin.unsqueeze(-2)
+ expected = split_x * cos_u
+ expected[..., :1, :].addcmul_(-sin_u, second_x)
+ expected[..., 1:, :].addcmul_(sin_u, first_x)
+ expected = expected.reshape_as(x)
+
+ actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
def test_ltx_rms_norm_no_affine_matches_official_torch_module():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
@@ -29,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
assert "weight" in norm.state_dict()
-def test_ltx_transformer_has_separate_cfg_cache_dit_config():
+def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
+ calls = []
+
+ class UpcastingNorm(nn.Module):
+ def forward(self, tensor):
+ return tensor.float()
+
+ class RecordingAttention:
+ def __call__(self, query, key, value, metadata):
+ calls.append((query, key, value, metadata))
+ return query
+
+ attention = SimpleNamespace(
+ heads=2,
+ head_dim=4,
+ to_qkv=None,
+ to_q=nn.Identity(),
+ to_k=nn.Identity(),
+ to_v=nn.Identity(),
+ norm_q=UpcastingNorm(),
+ norm_k=UpcastingNorm(),
+ to_gate_logits=None,
+ attn=RecordingAttention(),
+ to_out=(nn.Identity(), nn.Identity()),
+ rope_type="split",
+ )
+ hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+ rotary = (
+ torch.ones(2, 2, 3, 2, dtype=torch.float32),
+ torch.zeros(2, 2, 3, 2, dtype=torch.float32),
+ )
+
+ processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
+ processor._slice_rope_for_tp = lambda rope, _: rope
+ output = processor(attention, hidden_states, query_rotary_emb=rotary)
+
+ query, key, value, metadata = calls[0]
+ assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+ assert metadata is None
+ torch.testing.assert_close(output, hidden_states)
+
+
+def test_ltx_transformer_has_fused_cfg_cache_dit_config():
adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
assert isinstance(adapter_config, CacheDiTAdapterConfig)
- assert adapter_config.has_separate_cfg
+ assert not adapter_config.has_separate_cfg
def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -66,3 +151,332 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
+
+
+def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
+ processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
+ monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
+ processor = processor_cls()
+ hidden_states = torch.zeros(1, 3, 4)
+ encoder_hidden_states = torch.zeros(1, 4, 4)
+ key_padding_mask = torch.tensor([[True, True, True, False]])
+
+ actual = processor._prepare_attention_mask(
+ SimpleNamespace(),
+ hidden_states,
+ encoder_hidden_states,
+ key_padding_mask,
+ batch_size=1,
+ sequence_length=4,
+ )
+
+ torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
+
+
+def _fake_ltx_parallel_attention():
+ calls = []
+
+ class Backend:
+ @staticmethod
+ def get_name():
+ return "FLASH_ATTN"
+
+ @staticmethod
+ def supports_attention_mask():
+ return True
+
+ def run(name):
+ def forward(query, _key, _value, _metadata):
+ calls.append(name)
+ return torch.zeros_like(query)
+
+ return forward
+
+ layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
+ nn.Module.__init__(layer)
+ layer.attn_backend = Backend
+ layer.attention = SimpleNamespace(forward=run("native"))
+ layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
+ layer.backend_pref = "FLASH_ATTN"
+ return layer, calls
+
+
+@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
+def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
+ layer, calls = _fake_ltx_parallel_attention()
+ query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
+ key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
+ key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
+ metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+
+ layer._run_local_attention(query, key, value, metadata)
+
+ assert calls == ["sdpa"]
+
+
+def test_ltx_unmasked_attention_keeps_native_backend():
+ layer, calls = _fake_ltx_parallel_attention()
+ query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
+
+ layer._run_local_attention(query, key, value, None)
+
+ assert calls == ["native"]
+
+
+def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
+ class FakeParallelLinear(nn.Linear):
+ def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
+ super().__init__(dim_in, dim_out, bias=bias)
+
+ monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
+ monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
+
+ without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
+ with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
+
+ assert without_bias.net[0].proj.bias is None
+ assert without_bias.net[2].bias is None
+ assert with_bias.net[0].proj.bias is not None
+ assert with_bias.net[2].bias is not None
+
+
+def _build_tiny_ltx_transformer(monkeypatch, **overrides):
+ class FakeBlock(nn.Module):
+ def __init__(self, **kwargs):
+ super().__init__()
+ self.init_kwargs = kwargs
+
+ def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
+ return hidden_states, audio_hidden_states
+
+ monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
+ kwargs = {
+ "in_channels": 4,
+ "out_channels": 4,
+ "num_attention_heads": 2,
+ "attention_head_dim": 4,
+ "cross_attention_dim": 8,
+ "vae_scale_factors": (1, 1, 1),
+ "pos_embed_max_pos": 4,
+ "base_height": 8,
+ "base_width": 8,
+ "audio_in_channels": 4,
+ "audio_out_channels": 4,
+ "audio_num_attention_heads": 2,
+ "audio_attention_head_dim": 4,
+ "audio_cross_attention_dim": 8,
+ "audio_scale_factor": 1,
+ "audio_pos_embed_max_pos": 4,
+ "audio_sampling_rate": 16,
+ "audio_hop_length": 4,
+ "num_layers": 1,
+ "caption_channels": 8,
+ "use_prompt_embeddings": False,
+ "cross_attn_mod": True,
+ "audio_cross_attn_mod": True,
+ }
+ kwargs.update(overrides)
+ return LTX2VideoTransformer3DModel(**kwargs)
+
+
+def test_ltx25_keyframe_absolute_embedding_matches_official():
+ hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
+ keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
+ embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
+
+ actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
+ expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
+
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
+
+
+def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
+ legacy_model = _build_tiny_ltx_transformer(monkeypatch)
+
+ assert legacy_model.config.ff_bias is True
+ assert legacy_model.config.audio_ff_bias is True
+ assert legacy_model.config.use_prompt_adaln_single is True
+ assert legacy_model.config.use_keyframes_abs_pos_embedding is False
+ assert hasattr(legacy_model, "prompt_adaln")
+ assert hasattr(legacy_model, "audio_prompt_adaln")
+ assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
+ assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
+ assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+ ltx25_model = _build_tiny_ltx_transformer(
+ monkeypatch,
+ ff_bias=False,
+ audio_ff_bias=True,
+ use_prompt_adaln_single=False,
+ use_keyframes_abs_pos_embedding=True,
+ )
+
+ assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
+ torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
+ assert not hasattr(ltx25_model, "prompt_adaln")
+ assert not hasattr(ltx25_model, "audio_prompt_adaln")
+ assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
+ assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+ monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
+ expected = torch.arange(8, dtype=torch.float32).view(1, 8)
+ loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
+
+ assert loaded == {"keyframes_abs_pos_embedding"}
+ torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
+
+
+def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
+ class CaptureAda(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.calls = []
+
+ def forward(self, timestep, *, batch_size, hidden_dtype):
+ self.calls.append(timestep.detach().clone())
+ embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
+ return embedding, embedding
+
+ class CaptureRope(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.output_dtypes = []
+
+ def forward(self, _coords, *, device, out_dtype=None):
+ self.output_dtypes.append(out_dtype)
+ empty = torch.empty(0, device=device, dtype=out_dtype)
+ return empty, empty
+
+ model = _build_tiny_ltx_transformer(monkeypatch)
+ rope = CaptureRope()
+ model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
+ model.prompt_modulation = False
+ video_scale_shift = CaptureAda()
+ video_gate = CaptureAda()
+ audio_scale_shift = CaptureAda()
+ audio_gate = CaptureAda()
+ model.av_cross_attn_video_scale_shift = video_scale_shift
+ model.av_cross_attn_video_a2v_gate = video_gate
+ model.av_cross_attn_audio_scale_shift = audio_scale_shift
+ model.av_cross_attn_audio_v2a_gate = audio_gate
+
+ model(
+ hidden_states=torch.zeros(1, 2, 4),
+ audio_hidden_states=torch.zeros(1, 1, 4),
+ encoder_hidden_states=torch.zeros(1, 1, 8),
+ audio_encoder_hidden_states=torch.zeros(1, 1, 8),
+ timestep=torch.tensor([[0.0, 2.0]]),
+ audio_timestep=torch.tensor([[5.0]]),
+ sigma=torch.tensor([2.0]),
+ audio_sigma=torch.tensor([5.0]),
+ video_coords=torch.zeros(1, 3, 2, 2),
+ audio_coords=torch.zeros(1, 1, 1, 2),
+ )
+
+ torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
+ torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
+ torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
+ torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
+ assert len(rope.output_dtypes) == 4
+ assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
+
+
+def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
+ class CaptureAttention(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.encoder_hidden_states = None
+
+ def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+ if encoder_hidden_states is not None:
+ self.encoder_hidden_states = encoder_hidden_states.detach().clone()
+ return torch.zeros_like(hidden_states)
+
+ class ZeroOutput(nn.Module):
+ def forward(self, hidden_states):
+ return torch.zeros_like(hidden_states)
+
+ block = object.__new__(LTX2VideoTransformerBlock)
+ nn.Module.__init__(block)
+ block.video_cross_attn_adaln = True
+ block.audio_cross_attn_adaln = True
+ block.cross_attn_adaln = True
+ block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+ block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+ block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+ block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
+ block.norm1 = nn.Identity()
+ block.audio_norm1 = nn.Identity()
+ block.norm2 = nn.Identity()
+ block.audio_norm2 = nn.Identity()
+ block.norm3 = nn.Identity()
+ block.audio_norm3 = nn.Identity()
+ block.attn1 = CaptureAttention()
+ block.audio_attn1 = CaptureAttention()
+ block.attn2 = CaptureAttention()
+ block.audio_attn2 = CaptureAttention()
+ block.ff = ZeroOutput()
+ block.audio_ff = ZeroOutput()
+
+ hidden_states = torch.ones(1, 1, 2)
+ audio_hidden_states = torch.ones(1, 1, 2)
+ encoder_hidden_states = torch.ones(1, 1, 2)
+ audio_encoder_hidden_states = torch.ones(1, 1, 2)
+ block(
+ hidden_states,
+ audio_hidden_states,
+ encoder_hidden_states,
+ audio_encoder_hidden_states,
+ temb=torch.zeros(1, 1, 18),
+ temb_audio=torch.zeros(1, 1, 18),
+ temb_ca_scale_shift=torch.empty(0),
+ temb_ca_audio_scale_shift=torch.empty(0),
+ temb_ca_gate=torch.empty(0),
+ temb_ca_audio_gate=torch.empty(0),
+ temb_prompt=None,
+ temb_prompt_audio=None,
+ use_a2v_cross_attention=False,
+ use_v2a_cross_attention=False,
+ )
+
+ torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
+ torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
+
+
+def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
+ column_calls = []
+
+ class FakeColumnParallelLinear(nn.Linear):
+ def __init__(self, input_size, output_size, *, bias=True, **kwargs):
+ column_calls.append((output_size, kwargs))
+ super().__init__(input_size, output_size, bias=bias)
+
+ class FakeRowParallelLinear(nn.Linear):
+ def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
+ super().__init__(input_size, output_size, bias=bias)
+
+ class FakeParallelAttention(nn.Module):
+ def __init__(self, **_kwargs):
+ super().__init__()
+
+ monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
+ monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
+ monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
+ monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+ attention = ltx2_transformer.LTX2Attention(
+ query_dim=8,
+ heads=4,
+ kv_heads=4,
+ dim_head=2,
+ qk_norm="rms_norm",
+ apply_gated_attention=True,
+ pack_qkv=False,
+ )
+
+ gate_output_size, gate_kwargs = column_calls[-1]
+ assert gate_output_size == 4
+ assert gate_kwargs["gather_output"] is False
+ assert gate_kwargs["return_bias"] is False
+ assert attention.query_num_heads == 2
diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py
index 108bb960d..964bb7d06 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_vae.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py
@@ -1,9 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+"""Unit tests for LTX video VAE tiling and distributed decode behavior."""
from types import SimpleNamespace
+from typing import Any
import pytest
import torch
@@ -11,6 +12,466 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+
+ class FakeBaseVocoder(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.input_dtype = None
+
+ def forward(self, value):
+ self.input_dtype = value.dtype
+ return value
+
+ monkeypatch.setattr(
+ ltx_runtime.torch,
+ "autocast",
+ lambda **_kwargs: (_ for _ in ()).throw(AssertionError("base vocoder must not use autocast")),
+ )
+ vocoder = FakeBaseVocoder()
+
+ output = ltx_runtime._run_ltx_vocoder(vocoder, torch.ones(1, dtype=torch.bfloat16))
+
+ assert vocoder.input_dtype == torch.bfloat16
+ assert output.dtype == torch.bfloat16
+
+
+class TestLTXDiffusionDecoder:
+ @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
+ def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+ DistributedLTX2VideoDiffusionDecoderModel,
+ )
+
+ model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+ torch.nn.Module.__init__(model)
+ model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None)
+
+ with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"):
+ model.set_parallel_size(2, mode=mode)
+
+ def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+ DistributedLTX2VideoDiffusionDecoderModel,
+ )
+
+ calls = []
+ model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+ torch.nn.Module.__init__(model)
+ model.distributed_executor = SimpleNamespace(
+ set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode))
+ )
+
+ model.set_parallel_size(4, mode="tile")
+
+ assert calls == [(4, "tile")]
+
+ def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+ LTX2VideoDiffusionDecoder3d,
+ LTX2VideoDiffusionDecoderModel,
+ )
+
+ class IdentityAttention:
+ def build_block_mask(self, _hidden_states):
+ return None
+
+ class IdentityBlock(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.attn = IdentityAttention()
+
+ def forward(self, hidden_states, _block_mask):
+ return hidden_states
+
+ class IdentityUpsample(torch.nn.Module):
+ def forward(self, hidden_states, *, drop_leading_frame):
+ assert drop_leading_frame
+ return hidden_states
+
+ decoder = object.__new__(LTX2VideoDiffusionDecoder3d)
+ torch.nn.Module.__init__(decoder)
+ decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])])
+ decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()])
+ decoder.trailing_pad_latent_frames = 2
+ decoder.temporal_compression_ratio = 8
+ decoder.stage5_kernel = (11, 3, 3)
+
+ # 17 frames contain one real frame followed by 16 propagated ghost
+ # frames. The reference keeps 11 for stage-5 NATTEN instead of
+ # cropping the tensor immediately back to one.
+ context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4))
+ assert context.shape == (1, 11, 3, 3, 4)
+
+ class FakeDecoder(torch.nn.Module):
+ def forward(self, z, *, generator, num_inference_steps):
+ assert generator is None
+ assert num_inference_steps is None
+ return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2)
+
+ model = object.__new__(LTX2VideoDiffusionDecoderModel)
+ torch.nn.Module.__init__(model)
+ model.decoder = FakeDecoder()
+ model.use_tiling = False
+ model.tile_sample_min_num_frames = 80
+ model.tile_sample_min_height = 768
+ model.tile_sample_min_width = 768
+ model.temporal_compression_ratio = 8
+ model.spatial_compression_ratio = 2
+
+ output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0]
+ assert output.shape == (1, 3, 1, 4, 6)
+
+ @pytest.mark.parametrize(
+ ("model_version", "extras", "expected"),
+ [
+ ("2.5", {}, True),
+ ("2.5", {"ltx2_use_conv_vae": True}, False),
+ ("2.3", {}, False),
+ ],
+ )
+ def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected):
+ from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+ assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected
+
+ def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+ with pytest.raises(TypeError, match="ltx2_use_conv_vae"):
+ _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
+
+ def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+ convert_ltx25_native_diffusion_decoder_state_dict,
+ )
+
+ qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
+ qkv_bias = torch.arange(6, dtype=torch.float32)
+ projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+ projection_bias = torch.ones(2)
+ gate = torch.tensor([2.0, 3.0])
+ mean = torch.tensor([0.1, 0.2])
+ std = torch.tensor([0.3, 0.4])
+ native = {
+ "encoder.ignored": torch.ones(1),
+ "per_channel_statistics.mean-of-means": mean,
+ "vae.per_channel_statistics.std-of-means": std,
+ "decoder.type_emb": torch.ones(2),
+ "decoder.coarse_head.weight": torch.ones(1),
+ "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight,
+ "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias,
+ "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]),
+ "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]),
+ "decoder.diff_blocks.0.attn.proj.weight": projection_weight,
+ "decoder.diff_blocks.0.attn.proj.bias": projection_bias,
+ "decoder.diff_blocks.0.gate_msa": gate,
+ "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2),
+ }
+
+ converted = convert_ltx25_native_diffusion_decoder_state_dict(native)
+
+ torch.testing.assert_close(converted["latents_mean"], mean)
+ torch.testing.assert_close(converted["latents_std"], std)
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:])
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0]))
+ torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0]))
+ torch.testing.assert_close(
+ converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight
+ )
+ torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias)
+ assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted
+ assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
+
+ def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+ convert_ltx25_native_diffusion_decoder_state_dict,
+ )
+
+ with pytest.raises(ValueError, match="not divisible by 3"):
+ convert_ltx25_native_diffusion_decoder_state_dict(
+ {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)}
+ )
+
+ def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self):
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ seen = {}
+
+ class ConvVae:
+ dtype = torch.float32
+ config = SimpleNamespace(timestep_conditioning=True)
+
+ def decode(self, *_args, **_kwargs):
+ raise AssertionError("the convolutional VAE decoder must not run")
+
+ class DiffusionDecoder:
+ dtype = torch.float32
+
+ def decode(self, latents, *, generator, return_dict):
+ seen["latents"] = latents
+ seen["generator"] = generator
+ seen["return_dict"] = return_dict
+ return (latents + 1,)
+
+ class AudioVae:
+ dtype = torch.float32
+
+ def decode(self, audio_latents, *, return_dict):
+ assert return_dict is False
+ return (audio_latents + 2,)
+
+ class VideoProcessor:
+ def postprocess_video(self, video, *, output_type):
+ seen["output_type"] = output_type
+ return video
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.use_diffusion_decoder = True
+ pipe.vae = ConvVae()
+ pipe.diffusion_decoder = DiffusionDecoder()
+ pipe.audio_vae = AudioVae()
+ pipe.vocoder = torch.nn.Identity()
+ pipe.video_processor = VideoProcessor()
+
+ generator = torch.Generator().manual_seed(123)
+ latents = torch.ones(1, 1)
+ output = pipe._decode_output(
+ latents=latents,
+ audio_latents=torch.ones(1, 1),
+ output_type="pt",
+ connector_prompt_embeds=torch.ones(1, 1),
+ generator=generator,
+ device=torch.device("cpu"),
+ decode_timestep=0.7,
+ decode_noise_scale=0.5,
+ prompt_batch_size=1,
+ )
+
+ torch.testing.assert_close(seen["latents"], latents)
+ assert seen["generator"] is generator
+ assert seen["return_dict"] is False
+ assert seen["output_type"] == "pt"
+ torch.testing.assert_close(output.output[0], latents + 1)
+ torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0))
+
+ def test_diffusion_decode_denormalizes_with_decoder_statistics(self):
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.use_diffusion_decoder = True
+ pipe.transformer_spatial_patch_size = 1
+ pipe.transformer_temporal_patch_size = 1
+ pipe.vae = SimpleNamespace(
+ latents_mean=torch.full((2,), 100.0),
+ latents_std=torch.full((2,), 10.0),
+ config=SimpleNamespace(scaling_factor=1.0),
+ )
+ pipe.diffusion_decoder = SimpleNamespace(
+ latents_mean=torch.tensor([10.0, 20.0]),
+ latents_std=torch.tensor([2.0, 4.0]),
+ config=SimpleNamespace(scaling_factor=2.0),
+ )
+ pipe.audio_vae = SimpleNamespace(
+ latents_mean=torch.tensor(0.0),
+ latents_std=torch.tensor(1.0),
+ )
+ forward_ctx = SimpleNamespace(
+ latent_num_frames=1,
+ latent_height=1,
+ latent_width=1,
+ original_audio_num_frames=1,
+ latent_mel_bins=2,
+ )
+
+ video, audio = pipe._unpack_and_denormalize_stage(
+ forward_ctx,
+ torch.tensor([[[1.0, 2.0]]]),
+ torch.zeros(1, 1, 2),
+ )
+
+ torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0]))
+ torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2))
+
+ def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ class ConvVae:
+ config = SimpleNamespace(timestep_conditioning=False)
+
+ def is_distributed_enabled(self):
+ raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives")
+
+ class DiffusionDecoder:
+ dtype = torch.float32
+
+ def decode(self, *_args, **_kwargs):
+ raise AssertionError("only the output rank may run DiffVAE decode")
+
+ monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+ monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.use_diffusion_decoder = True
+ pipe.vae = ConvVae()
+ pipe.diffusion_decoder = DiffusionDecoder()
+
+ output = pipe._decode_output(
+ latents=torch.ones(1, 1),
+ audio_latents=torch.ones(1, 1),
+ output_type="pt",
+ connector_prompt_embeds=torch.ones(1, 1),
+ generator=None,
+ device=torch.device("cpu"),
+ decode_timestep=0.0,
+ decode_noise_scale=None,
+ prompt_batch_size=1,
+ )
+
+ assert output.output[0].numel() == 0
+ assert output.output[1].numel() == 0
+
+ def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch):
+ import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+ class ConvVae:
+ config = SimpleNamespace(timestep_conditioning=False)
+
+ def is_distributed_enabled(self):
+ raise AssertionError("DiffVAE decode must query the active decoder")
+
+ class DiffusionDecoder:
+ dtype = torch.float32
+
+ def __init__(self):
+ self.decode_calls = 0
+
+ def is_distributed_enabled(self):
+ return True
+
+ def decode(self, *_args, **_kwargs):
+ self.decode_calls += 1
+ return (torch.empty(0, 3, 0, 0, 0),)
+
+ monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+ monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+ pipe = object.__new__(LTX2Pipeline)
+ torch.nn.Module.__init__(pipe)
+ pipe.use_diffusion_decoder = True
+ pipe.vae = ConvVae()
+ pipe.diffusion_decoder = DiffusionDecoder()
+
+ output = pipe._decode_output(
+ latents=torch.ones(1, 1),
+ audio_latents=torch.ones(1, 1),
+ output_type="pt",
+ connector_prompt_embeds=torch.ones(1, 1),
+ generator=None,
+ device=torch.device("cpu"),
+ decode_timestep=0.0,
+ decode_noise_scale=None,
+ prompt_batch_size=1,
+ )
+
+ assert pipe.diffusion_decoder.decode_calls == 1
+ assert output.output[0].numel() == 0
+ assert output.output[1].numel() == 0
+
+ def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+ DistributedLTX2VideoDiffusionDecoderModel,
+ )
+
+ expected = torch.ones(1, 3, 1, 2, 2)
+ seen = {}
+
+ def native_tiled_decode(self, z, generator=None, num_inference_steps=None):
+ seen["z"] = z
+ seen["generator"] = generator
+ seen["num_inference_steps"] = num_inference_steps
+ return expected
+
+ monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode)
+ model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+ torch.nn.Module.__init__(model)
+ model.is_distributed_enabled = lambda: False
+ generator = torch.Generator().manual_seed(17)
+ z = torch.zeros(1, 1, 1, 1, 1)
+
+ result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode(
+ model,
+ z,
+ generator=generator,
+ num_inference_steps=1,
+ )
+
+ assert result is expected
+ assert seen["z"] is z
+ assert seen["generator"] is generator
+ assert seen["num_inference_steps"] == 1
+
+ def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
+ from diffusers.utils.torch_utils import randn_tensor
+
+ from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+ DistributedLTX2VideoDiffusionDecoderModel,
+ )
+
+ model = DistributedLTX2VideoDiffusionDecoderModel(
+ out_channels=1,
+ latent_channels=1,
+ patch_size=1,
+ decoder_head_dim=8,
+ decoder_stage_channels=(8, 8, 8, 8, 8),
+ decoder_stage_depths=(1, 1, 1, 1, 1),
+ decoder_stage_kernels=((1, 1, 1),) * 4,
+ decoder_upsample_strides=((1, 1, 1),) * 4,
+ decoder_upsample_channel_reductions=(1, 1, 1, 1),
+ decoder_stage5_kernel=(1, 1, 1),
+ spatial_compression_ratio=1,
+ temporal_compression_ratio=1,
+ )
+ model.tile_sample_min_num_frames = 2
+ model.tile_sample_stride_num_frames = 1
+ model.tile_sample_min_height = 2
+ model.tile_sample_stride_height = 1
+ model.tile_sample_min_width = 2
+ model.tile_sample_stride_width = 1
+ model.distributed_executor = SimpleNamespace(group=None)
+ model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1)
+ monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None)
+
+ z = torch.zeros(1, 1, 2, 3, 3)
+ distributed_generator = torch.Generator().manual_seed(123)
+ reference_generator = torch.Generator().manual_seed(123)
+ tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1)
+
+ assert grid_spec.grid_shape == (1, 3, 3)
+ assert grid_spec.split_dims == (2, 3, 4)
+ for task in tasks:
+ tile_shape = model._tiled_pixel_shape_from_features(
+ task.tensor,
+ drop_leading_frame=task.drop_leading_frame,
+ crop_trailing_ghost=task.crop_trailing_ghost,
+ )
+ expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype)
+ actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype)
+ torch.testing.assert_close(actual, expected)
+
+ assert torch.equal(distributed_generator.get_state(), reference_generator.get_state())
+
+
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
@@ -113,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
},
output_dtype=torch.float32,
)
- seen = {}
+ seen: dict[str, Any] = {}
def exec_tile(task):
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
diff --git a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py
new file mode 100644
index 000000000..d3d885d57
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py
@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CUDA regression tests for LTX-2 BWE vocoder precision."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+class _BWEConvVocoder(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.conv = torch.nn.Conv1d(
+ 1,
+ 1,
+ kernel_size=1,
+ bias=False,
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ self.bwe_generator = torch.nn.Identity()
+ self.input_dtype = None
+ self.conv_output_dtype = None
+
+ def forward(self, value):
+ self.input_dtype = value.dtype
+ output = self.conv(value)
+ self.conv_output_dtype = output.dtype
+ return output
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32():
+ from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder
+
+ vocoder = _BWEConvVocoder()
+ generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16)
+
+ output = _run_ltx_vocoder(vocoder, generated_mel)
+
+ assert next(vocoder.parameters()).dtype == torch.bfloat16
+ assert vocoder.input_dtype == torch.float32
+ assert vocoder.conv_output_dtype == torch.float32
+ assert output.dtype == torch.bfloat16
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
index 60ed7090f..382ecf9aa 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import json
+import sys
from multiprocessing.reduction import ForkingPickler
from types import SimpleNamespace
-from unittest.mock import Mock
+from typing import Any, TypeVar
+from unittest.mock import Mock, patch
import numpy as np
import pytest
@@ -13,6 +15,14 @@ import torch.nn as nn
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+_ItemT = TypeVar("_ItemT")
+_ResultT = TypeVar("_ResultT")
+
+
+def _append_and_return(items: list[_ItemT], item: _ItemT, result: _ResultT) -> _ResultT:
+ items.append(item)
+ return result
+
def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -28,6 +38,7 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
+ pipeline.text_encoder = object()
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
@@ -46,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
- pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+ decoded_audio = torch.zeros(1)
+ pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio))
sampling = OmniDiffusionSamplingParams(
quality="high",
width=1344,
@@ -75,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
num_inference_steps=50,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
- assert output.output == pipeline.decode.return_value
+ assert output.output[0].shape == (1, 1, 1, 1, 3)
+ assert output.output[0].dtype == torch.uint8
+ assert output.output[1] is decoded_audio
def test_pipeline_import_registry_and_component_discovery():
@@ -100,6 +114,40 @@ def test_pipeline_import_registry_and_component_discovery():
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
+def test_encoder_free_stage_skips_text_encoder_during_dlo_discovery():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.transformer = torch.nn.Linear(1, 1)
+ pipeline.video_vae = torch.nn.Linear(1, 1)
+ pipeline.audio_vae = torch.nn.Linear(1, 1)
+ pipeline.text_encoder = None
+ pipeline._dit_modules = ["transformer"]
+ pipeline._encoder_modules = []
+ pipeline._vae_modules = ["video_vae", "audio_vae"]
+
+ discovered = ModuleDiscovery.discover(pipeline)
+
+ assert discovered.dit_names == ["transformer"]
+ assert discovered.encoder_names == []
+ assert discovered.vae_names == ["video_vae", "audio_vae"]
+
+
+def test_encoder_free_stage_skips_missing_component_during_dlo_registration():
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+ cache = Mock()
+ video_vae = Mock()
+ audio_vae = Mock()
+
+ pipeline_module._register_dlo_component_cache(cache, None, video_vae, audio_vae)
+
+ video_vae.set_omni_component_cache.assert_called_once_with(cache)
+ audio_vae.set_omni_component_cache.assert_called_once_with(cache)
+
+
def _write_partition_index(path, *, partition, tasks):
path.mkdir(parents=True)
(path / "model_index.json").write_text(
@@ -184,6 +232,7 @@ def test_startup_task_rejects_unsupported_value():
def test_combined_task_inference_and_transformer_routing():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.errors import OmniClientError
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
@@ -195,6 +244,9 @@ def test_combined_task_inference_and_transformer_routing():
assert pipeline._resolve_task(None, {"image": object()}) == "fl2va"
assert pipeline._resolve_task(None, {"audio": object()}) == "ref2va"
assert pipeline._resolve_task(None, {"video": object()}) == "ref2va"
+ assert pipeline._resolve_task("fl2va", {"image": object()}, has_turbo_lora=True) == "fl2va"
+ with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+ pipeline._resolve_task("ref2va", {}, has_turbo_lora=True)
pipeline.partition = "ref2va"
pipeline.supported_tasks = frozenset({"ref2va"})
assert pipeline._resolve_task(None, {"image": object()}) == "ref2va"
@@ -259,7 +311,7 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
):
(tmp_path / partition_name / component).mkdir()
- created = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
+ created: dict[str, list[Any]] = {"dit": [], "text_encoder": [], "video_vae": [], "audio_vae": []}
class FakeModule(torch.nn.Module):
def __init__(self, *args, **kwargs):
@@ -277,15 +329,15 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
return create
- tokenizer_calls = []
- processor_calls = []
+ tokenizer_calls: list[Any] = []
+ processor_calls: list[Any] = []
tokenizer_cls = Mock(spec=pipeline_module.Qwen2TokenizerFast)
- tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
- tokenizer_calls.append((args, kwargs)) or object()
+ tokenizer_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+ tokenizer_calls, (args, kwargs), object()
)
processor_cls = Mock(spec=pipeline_module.Qwen3VLProcessor)
- processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: (
- processor_calls.append((args, kwargs)) or object()
+ processor_cls.from_pretrained.side_effect = lambda *args, **kwargs: _append_and_return(
+ processor_calls, (args, kwargs), object()
)
monkeypatch.setattr(pipeline_module, "MiniMaxH3DiTModel", FakeDiT)
monkeypatch.setattr(
@@ -362,8 +414,18 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
if expected_partition == "ref2va":
assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
assert [source.model_or_path for source in pipeline.weights_sources] == [
- str(tmp_path / partition_name) for partition_name in source_partitions
+ *(str(tmp_path / partition_name) for partition_name in source_partitions),
+ str(component_path),
]
+ assert [source.subfolder for source in pipeline.weights_sources] == [
+ *("transformer" for _ in source_partitions),
+ "text_encoder",
+ ]
+ expected_prefixes = ["transformer."]
+ if expected_dits == 2:
+ expected_prefixes.append("transformers_ref.")
+ expected_prefixes.append("text_encoder.")
+ assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
expected_patterns = (
pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
if expected_partition == "combined"
@@ -434,18 +496,41 @@ def test_joint_postprocess_is_multiprocessing_picklable():
postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
- video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
+ # decode() hands over quantised frames already laid out as (B, T, H, W, C).
+ video = torch.arange(2 * 2 * 4 * 5 * 3, dtype=torch.uint8).reshape(2, 2, 4, 5, 3)
audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
result = postprocess((video, audio), output_type="np")
assert isinstance(result["video"], list)
+ assert result["video"][0].dtype == np.uint8
assert result["video"][0].shape == (2, 4, 5, 3)
+ np.testing.assert_array_equal(result["video"][0], video[0].numpy())
np.testing.assert_array_equal(result["audio"], audio.numpy())
assert result["audio_sample_rate"] == 32000
assert result["fps"] == 24
+def test_joint_video_output_is_quantized_before_transfer():
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _minimax_h3_post_process,
+ _prepare_minimax_h3_video_output,
+ )
+
+ video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
+ expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
+ prepared = _prepare_minimax_h3_video_output(video)
+
+ assert prepared.dtype == torch.uint8
+ assert prepared.is_contiguous()
+ assert prepared.shape == (2, 2, 4, 5, 3)
+ torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
+
+ result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
+ assert result["video"][0].dtype == np.uint8
+ np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
+
+
def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -513,7 +598,7 @@ def test_base_schedule_overrides_the_uniform_sigma_positions():
@pytest.mark.parametrize(
"base_schedule",
[
- [0.9, 0.5, 0.0],
+ [1.5, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.5, 0.5, 0.0],
[1.0],
@@ -551,6 +636,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
+ pipeline.text_encoder = Mock()
pipeline._base_schedule_by_partition = schedules
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = SimpleNamespace(cache_dit=None)
@@ -567,7 +653,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
- pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+ pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1)))
return pipeline
@@ -615,7 +701,7 @@ def test_partially_constructed_pipeline_falls_back_to_the_uniform_schedule():
def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
- diffuse_calls = []
+ diffuse_calls: list[dict[str, Any]] = []
base_schedule = [1.0, 0.7, 0.4, 0.15, 0.0]
pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": base_schedule, "ref2va": None})
@@ -632,7 +718,7 @@ def test_distilled_forward_reports_denoising_steps_not_sigma_boundaries():
def test_distilled_forward_accepts_the_matching_explicit_step_count():
- diffuse_calls = []
+ diffuse_calls: list[dict[str, Any]] = []
pipeline = _distilled_pipeline(diffuse_calls, {"fl2va": [1.0, 0.7, 0.4, 0.15, 0.0], "ref2va": None})
pipeline.forward(_t2va_batch(num_inference_steps=4))
@@ -702,8 +788,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
)
assert output.shape == q.shape
- assert attention.attention.metadata.attn_mask is None
- assert attention.attention.metadata.extra["valid_kv_length"] == 5
+ metadata = attention.attention.metadata
+ assert metadata is not None
+ assert metadata.attn_mask is None
+ assert metadata.extra["valid_kv_length"] == 5
def test_packed_attention_skips_mask_for_packed_mask_free_backend():
@@ -743,9 +831,17 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
packed_total=8,
)
- assert attention.attention.metadata.attn_mask is None
- assert attention.attention.metadata.extra["valid_kv_length"] == 5
- assert attention.attention.metadata.extra["npu_attn_varlen"] is True
+ metadata = attention.attention.metadata
+ assert metadata is not None
+ assert metadata.attn_mask is None
+ assert metadata.extra["valid_kv_length"] == 5
+ assert metadata.extra["npu_attn_varlen"] is True
+ packed_padding = metadata.packed_padding
+ assert packed_padding is not None
+ assert packed_padding.q_length == 5
+ assert packed_padding.kv_length == 5
+ assert packed_padding.cu_seqlens_q.tolist() == [0, 5]
+ assert packed_padding.cu_seqlens_k.tolist() == [0, 5]
def test_packed_attention_keeps_padding_mask_for_other_backends():
@@ -785,12 +881,90 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
packed_total=8,
)
+ metadata = attention.attention.metadata
+ assert metadata is not None
assert torch.equal(
- attention.attention.metadata.attn_mask,
+ metadata.attn_mask,
torch.tensor([[True, True, True, True, True, False, False, False]]),
)
+def _fake_packed_attention(
+ backend_name: str,
+ *,
+ prefix_kv_slicing: bool = False,
+ supports_multi_doc: bool = False,
+):
+ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+ MiniMaxH3Attention,
+ )
+
+ class FakeBackend:
+ supports_prefix_kv_slicing = prefix_kv_slicing
+
+ @staticmethod
+ def get_name() -> str:
+ return backend_name
+
+ @classmethod
+ def supports_multi_doc_packed_varlen(cls) -> bool:
+ return supports_multi_doc
+
+ class FakeAttention(torch.nn.Module):
+ attn_backend = FakeBackend
+
+ def __init__(self):
+ super().__init__()
+ self.metadata = None
+
+ def forward(self, query, key, value, metadata):
+ self.metadata = metadata
+ return query
+
+ attention = object.__new__(MiniMaxH3Attention)
+ torch.nn.Module.__init__(attention)
+ attention.attention = FakeAttention()
+ return attention
+
+
+def test_packed_attention_drops_prefix_semantics_for_a_co_batched_layout():
+ """A batch's valid rows are block-diagonal, so no prefix can describe them."""
+ attention = _fake_packed_attention("FLASH_ATTN", supports_multi_doc=True)
+ q = torch.randn(12, 2, 4)
+
+ attention._run_packed_attention(
+ q,
+ q,
+ q,
+ # Two requests: rows + padding tail each.
+ cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
+ max_seqlen=5,
+ packed_total=12,
+ num_requests=2,
+ )
+
+ assert attention.attention.metadata.attn_mask is None
+ assert attention.attention.metadata.extra["valid_kv_length"] == 12
+
+
+@pytest.mark.parametrize("backend_name", ["TORCH_SDPA", "FLASH_ATTN"])
+def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_name):
+ """Backend names alone must not allow cross-request attention."""
+ attention = _fake_packed_attention(backend_name, supports_multi_doc=False)
+ q = torch.randn(12, 2, 4)
+
+ with pytest.raises(ValueError, match=backend_name):
+ attention._run_packed_attention(
+ q,
+ q,
+ q,
+ cu_seqlens=torch.tensor([0, 5, 6, 11, 12], dtype=torch.int32),
+ max_seqlen=5,
+ packed_total=12,
+ num_requests=2,
+ )
+
+
def test_reference_image_resize_contract():
from PIL import Image
@@ -852,6 +1026,80 @@ def test_minimax_h3_advertises_the_official_ref2va_image_limit():
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
+def test_text_attention_routes_local_gqa_heads_through_sdpa_helper(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+ class FakeEncoderGroup:
+ rank_in_group = 0
+ world_size = 2
+
+ def all_reduce(self, tensor):
+ del tensor
+
+ config = SimpleNamespace(
+ hidden_size=8,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ head_dim=2,
+ rms_norm_eps=1e-6,
+ )
+ with (
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+ ):
+ attention = encoder_module.MiniMaxH3Qwen3VLTextAttention(FakeEncoderGroup(), config, torch.float32)
+ # This CPU contract test only needs projection shapes and head reshaping.
+ # Avoid dispatching platform-specific norm and linear kernels (for example,
+ # rocm_unquantized_gemm when the test suite is collected on ROCm).
+ attention.q_norm = nn.Identity()
+ attention.k_norm = nn.Identity()
+ with torch.no_grad():
+ attention.qkv_proj.weight.zero_()
+ attention.o_proj = nn.Linear(
+ attention.qkv_proj.local_num_heads * attention.head_dim,
+ config.hidden_size,
+ bias=False,
+ )
+ attn_call: dict[str, object] = {}
+
+ def fake_attention(query, key, value):
+ attn_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape)
+ return query
+
+ monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", fake_attention)
+ hidden_states = torch.randn(1, 3, config.hidden_size)
+ cos = torch.ones(1, 3, config.head_dim)
+ sin = torch.zeros_like(cos)
+
+ output = attention(hidden_states, (cos, sin))
+
+ assert attn_call["query_shape"] == (1, config.num_attention_heads // 2, 3, config.head_dim)
+ assert attn_call["key_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
+ assert attn_call["value_shape"] == (1, config.num_key_value_heads // 2, 3, config.head_dim)
+ assert output.shape == hidden_states.shape
+
+
+def test_text_attention_sdpa_helper_preserves_expanded_kv_fallback(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+ sdpa_call: dict[str, object] = {}
+
+ def fake_sdpa(query, key, value, **kwargs):
+ sdpa_call.update(query_shape=query.shape, key_shape=key.shape, value_shape=value.shape, kwargs=kwargs)
+ return query
+
+ monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+ output = encoder_module._scaled_dot_product_attention(
+ torch.randn(1, 4, 3, 8),
+ torch.randn(1, 2, 3, 8),
+ torch.randn(1, 2, 3, 8),
+ )
+
+ assert sdpa_call["query_shape"] == sdpa_call["key_shape"] == sdpa_call["value_shape"] == (1, 4, 3, 8)
+ assert sdpa_call["kwargs"] == {"dropout_p": 0.0, "is_causal": True}
+ assert output.shape == (1, 4, 3, 8)
+
+
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
from vllm_omni.diffusion.models.minimax_h3.encoder import (
MiniMaxH3Qwen3VLEncoder,
@@ -909,7 +1157,7 @@ def test_encoder_forward_forwards_video_inputs():
def test_reference_video_shape_uses_h3_adapt_shape_policy():
- from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+ from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_reference_video_shape,
)
@@ -935,6 +1183,66 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
assert list(encoder.named_parameters()) == []
+def test_global_quant_config_is_shared_by_dit_and_encoder():
+ from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+ global_config = build_quant_config("fp8")
+
+ assert resolve_component_quant_config(global_config, "transformer") is global_config
+ assert resolve_component_quant_config(global_config, "text_encoder") is global_config
+
+
+def test_minimax_h3_fp8_component_selection():
+ from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+ dit_only = build_quant_config({"transformer": {"method": "fp8"}})
+ encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}})
+
+ assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8"
+ assert resolve_component_quant_config(dit_only, "text_encoder") is None
+ assert resolve_component_quant_config(encoder_only, "transformer") is None
+ assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8"
+
+
+def test_text_encoder_rejects_serialized_fp8():
+ from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config
+
+ config = SimpleNamespace(
+ get_name=lambda: "fp8",
+ is_checkpoint_fp8_serialized=True,
+ )
+ with pytest.raises(ValueError, match="online FP8 only"):
+ _validate_encoder_quant_config(config)
+
+
+def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module
+
+ group = SimpleNamespace(rank_in_group=0, world_size=1)
+ method = UnquantizedLinearMethod()
+ quant_config = Mock()
+ quant_config.get_quant_method.return_value = method
+ prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj"
+
+ with (
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+ ):
+ linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+ group,
+ input_size=4,
+ intermediate_size=8,
+ dtype=torch.bfloat16,
+ quant_config=quant_config,
+ prefix=prefix,
+ )
+
+ assert linear.quant_method is method
+ quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
+
+
def test_no_offload_keeps_text_encoder_resident():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
@@ -967,6 +1275,7 @@ def test_model_offload_uses_hooked_text_encoder_call():
)
expected = torch.ones(2, 3)
pipeline.text_encoder = Mock(return_value=expected)
+ pipeline._model_cpu_offload_modules = [pipeline.text_encoder]
input_ids = torch.tensor([1, 2])
vision_kwargs = {"pixel_values": torch.ones(1, 4)}
@@ -977,15 +1286,21 @@ def test_model_offload_uses_hooked_text_encoder_call():
pipeline.text_encoder.load_to_device.assert_not_called()
-def test_layerwise_offload_releases_text_encoder():
+@pytest.mark.parametrize(
+ "offload_flag",
+ ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
+)
+def test_layerwise_offload_releases_text_encoder(offload_flag):
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.od_config = SimpleNamespace(
enable_cpu_offload=False,
- enable_layerwise_offload=True,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
)
+ setattr(pipeline.od_config, offload_flag, True)
pipeline.text_encoder = Mock()
expected = torch.ones(2, 3)
pipeline.text_encoder.encode_ids.return_value = expected
@@ -1324,7 +1639,7 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
lambda path: f"frames:{path}",
)
- rows, shapes = pipeline._encode_video_conditions(None, count=2)
+ rows, shapes = pipeline._encode_video_conditions_resident(None, count=2)
assert pipeline.video_vae.calls == [
"frames:video-1.mp4",
@@ -1406,12 +1721,12 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
"audio_codecs": (),
"file_size": 1024,
}
- transcode_calls = []
+ transcode_calls: list[Any] = []
monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
monkeypatch.setattr(
reference_video_module,
"_transcode_reference_video",
- lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
+ lambda source, **kwargs: _append_and_return(transcode_calls, (source, kwargs), "prepared.mp4"),
)
prepared = reference_video_module.prepare_reference_videos(
@@ -1428,34 +1743,305 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
assert transcode_calls[0][1]["target_frame_count"] == 209
-def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch):
+def test_ref2va_transcode_zero_frame_count_keeps_video_stream(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
- decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3)
- load_calls = []
+ commands: list[list[str]] = []
+ monkeypatch.setattr(
+ reference_video_module.subprocess,
+ "run",
+ lambda command, **_kwargs: commands.append(command),
+ )
+
+ reference_video_module._transcode_reference_video(
+ "input.mp4",
+ target_width=768,
+ target_height=448,
+ target_frame_count=0,
+ workdir=str(tmp_path),
+ duration_seconds=4.0,
+ )
+
+ assert "-frames:v" not in commands[0]
+
+
+def test_prepared_reference_video_descriptor_round_trip(tmp_path):
+ from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+ deserialize_prepared_reference_videos,
+ serialize_prepared_reference_videos,
+ )
+
+ item = {
+ "original_path": str(tmp_path / "source.mp4"),
+ "prepared_path": str(tmp_path / "prepared.mp4"),
+ "input_has_audio": True,
+ "width": 1344,
+ "height": 768,
+ "start_time_seconds": 0.0,
+ "duration_seconds": 5.2,
+ "audio_duration_seconds": 5.2,
+ }
+
+ descriptor = serialize_prepared_reference_videos([item], str(tmp_path))
+
+ assert deserialize_prepared_reference_videos(descriptor) == (str(tmp_path), [item])
+
+
+def test_prepared_reference_video_descriptor_rejects_unknown_fields():
+ from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+ deserialize_prepared_reference_videos,
+ )
+
+ with pytest.raises(ValueError, match="invalid MiniMax H3 prepared-reference-video descriptor"):
+ deserialize_prepared_reference_videos('{"artifact_dir":"/tmp","videos":[{"path":"x"}]}')
+
+
+def test_stage_one_reuses_existing_prepared_reference_video(tmp_path):
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _reuse_prepared_reference_videos,
+ )
+
+ prepared_path = tmp_path / "prepared.mp4"
+ prepared_path.touch()
+ prepared = [{"prepared_path": str(prepared_path)}]
+
+ assert _reuse_prepared_reference_videos(prepared, expected_count=1) is prepared
+
+
+def test_stage_one_rejects_invalid_prepared_reference_video(tmp_path):
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _reuse_prepared_reference_videos,
+ )
+
+ with pytest.raises(ValueError, match="count does not match"):
+ _reuse_prepared_reference_videos([], expected_count=1)
+ with pytest.raises(ValueError, match="prepared reference video is unavailable"):
+ _reuse_prepared_reference_videos(
+ [{"prepared_path": str(tmp_path / "missing.mp4")}],
+ expected_count=1,
+ )
+
+
+def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+ decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
+ decode_calls: list[Any] = []
monkeypatch.setattr(
reference_video_module,
"_probe_video",
- lambda _path: {"frame_count": len(decoded)},
+ lambda _path: {
+ "frame_count": 25,
+ "width": 2,
+ "height": 2,
+ },
)
monkeypatch.setattr(
reference_video_module,
- "load_video_frames",
- lambda path: load_calls.append(path) or decoded,
+ "_decode_video_frames_ffmpeg",
+ lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), decoded),
)
sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")
- assert load_calls == ["prepared.mp4"]
- for frame, expected in zip(
- sampled["frames"],
- (decoded[0], decoded[12], decoded[24]),
- strict=True,
- ):
+ assert decode_calls == [
+ (
+ "prepared.mp4",
+ {
+ "frame_count": 25,
+ "indices": [0, 12, 24],
+ "width": 2,
+ "height": 2,
+ },
+ )
+ ]
+ for frame, expected in zip(sampled["frames"], decoded, strict=True):
np.testing.assert_array_equal(frame, expected)
assert sampled["block_timestamps"] == [0.25, 1.0]
+@pytest.mark.parametrize(
+ ("indices", "frame_count"),
+ [([0, 12], 25), (None, 2)],
+)
+def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+ expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3)
+ observed = []
+
+ def fake_run(command, **kwargs):
+ observed.append((command, kwargs))
+ return SimpleNamespace(stdout=expected.tobytes())
+
+ monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run)
+
+ actual = reference_video_module._decode_video_frames_ffmpeg(
+ "prepared.mp4",
+ frame_count=frame_count,
+ indices=indices,
+ width=3,
+ height=2,
+ )
+
+ np.testing.assert_array_equal(actual, expected)
+ command, kwargs = observed[0]
+ assert command[0] == "ffmpeg"
+ assert command[command.index("-frames:v") + 1] == "2"
+ if indices is None:
+ assert "-vf" not in command
+ else:
+ assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)"
+ assert command[-1] == "pipe:1"
+ assert kwargs == {"check": True, "capture_output": True}
+
+
+def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+ expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
+ decode_calls: list[Any] = []
+ monkeypatch.setitem(sys.modules, "decord", None)
+ monkeypatch.setattr(
+ reference_video_module,
+ "_probe_video",
+ lambda _path: {
+ "frame_count": 2,
+ "width": 3,
+ "height": 2,
+ },
+ )
+ monkeypatch.setattr(
+ reference_video_module,
+ "_decode_video_frames_ffmpeg",
+ lambda path, **kwargs: _append_and_return(decode_calls, (path, kwargs), expected),
+ )
+
+ actual = reference_video_module.load_video_frames("prepared.mp4")
+
+ assert actual is expected
+ assert decode_calls == [
+ (
+ "prepared.mp4",
+ {
+ "frame_count": 2,
+ "width": 3,
+ "height": 2,
+ },
+ )
+ ]
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "message"),
+ [
+ ({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"),
+ ({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"),
+ (
+ {"frame_count": 1, "indices": [], "width": 3, "height": 2},
+ "invalid frame indices",
+ ),
+ ],
+)
+def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+ from vllm_omni.errors import OmniClientError
+
+ with pytest.raises(OmniClientError, match=message):
+ reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs)
+
+
+def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+ calls: list[Any] = []
+ probe = {
+ "streams": [
+ {
+ "codec_type": "video",
+ "width": 4,
+ "height": 2,
+ "r_frame_rate": "24/1",
+ "nb_frames": "3",
+ "duration": "0.125",
+ "sample_aspect_ratio": "1:1",
+ }
+ ],
+ "format": {"duration": "0.125", "size": "123"},
+ }
+ monkeypatch.setattr(
+ reference_video_module.subprocess,
+ "run",
+ lambda command, **kwargs: _append_and_return(
+ calls,
+ (command, kwargs),
+ SimpleNamespace(stdout=json.dumps(probe)),
+ ),
+ )
+
+ metadata = reference_video_module._probe_video("prepared.mp4")
+
+ assert metadata["frame_count"] == 3
+ assert len(calls) == 1
+ assert "-count_frames" not in calls[0][0]
+
+
+def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+
+ calls: list[Any] = []
+ stream = {
+ "codec_type": "video",
+ "width": 4,
+ "height": 2,
+ "r_frame_rate": "24/1",
+ "duration": "0.125",
+ "sample_aspect_ratio": "1:1",
+ }
+ first_probe = {
+ "streams": [stream],
+ "format": {"duration": "0.125", "size": "123"},
+ }
+ second_probe = {
+ "streams": [{**stream, "nb_read_frames": "3"}],
+ "format": {"duration": "0.125", "size": "123"},
+ }
+ monkeypatch.setattr(
+ reference_video_module.subprocess,
+ "run",
+ lambda command, **kwargs: _append_and_return(
+ calls,
+ (command, kwargs),
+ SimpleNamespace(stdout=json.dumps(first_probe if not calls else second_probe)),
+ ),
+ )
+
+ metadata = reference_video_module._probe_video("prepared.mp4")
+
+ assert metadata["frame_count"] == 3
+ assert len(calls) == 2
+ assert "-count_frames" in calls[1][0]
+
+
+def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
+ from vllm_omni.errors import OmniClientError
+
+ monkeypatch.setattr(
+ reference_video_module.subprocess,
+ "run",
+ lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"),
+ )
+
+ with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"):
+ reference_video_module._decode_video_frames_ffmpeg(
+ "prepared.mp4",
+ frame_count=1,
+ width=3,
+ height=2,
+ )
+
+
def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
@@ -1487,12 +2073,16 @@ def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_p
"file_size": 1024,
},
}
- transcode_calls = []
+ transcode_calls: list[Any] = []
monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
monkeypatch.setattr(
reference_video_module,
"_transcode_reference_video",
- lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
+ lambda source, **kwargs: _append_and_return(
+ transcode_calls,
+ (source, kwargs),
+ f"{source}.prepared.mp4",
+ ),
)
prepared = reference_video_module.prepare_reference_videos(
@@ -1671,7 +2261,7 @@ def test_g4_reference_image_file_format_and_size_contract(tmp_path):
def test_g4_standalone_audio_duration_and_total_duration_contract():
- from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+ from vllm_omni.model_executor.models.minimax_h3.reference_video import (
validate_reference_audio_waveforms,
)
@@ -1702,7 +2292,7 @@ def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch):
waveform = torch.zeros(1, 10)
with pytest.raises(ValueError, match="max_duration_seconds must be positive"):
- pipeline._encode_audio_conditions(
+ pipeline._encode_audio_conditions_resident(
[(waveform, 10)],
max_duration_seconds=0,
)
@@ -1723,7 +2313,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
pipeline.audio_vae = FakeAudioVAE()
waveform = torch.arange(40, dtype=torch.float32).reshape(1, 40)
- encoded, lengths = pipeline._encode_audio_conditions(
+ encoded, lengths = pipeline._encode_audio_conditions_resident(
[(waveform, 10)],
max_duration_seconds=2.5,
)
@@ -1746,7 +2336,7 @@ def test_ref2va_standalone_audio_condition_is_bounded_to_output_duration():
],
)
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
- from vllm_omni.diffusion.models.minimax_h3.reference_video import (
+ from vllm_omni.model_executor.models.minimax_h3.reference_video import (
_validate_reference_video_metadata,
)
@@ -1802,21 +2392,25 @@ def _one_layer_text_encoder():
group = _SingleRankEncoderGroup(0)
layer = nn.Module()
layer.self_attn = nn.Module()
- layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
- group,
- hidden_size=_ENCODER_HIDDEN,
- num_heads=_ENCODER_NUM_HEADS,
- num_kv_heads=_ENCODER_NUM_KV_HEADS,
- head_dim=_ENCODER_HEAD_DIM,
- dtype=torch.float32,
- )
layer.mlp = nn.Module()
- layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
- group,
- input_size=_ENCODER_HIDDEN,
- intermediate_size=_ENCODER_INTERMEDIATE,
- dtype=torch.float32,
- )
+ with (
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+ ):
+ layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
+ group,
+ hidden_size=_ENCODER_HIDDEN,
+ num_heads=_ENCODER_NUM_HEADS,
+ num_kv_heads=_ENCODER_NUM_KV_HEADS,
+ head_dim=_ENCODER_HEAD_DIM,
+ dtype=torch.float32,
+ )
+ layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+ group,
+ input_size=_ENCODER_HIDDEN,
+ intermediate_size=_ENCODER_INTERMEDIATE,
+ dtype=torch.float32,
+ )
layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)
encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
@@ -1849,6 +2443,7 @@ def _write_text_encoder_checkpoint(path, *, omit=()):
safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
(path / "model.safetensors.index.json").write_text(json.dumps(index))
+ return tensors
@pytest.mark.parametrize(
@@ -1870,10 +2465,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
tmp_path, omitted_sources, expected_detail, unreported_target
):
encoder = _one_layer_text_encoder()
- _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
+ weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
with pytest.raises(RuntimeError) as excinfo:
- encoder._load_weights(str(tmp_path))
+ encoder.load_weights(weights.items())
message = str(excinfo.value)
assert expected_detail in message
@@ -1882,17 +2477,17 @@ def test_encoder_load_reports_missing_fused_source_shards(
def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
encoder = _one_layer_text_encoder()
- _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
+ weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
- encoder._load_weights(str(tmp_path))
+ encoder.load_weights(weights.items())
def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
encoder = _one_layer_text_encoder()
- _write_text_encoder_checkpoint(tmp_path)
+ weights = _write_text_encoder_checkpoint(tmp_path)
- encoder._load_weights(str(tmp_path))
+ encoder.load_weights(weights.items())
params = dict(encoder.named_parameters())
q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py
new file mode 100644
index 000000000..b0ebf1038
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_dlo_lifecycle.py
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _OffloadAbort(BaseException):
+ pass
+
+
+def test_encoder_non_block_children_use_one_shared_snapshot_stager(monkeypatch, mocker):
+ from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+
+ class VisionStack(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.patch_embed = torch.nn.Linear(2, 2)
+ self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+ class TextStack(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.embed_tokens = torch.nn.Embedding(4, 2)
+ self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)])
+
+ encoder = object.__new__(encoder_module.MiniMaxH3Qwen3VLEncoder)
+ torch.nn.Module.__init__(encoder)
+ encoder.device_target = torch.device("cpu")
+ encoder.vision = VisionStack()
+ encoder.text_model = TextStack()
+ hook = mocker.Mock()
+ hook.pin_memory = False
+ encoder._omni_layerwise_hooks = [hook]
+ encoder._omni_layerwise_enabled = True
+ cache = mocker.Mock()
+ encoder.set_omni_component_cache(cache)
+ stager = mocker.Mock()
+ stager_cls = mocker.Mock(return_value=stager)
+ monkeypatch.setattr(encoder_module, "PinnedModuleStager", stager_cls)
+
+ encoder.load_to_device()
+ encoder.offload_to_cpu()
+
+ stager_cls.assert_called_once_with(
+ [encoder.vision.patch_embed, encoder.text_model.embed_tokens],
+ torch.device("cpu"),
+ pin_memory=False,
+ cache_retention=cache,
+ )
+ stager.load.assert_called_once_with()
+ stager.offload.assert_called_once_with()
+ hook.offload_layer.assert_called_once_with()
+
+
+def test_manual_component_failure_forces_retained_cache_release(mocker):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.od_config = mocker.Mock()
+ pipeline.od_config.enable_layerwise_offload = False
+ pipeline.od_config.enable_distributed_layerwise_offload = True
+ pipeline._model_cpu_offload_modules = []
+ pipeline._dlo_component_cache = mocker.Mock()
+ component = mocker.Mock()
+ component.offload_to_cpu.side_effect = _OffloadAbort("offload failed")
+
+ with pytest.raises(RuntimeError, match="component failed"):
+ with pipeline._component_on_device(component):
+ raise RuntimeError("component failed")
+
+ component.load_to_device.assert_called_once_with()
+ component.offload_to_cpu.assert_called_once_with()
+ pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_manual_component_offload_failure_forces_retained_cache_release(mocker):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.od_config = mocker.Mock()
+ pipeline.od_config.enable_layerwise_offload = False
+ pipeline.od_config.enable_distributed_layerwise_offload = True
+ pipeline._model_cpu_offload_modules = []
+ pipeline._dlo_component_cache = mocker.Mock()
+ component = mocker.Mock()
+ component.offload_to_cpu.side_effect = [None, _OffloadAbort("offload failed")]
+
+ with pipeline._component_on_device(component):
+ pass
+ with pytest.raises(_OffloadAbort, match="offload failed"):
+ with pipeline._component_on_device(component):
+ pass
+
+ assert component.load_to_device.call_count == 2
+ assert component.offload_to_cpu.call_count == 2
+ pipeline._dlo_component_cache.release_if_needed.assert_called_once_with(force=True)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py
new file mode 100644
index 000000000..b169d918e
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_sdpa.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections.abc import Iterator
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLEncoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def restore_cudnn_sdp_state() -> Iterator[None]:
+ original = torch.backends.cuda.cudnn_sdp_enabled()
+ try:
+ yield
+ finally:
+ torch.backends.cuda.enable_cudnn_sdp(original)
+
+
+def _encoder_without_weights() -> MiniMaxH3Qwen3VLEncoder:
+ encoder = MiniMaxH3Qwen3VLEncoder("unused", device=torch.device("cpu"), load_model=False)
+ encoder.text_model = nn.Linear(1, 1)
+ return encoder
+
+
+@pytest.mark.parametrize("initial_state", [False, True])
+def test_encode_ids_restores_cudnn_sdp_after_success(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
+ encoder = _encoder_without_weights()
+ expected = torch.tensor([1.0])
+
+ def encode(*args, **kwargs):
+ assert torch.backends.cuda.cudnn_sdp_enabled()
+ return expected
+
+ monkeypatch.setattr(encoder, "_encode", encode)
+ torch.backends.cuda.enable_cudnn_sdp(initial_state)
+
+ actual = encoder.encode_ids(torch.tensor([1]))
+
+ torch.testing.assert_close(actual, expected)
+ assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
+
+
+@pytest.mark.parametrize("initial_state", [False, True])
+def test_encode_ids_restores_cudnn_sdp_after_failure(monkeypatch: pytest.MonkeyPatch, initial_state: bool) -> None:
+ encoder = _encoder_without_weights()
+
+ def fail(*args, **kwargs):
+ assert torch.backends.cuda.cudnn_sdp_enabled()
+ raise RuntimeError("encode failed")
+
+ monkeypatch.setattr(encoder, "_encode", fail)
+ torch.backends.cuda.enable_cudnn_sdp(initial_state)
+
+ with pytest.raises(RuntimeError, match="encode failed"):
+ encoder.encode_ids(torch.tensor([1]))
+
+ assert torch.backends.cuda.cudnn_sdp_enabled() is initial_state
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
new file mode 100644
index 000000000..8945b9b47
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py
@@ -0,0 +1,498 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.minimax_h3.fasth3 import (
+ FASTH3_BASE_MODEL,
+ FASTH3_BASE_SCHEDULE,
+ FASTH3_DENOISE_STEPS,
+ FASTH3_FORMAT,
+ FastH3AdapterError,
+ FastH3WeightFusion,
+ resolve_fasth3_fusion,
+)
+from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
+from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
+from vllm_omni.errors import OmniClientError
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_RANK = 2
+_HIDDEN = 4
+_HEAD_DIM = 2
+_HEADS = 3
+_INNER = _HEAD_DIM * _HEADS # attention inner size
+_FFN = 5
+
+
+def _factors(out_dim: int, in_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+ """A rank-2 pair whose product is reproducible and non-symmetric."""
+ a = torch.arange(_RANK * in_dim, dtype=torch.float32).reshape(_RANK, in_dim) / (in_dim * _RANK)
+ b = torch.arange(out_dim * _RANK, dtype=torch.float32).reshape(out_dim, _RANK) / (out_dim * _RANK)
+ return a, b
+
+
+def _write_adapter(
+ path,
+ *,
+ tensors=None,
+ drop: str | None = None,
+ blocks: int = 1,
+ metadata: dict[str, str | None] | None = None,
+) -> None:
+ """Write an artifact in the published ``fastvideo-lora-v2`` shape."""
+ payload: dict[str, torch.Tensor] = {}
+ for block in range(blocks):
+ for suffix, (out_dim, in_dim) in {
+ "attn.to_q": (_INNER, _HIDDEN),
+ "attn.to_k": (_INNER, _HIDDEN),
+ "attn.to_v": (_INNER, _HIDDEN),
+ "attn.to_out.0": (_HIDDEN, _INNER),
+ "ff.net.0.proj": (2 * _FFN, _HIDDEN),
+ "ff.net.2": (_HIDDEN, _FFN),
+ }.items():
+ a, b = _factors(out_dim, in_dim)
+ payload[f"transformer_blocks.{block}.{suffix}.lora_A.weight"] = a
+ payload[f"transformer_blocks.{block}.{suffix}.lora_B.weight"] = b
+ payload.update(tensors or {})
+ if drop is not None:
+ payload.pop(drop, None)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ # The published writer records what it emitted; mirror it so the fixtures
+ # are as self-describing as the real artifact.
+ declared = {
+ "format": FASTH3_FORMAT,
+ "finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1",
+ "base_model": FASTH3_BASE_MODEL,
+ "rank": str(_RANK),
+ "low_rank_tensors": str(sum(1 for key in payload if key.endswith((".lora_A.weight", ".lora_B.weight")))),
+ "diff_tensors": str(sum(1 for key in payload if key.endswith((".diff", ".diff_b")))),
+ "set_weight_tensors": str(sum(1 for key in payload if key.endswith(".set_weight"))),
+ }
+ # safetensors metadata is string-only, so a ``None`` override drops the key.
+ declared.update(metadata or {})
+ save_file(payload, str(path), metadata={key: value for key, value in declared.items() if value is not None})
+
+
+def _claim(path, **kwargs) -> FastH3WeightFusion | None:
+ """Whatever from_path decides, including declining the artifact."""
+ kwargs.setdefault("head_dim", _HEAD_DIM)
+ kwargs.setdefault("num_blocks", 1)
+ kwargs.setdefault("num_refiner_blocks", 0)
+ return FastH3WeightFusion.from_path(path, **kwargs)
+
+
+def _load(path, **kwargs) -> FastH3WeightFusion:
+ """The same, for the tests that only make sense on a claimed adapter."""
+ fusion = _claim(path, **kwargs)
+ assert fusion is not None, f"{path} was not claimed as a FastH3 adapter"
+ return fusion
+
+
+def test_only_an_artifact_carrying_the_release_identity_is_claimed(tmp_path):
+ plain = tmp_path / "peft" / "adapter_model.safetensors"
+ plain.parent.mkdir(parents=True)
+ save_file({"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(plain))
+ # No fastvideo-lora-v2 metadata: this is somebody else's LoRA and has to
+ # stay on the dynamic route rather than being fused.
+ assert _claim(plain.parent) is None
+ assert _claim(tmp_path / "missing") is None
+
+ claimed = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(claimed)
+ assert _claim(claimed.parent) is not None
+
+
+@pytest.mark.parametrize(
+ "identity",
+ [
+ # An ordinary H3 adapter out of FastVideo's own extraction tools.
+ {"finetuned_model": "someone/minimax-h3-style-lora"},
+ # Somebody else's adapter that merely names the student it imitates.
+ {"finetuned_model": "someone/fasth3-style-lora-for-h3"},
+ {"finetuned_model": ""},
+ {"finetuned_model": None},
+ # A FastH3 name over a different base model.
+ {"base_model": "Wan-AI/Wan2.2-TI2V-5B"},
+ ],
+)
+def test_a_generic_fastvideo_h3_adapter_stays_on_the_dynamic_route(identity, tmp_path):
+ """The fixture edits every block, so only the identity separates it from FastH3."""
+ path = tmp_path / "generic" / "adapter_model.safetensors"
+ _write_adapter(path, metadata=identity)
+
+ assert _claim(path.parent) is None
+
+
+def test_a_claimed_artifact_must_declare_every_tensor_count(tmp_path):
+ """Declining is for other people's adapters; a FastH3 file is held to its word."""
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, metadata={"diff_tensors": None})
+
+ with pytest.raises(FastH3AdapterError, match="omits diff_tensors"):
+ _claim(path.parent)
+
+
+def test_the_published_bundle_root_is_refused_rather_than_guessed(tmp_path):
+ root = tmp_path / "FastVideo-FastH3-4-step-Preview-v1-LoRA"
+ for slug in ("dense-datafree", "vsa-datafree"):
+ _write_adapter(root / slug / "adapter_model.safetensors")
+ (root / "adapter_manifest.json").write_text(json.dumps({"schema_version": "fasth3-lora-bundle-v1"}))
+
+ with pytest.raises(FastH3AdapterError, match="point --lora-path at one variant"):
+ _load(root)
+ # One variant inside it is unambiguous.
+ assert _claim(root / "dense-datafree") is not None
+
+
+def test_a_multi_shard_non_fasth3_adapter_stays_on_the_dynamic_route(tmp_path):
+ # A plain PEFT LoRA saved as several shards must not hard-fail startup with a
+ # FastH3-specific message; it has to fall through to the dynamic LoRA route.
+ directory = tmp_path / "peft-sharded"
+ directory.mkdir()
+ for shard in ("adapter_model-00001-of-00002.safetensors", "adapter_model-00002-of-00002.safetensors"):
+ save_file(
+ {"transformer_blocks.0.attn.to_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))}, str(directory / shard)
+ )
+
+ assert _claim(directory) is None
+
+
+def test_low_rank_factors_reach_the_fused_projections(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path)
+ fusion = _load(path.parent)
+
+ qkv = torch.zeros((3 * _INNER, _HIDDEN), dtype=torch.float32)
+ fused = fusion.fuse("blocks.0.attn.qkv_proj.weight", qkv).cpu()
+ # The checkpoint stores one head group at a time as [q, k, v], so the delta
+ # has to survive the loader's own unpacking as three separate projections.
+ q, k, v = torch.split(
+ _reorder_grouped_qkv_to_qkv(fused, num_query_groups=_HEADS, heads_per_group=1, head_dim=_HEAD_DIM),
+ [_INNER, _INNER, _INNER],
+ )
+ a, b = _factors(_INNER, _HIDDEN)
+ expected = b @ a
+ for got in (q, k, v):
+ assert torch.allclose(got, expected, atol=1e-5)
+
+
+def test_the_fused_mlp_delta_is_swapped_into_gate_first_order(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path)
+ fusion = _load(path.parent)
+
+ fused = fusion.fuse("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))).cpu()
+ a, b = _factors(2 * _FFN, _HIDDEN)
+ value, gate = (b @ a).chunk(2, dim=0)
+ # The diffusers export is value-first; H3's fc1 is gate-first.
+ assert torch.allclose(fused, torch.cat((gate, value), dim=0), atol=1e-5)
+
+
+def test_diff_and_diff_b_edit_weights_and_biases(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(
+ path,
+ tensors={
+ "transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.25),
+ "transformer_blocks.0.adaln_proj.linear.diff_b": torch.full((_HIDDEN,), -0.5),
+ "context_embedder.diff": torch.full((_HIDDEN, _HIDDEN), 0.125),
+ "context_embedder.diff_b": torch.full((_HIDDEN,), 2.0),
+ },
+ )
+ fusion = _load(path.parent)
+
+ # A full-rank delta on an RMSNorm vector is exactly what a LoRA layer
+ # cannot express, and why this release is fused instead of switched.
+ assert torch.allclose(fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 1.25))
+ assert torch.allclose(
+ fusion.fuse("blocks.0.adaln_proj.linear.bias", torch.ones(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 0.5)
+ )
+ assert torch.allclose(
+ fusion.fuse("condition_proj.weight", torch.zeros((_HIDDEN, _HIDDEN))).cpu(),
+ torch.full((_HIDDEN, _HIDDEN), 0.125),
+ )
+ assert torch.allclose(fusion.fuse("condition_proj.bias", torch.zeros(_HIDDEN)).cpu(), torch.full((_HIDDEN,), 2.0))
+
+
+def test_a_streamed_checkpoint_keeps_the_parameters_the_adapter_never_edits(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 3.0)})
+ fusion = _load(path.parent)
+
+ untouched = torch.arange(6, dtype=torch.bfloat16).reshape(2, 3)
+ streamed = dict(
+ fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN)), ("blocks.0.attn.q_norm.weight", untouched)])
+ )
+ assert torch.allclose(streamed["blocks.0.norm1.weight"].cpu(), torch.full((_HIDDEN,), 3.0))
+ assert streamed["blocks.0.attn.q_norm.weight"] is untouched
+
+ # validate_fully_applied() releases the deltas, so a second stream would
+ # fuse nothing at all; it has to fail rather than serve base H3 weights.
+ with pytest.raises(FastH3AdapterError, match="already been fused"):
+ next(fusion.apply([("blocks.0.norm1.weight", torch.zeros(_HIDDEN))]))
+
+
+def test_the_fused_result_keeps_the_parameter_dtype(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, tensors={"transformer_blocks.0.norm1.diff": torch.full((_HIDDEN,), 0.5)})
+ fusion = _load(path.parent)
+
+ fused = fusion.fuse("blocks.0.norm1.weight", torch.ones(_HIDDEN, dtype=torch.bfloat16))
+ assert fused.dtype == torch.bfloat16
+
+
+def test_a_vsa_variant_is_recognised_by_its_compression_gates(tmp_path):
+ dense = tmp_path / "dense" / "adapter_model.safetensors"
+ _write_adapter(dense)
+ assert _load(dense.parent).requires_vsa is False
+
+ sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+ _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+ assert _load(sparse.parent).requires_vsa is True
+
+
+def test_a_tensor_naming_no_h3_parameter_is_an_error(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, tensors={"transformer_blocks.0.attn.norm_q.lora_A.weight": torch.ones((_RANK, _HIDDEN))})
+ # Dropping it silently would load a model that is not the distilled student.
+ with pytest.raises(FastH3AdapterError, match="name no known"):
+ _load(path.parent)
+
+
+def test_an_unpaired_low_rank_factor_is_an_error(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, drop="transformer_blocks.0.attn.to_q.lora_B.weight")
+
+ with pytest.raises(FastH3AdapterError, match="unpaired factor"):
+ _load(path.parent)
+
+
+def test_a_delta_the_checkpoint_never_offered_is_reported(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path)
+ fusion = _load(path.parent)
+
+ fusion.fuse("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN)))
+ # An unapplied delta is the failure that matters: the model would load and
+ # generate, just not as the distilled student.
+ with pytest.raises(FastH3AdapterError, match="never provided"):
+ fusion.validate_fully_applied()
+
+ for name, shape in (
+ ("blocks.0.attn.out_proj.weight", (_HIDDEN, _INNER)),
+ ("blocks.0.mlp.fc1.weight", (2 * _FFN, _HIDDEN)),
+ ("blocks.0.mlp.fc2.weight", (_HIDDEN, _FFN)),
+ ):
+ fusion.fuse(name, torch.zeros(shape))
+ fusion.validate_fully_applied()
+
+
+def _pipeline_stub(fusion=None):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "combined"
+ pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._fasth3 = fusion
+ pipeline.od_config = SimpleNamespace()
+ pipeline.default_video_shift, pipeline.default_audio_shift = 12.0, 3.0
+ return pipeline
+
+
+def _sampling(steps=FASTH3_DENOISE_STEPS, **kwargs) -> SimpleNamespace:
+ fields = {"num_inference_steps": steps, "extra_args": {}, "lora_request": None}
+ return SimpleNamespace(**{**fields, **kwargs})
+
+
+def _check_request(fusion, sampling) -> None:
+ fusion.check_request(sampling, video_shift=12.0, audio_shift=3.0)
+
+
+def _fusion(tmp_path, **kwargs) -> FastH3WeightFusion:
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, **kwargs)
+ return _load(path.parent)
+
+
+@pytest.mark.parametrize("task", ["fl2va", "ref2va"])
+def test_an_active_fasth3_fusion_restricts_requests_to_t2va(task, tmp_path):
+ pipeline = _pipeline_stub(_fusion(tmp_path))
+
+ with pytest.raises(OmniClientError, match="distills \\['t2va'\\] only"):
+ pipeline._resolve_task(task, {})
+ assert pipeline._resolve_task("t2va", {}) == "t2va"
+ assert _pipeline_stub()._resolve_task(task, {}) == task
+
+
+@pytest.mark.parametrize("num_inference_steps", [None, 5, 50])
+def test_fasth3_requires_one_forward_per_sigma_interval(num_inference_steps, tmp_path):
+ fusion = _fusion(tmp_path)
+
+ # The five sigma points bound four forwards, and a request states forwards:
+ # the interval contract the pinned schedules use, and the count the step
+ # scheduler admits a request on before any pipeline hook runs.
+ with pytest.raises(OmniClientError, match=f"num_inference_steps={FASTH3_DENOISE_STEPS}"):
+ _check_request(fusion, _sampling(num_inference_steps))
+ _check_request(fusion, _sampling())
+ # The pinned branch coerces the string form; this one cannot disagree.
+ _check_request(fusion, _sampling(str(FASTH3_DENOISE_STEPS)))
+
+
+def test_the_request_denoises_on_the_release_ladder(tmp_path):
+ pipeline = _pipeline_stub(_fusion(tmp_path))
+
+ positions, num_steps = pipeline._resolve_sigma_positions("t2va", _sampling())
+
+ # dmd_denoising_steps [999, 749, 500, 250] out of 1000, closed with the
+ # terminal 0.0: the pre-shift positions the student was distilled at, not
+ # the uniform ladder num_inference_steps would otherwise derive.
+ assert positions == (0.999, 0.749, 0.5, 0.25, 0.0)
+ # Five sigma points bound four transformer forwards, and forwards is the
+ # unit num_steps carries downstream (Cache-DiT, quality hints).
+ assert num_steps == FASTH3_DENOISE_STEPS == len(positions) - 1
+ # H3's own per-modality shifts still apply on top, and are what turn those
+ # positions into the noise levels the four forwards run at.
+ assert minimax_h3_time_shift_sigmas(shift_scale=12.0, base_schedule=positions) == pytest.approx(
+ [0.999917, 0.972833, 0.923077, 0.8, 0.0], abs=1e-6
+ )
+ assert minimax_h3_time_shift_sigmas(shift_scale=3.0, base_schedule=positions) == pytest.approx(
+ [0.999666, 0.899520, 0.75, 0.5, 0.0], abs=1e-6
+ )
+ # Without a fused adapter the same request keeps the undistilled ladder.
+ assert _pipeline_stub()._resolve_sigma_positions("t2va", _sampling()) == (None, FASTH3_DENOISE_STEPS)
+
+
+def test_the_contract_leaves_the_checkpoint_shifts_alone(tmp_path):
+ pipeline = _pipeline_stub(_fusion(tmp_path))
+ pipeline.partition = "fl2va"
+
+ pipeline._fasth3.check_serving_contract(
+ partition=pipeline.partition,
+ od_config=pipeline.od_config,
+ video_shift=pipeline.default_video_shift,
+ audio_shift=pipeline.default_audio_shift,
+ )
+
+ # The ladder is stated pre-shift, so H3's 12/3 shifts are what place the
+ # student on the levels it was distilled at; the contract must not touch them.
+ assert pipeline.default_video_shift == 12.0
+ assert pipeline.default_audio_shift == 3.0
+
+
+@pytest.mark.parametrize(
+ "extra",
+ [{"flow_shift": 1.0}, {"audio_flow_shift": 1.0}, {"flow_shift": "bad"}],
+)
+def test_a_request_may_not_move_the_student_off_its_rungs(extra, tmp_path):
+ fusion = _fusion(tmp_path)
+
+ with pytest.raises(OmniClientError, match="FastH3 requires"):
+ _check_request(fusion, _sampling(extra_args=extra))
+
+ _check_request(fusion, _sampling(extra_args={"flow_shift": 12.0, "audio_flow_shift": 3.0}))
+
+
+def test_a_request_may_not_carry_a_lora_on_a_fused_server(tmp_path):
+ fusion = _fusion(tmp_path)
+
+ # The dynamic LoRA manager is skipped for a fused adapter, so an adapter
+ # asked for per request would be neither applied nor reported.
+ with pytest.raises(OmniClientError, match="per-request lora is unavailable"):
+ _check_request(fusion, _sampling(lora_request=SimpleNamespace(lora_int_id=1)))
+
+
+def test_offload_is_refused_because_it_bypasses_the_fusion(tmp_path):
+ fusion = _fusion(tmp_path)
+ for flag in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload"):
+ # A host-weight plan installs the transformer without load_weights(),
+ # so the fusion and its completeness check would both be skipped.
+ with pytest.raises(ValueError, match="cannot be combined with"):
+ fusion.check_serving_contract(
+ partition="fl2va",
+ od_config=SimpleNamespace(**{flag: True}),
+ video_shift=12.0,
+ audio_shift=3.0,
+ )
+
+
+def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
+ sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+ _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+ contract = {"od_config": SimpleNamespace(), "video_shift": 12.0, "audio_shift": 3.0}
+ with pytest.raises(ValueError, match="Video Sparse Attention variant"):
+ _load(sparse.parent).check_serving_contract(partition="fl2va", **contract)
+
+ with pytest.raises(ValueError, match="cannot serve a Ref2VA partition"):
+ _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
+
+
+def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
+ # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
+ # so a .diff aimed at one would otherwise be added transposed.
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, tensors={"transformer_blocks.0.attn.to_q.diff": torch.ones((_INNER, _HIDDEN))})
+ with pytest.raises(FastH3AdapterError, match="fused layout"):
+ _claim(path.parent)
+
+
+def test_an_adapter_that_leaves_blocks_untouched_is_refused(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, blocks=2)
+
+ # Every tensor in this file is well formed and pairs up, so the fusion's
+ # own completeness check passes; only the model's depth reveals that a
+ # third of it would have been served as base H3 weights.
+ with pytest.raises(FastH3AdapterError, match="missing=\\[2\\]"):
+ _claim(path.parent, num_blocks=3)
+ with pytest.raises(FastH3AdapterError, match="unknown=\\[1\\]"):
+ _claim(path.parent, num_blocks=1)
+ assert _claim(path.parent, num_blocks=2) is not None
+
+
+def test_an_adapter_shorter_than_it_declares_is_refused(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path, metadata={"low_rank_tensors": "724", "diff_tensors": "85"})
+
+ # The writer records what it emitted, which is the file's one statement
+ # about its own completeness.
+ with pytest.raises(FastH3AdapterError, match="declares low_rank_tensors=724 but carries 12"):
+ _claim(path.parent)
+
+
+def test_only_a_configured_adapter_path_is_read(tmp_path):
+ path = tmp_path / "fasth3" / "adapter_model.safetensors"
+ _write_adapter(path)
+ transformer = SimpleNamespace(
+ arch=SimpleNamespace(attention_head_dim=_HEAD_DIM, num_layers=1, token_refiner_num_layers=0)
+ )
+ unbuilt = object() # a transformer whose arch must not be touched
+
+ assert resolve_fasth3_fusion(SimpleNamespace(lora_path=str(path)), transformer) is not None
+ assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path)]), transformer) is not None
+ # Nothing configured, or more than one artifact, stays on the dynamic route
+ # without reading anything off the model.
+ assert resolve_fasth3_fusion(SimpleNamespace(lora_path=None), unbuilt) is None
+ assert resolve_fasth3_fusion(SimpleNamespace(), unbuilt) is None
+ assert resolve_fasth3_fusion(SimpleNamespace(lora_path=[str(path), str(path)]), unbuilt) is None
+
+
+def test_a_pipeline_that_fused_its_adapter_needs_no_lora_manager(tmp_path):
+ assert _pipeline_stub(_fusion(tmp_path)).lora_is_fused is True
+ assert _pipeline_stub().lora_is_fused is False
+
+
+def test_the_ladder_is_the_one_the_release_publishes():
+ # FASTH3_DENOISE_STEPS is derived from it, so the two cannot drift apart.
+ assert FASTH3_BASE_SCHEDULE.base_schedule == (0.999, 0.749, 0.5, 0.25, 0.0)
+ assert FASTH3_BASE_SCHEDULE.num_inference_steps == FASTH3_DENOISE_STEPS == 4
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py
new file mode 100644
index 000000000..5f98a88d5
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_lora.py
@@ -0,0 +1,416 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from vllm.lora.lora_weights import PackedLoRALayerWeights
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.models.minimax_h3 import lora as lora_module
+from vllm_omni.diffusion.models.minimax_h3.lora import (
+ load_minimax_h3_turbo_lora,
+)
+from vllm_omni.errors import OmniClientError
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_TINY_TARGET_DIMS = {
+ "attn.to_q": (2, 3),
+ "attn.to_k": (2, 3),
+ "attn.to_v": (2, 3),
+ "attn.to_out.0": (3, 2),
+ "ff.net.0.proj": (2, 4),
+ "ff.net.2": (2, 2),
+}
+
+
+@pytest.fixture(autouse=True)
+def _use_tiny_h3_dimensions(monkeypatch):
+ monkeypatch.setattr(lora_module, "_TURBO_TARGET_DIMS", _TINY_TARGET_DIMS)
+
+
+def _request(path) -> LoRARequest:
+ return LoRARequest(
+ lora_name="turbo",
+ lora_int_id=1,
+ lora_path=str(path),
+ )
+
+
+def _write_tiny_turbo(
+ path,
+ *,
+ alpha: str = "128",
+ key_format: str = "minimax-h3-diffusers",
+ omit_target: str | None = None,
+ shape_overrides: dict[str, tuple[int, int]] | None = None,
+) -> None:
+ rank = 128
+ tensors = {}
+ overrides = shape_overrides or {}
+ target_suffixes = (
+ "attn.to_q",
+ "attn.to_k",
+ "attn.to_v",
+ "attn.to_out.0",
+ "ff.net.0.proj",
+ "ff.net.2",
+ )
+ for prefix, block_count in (
+ ("transformer_blocks", 50),
+ ("token_refiner.refiner_blocks", 2),
+ ):
+ for block_index in range(block_count):
+ for suffix in target_suffixes:
+ target = f"{prefix}.{block_index}.{suffix}"
+ if target == omit_target:
+ continue
+ input_dim, output_dim = _TINY_TARGET_DIMS[suffix]
+ a_name = f"{target}.lora_A.default.weight"
+ b_name = f"{target}.lora_B.default.weight"
+ tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
+ if suffix == "ff.net.0.proj":
+ if b_name in overrides:
+ tensors[b_name] = torch.ones(overrides[b_name])
+ else:
+ tensors[b_name] = torch.cat(
+ (
+ torch.ones(output_dim // 2, rank),
+ torch.full((output_dim // 2, rank), 2.0),
+ ),
+ dim=0,
+ )
+ else:
+ tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
+ save_file(
+ tensors,
+ str(path),
+ metadata={"alpha": alpha, "key_format": key_format},
+ )
+
+
+def test_h3_turbo_loads_through_legacy_lora_model_and_swaps_ffn(tmp_path):
+ path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(path)
+
+ loaded = load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+ assert loaded is not None
+ lora_model, peft_helper = loaded
+ assert peft_helper.r == 128
+ assert peft_helper.lora_alpha == 128
+ assert len(lora_model.loras) == 312
+ assert "blocks.0.attn.to_q" in lora_model.loras
+ assert "blocks.0.mlp.fc1" in lora_model.loras
+ fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
+ assert isinstance(fc1, PackedLoRALayerWeights)
+ assert len(fc1.lora_a) == 2
+ assert len(fc1.lora_b) == 2
+ torch.testing.assert_close(fc1.lora_a[0], torch.ones(128, 2))
+ torch.testing.assert_close(fc1.lora_a[1], torch.ones(128, 2))
+ torch.testing.assert_close(fc1.lora_b[0], torch.full((2, 128), 2.0))
+ torch.testing.assert_close(fc1.lora_b[1], torch.ones(2, 128))
+
+
+@pytest.mark.parametrize(
+ ("tensor_suffix", "bad_shape"),
+ [
+ ("attn.to_q.lora_A.default.weight", (128, 1)),
+ ("attn.to_q.lora_A.default.weight", (128, 3)),
+ ("attn.to_q.lora_B.default.weight", (2, 128)),
+ ("attn.to_q.lora_B.default.weight", (4, 128)),
+ ],
+)
+def test_h3_turbo_rejects_under_and_oversized_global_shapes(tmp_path, tensor_suffix, bad_shape):
+ path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ tensor_name = f"transformer_blocks.0.{tensor_suffix}"
+ _write_tiny_turbo(path, shape_overrides={tensor_name: bad_shape})
+
+ with pytest.raises(ValueError, match="invalid global shape"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+def test_legacy_manager_uses_the_h3_model_loader_without_changing_its_interface(tmp_path):
+ path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(path)
+
+ class _Pipeline:
+ def _load_diffusion_lora_adapter(self, **kwargs):
+ return load_minimax_h3_turbo_lora(partition="fl2va", **kwargs)
+
+ manager = object.__new__(DiffusionLoRAManager)
+ manager.pipeline = _Pipeline()
+ manager.dtype = torch.float32
+ manager._expected_lora_modules = {"to_q", "fc1"}
+
+ lora_model, peft_helper = manager._load_adapter(_request(path))
+
+ assert lora_model.id == 1
+ assert peft_helper.lora_alpha == 128
+ assert len(lora_model.loras) == 312
+ assert "blocks.0.attn.to_q" in lora_model.loras
+ assert "blocks.0.mlp.fc1" in lora_model.loras
+
+
+def test_h3_turbo_rejects_wrong_alpha_and_ref2va(tmp_path):
+ wrong_alpha = tmp_path / "wrong_alpha" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ wrong_alpha.parent.mkdir()
+ _write_tiny_turbo(wrong_alpha, alpha="8")
+ with pytest.raises(ValueError, match="requires alpha=128"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(wrong_alpha),
+ lora_path=wrong_alpha,
+ dtype=torch.float32,
+ )
+
+ valid = tmp_path / "valid" / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ valid.parent.mkdir()
+ _write_tiny_turbo(valid)
+ with pytest.raises(ValueError, match="supports FL2VA/T2VA only"):
+ load_minimax_h3_turbo_lora(
+ partition="ref2va",
+ lora_request=_request(valid),
+ lora_path=valid,
+ dtype=torch.float32,
+ )
+
+
+@pytest.mark.parametrize(
+ "offload_mode",
+ [
+ "model-level CPU offload (--enable-cpu-offload)",
+ "layerwise offload (--enable-layerwise-offload)",
+ ],
+)
+def test_h3_turbo_rejects_offload_modes(tmp_path, offload_mode):
+ path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(path)
+
+ with pytest.raises(ValueError, match="does not support"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ unsupported_offload_mode=offload_mode,
+ )
+
+
+def test_h3_turbo_allows_distributed_layerwise_offload(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import (
+ pipeline_minimax_h3 as pipeline_module,
+ )
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=True,
+ )
+ captured = {}
+
+ def load_turbo(**kwargs):
+ captured.update(kwargs)
+ return object(), object()
+
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", load_turbo)
+ loaded = pipeline._load_diffusion_lora_adapter(
+ lora_request=_request("turbo"),
+ lora_path="turbo",
+ dtype=torch.bfloat16,
+ )
+
+ assert loaded is not None
+ assert captured["unsupported_offload_mode"] is None
+
+
+def test_h3_turbo_accepts_only_the_declared_v1_artifact(tmp_path):
+ unsupported = tmp_path / "minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors"
+ _write_tiny_turbo(unsupported)
+
+ with pytest.raises(ValueError, match="supports only"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(unsupported),
+ lora_path=unsupported,
+ dtype=torch.float32,
+ )
+
+ supported = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(supported)
+ # Directory resolution selects the declared artifact even when another
+ # v1.0 checkpoint is present beside it.
+ assert (
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(tmp_path),
+ lora_path=tmp_path,
+ dtype=torch.float32,
+ )
+ is not None
+ )
+
+
+def test_h3_turbo_rejects_a_truncated_declared_artifact(tmp_path):
+ path = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(path, omit_target="transformer_blocks.49.ff.net.2")
+
+ with pytest.raises(ValueError, match="target set does not match"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+def test_non_h3_checkpoint_falls_back_to_the_generic_peft_loader(tmp_path):
+ path = tmp_path / "other.safetensors"
+ _write_tiny_turbo(path, key_format="other")
+
+ assert (
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+ is None
+ )
+
+ declared = tmp_path / "minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_turbo(declared, key_format="other")
+ with pytest.raises(ValueError, match="requires safetensors metadata"):
+ load_minimax_h3_turbo_lora(
+ partition="fl2va",
+ lora_request=_request(declared),
+ lora_path=declared,
+ dtype=torch.float32,
+ )
+
+
+def test_only_an_active_recognized_turbo_adapter_restricts_ref2va(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import (
+ pipeline_minimax_h3 as pipeline_module,
+ )
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "combined"
+ pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+ request = _request("generic-peft")
+
+ def load():
+ return pipeline._load_diffusion_lora_adapter(
+ lora_request=request,
+ lora_path=request.lora_path,
+ dtype=torch.float32,
+ )
+
+ def resolve(scale):
+ sampling = SimpleNamespace(lora_request=request, lora_scale=scale)
+ return pipeline._resolve_task(
+ "ref2va",
+ {},
+ has_turbo_lora=pipeline._has_active_turbo_lora(sampling),
+ )
+
+ recognized = (object(), object())
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: recognized)
+ assert load() is recognized
+ assert resolve(0.0) == "ref2va"
+ with pytest.raises(OmniClientError, match="supports T2VA/FL2VA requests only"):
+ resolve(1.0)
+
+ # Simulate manager eviction followed by a generic PEFT adapter reusing the
+ # same client-supplied ID. A real reload must replace the stale kind.
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+ assert load() is None
+ assert resolve(1.0) == "ref2va"
+
+
+def test_h3_turbo_requires_all_loaded_targets_to_bind():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline._turbo_lora_adapter_ids = {1}
+ lora_model = SimpleNamespace(id=1, loras={"to_q": object(), "to_k": object()})
+
+ pipeline._validate_diffusion_lora_binding(
+ lora_model=lora_model,
+ bound_lora_names=frozenset(lora_model.loras),
+ )
+ with pytest.raises(ValueError, match="bound=1/2"):
+ pipeline._validate_diffusion_lora_binding(
+ lora_model=lora_model,
+ bound_lora_names=frozenset({"to_q"}),
+ )
+
+
+@pytest.mark.parametrize(
+ ("num_inference_steps", "extra_args", "error"),
+ [
+ (None, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+ (4, {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+ ("5", {"flow_shift": 6.0, "audio_flow_shift": 3.0}, "num_inference_steps=5"),
+ (5, {"flow_shift": 7.0, "audio_flow_shift": 3.0}, "flow_shift=6"),
+ (5, {"flow_shift": "bad", "audio_flow_shift": 3.0}, "flow_shift=6"),
+ (5, {"flow_shift": 6.0, "audio_flow_shift": 4.0}, "audio_flow_shift=3"),
+ (5, {"flow_shift": 6.0, "audio_flow_shift": []}, "audio_flow_shift=3"),
+ ],
+)
+def test_h3_turbo_rejects_unsupported_sampling(num_inference_steps, extra_args, error):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ pipeline.default_video_shift = 12.0
+ pipeline.default_audio_shift = 3.0
+ sampling = SimpleNamespace(
+ num_inference_steps=num_inference_steps,
+ extra_args=extra_args,
+ )
+
+ with pytest.raises(OmniClientError, match=error):
+ pipeline._validate_turbo_sampling(sampling)
+
+
+def test_h3_turbo_accepts_five_sigma_points_for_four_nfe():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ pipeline.default_video_shift = 12.0
+ pipeline.default_audio_shift = 3.0
+ pipeline._validate_turbo_sampling(
+ SimpleNamespace(
+ num_inference_steps=5,
+ extra_args={"flow_shift": 6.0, "audio_flow_shift": 3.0},
+ )
+ )
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py
new file mode 100644
index 000000000..c4e581381
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py
@@ -0,0 +1,864 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+from vllm.lora.lora_weights import PackedLoRALayerWeights
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.models.minimax_h3.npu import lora as lora_module
+from vllm_omni.diffusion.models.minimax_h3.npu.lora import load_minimax_h3_native_lora
+from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule
+from vllm_omni.errors import OmniClientError
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_TINY_HIDDEN = 2
+_TINY_FFN = 2
+_TINY_TIME = 2
+_TINY_BLOCK_ADALN = 4
+_TINY_FINAL_ADALN = 3
+_TINY_RANK = 4
+_TINY_NUM_QUERY_GROUPS = 2
+_TINY_HEADS_PER_GROUP = 1
+_TINY_HEAD_DIM = 1
+_TINY_QKV_SLICE = _TINY_NUM_QUERY_GROUPS * _TINY_HEADS_PER_GROUP * _TINY_HEAD_DIM
+# Grouped qkv stores (heads_per_group + 2) * head_dim rows per query group, so
+# the attention inner size follows from the head layout and cannot be picked
+# independently: 3 * inner must equal the grouped row count.
+_TINY_INNER = _TINY_QKV_SLICE
+
+
+@pytest.fixture(autouse=True)
+def _use_tiny_native_dimensions(monkeypatch):
+ monkeypatch.setattr(lora_module, "_NATIVE_RANK", _TINY_RANK)
+ monkeypatch.setattr(lora_module, "_NATIVE_ALPHA", _TINY_RANK)
+ monkeypatch.setattr(lora_module, "_NATIVE_HIDDEN_SIZE", _TINY_HIDDEN)
+ monkeypatch.setattr(lora_module, "_NATIVE_ATTENTION_INNER_SIZE", _TINY_INNER)
+ monkeypatch.setattr(lora_module, "_NATIVE_FFN_HIDDEN_SIZE", _TINY_FFN)
+ monkeypatch.setattr(lora_module, "_NATIVE_TIME_EMBED_DIM", _TINY_TIME)
+ monkeypatch.setattr(lora_module, "_NATIVE_BLOCK_ADALN_OUT", _TINY_BLOCK_ADALN)
+ monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_OUT", _TINY_FINAL_ADALN)
+ monkeypatch.setattr(lora_module, "_NATIVE_NUM_QUERY_GROUPS", _TINY_NUM_QUERY_GROUPS)
+ monkeypatch.setattr(lora_module, "_NATIVE_HEADS_PER_GROUP", _TINY_HEADS_PER_GROUP)
+ monkeypatch.setattr(lora_module, "_NATIVE_HEAD_DIM", _TINY_HEAD_DIM)
+ monkeypatch.setattr(lora_module, "_NATIVE_QKV_SLICE", _TINY_QKV_SLICE)
+ monkeypatch.setattr(
+ lora_module,
+ "_NATIVE_TARGET_DIMS",
+ {
+ "attn.qkv_proj": (_TINY_HIDDEN, 3 * _TINY_INNER),
+ "attn.out_proj": (_TINY_INNER, _TINY_HIDDEN),
+ "mlp.fc1": (_TINY_HIDDEN, 2 * _TINY_FFN),
+ "mlp.fc2": (_TINY_FFN, _TINY_HIDDEN),
+ "adaln_proj.linear": (_TINY_TIME, _TINY_BLOCK_ADALN),
+ },
+ )
+ monkeypatch.setattr(lora_module, "_NATIVE_FINAL_ADALN_DIMS", (_TINY_TIME, _TINY_FINAL_ADALN))
+ monkeypatch.setattr(
+ lora_module,
+ "_NATIVE_EXPECTED_TARGETS",
+ frozenset(
+ [
+ *(
+ f"blocks.{block_index}.{suffix}"
+ for block_index in range(50)
+ for suffix in lora_module._NATIVE_TARGET_SUFFIXES
+ )
+ ]
+ + [
+ *(
+ f"token_refiner.blocks.{block_index}.{suffix}"
+ for block_index in range(2)
+ for suffix in lora_module._NATIVE_TOKEN_REFINER_SUFFIXES
+ )
+ ]
+ + ["final_layer.adaln_proj.linear"]
+ ),
+ )
+
+
+def _request(path) -> LoRARequest:
+ return LoRARequest(
+ lora_name="flashgen",
+ lora_int_id=7,
+ lora_path=str(path),
+ )
+
+
+def _write_tiny_native(
+ path,
+ *,
+ omit_target: str | None = None,
+ shape_overrides: dict[str, tuple[int, int]] | None = None,
+ metadata: dict[str, str] | None = None,
+) -> None:
+ rank = _TINY_RANK
+ tensors = {}
+ overrides = shape_overrides or {}
+ for target in sorted(lora_module._NATIVE_EXPECTED_TARGETS):
+ if target == omit_target:
+ continue
+ input_dim, output_dim = lora_module._native_target_dims(target)
+ raw_target = f"transformer.{target}"
+ a_name = f"{raw_target}.lora_A.default.weight"
+ b_name = f"{raw_target}.lora_B.default.weight"
+ tensors[a_name] = torch.ones(overrides.get(a_name, (rank, input_dim)))
+ if target.endswith(".attn.qkv_proj"):
+ grouped_rows = []
+ for group in range(_TINY_NUM_QUERY_GROUPS):
+ grouped_rows.extend(
+ [
+ torch.full((1, rank), float(group * 10 + 1)),
+ torch.full((1, rank), float(group * 10 + 2)),
+ torch.full((1, rank), float(group * 10 + 3)),
+ ]
+ )
+ tensors[b_name] = torch.cat(grouped_rows, dim=0)
+ elif target.endswith(".mlp.fc1"):
+ tensors[b_name] = torch.cat(
+ (
+ torch.full((output_dim // 2, rank), 2.0),
+ torch.full((output_dim // 2, rank), 1.0),
+ ),
+ dim=0,
+ )
+ else:
+ tensors[b_name] = torch.ones(overrides.get(b_name, (output_dim, rank)))
+ save_file(
+ tensors,
+ str(path),
+ metadata=metadata
+ or {
+ "format": "pt",
+ "key_format": "minimax-h3-native",
+ "qkv_layout": "grouped",
+ "lora_rank": str(rank),
+ "lora_alpha": str(rank),
+ "base_schedule": "1.0,0.7,0.4,0.15,0.0",
+ "tasks": "t2va",
+ },
+ )
+
+
+def test_h3_native_loads_and_packs_qkv_and_fc1(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+
+ loaded = load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+ assert loaded is not None
+ lora_model, peft_helper, schedule = loaded
+ assert peft_helper.r == _TINY_RANK
+ assert schedule.num_inference_steps == 4
+ assert len(lora_model.loras) == 259
+ assert "blocks.0.adaln_proj.linear" in lora_model.loras
+ assert "final_layer.adaln_proj.linear" in lora_model.loras
+
+ # The writer tags group g with q=g*10+1, k=g*10+2, v=g*10+3, so a correct
+ # grouped -> q/k/v permutation yields slices that are uniform per role.
+ qkv = lora_model.get_lora("blocks.0.attn.qkv_proj")
+ assert isinstance(qkv, PackedLoRALayerWeights)
+ torch.testing.assert_close(qkv.lora_b[0], torch.tensor([[1.0] * _TINY_RANK, [11.0] * _TINY_RANK]))
+ torch.testing.assert_close(qkv.lora_b[1], torch.tensor([[2.0] * _TINY_RANK, [12.0] * _TINY_RANK]))
+ torch.testing.assert_close(qkv.lora_b[2], torch.tensor([[3.0] * _TINY_RANK, [13.0] * _TINY_RANK]))
+
+ fc1 = lora_model.get_lora("blocks.0.mlp.fc1")
+ assert isinstance(fc1, PackedLoRALayerWeights)
+ torch.testing.assert_close(fc1.lora_b[0], torch.full((_TINY_FFN, _TINY_RANK), 2.0))
+ torch.testing.assert_close(fc1.lora_b[1], torch.full((_TINY_FFN, _TINY_RANK), 1.0))
+
+
+def test_h3_native_rejects_bad_metadata_and_ref2va(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path, metadata={"key_format": "minimax-h3-native", "qkv_layout": "runtime"})
+ with pytest.raises(ValueError, match="qkv_layout='grouped'"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+ # The exact-filename guard runs before the partition check, so the artifact
+ # has to keep the published name for this to reach the Ref2VA rejection.
+ valid_dir = tmp_path / "valid"
+ valid_dir.mkdir()
+ valid = valid_dir / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(valid)
+ with pytest.raises(ValueError, match="supports T2VA only"):
+ load_minimax_h3_native_lora(
+ partition="ref2va",
+ lora_request=_request(valid),
+ lora_path=valid,
+ dtype=torch.float32,
+ )
+
+ renamed = tmp_path / "renamed.safetensors"
+ _write_tiny_native(renamed)
+ with pytest.raises(ValueError, match="supports only"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(renamed),
+ lora_path=renamed,
+ dtype=torch.float32,
+ )
+
+
+def test_h3_native_rejects_truncated_artifact(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path, omit_target="blocks.49.mlp.fc2")
+ with pytest.raises(ValueError, match="target set does not match"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+def test_h3_native_rejects_schedule_with_wrong_interval_count(tmp_path):
+ """A mislabeled schedule must fail closed instead of silently changing steps."""
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(
+ path,
+ metadata={
+ "format": "pt",
+ "key_format": "minimax-h3-native",
+ "qkv_layout": "grouped",
+ "lora_rank": str(_TINY_RANK),
+ "lora_alpha": str(_TINY_RANK),
+ "base_schedule": "1.0,0.5,0.0",
+ "tasks": "t2va",
+ },
+ )
+ with pytest.raises(ValueError, match="4-interval base_schedule"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+@pytest.mark.parametrize(
+ ("raw_schedule", "message"),
+ [
+ (None, "requires safetensors metadata base_schedule"),
+ ("", "must not be empty"),
+ (",,", "must not be empty"),
+ ("1.0,oops,0.15,0.0", "is malformed"),
+ # A capped opening rung below 1.0 is a legal DMD2 schedule, so the
+ # out-of-range cases are a first position above 1.0 and a tail that
+ # never reaches clean latents.
+ ("1.5,0.7,0.4,0.15,0.0", "is malformed"),
+ ("1.0,0.7,0.4,0.15", "is malformed"),
+ ],
+)
+def test_h3_native_rejects_malformed_schedule_metadata(tmp_path, raw_schedule, message):
+ """safetensors metadata is string-only, so the loader owns the parse contract."""
+ metadata = {
+ "format": "pt",
+ "key_format": "minimax-h3-native",
+ "qkv_layout": "grouped",
+ "lora_rank": str(_TINY_RANK),
+ "lora_alpha": str(_TINY_RANK),
+ "tasks": "t2va",
+ }
+ if raw_schedule is not None:
+ metadata["base_schedule"] = raw_schedule
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path, metadata=metadata)
+
+ with pytest.raises(ValueError, match=message):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+def test_h3_native_declared_file_fails_closed_on_invalid_metadata(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path, metadata={"key_format": "other"})
+ with pytest.raises(ValueError, match="requires safetensors metadata"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+
+def test_pipeline_native_schedule_and_task_validation(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = {7}
+ pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
+ pipeline._base_schedule_by_partition = {"fl2va": None}
+
+ sampling = SimpleNamespace(
+ lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
+ lora_scale=1.0,
+ num_inference_steps=4,
+ extra_args={},
+ )
+ assert pipeline._sigma_schedule_for_request(sampling, "t2va").num_inference_steps == 4
+ with pytest.raises(OmniClientError, match="num_inference_steps must be 4"):
+ pipeline._validate_native_sampling(
+ SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
+ task="t2va",
+ )
+ with pytest.raises(OmniClientError, match="supports T2VA requests only"):
+ pipeline._resolve_task(
+ "fl2va",
+ {},
+ has_turbo_lora=False,
+ has_native_lora=True,
+ )
+
+ pipeline._base_schedule_by_partition = {"fl2va": DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
+ with pytest.raises(OmniClientError, match="already pins base_schedule"):
+ pipeline._sigma_schedule_for_request(sampling, "t2va")
+
+
+def test_pipeline_native_step_count_follows_adapter_schedule():
+ """The accepted step count is the adapter's interval count, never a literal."""
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline._native_lora_adapter_ids = {7}
+ pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.5, 0.0])}
+ request = LoRARequest("flashgen", 7, "/tmp/native.safetensors")
+
+ pipeline._validate_native_sampling(
+ SimpleNamespace(lora_request=request, num_inference_steps=2),
+ task="t2va",
+ )
+ with pytest.raises(OmniClientError, match="num_inference_steps must be 2 or omitted, not 3"):
+ pipeline._validate_native_sampling(
+ SimpleNamespace(lora_request=request, num_inference_steps=3),
+ task="t2va",
+ )
+ with pytest.raises(OmniClientError, match="requires num_inference_steps=2"):
+ pipeline._validate_native_sampling(
+ SimpleNamespace(lora_request=request, num_inference_steps=4),
+ task="t2va",
+ )
+
+
+@pytest.mark.parametrize("step_execution", [False, True])
+def test_pipeline_native_omitted_step_count_follows_execution_mode(step_execution):
+ """Only request mode can resolve an omitted count from the adapter schedule.
+
+ ``StepScheduler`` reads ``num_inference_steps`` off the request at admission,
+ before this pipeline sees it, so under step execution an omitted value can
+ only disagree with the denoise loop. The advertised contract has to match.
+ """
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline._native_lora_adapter_ids = {7}
+ pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
+ pipeline.od_config = SimpleNamespace(step_execution=step_execution)
+ sampling = SimpleNamespace(
+ lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
+ num_inference_steps=None,
+ )
+
+ if not step_execution:
+ # Request mode defaults to the adapter's own interval count.
+ pipeline._validate_native_sampling(sampling, task="t2va")
+ return
+
+ with pytest.raises(OmniClientError, match="explicit num_inference_steps=4 under step execution"):
+ pipeline._validate_native_sampling(sampling, task="t2va")
+ # The step-mode message must not offer omission it cannot accept.
+ with pytest.raises(OmniClientError, match=r"must be 4, not 5$"):
+ pipeline._validate_native_sampling(
+ SimpleNamespace(lora_request=sampling.lora_request, num_inference_steps=5),
+ task="t2va",
+ )
+
+
+def test_pipeline_replaces_native_classification_after_reload(monkeypatch, tmp_path):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+ request = _request(path)
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ )
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+
+ loaded = pipeline._load_diffusion_lora_adapter(
+ lora_request=request,
+ lora_path=path,
+ dtype=torch.float32,
+ )
+ assert loaded is not None
+ assert request.lora_int_id in pipeline._native_lora_adapter_ids
+ assert request.lora_int_id in pipeline._lora_sigma_schedules
+
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
+ assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
+ assert request.lora_int_id not in pipeline._native_lora_adapter_ids
+ assert request.lora_int_id not in pipeline._lora_sigma_schedules
+
+
+@pytest.mark.parametrize(
+ "offload_mode",
+ [
+ "model-level CPU offload (--enable-cpu-offload)",
+ "layerwise offload (--enable-layerwise-offload)",
+ ],
+)
+def test_h3_native_rejects_offload_modes(tmp_path, offload_mode):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+
+ with pytest.raises(ValueError, match="does not support"):
+ load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ unsupported_offload_mode=offload_mode,
+ )
+
+
+def test_h3_native_allows_distributed_layerwise_offload(monkeypatch):
+ """DLO keeps LoRA A/B buffers resident, so native must not fail closed."""
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=True,
+ )
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+ captured: dict[str, object] = {}
+ schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
+
+ def load_native(**kwargs):
+ captured.update(kwargs)
+ return object(), object(), schedule
+
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", load_native)
+
+ request = _request("flashgen")
+ loaded = pipeline._load_diffusion_lora_adapter(
+ lora_request=request,
+ lora_path="flashgen",
+ dtype=torch.bfloat16,
+ )
+
+ assert loaded is not None
+ assert captured["unsupported_offload_mode"] is None
+ assert request.lora_int_id in pipeline._native_lora_adapter_ids
+ assert pipeline._lora_sigma_schedules[request.lora_int_id] is schedule
+
+
+def test_h3_native_qkv_reorder_matches_base_loader_contract():
+ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import _reorder_grouped_qkv_to_qkv
+
+ num_query_groups = 56
+ heads_per_group = 1
+ head_dim = 128
+ grouped = torch.arange(num_query_groups * (heads_per_group + 2) * head_dim, dtype=torch.float32)
+ grouped = grouped.reshape(num_query_groups, (heads_per_group + 2) * head_dim)
+ reordered = _reorder_grouped_qkv_to_qkv(
+ grouped.reshape(-1, 1),
+ num_query_groups=num_query_groups,
+ heads_per_group=heads_per_group,
+ head_dim=head_dim,
+ ).reshape(-1)
+ q_size = num_query_groups * heads_per_group * head_dim
+ k_size = num_query_groups * head_dim
+ q, k, v = torch.split(reordered, [q_size, k_size, k_size])
+ assert q.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 0]
+ assert k.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, head_dim]
+ assert v.reshape(num_query_groups, head_dim)[0, 0] == grouped[0, 2 * head_dim]
+
+
+def test_native_target_pattern_matches_component_qualified_names():
+ """The manager matches ".", so the pattern needs the prefix."""
+ import regex as re
+
+ pattern = lora_module._NATIVE_TARGET_PATTERN
+ matched = [
+ target
+ for target in lora_module._NATIVE_EXPECTED_TARGETS
+ if re.search(pattern, f"transformer.{target}") is not None
+ ]
+ assert len(matched) == 259
+ assert re.search(pattern, "transformer.blocks.0.attn.q_proj") is None
+ assert re.search(pattern, "transformer.token_refiner.blocks.0.adaln_proj.linear") is None
+
+
+def test_replace_layers_binds_every_native_target(monkeypatch):
+ """Regression for bound=0/259: layer replacement must reach all 259 modules."""
+ from vllm.lora.layers import BaseLayerWithLoRA
+ from vllm.lora.peft_helper import PEFTHelper
+
+ from vllm_omni.diffusion.lora import manager as manager_module
+
+ class _StubLoRALayer(BaseLayerWithLoRA):
+ def __init__(self, base_layer):
+ super().__init__()
+ self.base_layer = base_layer
+
+ def _linear() -> torch.nn.Module:
+ return torch.nn.Linear(_TINY_HIDDEN, _TINY_HIDDEN)
+
+ def _block(*, with_adaln: bool) -> torch.nn.Module:
+ block = torch.nn.Module()
+ block.attn = torch.nn.Module()
+ block.attn.qkv_proj = _linear()
+ block.attn.out_proj = _linear()
+ block.mlp = torch.nn.Module()
+ block.mlp.fc1 = _linear()
+ block.mlp.fc2 = _linear()
+ if with_adaln:
+ block.adaln_proj = torch.nn.Module()
+ block.adaln_proj.linear = _linear()
+ return block
+
+ transformer = torch.nn.Module()
+ transformer.blocks = torch.nn.ModuleList([_block(with_adaln=True) for _ in range(50)])
+ transformer.token_refiner = torch.nn.Module()
+ transformer.token_refiner.blocks = torch.nn.ModuleList([_block(with_adaln=False) for _ in range(2)])
+ transformer.final_layer = torch.nn.Module()
+ transformer.final_layer.adaln_proj = torch.nn.Module()
+ transformer.final_layer.adaln_proj.linear = _linear()
+ pipeline = torch.nn.Module()
+ pipeline.transformer = transformer
+
+ monkeypatch.setattr(manager_module, "from_layer_diffusion", lambda *, layer, **_: _StubLoRALayer(layer))
+
+ manager = object.__new__(DiffusionLoRAManager)
+ manager.pipeline = pipeline
+ manager.dtype = torch.float32
+ manager.max_cached_adapters = 1
+ manager._lora_modules = {}
+ manager._max_lora_rank = 64
+ manager._resident_lora_device = None
+
+ peft_helper = PEFTHelper.from_dict(
+ {
+ "r": _TINY_RANK,
+ "lora_alpha": _TINY_RANK,
+ "target_modules": lora_module._NATIVE_TARGET_PATTERN,
+ }
+ )
+ manager._replace_layers_with_lora(peft_helper)
+
+ assert set(manager._lora_modules) == {f"transformer.{target}" for target in lora_module._NATIVE_EXPECTED_TARGETS}
+
+
+def test_legacy_manager_uses_native_loader(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+
+ class _Pipeline:
+ def _load_diffusion_lora_adapter(self, **kwargs):
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ MiniMaxH3Pipeline,
+ )
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ )
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+ return pipeline._load_diffusion_lora_adapter(**kwargs)
+
+ manager = object.__new__(DiffusionLoRAManager)
+ manager.pipeline = _Pipeline()
+ manager.dtype = torch.float32
+ manager._expected_lora_modules = {"qkv_proj", "fc1", "adaln_proj.linear"}
+
+ lora_model, peft_helper = manager._load_adapter(_request(path))
+ assert lora_model.id == 7
+ assert peft_helper.lora_alpha == _TINY_RANK
+ assert len(lora_model.loras) == 259
+
+
+def test_pipeline_schedule_inactive_when_scale_zero():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = {7}
+ pipeline._lora_sigma_schedules = {7: DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])}
+ pipeline._base_schedule_by_partition = {"fl2va": None}
+
+ sampling = SimpleNamespace(
+ lora_request=LoRARequest("flashgen", 7, "/tmp/native.safetensors"),
+ lora_scale=0.0,
+ num_inference_steps=4,
+ )
+ assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
+ assert not pipeline._has_active_native_lora(sampling)
+
+
+def test_pipeline_schedule_falls_back_after_eviction(monkeypatch, tmp_path):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module
+
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+ request = _request(path)
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.partition = "fl2va"
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ )
+ pipeline._turbo_lora_adapter_ids = set()
+ pipeline._native_lora_adapter_ids = set()
+ pipeline._lora_sigma_schedules = {}
+ pipeline._base_schedule_by_partition = {"fl2va": None}
+
+ loaded = pipeline._load_diffusion_lora_adapter(
+ lora_request=request,
+ lora_path=path,
+ dtype=torch.float32,
+ )
+ assert loaded is not None
+ assert request.lora_int_id in pipeline._lora_sigma_schedules
+
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_turbo_lora", lambda **_: None)
+ monkeypatch.setattr(pipeline_module, "load_minimax_h3_native_lora", lambda **_: None)
+ assert pipeline._load_diffusion_lora_adapter(lora_request=request, lora_path=path, dtype=torch.float32) is None
+ assert request.lora_int_id not in pipeline._lora_sigma_schedules
+
+ sampling = SimpleNamespace(
+ lora_request=request,
+ lora_scale=1.0,
+ num_inference_steps=4,
+ )
+ assert pipeline._sigma_schedule_for_request(sampling, "t2va") is None
+
+
+def test_native_packed_qkv_slices_are_tp2_divisible():
+ assert lora_module._NATIVE_QKV_SLICE % 2 == 0
+ assert lora_module._NATIVE_FFN_HIDDEN_SIZE % 2 == 0
+
+
+def test_lora_manager_activates_native_packed_qkv(tmp_path):
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+ loaded = load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+ assert loaded is not None
+ lora_model, _, _ = loaded
+ packed = lora_model.get_lora("blocks.0.attn.qkv_proj")
+ assert isinstance(packed, PackedLoRALayerWeights)
+ assert len(packed.lora_b) == 3
+ assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b)
+
+ class _DummyLoRALayer:
+ def __init__(self):
+ self.n_slices = 3
+ self.output_slices = (_TINY_QKV_SLICE, _TINY_QKV_SLICE, _TINY_QKV_SLICE)
+ self.set_calls: list[tuple[list, list]] = []
+ self.reset_calls = 0
+
+ def reset_lora(self, index: int):
+ self.reset_calls += 1
+
+ def set_lora(self, index: int, lora_a, lora_b):
+ self.set_calls.append((lora_a, lora_b))
+
+ layer = _DummyLoRALayer()
+ manager = object.__new__(DiffusionLoRAManager)
+ manager.pipeline = torch.nn.Module()
+ manager._lora_modules = {"blocks.0.attn.qkv_proj": layer}
+ manager._registered_adapters = {
+ lora_model.id: lora_model,
+ }
+ manager._active_adapter_id = None
+ manager._adapter_scales = {}
+ manager._activate_adapter(lora_model.id, scale=0.5)
+
+ assert manager._active_adapter_id == lora_model.id
+
+ assert len(layer.set_calls) == 1
+ lora_a_list, lora_b_list = layer.set_calls[0]
+ assert len(lora_a_list) == 3
+ assert len(lora_b_list) == 3
+ torch.testing.assert_close(lora_b_list[0], packed.lora_b[0] * 0.5)
+ torch.testing.assert_close(lora_b_list[1], packed.lora_b[1] * 0.5)
+ torch.testing.assert_close(lora_b_list[2], packed.lora_b[2] * 0.5)
+
+
+def test_lora_manager_tp2_splits_native_packed_qkv_per_rank():
+ """Each Q/K/V slice must be divisible by TP2 so rank-local layers can shard rows."""
+
+ tp_size = 2
+ slice_rows = lora_module._NATIVE_QKV_SLICE
+ assert slice_rows % tp_size == 0
+ local_rows = slice_rows // tp_size
+
+ class _TpQkvLoRALayer:
+ def __init__(self, tp_rank: int):
+ self.tp_rank = tp_rank
+ self.n_slices = 3
+ self.output_slices = (local_rows, local_rows, local_rows)
+ self.set_calls: list[tuple[list, list]] = []
+
+ def reset_lora(self, index: int):
+ return
+
+ def set_lora(self, index: int, lora_a, lora_b):
+ self.set_calls.append((lora_a, lora_b))
+
+ for tp_rank in range(tp_size):
+ layer = _TpQkvLoRALayer(tp_rank)
+ start = tp_rank * local_rows
+ end = start + local_rows
+ full_packed = PackedLoRALayerWeights(
+ module_name="blocks.0.attn.qkv_proj",
+ rank=lora_module._NATIVE_RANK,
+ lora_alphas=[64, 64, 64],
+ lora_a=[torch.ones(lora_module._NATIVE_RANK, lora_module._NATIVE_HIDDEN_SIZE)] * 3,
+ lora_b=[torch.full((slice_rows, lora_module._NATIVE_RANK), float(i + 1)) for i in range(3)],
+ scaling=[1.0, 1.0, 1.0],
+ )
+ local_b = [b[start:end].contiguous() for b in full_packed.lora_b]
+ layer.set_lora(0, full_packed.lora_a, local_b)
+ assert len(layer.set_calls) == 1
+ _, lora_b_list = layer.set_calls[0]
+ assert all(b.shape[0] == local_rows for b in lora_b_list)
+ assert lora_b_list[0][0, 0].item() == 1.0
+ assert lora_b_list[1][0, 0].item() == 2.0
+ assert lora_b_list[2][0, 0].item() == 3.0
+
+
+# 259 modules, of which 52 qkv expand to 3 slices and 52 fc1 expand to 2.
+_EXPECTED_NATIVE_SLICE_PAIRS = 415
+
+
+def _iter_native_slice_pairs(lora_model):
+ """Yield ``(module_name, lora_a, lora_b)`` per slice, flattening packed weights."""
+ for module_name, weights in lora_model.loras.items():
+ lora_a = weights.lora_a if isinstance(weights.lora_a, list) else [weights.lora_a]
+ lora_b = weights.lora_b if isinstance(weights.lora_b, list) else [weights.lora_b]
+ for a, b in zip(lora_a, lora_b, strict=True):
+ yield module_name, a, b
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_native_lora_load_is_platform_agnostic(tmp_path, dtype):
+ """The artifact is trained on Ascend, but loading it is not Ascend-only.
+
+ This whole module runs in the CPU job with no ``torch_npu`` present, so a
+ passing run is itself the evidence. The assertions pin the property that
+ makes it portable: tensors land on CPU in exactly the requested dtype and
+ the loader never picks a device, leaving placement to the pipeline.
+ """
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+
+ loaded = load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=dtype,
+ )
+
+ assert loaded is not None
+ lora_model, _, _ = loaded
+ tensors = {
+ f"{module_name}.{side}": tensor
+ for module_name, a, b in _iter_native_slice_pairs(lora_model)
+ for side, tensor in (("lora_a", a), ("lora_b", b))
+ }
+ assert sum(1 for _ in _iter_native_slice_pairs(lora_model)) == _EXPECTED_NATIVE_SLICE_PAIRS
+ assert {name: tensor.device.type for name, tensor in tensors.items() if tensor.device.type != "cpu"} == {}
+ assert {name: tensor.dtype for name, tensor in tensors.items() if tensor.dtype != dtype} == {}
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+def test_native_lora_delta_matches_between_cpu_and_cuda(tmp_path):
+ """The packed adapter produces the same delta weight on CUDA as on CPU.
+
+ Covers the GPU side of the same artifact: every ``lora_b @ lora_a`` product,
+ including the reordered q/k/v slices, is recomputed on CUDA and compared
+ against the CPU result, so a GPU deployment cannot silently diverge.
+ """
+ path = tmp_path / "minimax_h3_t2va_flashgen_4step_v1.0_768p_bf16.safetensors"
+ _write_tiny_native(path)
+
+ loaded = load_minimax_h3_native_lora(
+ partition="fl2va",
+ lora_request=_request(path),
+ lora_path=path,
+ dtype=torch.float32,
+ )
+
+ assert loaded is not None
+ lora_model, _, _ = loaded
+ checked = 0
+ for module_name, a, b in _iter_native_slice_pairs(lora_model):
+ cpu_delta = b @ a
+ cuda_delta = b.cuda() @ a.cuda()
+ assert cuda_delta.device.type == "cuda"
+ torch.testing.assert_close(cuda_delta.cpu(), cpu_delta, msg=f"delta mismatch for {module_name}")
+ checked += 1
+ assert checked == _EXPECTED_NATIVE_SLICE_PAIRS
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
new file mode 100644
index 000000000..c993de492
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py
@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_h3_profiler_targets_cover_aggregate_reference_encoders():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ targets = MiniMaxH3Pipeline._PROFILER_TARGETS
+ assert "_encode_visual_conditions" in targets
+ assert "_encode_reference_audio_conditions" in targets
+ assert "_encode_video_conditions" not in targets
+ assert "_encode_video_audio_conditions" not in targets
+ assert "_encode_audio_conditions" not in targets
+
+
+@pytest.mark.parametrize("distributed", [False, True])
+def test_manual_component_offload_wins_over_requested_model_offload(distributed):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=True,
+ enable_layerwise_offload=not distributed,
+ enable_distributed_layerwise_offload=distributed,
+ )
+ pipeline._model_cpu_offload_modules = []
+ pipeline.text_encoder = Mock()
+ expected = torch.ones(2, 3)
+ pipeline.text_encoder.encode_ids.return_value = expected
+ component = Mock()
+
+ actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+ with pipeline._component_on_device(component):
+ component.load_to_device.assert_called_once_with()
+ component.offload_to_cpu.assert_not_called()
+
+ assert actual is expected
+ pipeline.text_encoder.load_to_device.assert_called_once_with()
+ pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
+ component.offload_to_cpu.assert_called_once_with()
+
+
+def test_requested_model_offload_without_backend_uses_resident_components():
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=True,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=False,
+ )
+ pipeline._model_cpu_offload_modules = []
+ pipeline.text_encoder = Mock()
+ expected = torch.ones(2, 3)
+ pipeline.text_encoder.encode_ids.return_value = expected
+
+ actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
+
+ assert actual is expected
+ pipeline.text_encoder.load_to_device.assert_called_once_with()
+ pipeline.text_encoder.encode_ids.assert_called_once()
+
+
+def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.transformer = torch.nn.Linear(2, 2)
+ pipeline.transformers_ref = torch.nn.Linear(2, 2)
+ pipeline.text_encoder = torch.nn.Linear(2, 2)
+ pipeline.video_vae = torch.nn.Linear(2, 2)
+ pipeline.audio_vae = torch.nn.Linear(2, 2)
+ apply_offload = Mock()
+ remove_offload = Mock()
+ monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
+ monkeypatch.setattr(module, "remove_sequential_offload", remove_offload)
+
+ pipeline.enable_omni_model_cpu_offload(
+ device=torch.device("cpu"),
+ pin_memory=False,
+ use_hsdp=False,
+ )
+
+ dits = [pipeline.transformer, pipeline.transformers_ref]
+ stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
+ apply_offload.assert_called_once_with(
+ dit_modules=dits,
+ encoder_modules=stages,
+ device=torch.device("cpu"),
+ pin_memory=False,
+ use_hsdp=False,
+ offload_initial_dits=True,
+ )
+
+ pipeline.disable_omni_model_cpu_offload()
+
+ remove_offload.assert_called_once_with([*dits, *stages])
+
+
+@pytest.mark.parametrize("decode_fails", [False, True])
+def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ component = torch.nn.Linear(2, 2)
+ events = []
+
+ @contextmanager
+ def record_component(value):
+ events.append(("activate", value))
+ try:
+ yield
+ finally:
+ events.append(("offload", value))
+
+ monkeypatch.setattr(module, "sequential_offload_component", record_component)
+ pipeline._model_cpu_offload_modules = [component]
+
+ def decode():
+ with pipeline._component_on_device(component):
+ events.append(("decode", component))
+ if decode_fails:
+ raise RuntimeError("decode failed")
+
+ if decode_fails:
+ with pytest.raises(RuntimeError, match="decode failed"):
+ decode()
+ else:
+ decode()
+
+ assert events == [
+ ("activate", component),
+ ("decode", component),
+ ("offload", component),
+ ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.video_vae = Mock()
+ pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+ events = []
+
+ @contextmanager
+ def record_component(value):
+ events.append(("activate", value))
+ try:
+ yield
+ finally:
+ events.append(("offload", value))
+
+ def encode_image(image):
+ events.append(("encode", image.size))
+ if encode_fails and image.width == 32:
+ raise RuntimeError("encode failed")
+ return torch.ones(image.width // 16, 4)
+
+ pipeline.video_vae.encode_image.side_effect = encode_image
+ monkeypatch.setattr(module, "sequential_offload_component", record_component)
+ monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+ monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+ images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+
+ if encode_fails:
+ with pytest.raises(RuntimeError, match="encode failed"):
+ pipeline._encode_visual_conditions(images, None, video_count=0)
+ else:
+ rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
+ assert rows.shape == (3, 4)
+ assert shapes == [(1, 1, 1), (1, 1, 2)]
+
+ assert events == [
+ ("activate", pipeline.video_vae),
+ ("encode", (16, 16)),
+ ("encode", (32, 16)),
+ ("offload", pipeline.video_vae),
+ ]
+
+
+def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.video_vae = Mock()
+ pipeline.video_vae.is_distributed_enabled.return_value = False
+ pipeline._model_cpu_offload_modules = [pipeline.video_vae]
+ events = []
+
+ @contextmanager
+ def record_component(value):
+ events.append(("activate", value))
+ try:
+ yield
+ finally:
+ events.append(("offload", value))
+
+ def encode_image(image):
+ events.append(("image", image.size))
+ return torch.ones(1, 4)
+
+ def encode_video(frames):
+ events.append(("video", frames))
+ return torch.ones(3, 4), (3, 2, 2)
+
+ pipeline.video_vae.encode_image.side_effect = encode_image
+ pipeline.video_vae.encode_video.side_effect = encode_video
+ monkeypatch.setattr(module, "sequential_offload_component", record_component)
+ monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
+ monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
+ monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
+ images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
+ prepared_videos = [{"prepared_path": "reference.mp4"}]
+
+ rows, shapes = pipeline._encode_visual_conditions(
+ images,
+ prepared_videos,
+ video_count=1,
+ )
+
+ assert rows.shape == (5, 4)
+ assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
+ assert events == [
+ ("activate", pipeline.video_vae),
+ ("image", (16, 16)),
+ ("image", (32, 16)),
+ ("video", "frames:reference.mp4"),
+ ("offload", pipeline.video_vae),
+ ]
+
+
+@pytest.mark.parametrize("encode_fails", [False, True])
+def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails):
+ from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ torch.nn.Module.__init__(pipeline)
+ pipeline.audio_vae = Mock()
+ pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
+ events = []
+
+ @contextmanager
+ def record_component(value):
+ events.append(("activate", value))
+ try:
+ yield
+ finally:
+ events.append(("offload", value))
+
+ def encode_embedded(*args, **kwargs):
+ events.append(("embedded",))
+ if encode_fails:
+ raise RuntimeError("audio encode failed")
+ return torch.ones(1, 2), [1]
+
+ def encode_standalone(*args, **kwargs):
+ events.append(("standalone",))
+ return torch.ones(1, 2), [1]
+
+ pipeline._encode_video_audio_conditions_resident = encode_embedded
+ pipeline._encode_audio_conditions_resident = encode_standalone
+ monkeypatch.setattr(module, "sequential_offload_component", record_component)
+
+ def call():
+ return pipeline._encode_reference_audio_conditions(
+ [{"input_has_audio": True}],
+ has_audio=[True],
+ standalone_audios=[(torch.ones(1), 16_000)],
+ max_duration_seconds=1.0,
+ )
+
+ if encode_fails:
+ with pytest.raises(RuntimeError, match="audio encode failed"):
+ call()
+ else:
+ result = call()
+ assert result[0].shape == (1, 2)
+ assert result[2].shape == (1, 2)
+
+ assert events[0] == ("activate", pipeline.audio_vae)
+ assert events[-1] == ("offload", pipeline.audio_vae)
+ assert events.count(("activate", pipeline.audio_vae)) == 1
+ assert events.count(("offload", pipeline.audio_vae)) == 1
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py
index efe53f826..53348f579 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py
@@ -74,6 +74,8 @@ def test_t2va_and_fl2va_packing_keep_update_rows_separate():
assert fl2va["update_mask"].sum().item() == 12
assert (~fl2va["update_mask"]).sum().item() == 6
assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
+ assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
+ assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
def test_ref2va_packing_tracks_video_and_audio_update_masks():
@@ -101,6 +103,33 @@ def test_ref2va_packing_tracks_video_and_audio_update_masks():
assert packed["cu_seqlens"].tolist() == [0, 30, 64]
+def test_ref2va_packing_publishes_physical_video_spans_only():
+ from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
+ minimax_h3_packed_sequence_ref2va_blocks,
+ )
+
+ packed = minimax_h3_packed_sequence_ref2va_blocks(
+ text_len=4,
+ latent_t=2,
+ latent_h=4,
+ latent_w=6,
+ audio_t=3,
+ ref_blocks=[
+ {"kind": "image", "latent_h": 4, "latent_w": 4},
+ {"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
+ {"kind": "audio", "ref_audio_t": 1},
+ {"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
+ ],
+ )
+
+ # Image and all audio rows intentionally stay out of the sparse spans.
+ assert packed["video_spans"] == (
+ {"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
+ {"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
+ {"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
+ )
+
+
def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
minimax_h3_audio_cond_noise_aug_rows,
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py
index 2b168499e..45f9229d6 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
@@ -36,7 +36,13 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
- assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+ assert set(MiniMaxH3DiTModel._sp_plan) == {
+ "sp_prepare",
+ "local_sp_prepare",
+ "sp_gather",
+ }
+ assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
+ assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}
model = object.__new__(MiniMaxH3DiTModel)
nn.Module.__init__(model)
@@ -89,6 +95,419 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)
+def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
+ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+ _build_rope_table,
+ )
+
+ local_freqs_half = torch.randn(2, 48)
+ local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1)
+
+ actual = _build_rope_table(local_freqs)
+ expected = torch.cat(
+ (torch.cos(local_freqs_half), torch.sin(local_freqs_half)),
+ dim=-1,
+ ).to(torch.bfloat16)
+
+ assert actual.shape == (2, 96)
+ torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+def test_h3_prepared_rope_table_is_rank_local_and_validated(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ class Rope(nn.Module):
+ def forward(self, position_ids):
+ freqs_half = position_ids[0].float()
+ return torch.cat((freqs_half, freqs_half), dim=-1)
+
+ model = object.__new__(h3.MiniMaxH3DiTModel)
+ nn.Module.__init__(model)
+ model.arch = h3.MiniMaxH3DiTArchConfig(rope_inv_freq_len=1)
+ model.rope = Rope()
+ model.local_sp_prepare = nn.Identity()
+ monkeypatch.setattr(h3, "_sequence_parallel_local_span", lambda *args, **kwargs: (1, 2))
+
+ position_ids = torch.arange(12, dtype=torch.long).reshape(1, 4, 3)
+ actual = model.prepare_rope_table(position_ids, seq_len=4)
+ expected = h3._build_rope_table(model.rope(position_ids[:, 1:3]))
+
+ torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+ model._validate_prepared_rope_table(actual, local_len=2, device=torch.device("cpu"))
+ with pytest.raises(ValueError, match="local_seq_len"):
+ model._validate_prepared_rope_table(actual[:1], local_len=2, device=torch.device("cpu"))
+
+
+def test_denoise_branch_prepares_rope_table_once_per_npu_run(monkeypatch: pytest.MonkeyPatch):
+ from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
+ from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+ MiniMaxH3DenoiseBranch,
+ minimax_h3_denoise_loop,
+ )
+
+ monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: True)
+
+ packed = {
+ "seq_len": torch.tensor(3),
+ "img_pos": torch.tensor([1]),
+ "audio_pos": torch.tensor([2]),
+ "text_pos": torch.tensor([0]),
+ "update_mask": torch.tensor([True]),
+ "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
+ "img_position_ids": torch.zeros(3, 3, dtype=torch.long),
+ "latent_grid": torch.tensor([1, 1, 1]),
+ "video_row_start": torch.tensor(1),
+ }
+ branch = MiniMaxH3DenoiseBranch(
+ packed=packed,
+ text_embeddings=torch.zeros(1, 2),
+ token_tags=torch.zeros(3, dtype=torch.long),
+ device=torch.device("cpu"),
+ )
+
+ class Model:
+ def __init__(self):
+ self.rope_table = torch.zeros(3, 6, dtype=torch.bfloat16)
+ self.prepare_calls = 0
+ self.forward_calls = 0
+
+ def prepare_rope_table(self, img_position_ids, *, seq_len):
+ assert img_position_ids is branch.static_kwargs["img_position_ids"]
+ assert seq_len == 3
+ self.prepare_calls += 1
+ return self.rope_table
+
+ def __call__(self, **kwargs):
+ assert kwargs["rope_table"] is self.rope_table
+ self.forward_calls += 1
+ return torch.zeros(1, 96), torch.zeros(1, 32)
+
+ model = Model()
+ minimax_h3_denoise_loop(
+ model=model,
+ positive=branch,
+ initial_video_rows=torch.zeros(1, 96),
+ initial_audio_rows=torch.zeros(1, 32),
+ keyframe_cond_rows=None,
+ sigmas_video=[1.0, 0.5, 0.0],
+ sigmas_audio=[1.0, 0.5, 0.0],
+ device=torch.device("cpu"),
+ )
+
+ assert model.prepare_calls == 1
+ assert model.forward_calls == 2
+
+
+def test_denoise_branch_keeps_rope_construction_on_the_reference_path_off_npu(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ from vllm_omni.diffusion.models.minimax_h3 import denoise_loop
+ from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+ MiniMaxH3DenoiseBranch,
+ minimax_h3_denoise_loop,
+ )
+
+ monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: False)
+ packed = {
+ "seq_len": torch.tensor(3),
+ "img_pos": torch.tensor([1]),
+ "audio_pos": torch.tensor([2]),
+ "text_pos": torch.tensor([0]),
+ "update_mask": torch.tensor([True]),
+ "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32),
+ "img_position_ids": torch.zeros(3, 3, dtype=torch.long),
+ "latent_grid": torch.tensor([1, 1, 1]),
+ "video_row_start": torch.tensor(1),
+ }
+ branch = MiniMaxH3DenoiseBranch(
+ packed=packed,
+ text_embeddings=torch.zeros(1, 2),
+ token_tags=torch.zeros(3, dtype=torch.long),
+ device=torch.device("cpu"),
+ )
+
+ class Model:
+ def __init__(self):
+ self.forward_calls = 0
+
+ def prepare_rope_table(self, *_args, **_kwargs):
+ raise AssertionError("the reference path must not pre-build a rope_table")
+
+ def __call__(self, **kwargs):
+ assert "rope_table" not in kwargs
+ self.forward_calls += 1
+ return torch.zeros(1, 96), torch.zeros(1, 32)
+
+ model = Model()
+ minimax_h3_denoise_loop(
+ model=model,
+ positive=branch,
+ initial_video_rows=torch.zeros(1, 96),
+ initial_audio_rows=torch.zeros(1, 32),
+ keyframe_cond_rows=None,
+ sigmas_video=[1.0, 0.5, 0.0],
+ sigmas_audio=[1.0, 0.5, 0.0],
+ device=torch.device("cpu"),
+ )
+
+ assert model.forward_calls == 2
+
+
+def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
+ from vllm_omni.diffusion import forward_context
+ from vllm_omni.diffusion.distributed import parallel_state
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+ monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict")
+ monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+ monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+ monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8)
+ monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+ monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1)
+
+ assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2)
+
+
+@pytest.mark.parametrize(
+ ("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"),
+ [
+ (False, "strict", 16, 8, 1, 1),
+ (True, "advanced_uaa", 16, 8, 1, 1),
+ (True, "strict", 17, 8, 1, 1),
+ (True, "strict", 16, 4, 2, 1),
+ (True, "strict", 16, 1, 1, 8),
+ ],
+)
+def test_local_span_falls_back_when_local_embedding_is_unsafe(
+ monkeypatch,
+ hooks_applied,
+ mode,
+ seq_len,
+ ulysses,
+ ring,
+ allgather,
+):
+ from vllm_omni.diffusion import forward_context
+ from vllm_omni.diffusion.distributed import parallel_state
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+ monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode)
+ monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+ monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+ monkeypatch.setattr(
+ parallel_state,
+ "get_ulysses_parallel_world_size",
+ lambda: ulysses,
+ )
+ monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring)
+ monkeypatch.setattr(
+ parallel_state,
+ "get_allgather_parallel_world_size",
+ lambda: allgather,
+ )
+
+ assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == (
+ 0,
+ seq_len,
+ )
+
+
+def test_local_embedding_spans_reassemble_multimodal_rows():
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ class _IdentityProjection(nn.Module):
+ def forward(self, x):
+ return x, None
+
+ class _IdentityRefiner(nn.Module):
+ def forward(self, x, **kwargs):
+ return x
+
+ model = object.__new__(h3.MiniMaxH3DiTModel)
+ nn.Module.__init__(model)
+ model.hidden_size = 2
+ model.video_patch_proj = _IdentityProjection()
+ model.audio_patch_proj = _IdentityProjection()
+ model.condition_proj = _IdentityProjection()
+ model.token_refiner = _IdentityRefiner()
+ model.time_embedder = nn.Identity()
+
+ seq_len = 16
+ x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2)
+ audio_x = x + 100
+ text = torch.tensor(
+ [[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]],
+ dtype=torch.bfloat16,
+ )
+ common = {
+ "x": x,
+ "audio_x": audio_x,
+ "text_embeddings_selected": text,
+ "unique_timesteps": torch.tensor([0.2, 0.3]),
+ "img_pos": torch.tensor([1, 2, 6, 7, 11, 12]),
+ "audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]),
+ "text_pos": torch.tensor([0, 5, 10, 15]),
+ "refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32),
+ "refiner_max_seqlen": 4,
+ "seq_len": seq_len,
+ "device": torch.device("cpu"),
+ "local_span": (0, seq_len),
+ }
+
+ full, full_t_emb = model._embed(**common)
+ local_spans = []
+ for rank in range(8):
+ local_span = (rank * (seq_len // 8), seq_len // 8)
+ local, local_t_emb = model._embed(
+ **{
+ **common,
+ "local_span": local_span,
+ }
+ )
+ assert local.shape == (2, 2)
+ torch.testing.assert_close(local_t_emb, full_t_emb)
+ local_spans.append(local)
+
+ torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+ "local_projection",
+ [False, True],
+ ids=["fallback", "strict-local"],
+)
+def test_sp_output_boundary_preserves_logits_and_minimizes_gather(
+ monkeypatch,
+ local_projection,
+):
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ class _StaticRope(nn.Module):
+ input_rows = None
+
+ def forward(self, position_ids):
+ self.input_rows = position_ids.shape[1]
+ return torch.zeros(position_ids.shape[1], 6)
+
+ class _Prepare(nn.Module):
+ def __init__(self, shard):
+ super().__init__()
+ self.shard = shard
+
+ def forward(self, hidden, rope, combined):
+ local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0]
+ return hidden[:local_len], rope[:local_len], combined[:local_len]
+
+ class _Gather(nn.Module):
+ def __init__(self, remote_rows, replace):
+ super().__init__()
+ self.remote_rows = remote_rows
+ self.replace = replace
+ self.input_width = None
+
+ def forward(self, tensor):
+ self.input_width = tensor.shape[-1]
+ if self.replace:
+ return self.remote_rows
+ return torch.cat((tensor, self.remote_rows), dim=0)
+
+ class _FinalLayer(nn.Module):
+ def forward(self, hidden, *, t_emb, inverse_indices):
+ del t_emb
+ inverse = inverse_indices.to(torch.float32).unsqueeze(-1)
+ return (
+ hidden[:, :2].to(torch.float32) + inverse,
+ hidden[:, 2:3].to(torch.float32) - inverse,
+ )
+
+ seq_len = 8
+ full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4)
+ inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0])
+ final_layer = _FinalLayer()
+ remote_video, remote_audio = final_layer(
+ full_hidden[seq_len // 2 :],
+ t_emb=torch.empty(0),
+ inverse_indices=inverse_indices[seq_len // 2 :],
+ )
+ if local_projection:
+ remote_rows = torch.cat((remote_video, remote_audio), dim=-1)
+ else:
+ remote_rows = full_hidden
+ gather = _Gather(remote_rows, replace=not local_projection)
+
+ model = object.__new__(h3.MiniMaxH3DiTModel)
+ nn.Module.__init__(model)
+ model.arch = h3.MiniMaxH3DiTArchConfig(
+ latents_dim=2,
+ audio_latents_dim=1,
+ patch_size=(1, 1, 1),
+ )
+ rope = _StaticRope()
+ model.rope = rope
+ model.sp_prepare = _Prepare(shard=True)
+ model.local_sp_prepare = _Prepare(shard=False)
+ model.sp_gather = gather
+ model.blocks = nn.ModuleList()
+ model.final_layer = final_layer
+ monkeypatch.setattr(
+ model,
+ "_embed",
+ lambda **kwargs: (
+ full_hidden[: seq_len // 2] if local_projection else full_hidden,
+ torch.tensor([[0.5]], dtype=torch.float32),
+ ),
+ )
+
+ def _local_span(*args, **kwargs):
+ del args, kwargs
+ return (0, seq_len // 2) if local_projection else (0, seq_len)
+
+ monkeypatch.setattr(
+ h3,
+ "_sequence_parallel_local_span",
+ _local_span,
+ )
+
+ img_pos = torch.tensor([0, 2, 4, 6])
+ audio_pos = torch.tensor([1, 3, 5, 7])
+ video, audio = model(
+ x=torch.zeros(1, seq_len, 2),
+ audio_x=torch.zeros(1, seq_len, 1),
+ img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long),
+ unique_timesteps=torch.tensor([0.5]),
+ inverse_indices=inverse_indices,
+ update_mask=torch.ones(img_pos.numel()),
+ token_tags=torch.zeros(seq_len, dtype=torch.long),
+ prompt_embeds=torch.empty(0, 2),
+ img_pos_info={"position_ids": img_pos},
+ audio_pos_info={"position_ids": audio_pos},
+ text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)},
+ img_pos_for_infer_output_info={"position_ids": img_pos},
+ packed_seq_params={
+ "cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32),
+ "max_seqlen_q": seq_len,
+ },
+ refiner_packed_seq_params={
+ "cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32),
+ "max_seqlen_q": 0,
+ },
+ )
+
+ expected_video, expected_audio = final_layer(
+ full_hidden,
+ t_emb=torch.empty(0),
+ inverse_indices=inverse_indices,
+ )
+ torch.testing.assert_close(video, expected_video.index_select(0, img_pos))
+ torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos))
+ expected_gather_width = 3 if local_projection else full_hidden.shape[-1]
+ expected_rope_rows = seq_len // 2 if local_projection else seq_len
+ assert gather.input_width == expected_gather_width
+ assert rope.input_rows == expected_rope_rows
+
+
@pytest.mark.parametrize(
("tp_size", "message"),
[
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
index b9ae3335c..92bd09ac2 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py
@@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
]
+def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
+ from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+ get_direct_mmap_adapter,
+ )
+ from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+ monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
+ monkeypatch.setattr(h3, "Attention", _FakeAttention)
+
+ arch = h3.MiniMaxH3DiTArchConfig(
+ hidden_size=1,
+ num_attention_heads=2,
+ attention_head_dim=1,
+ rope_inv_freq_len=1,
+ )
+ attention = h3.MiniMaxH3Attention(
+ arch,
+ quant_config=None,
+ prefix="blocks.0.attn",
+ )
+ transformer = object.__new__(h3.MiniMaxH3DiTModel)
+ nn.Module.__init__(transformer)
+ transformer.blocks = nn.ModuleList([nn.Module()])
+ transformer.blocks[0].attn = attention
+ pipeline = nn.Module()
+ pipeline.transformer = transformer
+
+ adapter = get_direct_mmap_adapter(pipeline)
+ assert adapter is not None
+ policy = adapter.policy_for(
+ "transformer.blocks.0.attn.qkv_proj.weight",
+ attention.qkv_proj.weight,
+ )
+ assert policy is not None
+ assert policy.allow_custom_loader
+ assert policy.transform is not None
+ checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)
+
+ assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
+ assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
+
+
def test_pipeline_resolves_transformer_component_quant_config():
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
_resolve_component_quant_config,
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
index 80d99895f..0fb043a1e 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from pathlib import Path
@@ -142,12 +142,6 @@ def _generate_joint_output(omni, config: QualityTestConfig):
first = outputs[0]
worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
- if hasattr(first, "request_output") and isinstance(first.request_output, list):
- first = first.request_output[0]
- worker_peak_memory_mb = max(
- worker_peak_memory_mb,
- float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
- )
if not hasattr(first, "images") or not first.images:
raise ValueError("Could not extract video frames from H3 output")
frames = first.images[0]
@@ -193,6 +187,7 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
"model": model_ref,
"enforce_eager": True,
"tensor_parallel_size": 2,
+ # The fused BF16 baseline only fits on H100-80GB with encoder TP.
"text_encoder_tp_size": 2,
"vae_use_tiling": True,
}
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py
new file mode 100644
index 000000000..45a109b28
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_gqa.py
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_minimax_h3_npu_sdpa_patch_is_independent_and_idempotent(monkeypatch) -> None:
+ from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+ from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+ def sdpa_sentinel(query, key, value):
+ del key, value
+ return query
+
+ def rope_sentinel(q, k, cos, sin):
+ del cos, sin
+ return q, k
+
+ monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
+ monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
+ monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+ monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+
+ npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
+
+ assert encoder_module._scaled_dot_product_attention is npu_minimax_h3._scaled_dot_product_attention_npu
+ assert encoder_module._apply_rotary_pos_emb is rope_sentinel
+ assert not npu_minimax_h3._ROPE_PATCHED
+ assert npu_minimax_h3._SDPA_PATCHED
+
+ monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+ npu_minimax_h3.apply_minimax_h3_qwen3vl_sdpa_patch()
+
+ assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize(
+ ("num_heads", "num_kv_heads", "expected_enable_gqa"),
+ [(2, 2, False), (4, 2, True)],
+)
+def test_minimax_h3_npu_sdpa_preserves_compressed_heads_and_selects_gqa(
+ monkeypatch,
+ num_heads: int,
+ num_kv_heads: int,
+ expected_enable_gqa: bool,
+) -> None:
+ from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+ captured: dict[str, object] = {}
+
+ def fake_sdpa(query, key, value, **kwargs):
+ captured.update(query=query, key=key, value=value, kwargs=kwargs)
+ return query
+
+ monkeypatch.setattr(F, "scaled_dot_product_attention", fake_sdpa)
+ query = torch.randn(1, num_heads, 3, 8)
+ key = torch.randn(1, num_kv_heads, 3, 8)
+ value = torch.randn_like(key)
+
+ output = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+ assert captured["query"] is query
+ assert captured["key"] is key
+ assert captured["value"] is value
+ assert captured["kwargs"] == {
+ "dropout_p": 0.0,
+ "is_causal": True,
+ "enable_gqa": expected_enable_gqa,
+ }
+ assert output is query
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize(
+ ("num_heads", "num_key_heads", "num_value_heads", "message"),
+ [
+ (4, 2, 1, "key and value to have the same number of heads"),
+ (4, 0, 0, "at least one KV head"),
+ (3, 2, 2, "query heads to be a multiple of KV heads"),
+ ],
+)
+def test_minimax_h3_npu_sdpa_rejects_invalid_head_configurations(
+ num_heads: int,
+ num_key_heads: int,
+ num_value_heads: int,
+ message: str,
+) -> None:
+ from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+ query = torch.randn(1, num_heads, 3, 8)
+ key = torch.randn(1, num_key_heads, 3, 8)
+ value = torch.randn(1, num_value_heads, 3, 8)
+
+ with pytest.raises(ValueError, match=message):
+ npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_minimax_h3_npu_causal_gqa_matches_expanded_kv() -> None:
+ """Exercise the real torch-npu causal GQA operator."""
+ from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+ torch.manual_seed(0)
+ batch_size, seq_len = 1, 16
+ num_heads, num_kv_heads, head_size = 8, 2, 128
+ device = torch.device("npu")
+ dtype = torch.bfloat16
+
+ query = torch.randn(batch_size, num_heads, seq_len, head_size, device=device, dtype=dtype)
+ key = torch.randn(batch_size, num_kv_heads, seq_len, head_size, device=device, dtype=dtype)
+ value = torch.randn_like(key)
+
+ actual = npu_minimax_h3._scaled_dot_product_attention_npu(query, key, value)
+
+ repeat_num = num_heads // num_kv_heads
+ expected = F.scaled_dot_product_attention(
+ query,
+ key.repeat_interleave(repeat_num, dim=1),
+ value.repeat_interleave(repeat_num, dim=1),
+ dropout_p=0.0,
+ is_causal=True,
+ )
+
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py
new file mode 100644
index 000000000..7422c0b5e
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_rope.py
@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_qwen3_vl_npu_rope_patch_is_independent_and_idempotent(monkeypatch) -> None:
+ from vllm_omni.diffusion.models.minimax_h3 import encoder as encoder_module
+ from vllm_omni.platforms.npu.models import minimax_h3 as npu_minimax_h3
+
+ def rope_sentinel(q, k, cos, sin):
+ del cos, sin
+ return q, k
+
+ def sdpa_sentinel(query, key, value):
+ del key, value
+ return query
+
+ monkeypatch.setattr(npu_minimax_h3, "_ROPE_PATCHED", False)
+ monkeypatch.setattr(npu_minimax_h3, "_SDPA_PATCHED", False)
+ monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+ monkeypatch.setattr(encoder_module, "_scaled_dot_product_attention", sdpa_sentinel)
+
+ npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
+
+ assert encoder_module._apply_rotary_pos_emb is npu_minimax_h3._apply_rotary_pos_emb_npu
+ assert encoder_module._scaled_dot_product_attention is sdpa_sentinel
+ assert npu_minimax_h3._ROPE_PATCHED
+ assert not npu_minimax_h3._SDPA_PATCHED
+
+ monkeypatch.setattr(encoder_module, "_apply_rotary_pos_emb", rope_sentinel)
+ npu_minimax_h3.apply_minimax_h3_qwen3vl_patch()
+
+ assert encoder_module._apply_rotary_pos_emb is rope_sentinel
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)])
+def test_qwen3_vl_npu_rope_matches_bnsd_reference(
+ num_q_heads: int,
+ num_kv_heads: int,
+ seq_len: int,
+) -> None:
+ from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
+
+ device = torch.device("npu")
+ dtype = torch.bfloat16
+ q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype)
+ k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype)
+ cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
+ sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
+
+ q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
+ cos_bnsd = cos.unsqueeze(1)
+ sin_bnsd = sin.unsqueeze(1)
+
+ def rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = x.chunk(2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+ torch.testing.assert_close(
+ q_out,
+ q * cos_bnsd + rotate_half(q) * sin_bnsd,
+ atol=2e-2,
+ rtol=2e-2,
+ )
+ torch.testing.assert_close(
+ k_out,
+ k * cos_bnsd + rotate_half(k) * sin_bnsd,
+ atol=2e-2,
+ rtol=2e-2,
+ )
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None:
+ """Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path."""
+ from types import SimpleNamespace
+
+ from vllm_omni.diffusion.models.minimax_h3.encoder import (
+ MiniMaxH3Qwen3VLTextRotaryEmbedding,
+ )
+ from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
+
+ device = torch.device("npu")
+ dtype = torch.bfloat16
+ batch_size, seq_len, head_dim = 2, 16, 128
+ config = SimpleNamespace(
+ hidden_size=head_dim,
+ num_attention_heads=1,
+ rope_parameters={
+ "rope_theta": 1_000_000.0,
+ "mrope_section": [24, 20, 20],
+ },
+ )
+ rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device)
+ positions = torch.arange(seq_len, device=device)
+ position_ids = torch.stack(
+ (
+ torch.stack((positions, positions + 7)),
+ torch.stack((positions * 2, positions * 2 + 3)),
+ torch.stack((positions * 3, positions * 3 + 5)),
+ )
+ )
+ q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype)
+ k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype)
+ cos, sin = rope(q, position_ids)
+
+ q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
+ cos_bnsd = cos.unsqueeze(1)
+ sin_bnsd = sin.unsqueeze(1)
+
+ def rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = x.chunk(2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+ assert not torch.equal(cos[0], cos[1])
+ torch.testing.assert_close(
+ q_out,
+ q * cos_bnsd + rotate_half(q) * sin_bnsd,
+ atol=2e-2,
+ rtol=2e-2,
+ )
+ torch.testing.assert_close(
+ k_out,
+ k * cos_bnsd + rotate_half(k) * sin_bnsd,
+ atol=2e-2,
+ rtol=2e-2,
+ )
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
new file mode 100644
index 000000000..3ed7649c6
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [
+ pytest.mark.core_model,
+ pytest.mark.diffusion,
+ *hardware_marks(res={"npu": "A3"}, num_cards=1),
+]
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)])
+def test_qwen3_vl_npu_swiglu_matches_split_reference(
+ shape: tuple[int, int, int],
+) -> None:
+ from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
+
+ packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16)
+ gate, up = packed.chunk(2, dim=-1)
+
+ torch.testing.assert_close(
+ npu_swiglu_from_packed(packed),
+ F.silu(gate) * up,
+ atol=2e-2,
+ rtol=2e-2,
+ )
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
+def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None:
+ from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
+
+ x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16)
+ weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16)
+ packed = F.linear(x, weight)
+ gate = F.linear(x, weight[:128])
+ up = F.linear(x, weight[128:])
+
+ torch.testing.assert_close(
+ npu_swiglu_from_packed(packed),
+ F.silu(gate) * up,
+ atol=2e-2,
+ rtol=2e-2,
+ )
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py
index 551fd68a5..65a0221ca 100644
--- a/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_rmsnorm.py
@@ -1,10 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
import torch
+import torch.nn as nn
from vllm_omni.diffusion.layers.norm import RMSNorm
-from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
+from vllm_omni.diffusion.models.minimax_h3.encoder import (
+ MiniMaxH3Qwen3VLRMSNorm,
+ MiniMaxH3Qwen3VLTextDecoderLayer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
@@ -24,3 +34,43 @@ def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
assert norm.weight.dtype == torch.bfloat16
assert set(norm.state_dict()) == {"weight"}
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
+
+
+def test_qwen3_vl_decoder_preserves_post_attention_residual_contract() -> None:
+ class ConstantAttention(nn.Module):
+ def forward(self, hidden_states, position_embeddings):
+ del position_embeddings
+ return torch.full_like(hidden_states, 2.0)
+
+ config = SimpleNamespace(
+ hidden_size=8,
+ intermediate_size=16,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ head_dim=2,
+ rms_norm_eps=1e-6,
+ )
+ group = SimpleNamespace(rank_in_group=0, world_size=1)
+ with (
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
+ patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
+ ):
+ layer = MiniMaxH3Qwen3VLTextDecoderLayer(group, config, torch.float32)
+ layer.input_layernorm = nn.Identity()
+ layer.self_attn = ConstantAttention()
+ layer.mlp = nn.Identity()
+
+ hidden_states = torch.arange(24, dtype=torch.float32).reshape(1, 3, 8)
+ attention_output = torch.full_like(hidden_states, 2.0)
+ expected_residual = hidden_states + attention_output
+ expected_normalized = RMSNorm.forward_native(layer.post_attention_layernorm, expected_residual)
+
+ with patch.object(
+ layer.post_attention_layernorm,
+ "forward",
+ wraps=layer.post_attention_layernorm.forward_native,
+ ):
+ output = layer(hidden_states, (torch.empty(0), torch.empty(0)))
+
+ assert isinstance(layer.post_attention_layernorm, MiniMaxH3Qwen3VLRMSNorm)
+ torch.testing.assert_close(output, expected_residual + expected_normalized)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
new file mode 100644
index 000000000..63c3a42f1
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py
@@ -0,0 +1,473 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MiniMax H3 step-wise execution (continuous batching) contract tests.
+
+These run on CPU against a stand-in DiT, so they cover the packing, the
+per-request scheduler math, and the runner-facing state wiring without needing
+checkpoint weights.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_HIDDEN = 8
+
+
+class _SegmentMeanModel:
+ """Stand-in DiT whose output depends on packed-document boundaries.
+
+ Every row is mixed with the mean of its own ``cu_seqlens`` document and with
+ its own timestep, so a wrong document boundary, row offset, or timestep
+ assignment shows up as a numeric difference instead of passing silently.
+ """
+
+ def __call__(self, **kwargs):
+ x = kwargs["x"][0]
+ audio_x = kwargs["audio_x"][0]
+ bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist()
+ if kwargs["packed_seq_params"].get("num_requests", 1) > 1:
+ assert all(start < stop for start, stop in zip(bounds[:-1], bounds[1:]))
+ img_pos = kwargs["img_pos_info"]["position_ids"]
+ audio_pos = kwargs["audio_pos_info"]["position_ids"]
+
+ pooled_video = torch.zeros_like(x)
+ pooled_audio = torch.zeros_like(audio_x)
+ for start, stop in zip(bounds[:-1], bounds[1:]):
+ if stop <= start:
+ continue
+ pooled_video[start:stop] = x[start:stop].mean(dim=0, keepdim=True)
+ pooled_audio[start:stop] = audio_x[start:stop].mean(dim=0, keepdim=True)
+
+ row_timesteps = kwargs["unique_timesteps"][kwargs["inverse_indices"]].unsqueeze(-1)
+ video = (pooled_video + 0.5 * x + row_timesteps)[img_pos]
+ audio = (pooled_audio + 0.5 * audio_x + row_timesteps)[audio_pos]
+ if not kwargs.get("skip_mask_out_condition", False):
+ video = video * kwargs["update_mask"].view(-1).unsqueeze(-1)
+ return video, audio
+
+
+def _make_branch(*, text_len: int, latent_t: int, latent_h: int, latent_w: int, audio_t: int, seed: int):
+ from vllm_omni.diffusion.models.minimax_h3.denoise_loop import MiniMaxH3DenoiseBranch
+ from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence
+
+ packed = minimax_h3_packed_sequence(
+ text_len=text_len,
+ latent_t=latent_t,
+ latent_h=latent_h,
+ latent_w=latent_w,
+ audio_t=audio_t,
+ include_keyframe_cond=False,
+ )
+ generator = torch.Generator().manual_seed(seed)
+ text_embeddings = torch.randn(text_len, _HIDDEN, generator=generator, dtype=torch.float32)
+ branch = MiniMaxH3DenoiseBranch(
+ packed=packed,
+ text_embeddings=text_embeddings,
+ token_tags=packed["token_tags"],
+ device=torch.device("cpu"),
+ )
+ video_rows = torch.randn(int(branch.img_pos.shape[0]), 96, generator=generator, dtype=torch.float32)
+ audio_rows = torch.randn(int(branch.audio_pos.shape[0]), 32, generator=generator, dtype=torch.float32)
+ return branch, video_rows, audio_rows
+
+
+def _sigmas(num_points: int, shift: float) -> list[float]:
+ from vllm_omni.diffusion.models.minimax_h3.time_request import minimax_h3_time_shift_sigmas
+
+ return minimax_h3_time_shift_sigmas(num_steps=num_points, shift_scale=shift)
+
+
+def _make_state(request_id: str, model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio):
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+ from vllm_omni.diffusion.worker.utils import StepRequestState
+
+ state = StepRequestState(request_id=request_id, sampling=SimpleNamespace())
+ state.latents = video_rows.clone()
+ state.timesteps = torch.tensor([1.0 - sigma for sigma in sigmas_video[:-1]], dtype=torch.float32)
+ state.step_index = 0
+ state.extra = {
+ mod._STEP_BRANCH: branch,
+ # Co-batched requests must share one DiT instance, or denoise_step()
+ # treats the batch as mixed-task and falls back to one forward each.
+ mod._STEP_TRANSFORMER: model,
+ mod._STEP_AUDIO_ROWS: audio_rows.clone(),
+ mod._STEP_COND_ANCHOR: None,
+ mod._STEP_AUDIO_ANCHOR: None,
+ mod._STEP_SIGMAS_VIDEO: sigmas_video,
+ mod._STEP_SIGMAS_AUDIO: sigmas_audio,
+ }
+ return state
+
+
+def _step_pipeline(model, *, packed_batch_supported: bool = True):
+ """A pipeline instance carrying only what the step methods touch."""
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+ pipeline = object.__new__(MiniMaxH3Pipeline)
+ pipeline.transformer = model
+ pipeline.device = torch.device("cpu")
+ pipeline._transformer_for_task = lambda task: model
+ pipeline._packed_batch_supported = lambda transformer: packed_batch_supported
+ return pipeline
+
+
+def test_step_execution_matches_request_mode_denoise_loop():
+ """Stepping through the contract must reproduce the request-mode loop."""
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+ from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+
+ model = _SegmentMeanModel()
+ branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
+ sigmas_video = _sigmas(6, 12.0)
+ sigmas_audio = _sigmas(6, 3.0)
+
+ reference_video, reference_audio = minimax_h3_denoise_loop(
+ model=model,
+ positive=branch,
+ initial_video_rows=video_rows,
+ initial_audio_rows=audio_rows,
+ keyframe_cond_rows=None,
+ sigmas_video=sigmas_video,
+ sigmas_audio=sigmas_audio,
+ device=torch.device("cpu"),
+ )
+
+ pipeline = _step_pipeline(model)
+ state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+ input_batch = SimpleNamespace(states=(state,))
+
+ steps = 0
+ while not state.denoise_completed:
+ noise_pred = pipeline.denoise_step(input_batch, states=[state])
+ pipeline.step_scheduler(state, noise_pred)
+ steps += 1
+
+ assert steps == len(sigmas_video) - 1
+ assert state.total_steps == steps
+ torch.testing.assert_close(state.latents, reference_video)
+ torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
+
+
+def test_batched_step_execution_matches_independent_requests():
+ """Two co-batched requests must land where they would have landed alone."""
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+ model = _SegmentMeanModel()
+ specs = [
+ # Exactly 64 packed rows exercises the no-padding-tail boundary case.
+ dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7),
+ dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
+ ]
+ # Different step counts, so the batch composition changes mid-flight.
+ schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))]
+
+ pipeline = _step_pipeline(model)
+
+ alone: list[tuple[torch.Tensor, torch.Tensor]] = []
+ for spec, (sigmas_video, sigmas_audio) in zip(specs, schedules):
+ branch, video_rows, audio_rows = _make_branch(**spec)
+ state = _make_state("solo", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+ while not state.denoise_completed:
+ pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
+ alone.append((state.latents, state.extra[mod._STEP_AUDIO_ROWS]))
+
+ states = []
+ for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
+ branch, video_rows, audio_rows = _make_branch(**spec)
+ if index == 0:
+ assert branch.used_len == branch.seq_len
+ states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio))
+
+ active = list(states)
+ while active:
+ noise_pred = pipeline.denoise_step(SimpleNamespace(states=tuple(active)), states=active)
+ offset = 0
+ for state in active:
+ rows = state.latents.shape[0]
+ pipeline.step_scheduler(state, noise_pred[offset : offset + rows])
+ offset += rows
+ assert offset == noise_pred.shape[0]
+ # Finished requests leave the batch, exactly like the runner drops them.
+ active = [state for state in active if not state.denoise_completed]
+
+ for state, (expected_video, expected_audio) in zip(states, alone):
+ torch.testing.assert_close(state.latents, expected_video)
+ torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], expected_audio)
+
+
+def test_both_modes_publish_denoise_progress_for_gated_attention():
+ """TRTLLM's skip gate and RAINFUSION's warmup stay dense without these."""
+ from vllm_omni.diffusion import forward_context as fc
+ from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+
+ model = _SegmentMeanModel()
+ branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=14)
+ sigmas_video = _sigmas(4, 12.0)
+ sigmas_audio = _sigmas(4, 3.0)
+
+ published: list[tuple[int | None, float | None]] = []
+
+ class _Recorder:
+ def __init__(self):
+ self.denoise_step_idx = None
+ self.denoise_timestep = None
+
+ def __setattr__(self, name, value):
+ object.__setattr__(self, name, value)
+ if name == "denoise_timestep":
+ published.append((self.denoise_step_idx, value))
+
+ recorder = _Recorder()
+ published.clear() # drop the pair __init__ recorded
+ original = fc._forward_context
+ fc._forward_context = recorder
+ try:
+ minimax_h3_denoise_loop(
+ model=model,
+ positive=branch,
+ initial_video_rows=video_rows,
+ initial_audio_rows=audio_rows,
+ keyframe_cond_rows=None,
+ sigmas_video=sigmas_video,
+ sigmas_audio=sigmas_audio,
+ device=torch.device("cpu"),
+ )
+ request_mode = list(published)
+
+ published.clear()
+ pipeline = _step_pipeline(model)
+ state = _make_state("req-0", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio)
+ while not state.denoise_completed:
+ pipeline.step_scheduler(state, pipeline.denoise_step(SimpleNamespace(states=(state,)), states=[state]))
+ step_mode = list(published)
+ finally:
+ fc._forward_context = original
+
+ num_steps = len(sigmas_video) - 1
+ expected = [(step, sigmas_video[step]) for step in range(num_steps)]
+ assert request_mode == expected + [(None, None)]
+ # Step mode has no loop to close, so it publishes only the per-step pairs.
+ assert step_mode == expected
+
+
+def test_mixed_step_batch_leaves_gated_attention_dense():
+ """A batch spanning different steps has no single progress point to publish."""
+ from vllm_omni.diffusion import forward_context as fc
+
+ model = _SegmentMeanModel()
+ first, first_video, first_audio = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=15)
+ second, second_video, second_audio = _make_branch(
+ text_len=5, latent_t=3, latent_h=6, latent_w=4, audio_t=2, seed=16
+ )
+ sigmas_video, sigmas_audio = _sigmas(6, 12.0), _sigmas(6, 3.0)
+ states = [
+ _make_state("req-0", model, first, first_video, first_audio, sigmas_video, sigmas_audio),
+ _make_state("req-1", model, second, second_video, second_audio, sigmas_video, sigmas_audio),
+ ]
+ states[1].step_index = 2 # admitted later, so it trails the batch
+
+ recorder = SimpleNamespace(denoise_step_idx="unset", denoise_timestep="unset")
+ original = fc._forward_context
+ fc._forward_context = recorder
+ try:
+ _step_pipeline(model).denoise_step(SimpleNamespace(states=tuple(states)), states=states)
+ finally:
+ fc._forward_context = original
+
+ assert recorder.denoise_step_idx is None
+ assert recorder.denoise_timestep is None
+
+
+def test_prepare_encode_seeds_runner_visible_state(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+
+ branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8)
+ sigmas_video = _sigmas(6, 12.0)
+ sigmas_audio = _sigmas(6, 3.0)
+ context = {
+ "height": 96,
+ "width": 64,
+ "latent_t": 2,
+ "latent_h": 4,
+ "latent_w": 6,
+ "audio_t": 3,
+ **{key: None for key in mod._MINIMAX_H3_DENOISE_INPUT_KEYS},
+ }
+
+ pipeline = _step_pipeline(_SegmentMeanModel())
+ monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
+ monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
+ monkeypatch.setattr(
+ mod.MiniMaxH3Pipeline,
+ "_build_denoise_inputs",
+ lambda self, **_: {
+ "branch": branch,
+ "video_rows": video_rows,
+ "audio_rows": audio_rows,
+ "cond_anchor": None,
+ "audio_anchor": None,
+ "sigmas_video": sigmas_video,
+ "sigmas_audio": sigmas_audio,
+ },
+ )
+
+ from vllm_omni.diffusion.worker.utils import StepRequestState
+
+ state = StepRequestState(
+ request_id="req-0",
+ sampling=SimpleNamespace(num_outputs_per_prompt=1),
+ prompt="a prompt",
+ )
+ pipeline.prepare_encode(state)
+
+ # The runner slices the batched velocity by this row count.
+ assert state.latents.shape == (int(branch.img_pos.shape[0]), 96)
+ assert state.total_steps == len(sigmas_video) - 1
+ assert state.step_index == 0
+ assert state.do_true_cfg is False
+ torch.testing.assert_close(state.current_timestep, torch.tensor(1.0 - sigmas_video[0]))
+ assert state.extra[mod._STEP_BRANCH] is branch
+ assert state.extra[mod._STEP_SHAPE]["height"] == 96
+
+
+def test_prepare_encode_rejects_request_mode_only_features():
+ """Multi-output and DLO have no representation in the step contract."""
+ from vllm_omni.diffusion.worker.utils import StepRequestState
+ from vllm_omni.errors import OmniClientError
+
+ # A request state carries exactly one latent tensor.
+ multi_output = StepRequestState(
+ request_id="req-0",
+ sampling=SimpleNamespace(num_outputs_per_prompt=2),
+ prompt="a prompt",
+ )
+ with pytest.raises(OmniClientError, match="one output per request"):
+ _step_pipeline(_SegmentMeanModel()).prepare_encode(multi_output)
+
+ # Distributed layerwise offload streams the DiT around a whole denoise loop.
+ single_output = StepRequestState(
+ request_id="req-1",
+ sampling=SimpleNamespace(num_outputs_per_prompt=1),
+ prompt="a prompt",
+ )
+ dlo_pipeline = _step_pipeline(_SegmentMeanModel())
+ dlo_pipeline._dlo_residency_controller = object()
+ with pytest.raises(ValueError, match="distributed layerwise offload"):
+ dlo_pipeline.prepare_encode(single_output)
+
+
+def test_prepare_encode_rejects_high_quality_cache_dit():
+ """quality=high installs a transformer-scoped Cache-DiT profile that
+ would leak across interleaved or co-batched step-mode requests."""
+ from vllm_omni.diffusion.worker.utils import StepRequestState
+ from vllm_omni.errors import OmniClientError
+
+ state = StepRequestState(
+ request_id="req-hi",
+ sampling=SimpleNamespace(num_outputs_per_prompt=1, quality="high"),
+ prompt="a prompt",
+ )
+ with pytest.raises(OmniClientError, match="quality=high"):
+ _step_pipeline(_SegmentMeanModel()).prepare_encode(state)
+
+
+def _fake_attention_module(
+ *,
+ use_ring: bool,
+ backend: str = "FLASH_ATTN",
+ supports_multi_doc: bool = True,
+):
+ """Build a bare MiniMaxH3Attention with the attributes ``_packed_batch_supported`` reads.
+
+ ``supports_multi_doc`` models the platform-dependent capability probe on
+ ``AttentionBackend`` (e.g. FLASH_ATTN returns True on CUDA/ROCm/MUSA and
+ False on NPU/XPU); ``backend`` remains only for logging/back-compat.
+ """
+ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import MiniMaxH3Attention
+
+ attn = object.__new__(MiniMaxH3Attention)
+ attn.attention = SimpleNamespace(
+ attn_backend=SimpleNamespace(
+ get_name=lambda: backend,
+ supports_multi_doc_packed_varlen=lambda: supports_multi_doc,
+ ),
+ use_ring=use_ring,
+ )
+ return attn
+
+
+class _FakeTransformer:
+ """The minimal ``modules()`` protocol ``_packed_batch_supported`` walks."""
+
+ def __init__(self, attentions):
+ self._attentions = list(attentions)
+
+ def modules(self):
+ return iter([self, *self._attentions])
+
+
+@pytest.mark.parametrize(
+ "attention",
+ [
+ _fake_attention_module(use_ring=True),
+ _fake_attention_module(use_ring=False, backend="XFORMERS", supports_multi_doc=False),
+ _fake_attention_module(use_ring=False, backend="FLASH_ATTN", supports_multi_doc=False),
+ ],
+)
+def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
+
+ assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
+
+
+def test_broadcast_rank0_exception_single_rank_reraises():
+ """Single-rank execution has no group; the helper just reraises."""
+ from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _broadcast_rank0_exception,
+ )
+ from vllm_omni.errors import OmniClientError
+
+ _broadcast_rank0_exception(None)
+ with pytest.raises(OmniClientError, match="bad ref"):
+ _broadcast_rank0_exception(OmniClientError("bad ref"))
+
+
+def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
+ """A rank-0 error becomes a matching client error on every other DiT rank."""
+ from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
+ from vllm_omni.errors import OmniClientError
+
+ def fake_rank_world(rank):
+ return lambda: (object(), rank, 4)
+
+ def make_broadcast(rank0_payload):
+ def fake_broadcast(payload_list, *, src, group):
+ payload_list[0] = rank0_payload
+
+ return fake_broadcast
+
+ monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
+ err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
+ rank0_payload = {
+ "type": type(err).__name__,
+ "message": str(err),
+ "status_code": err.status_code,
+ "error_type": err.error_type,
+ }
+ monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
+ with pytest.raises(OmniClientError) as rank0_info:
+ mod._broadcast_rank0_exception(err)
+ assert rank0_info.value is err
+
+ monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
+ with pytest.raises(OmniClientError) as rank2_info:
+ mod._broadcast_rank0_exception(None)
+ assert rank2_info.value.status_code == 422
+ assert rank2_info.value.error_type == "UnprocessableEntityError"
+ assert "invalid reference-video file" in str(rank2_info.value)
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
new file mode 100644
index 000000000..7a1c54f7b
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py
@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.gpu, pytest.mark.diffusion]
+
+
+def _selected_operators():
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import (
+ resolve_h3_vae_operators,
+ )
+
+ if not current_omni_platform.is_available():
+ pytest.skip("No accelerator is available")
+ device = current_omni_platform.get_torch_device()
+ operators = resolve_h3_vae_operators(device)
+ if operators is None:
+ pytest.skip("No H3 VAE operator implementation is registered for this target")
+ return device, operators
+
+
+def _operator_set(*, supports=lambda _device: True):
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae.dispatch import H3VAEOperatorSet
+
+ return H3VAEOperatorSet(
+ supports=supports,
+ qk_norm_rope=lambda *_args: None,
+ scaled_residual=lambda *_args: None,
+ )
+
+
+def _qk_reference(x, cos, sin):
+ normalized = F.rms_norm(x.float(), (64,), None, 1e-5).to(x.dtype)
+ rotary = normalized[..., :48]
+ first, second = rotary.chunk(2, dim=-1)
+ rotated = torch.cat((-second, first), dim=-1)
+ return torch.cat(
+ (rotary * cos + rotated * sin, normalized[..., 48:]),
+ dim=-1,
+ )
+
+
+def _vit_norm_input(_module, hidden_states):
+ return hidden_states.float()
+
+
+def _same_dtype_norm_input(_module, hidden_states):
+ return hidden_states
+
+
+def _failing_norm_input(_module, _hidden_states):
+ raise RuntimeError("unsupported remote normalization semantics")
+
+
+@pytest.mark.parametrize(("batch", "sequence"), [(1, 1), (1, 195), (2, 1797)])
+def test_h3_vae_qk_norm_rope_is_bit_exact(batch, sequence):
+ device, operators = _selected_operators()
+
+ torch.manual_seed(17)
+ qkv = torch.randn(
+ batch,
+ sequence,
+ 32,
+ 192,
+ device=device,
+ dtype=torch.float16,
+ )
+ q, k, _ = qkv.chunk(3, dim=-1)
+ cos = torch.randn(batch, sequence, 1, 48, device=device, dtype=torch.float16)
+ sin = torch.randn_like(cos)
+
+ expected_q = _qk_reference(q, cos, sin)
+ expected_k = _qk_reference(k, cos, sin)
+ with torch.inference_mode():
+ actual = operators.qk_norm_rope(q, k, (cos, sin), 1e-5)
+
+ assert actual is not None
+ actual_q, actual_k = actual
+ assert torch.equal(actual_q, expected_q)
+ assert torch.equal(actual_k, expected_k)
+
+
+def test_h3_vae_scaled_residual_is_bit_exact():
+ device, operators = _selected_operators()
+
+ torch.manual_seed(29)
+ residual = torch.randn(195, 2048, device=device, dtype=torch.float32)
+ branch = torch.randn(195, 2048, device=device, dtype=torch.float16)
+ scale = torch.randn(2048, device=device, dtype=torch.float32)
+ expected = residual + branch * scale
+
+ with torch.inference_mode():
+ actual = operators.scaled_residual(residual, branch, scale)
+
+ assert actual is not None
+ assert torch.equal(actual, expected)
+
+
+def test_h3_vae_exact_ops_reject_unsupported_inputs():
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae.qk_norm_rope import (
+ try_qk_norm_rope_exact,
+ )
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae.scaled_residual import try_scaled_residual_exact
+
+ q = torch.randn(1, 2, 32, 64, dtype=torch.float32)
+ cos = torch.randn(1, 2, 1, 48, dtype=torch.float32)
+ residual = torch.randn(2, 2048)
+ branch = torch.randn(2, 2048, dtype=torch.float16)
+ scale = torch.randn(2048)
+ with torch.inference_mode():
+ assert try_qk_norm_rope_exact(q, q, (cos, cos), 1e-5) is None
+ assert try_scaled_residual_exact(residual, branch, scale) is None
+
+
+def _make_decoder():
+ class Attention(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.to_qkv = nn.Linear(8, 24)
+ self.to_out = nn.Linear(8, 8)
+ self.norm_q = nn.RMSNorm(8, elementwise_affine=False)
+ self.norm_k = nn.RMSNorm(8, elementwise_affine=False)
+ self.spatial_parallel = False
+ self.dim_head = 8
+
+ def perform_attention(self, query, _key, _value, _pack_info):
+ return query
+
+ def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
+ return hidden_states
+
+ class FeedForward(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.w1 = nn.Linear(8, 32)
+ self.w2 = nn.Linear(16, 8)
+ self.use_gated = True
+ self.act_fn = nn.SiLU()
+ self._compile_forward_enabled = False
+ self._compile_forward_fatal = False
+ self._compiled_forward = None
+
+ def forward(self, hidden_states):
+ hidden_states = self.w1(hidden_states)
+ gate, hidden_states = hidden_states.chunk(2, dim=-1)
+ return self.w2(self.act_fn(gate) * hidden_states)
+
+ class Block(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.attn = Attention()
+ self.ff = FeedForward()
+ self.norm1 = nn.RMSNorm(8)
+ self.norm2 = nn.RMSNorm(8)
+ self.scale1 = nn.Parameter(torch.zeros(8))
+ self.scale2 = nn.Parameter(torch.zeros(8))
+ self.use_scale = True
+
+ def forward(self, hidden_states, rotary_pos_emb=None, pack_info=None):
+ return hidden_states
+
+ decoder = nn.Module()
+ decoder.transformer_blocks = nn.ModuleList([Block(), Block()])
+ decoder.proj_out = nn.Linear(8, 8)
+ return decoder
+
+
+def test_video_vae_installs_exact_optimizations(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
+
+ class FakeModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.decoder = torch.nn.Module()
+
+ class FakeRemote(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.model = FakeModel()
+
+ monkeypatch.setattr(
+ vae_module,
+ "_load_component_config",
+ lambda _path: {
+ "latent_channels": 1,
+ "latents_mean": [0.0],
+ "latents_std": [1.0],
+ },
+ )
+ monkeypatch.setattr(
+ vae_module,
+ "_load_remote_component",
+ lambda _path, _config: FakeRemote(),
+ )
+ install = Mock(return_value=True)
+ monkeypatch.setattr(vae_module, "install_h3_vae_optimizations", install)
+ monkeypatch.setattr(vae_module, "PinnedModuleStager", Mock())
+
+ video_vae = vae_module.MiniMaxH3VideoVAE(
+ "unused",
+ device=torch.device("cuda"),
+ load_device=torch.device("cpu"),
+ )
+
+ install.assert_called_once_with(
+ video_vae.model.decoder,
+ device=torch.device("cuda"),
+ )
+
+
+def test_h3_vae_install_precasts_only_block_linears(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ operators = _operator_set()
+ monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
+ decoder = _make_decoder()
+
+ assert vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+
+ for block in decoder.transformer_blocks:
+ assert block.attn.to_qkv.weight.dtype == torch.float16
+ assert block.attn.to_out.weight.dtype == torch.float16
+ assert block.ff.w1.weight.dtype == torch.float16
+ assert block.ff.w2.weight.dtype == torch.float16
+ assert block.forward.__func__.__name__ == "_optimized_transformer_block"
+ assert block.attn.forward.__func__.__name__ == "_optimized_attention"
+ assert block.ff.forward.__func__.__name__ == "_optimized_feed_forward"
+ assert block.attn._omni_qk_norm_rope is operators.qk_norm_rope
+ assert block._omni_scaled_residual is operators.scaled_residual
+ assert decoder.proj_out.weight.dtype == torch.float32
+
+ # Repeated installation is idempotent.
+ assert vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+
+
+def test_h3_vae_install_accepts_remote_integer_parallel_flag(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ monkeypatch.setattr(
+ vae_ops,
+ "resolve_h3_vae_operators",
+ lambda _device: _operator_set(),
+ )
+ decoder = _make_decoder()
+ for block in decoder.transformer_blocks:
+ # The official checkpoint loads this JSON boolean-like field as int 0.
+ block.attn.spatial_parallel = 0
+
+ assert vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+ assert getattr(decoder, "_omni_h3_vae_optimizations_installed", False)
+
+
+def test_h3_vae_swiglu_uses_post_linear_fp16_output(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ device, operators = _selected_operators()
+ monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: operators)
+ decoder = _make_decoder().to(device)
+ feed_forward = decoder.transformer_blocks[0].ff
+ reference_forward = type(feed_forward).forward
+ assert vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=device,
+ )
+
+ hidden_states = torch.randn(4, 8, device=device, dtype=torch.float32)
+ with torch.inference_mode(), torch.autocast(device.type, dtype=torch.float16):
+ expected = reference_forward(feed_forward, hidden_states)
+ actual = feed_forward(hidden_states)
+
+ assert expected.dtype == torch.float16
+ assert torch.equal(actual, expected)
+
+
+def test_h3_vae_install_leaves_unsupported_target_untouched(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ monkeypatch.setattr(vae_ops, "resolve_h3_vae_operators", lambda _device: None)
+ decoder = _make_decoder()
+
+ assert not vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+ assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
+ for block in decoder.transformer_blocks:
+ assert block.attn.to_qkv.weight.dtype == torch.float32
+ assert block.ff.w1.weight.dtype == torch.float32
+ assert block.forward.__func__.__name__ == "forward"
+ assert block.attn.forward.__func__.__name__ == "forward"
+
+
+@pytest.mark.parametrize("norm_input", [None, _same_dtype_norm_input, _failing_norm_input])
+def test_h3_vae_install_requires_fp32_attention_norm_semantics(monkeypatch, norm_input):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ monkeypatch.setattr(
+ vae_ops,
+ "resolve_h3_vae_operators",
+ lambda _device: _operator_set(),
+ )
+ decoder = _make_decoder()
+ forward_globals = type(decoder.transformer_blocks[0].attn).forward.__globals__
+ monkeypatch.setitem(forward_globals, "_vit_norm_input", norm_input)
+
+ assert not vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+ assert not hasattr(decoder, "_omni_h3_vae_optimizations_installed")
+ assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
+
+
+@pytest.mark.parametrize(
+ "break_contract",
+ [
+ lambda block: delattr(block.attn, "spatial_parallel"),
+ lambda block: setattr(block.attn, "spatial_parallel", 2),
+ lambda block: setattr(block.attn, "perform_attention", None),
+ lambda block: setattr(block, "scale1", nn.Parameter(torch.zeros(7))),
+ lambda block: setattr(block.ff, "_compile_forward_enabled", True),
+ ],
+)
+def test_h3_vae_install_rejects_incompatible_remote_contract(monkeypatch, break_contract):
+ from vllm_omni.diffusion.models.minimax_h3.ops import vae as vae_ops
+
+ monkeypatch.setattr(
+ vae_ops,
+ "resolve_h3_vae_operators",
+ lambda _device: _operator_set(),
+ )
+ decoder = _make_decoder()
+ break_contract(decoder.transformer_blocks[0])
+
+ assert not vae_ops.install_h3_vae_optimizations(
+ decoder,
+ device=torch.device("meta"),
+ )
+ assert decoder.transformer_blocks[0].attn.to_qkv.weight.dtype == torch.float32
+
+
+def test_h3_vae_dispatch_is_extended_by_adding_an_operator_set(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
+
+ first = _operator_set(supports=lambda _device: False)
+ added = _operator_set(supports=lambda device: device.type == "meta")
+ monkeypatch.setattr(dispatch, "H3_VAE_OPERATOR_TABLE", (first, added))
+
+ assert dispatch.resolve_h3_vae_operators(torch.device("meta")) is added
+ assert dispatch.resolve_h3_vae_operators(torch.device("cpu")) is None
+
+
+def test_h3_vae_dispatch_selects_supported_cuda_capabilities(monkeypatch):
+ from vllm_omni.diffusion.models.minimax_h3.ops.vae import dispatch
+
+ platform = Mock()
+ platform.is_cuda.return_value = True
+ platform.is_available.return_value = True
+ monkeypatch.setattr(dispatch, "HAS_TRITON", True)
+ monkeypatch.setattr(dispatch, "current_omni_platform", platform)
+
+ for capability in (90, 100, 103):
+ platform.get_device_capability.return_value.to_int.return_value = capability
+ assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is not None
+
+ for capability in (89, 101, 110):
+ platform.get_device_capability.return_value.to_int.return_value = capability
+ assert dispatch.resolve_h3_vae_operators(torch.device("cuda:0")) is None
diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py
new file mode 100644
index 000000000..ef04019f6
--- /dev/null
+++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CPU tests for the decoder tile-shortage guard in the MiniMax-H3 video VAE."""
+
+import sys
+import types
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+RATIO = 16
+# The checkpoint keeps two independent grids -- ``split_tiles`` selects
+# ``decoder_tile_size``/``decoder_tile_overlap_min`` when ``is_decoder`` is set
+# and ``tile_size``/``tile_overlap_min`` otherwise. The shipped H3 config leaves
+# both at 256/64, so the two happen to coincide there, but that is a config
+# coincidence. The double deliberately gives them different values so that
+# passing the wrong flag changes the tile count and fails a test.
+TILE_SIZE, OVERLAP = 256, 64 # decoder, matching the shipped config
+ENCODER_TILE_SIZE, ENCODER_OVERLAP = 128, 32
+
+
+class _FakeCheckpointModel:
+ """The parts of the checkpoint's AutoencoderKL the guard touches."""
+
+ def __init__(self):
+ self.vae_ratio = RATIO
+ self.parallel_tiling = True
+ self.decoded_with = []
+
+ def split_tiles(self, input_len, is_decoder=False):
+ """Same grid arithmetic as the checkpoint, on pixel dimensions."""
+ tile = TILE_SIZE if is_decoder else ENCODER_TILE_SIZE
+ overlap = OVERLAP if is_decoder else ENCODER_OVERLAP
+ if tile >= input_len:
+ return [0], [input_len], []
+ n = -(-input_len // tile)
+ while tile * n - overlap * (n - 1) - input_len < 0:
+ n += 1
+ return list(range(n)), [tile] * n, [overlap] * (n - 1)
+
+
+def _vae(parallel_size, state):
+ """A stand-in instance: the guard only reads .model/.remote/.parallel_size."""
+ module = types.ModuleType("fake_ckpt.parallel")
+ module.get_parallel_state = lambda: state
+ sys.modules.setdefault("fake_ckpt", types.ModuleType("fake_ckpt"))
+ sys.modules["fake_ckpt.parallel"] = module
+
+ remote = type("Remote", (), {"__module__": "fake_ckpt.klvae"})()
+
+ vae = object.__new__(MiniMaxH3VideoVAE)
+ vae.model = _FakeCheckpointModel()
+ vae.remote = remote
+ vae.parallel_size = parallel_size
+ return vae
+
+
+def _latent(h, w):
+ return torch.zeros(1, 24, 9, h, w)
+
+
+@pytest.mark.parametrize(
+ ("h", "w", "expected"),
+ [
+ (16, 16, 1), # 256x256 px -> single tile
+ (24, 16, 2), # 384x256 px
+ (24, 24, 4), # 384x384 px
+ (48, 84, 28), # 768x1344 px, the shipped H3 request
+ (96, 168, 112), # 1536x2688 px, a larger canvas the request path also allows
+ ],
+)
+def test_tile_count_matches_the_checkpoint_grid(h, w, expected):
+ assert MiniMaxH3VideoVAE._decoder_tile_count(_vae(1, {}), _latent(h, w)) == expected
+
+
+@pytest.mark.parametrize(("h", "w", "decoder_tiles", "encoder_tiles"), [(16, 16, 1, 9), (24, 24, 4, 16)])
+def test_tile_count_uses_the_decoder_grid_not_the_encoder_one(h, w, decoder_tiles, encoder_tiles):
+ """Passing ``is_decoder=False`` would silently change the count, so pin it."""
+ vae = _vae(1, {})
+ assert MiniMaxH3VideoVAE._decoder_tile_count(vae, _latent(h, w)) == decoder_tiles
+
+ px_h, px_w = h * RATIO, w * RATIO
+ as_encoder = len(vae.model.split_tiles(px_h, False)[0]) * len(vae.model.split_tiles(px_w, False)[0])
+ assert as_encoder == encoder_tiles != decoder_tiles
+
+
+def test_rank_local_tiling_restores_the_group_state():
+ state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"}
+ vae = _vae(4, state)
+
+ with vae._rank_local_tiling():
+ assert state["sp_size"] == 1
+ assert state["sp_rank"] == 0
+ assert state["sp_enabled"] is False
+ assert state["sp_process_group"] is None
+ assert vae.model.parallel_tiling is False
+
+ assert state == {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"}
+ assert vae.model.parallel_tiling is True
+
+
+def test_rank_local_tiling_restores_after_an_exception():
+ state = {"sp_size": 4, "sp_rank": 1, "sp_enabled": True, "sp_process_group": "pg"}
+ vae = _vae(4, state)
+
+ with pytest.raises(RuntimeError, match="decode failed"):
+ with vae._rank_local_tiling():
+ raise RuntimeError("decode failed")
+
+ assert state["sp_size"] == 4
+ assert vae.model.parallel_tiling is True
+
+
+# ---------------------------------------------------------------------------
+# Dispatch: every one of these calls ``decode_latent``, so deleting the guard
+# turns the whole matrix red rather than leaving it green.
+# ---------------------------------------------------------------------------
+
+
+class _RecordingProcessor:
+ @staticmethod
+ def revert_tensor(decoded):
+ return decoded
+
+
+def _dispatch_vae(parallel_size):
+ """A VAE whose ``decode_base`` records the parallel state it ran under."""
+ state = {
+ "group_size": parallel_size,
+ "group_rank": 0,
+ "local_process_group": object(),
+ "sp_size": parallel_size,
+ "sp_rank": 0,
+ "sp_enabled": True,
+ "sp_process_group": object(),
+ "tp_size": 1,
+ "tp_rank": 0,
+ }
+ vae = _vae(parallel_size, state)
+ vae.config_dict = {
+ "latent_channels": 24,
+ "latents_mean": [0.0] * 24,
+ "latents_std": [1.0] * 24,
+ }
+ vae.model.processor = _RecordingProcessor()
+
+ seen = {}
+
+ def decode_base(sample):
+ seen["sp_size"] = state["sp_size"]
+ seen["sp_enabled"] = state["sp_enabled"]
+ seen["sp_process_group"] = state["sp_process_group"]
+ seen["parallel_tiling"] = vae.model.parallel_tiling
+ # the guard never inspects the decoded shape, so keep it small
+ return torch.zeros(1, 3, 2, 4, 4)
+
+ vae.model.decode_base = decode_base
+ return vae, state, seen
+
+
+@pytest.mark.parametrize(
+ ("parallel_size", "h", "w", "falls_back"),
+ [
+ (4, 24, 16, True), # 2 tiles, ranks 2 and 3 would get none
+ (4, 16, 16, True), # 1 tile
+ (4, 24, 24, False), # 4 tiles, exactly enough
+ (4, 48, 84, False), # 28 tiles, the shipped 768x1344 request
+ (4, 96, 168, False), # 112 tiles, 1536x2688
+ (2, 24, 16, False), # 2 tiles across 2 ranks
+ (1, 16, 16, False), # never parallel, nothing to guard
+ ],
+)
+def test_decode_latent_falls_back_exactly_when_tiles_are_short(parallel_size, h, w, falls_back):
+ """Fewer tiles than ranks is the hang condition; equal or more is fine."""
+ vae, _, seen = _dispatch_vae(parallel_size)
+
+ MiniMaxH3VideoVAE.decode_latent(vae, _latent(h, w))
+
+ # only the guard clears ``parallel_tiling``, so it is the signal that it fired
+ assert (seen["parallel_tiling"] is False) is falls_back
+ assert seen["sp_size"] == (1 if falls_back else parallel_size)
+
+
+def test_decode_latent_isolates_and_restores_the_whole_group_state():
+ """1 tile, 4 ranks: ``decode_base`` must observe rank-local tiling."""
+ vae, state, seen = _dispatch_vae(4)
+
+ frames = MiniMaxH3VideoVAE.decode_latent(vae, _latent(16, 16))
+
+ assert seen["sp_size"] == 1
+ assert seen["sp_enabled"] is False
+ assert seen["sp_process_group"] is None
+ assert seen["parallel_tiling"] is False
+ assert frames.ndim == 5 and frames.dtype is torch.float32
+ # the group state is restored once decode returns
+ assert state["sp_size"] == 4 and state["sp_enabled"] is True
+ assert vae.model.parallel_tiling is True
+
+
+def test_decode_latent_leaves_the_group_state_alone_when_the_grid_is_large_enough():
+ """The shipped 768x1344 request is 28 tiles: the guard must not fire."""
+ vae, state, seen = _dispatch_vae(4)
+
+ MiniMaxH3VideoVAE.decode_latent(vae, _latent(48, 84))
+
+ assert seen["sp_size"] == 4
+ assert seen["sp_enabled"] is True
+ assert seen["sp_process_group"] is state["sp_process_group"]
+ assert seen["parallel_tiling"] is True
diff --git a/vllm_omni/model_executor/models/ming_tts/fm/__init__.py b/tests/diffusion/models/pi0/__init__.py
similarity index 100%
rename from vllm_omni/model_executor/models/ming_tts/fm/__init__.py
rename to tests/diffusion/models/pi0/__init__.py
diff --git a/tests/diffusion/models/pi0/test_pi0_parity.py b/tests/diffusion/models/pi0/test_pi0_parity.py
new file mode 100644
index 000000000..3b6358b50
--- /dev/null
+++ b/tests/diffusion/models/pi0/test_pi0_parity.py
@@ -0,0 +1,366 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle.
+
+Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit
+matching action chunks with LeRobot's ``PI0Policy`` when fed:
+ * the same weights (``lerobot/pi0_base``)
+ * the same pre-processed inputs (images, masks, tokens, state)
+ * the same initial noise tensor
+
+π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed
+``num_steps``), so the output is deterministic once the noise is fixed and
+``torch.allclose`` on the final action chunk is a valid oracle. This is the
+authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4).
+
+The OpenPI websocket online-serving e2e lives separately in
+``tests/e2e/online_serving/test_pi0_expansion.py``.
+
+Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni
+env), with the vllm-omni ``pi0`` package importable::
+
+ python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s
+
+Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env).
+The run is CPU/float32 with fixed defaults; the only override is
+``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the
+HF download; defaults to ``lerobot/pi0_base``).
+"""
+
+from __future__ import annotations
+
+import copy
+import importlib.util
+import os
+
+import pytest
+import torch
+
+# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it
+# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md).
+# Additionally gated on lerobot being importable.
+_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+
+# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ──
+DEVICE = "cpu"
+DTYPE_STR = "float32"
+ATOL = 1e-4
+NUM_STEPS = 10
+BATCH_SIZE = 2
+ACTION_DIM = 32
+STATE_DIM = 32
+ACTION_HORIZON = 50
+MAX_TOKEN_LEN = 48
+
+# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF
+# download; defaults to the HF repo id.
+MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base")
+
+
+def _resolve_checkpoint_dir() -> str:
+ """Return a local dir containing the pi0_base checkpoint (download if needed)."""
+ if os.path.isdir(MODEL_PATH):
+ return MODEL_PATH
+ from huggingface_hub import snapshot_download
+
+ return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
+
+
+# ─── Dummy dataset stats (identity) ──────────────────────────────────
+def _dummy_dataset_stats() -> dict:
+ return {
+ "observation.state": {
+ "mean": torch.zeros(STATE_DIM),
+ "std": torch.ones(STATE_DIM),
+ "q01": torch.zeros(STATE_DIM),
+ "q99": torch.ones(STATE_DIM),
+ },
+ "action": {
+ "mean": torch.zeros(ACTION_DIM),
+ "std": torch.ones(ACTION_DIM),
+ "q01": torch.zeros(ACTION_DIM),
+ "q99": torch.ones(ACTION_DIM),
+ },
+ "images": {
+ cam: {
+ "mean": torch.zeros(3, 224, 224),
+ "std": torch.ones(3, 224, 224),
+ "q01": torch.zeros(3, 224, 224),
+ "q99": torch.ones(3, 224, 224),
+ }
+ for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+ },
+ }
+
+
+def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict:
+ """Reproducible dummy inputs — identical across both implementations."""
+ g = torch.Generator(device="cpu").manual_seed(0)
+ prompt = "Pick up the red block and place it in the bin"
+ return {
+ "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
+ "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
+ "observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+ device
+ ),
+ "observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+ device
+ ),
+ "observation.images.right_wrist_0_rgb": torch.rand(
+ batch_size, 3, 224, 224, generator=g, dtype=torch.float32
+ ).to(device),
+ "task": [prompt for _ in range(batch_size)],
+ }
+
+
+# ─── LeRobot instantiation ────────────────────────────────────────────
+def _instantiate_lerobot():
+ from lerobot.policies.pi0 import PI0Policy
+ from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+
+ policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True)
+ policy.to(DEVICE)
+ policy.config.device = DEVICE
+ policy.eval()
+
+ pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
+ return policy, pre, post
+
+
+# ─── vllm-omni instantiation ──────────────────────────────────────────
+def _instantiate_vllm_omni():
+ """Build the vllm-omni π0 model in isolation (no pipeline, no engine)."""
+ from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction
+
+ cfg = Pi0Config(
+ max_action_dim=ACTION_DIM,
+ max_state_dim=STATE_DIM,
+ chunk_size=ACTION_HORIZON,
+ num_inference_steps=NUM_STEPS,
+ dtype=DTYPE_STR,
+ )
+ model = Pi0ForActionPrediction(cfg)
+ model.to(DEVICE).eval()
+ _load_lerobot_weights(model)
+ return model
+
+
+def _load_lerobot_weights(model):
+ """Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model.
+
+ ``Pi0ForActionPrediction.load_weights`` handles the leading ``model.``
+ prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass
+ the raw checkpoint dict.
+ """
+ import safetensors.torch
+
+ cache_dir = _resolve_checkpoint_dir()
+ path = os.path.join(cache_dir, "model.safetensors")
+ state = safetensors.torch.load_file(path)
+ model.load_weights(list(state.items()))
+
+
+# ─── Helpers to extract LeRobot's pre-processed inputs ────────────────
+def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
+ """Mimic what ``PI0Policy.predict_action_chunk`` feeds into
+ ``self.model.sample_actions``. We use these *exact* tensors for vllm-omni
+ so any divergence must come from the core model, not preprocessing.
+ """
+ images, img_masks = lerobot_policy._preprocess_images(processed_batch)
+ from lerobot.utils.constants import (
+ OBS_LANGUAGE_ATTENTION_MASK,
+ OBS_LANGUAGE_TOKENS,
+ )
+
+ lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS]
+ lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
+ state = lerobot_policy.prepare_state(processed_batch)
+ return images, img_masks, lang_tokens, lang_masks, state
+
+
+# ─── Shared fixed-noise sampler ───────────────────────────────────────
+def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
+ g = torch.Generator(device="cpu").manual_seed(42)
+ return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
+
+
+# ─── Main test ────────────────────────────────────────────────────────
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi0_vllm_omni_vs_lerobot():
+ print("\n[parity] Instantiating LeRobot…")
+ lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
+
+ print("[parity] Instantiating vllm-omni…")
+ omni_model = _instantiate_vllm_omni()
+
+ print("[parity] Preparing shared inputs…")
+ raw_batch = _create_dummy_batch()
+ processed_batch = lerobot_pre(copy.deepcopy(raw_batch))
+ images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch)
+ noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+ print(f"[parity] state.shape={state.shape} lang_tokens.shape={lang_tokens.shape}")
+ print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})")
+ print(f"[parity] noise.shape={noise.shape} dtype={noise.dtype}")
+
+ # ── LeRobot forward ──
+ print("[parity] Running LeRobot sample_actions…")
+ with torch.no_grad():
+ lerobot_actions = lerobot_policy.model.sample_actions(
+ images,
+ img_masks,
+ lang_tokens,
+ lang_masks,
+ state,
+ noise=noise,
+ num_steps=NUM_STEPS,
+ )
+ print(
+ f"[parity] LeRobot actions: shape={lerobot_actions.shape} "
+ f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}"
+ )
+
+ # ── vllm-omni forward ──
+ print("[parity] Running vllm-omni sample_actions…")
+ with torch.no_grad():
+ omni_actions = omni_model.sample_actions(
+ images=images,
+ image_masks=img_masks,
+ lang_tokens=lang_tokens,
+ lang_masks=lang_masks,
+ state=state,
+ noise=noise,
+ num_steps=NUM_STEPS,
+ )
+ print(
+ f"[parity] vllm-omni actions: shape={omni_actions.shape} "
+ f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}"
+ )
+
+ # ── Compare ──
+ diff = (lerobot_actions.float() - omni_actions.float()).abs()
+ max_diff = diff.max().item()
+ mean_diff = diff.mean().item()
+ print(f"[parity] |Δ| max={max_diff:.2e} mean={mean_diff:.2e} atol={ATOL:.1e}")
+ close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL)
+ print(f"[parity] torch.allclose(atol={ATOL}): {close}")
+
+ if not close:
+ print("\n[parity] ⚠️ Outputs diverge — running per-stage diagnostics…")
+ _diagnose_divergence(
+ lerobot_policy.model,
+ omni_model,
+ images,
+ img_masks,
+ lang_tokens,
+ lang_masks,
+ state,
+ noise,
+ )
+
+ assert close, (
+ f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e} mean_diff={mean_diff:.2e}"
+ )
+
+
+# ─── Per-stage divergence diagnostics ─────────────────────────────────
+@torch.no_grad()
+def _diagnose_divergence(
+ lerobot_flow_model,
+ omni_model,
+ images,
+ img_masks,
+ lang_tokens,
+ lang_masks,
+ state,
+ noise,
+):
+ """Localize a numerical mismatch to a specific pipeline stage:
+ 1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
+ 2. Prefix KV cache layer 0 — PaliGemma LM attention.
+ 3. A single denoise_step velocity at t=1.0 — action expert forward.
+ """
+ from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+ make_att_2d_masks,
+ prepare_attention_masks_4d,
+ )
+
+ # ── Stage 1: prefix embeddings ──
+ lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix(
+ images, img_masks, lang_tokens, lang_masks
+ )
+ sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+ total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item()
+ print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e} (shape={tuple(sg_prefix_embs.shape)})")
+ print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}")
+ print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}")
+
+ num_cams = len(images)
+ img_len = 256 * num_cams
+ lang_len = lr_prefix_embs.shape[1] - img_len
+ img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item()
+ lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item()
+ print(f"[diag] image slice [:{img_len}] max|Δ| = {img_diff:.2e} (num_cams={num_cams})")
+ print(f"[diag] lang slice [{img_len}:] max|Δ| = {lang_diff:.2e} (len={lang_len})")
+
+ def _stats(name, t):
+ t = t.float()
+ print(
+ f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} "
+ f"std={t.std().item():.4f} min={t.min().item():+.2f} "
+ f"max={t.max().item():+.2f}"
+ )
+
+ _stats("LeRobot ", lr_prefix_embs)
+ _stats("vllm-omni", sg_prefix_embs)
+
+ # ── Stage 2: prefix KV cache ──
+ prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att)
+ prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1
+ prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d)
+
+ lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager"
+ _, lr_kv = lerobot_flow_model.paligemma_with_expert.forward(
+ attention_mask=prefix_att_4d,
+ position_ids=prefix_pos,
+ past_key_values=None,
+ inputs_embeds=[lr_prefix_embs, None],
+ use_cache=True,
+ )
+ _, sg_kv = omni_model.paligemma_with_expert.forward(
+ attention_mask=prefix_att_4d,
+ position_ids=prefix_pos,
+ past_key_values=None,
+ inputs_embeds=[sg_prefix_embs, None],
+ use_cache=True,
+ )
+
+ def _layer0_kv(cache):
+ if isinstance(cache, list):
+ return cache[0]
+ layer0 = cache.layers[0]
+ return layer0.keys, layer0.values
+
+ try:
+ lr_k, lr_v = _layer0_kv(lr_kv)
+ sg_k, sg_v = _layer0_kv(sg_kv)
+ dk = (lr_k.float() - sg_k.float()).abs().max().item()
+ dv = (lr_v.float() - sg_v.float()).abs().max().item()
+ print(f"[diag] prefix KV layer0 K max|Δ|={dk:.2e} V max|Δ|={dv:.2e}")
+ except Exception as e: # noqa: BLE001
+ print(f"[diag] could not extract prefix KV for comparison: {e}")
+
+ # ── Stage 3: a single denoise_step at t=1.0 ──
+ bsize = state.shape[0]
+ t = torch.ones(bsize, dtype=torch.float32, device=state.device)
+ lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t)
+ sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t)
+ print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}")
+
+
+if __name__ == "__main__":
+ test_pi0_vllm_omni_vs_lerobot()
diff --git a/tests/diffusion/models/pi0/test_pi0_units.py b/tests/diffusion/models/pi0/test_pi0_units.py
new file mode 100644
index 000000000..a1fa18f1a
--- /dev/null
+++ b/tests/diffusion/models/pi0/test_pi0_units.py
@@ -0,0 +1,585 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU unit tests for the π0 (Pi-Zero) VLA model.
+
+Deliberately lightweight: pure utility code (config plumbing, mask building,
+sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
+that runs on CPU with no weights and no lerobot. The full LeRobot parity check
+lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
+
+ pytest tests/diffusion/models/pi0/test_pi0_units.py -v
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+ OPENPI_ATTENTION_MASK_VALUE,
+ Pi0ForActionPrediction,
+ _apply_norm,
+ _build_norm_buffers,
+ create_sinusoidal_pos_embedding,
+ get_gemma_config,
+ make_att_2d_masks,
+ prepare_attention_masks_4d,
+)
+from vllm_omni.diffusion.models.pi0.processor_pi0 import (
+ Pi0ImageProcessor,
+ build_model_inputs,
+ pil_image_to_tensor,
+ resize_with_pad,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ----------------------------------------------------------------------------
+# Pi0Config dataclass + checkpoint resolver
+# ----------------------------------------------------------------------------
+_LEROBOT_CFG = {
+ "type": "pi0",
+ "paligemma_variant": "gemma_2b",
+ "action_expert_variant": "gemma_300m",
+ "chunk_size": 50,
+ "max_action_dim": 32,
+ "max_state_dim": 32,
+ "num_inference_steps": 10,
+ "image_resolution": [224, 224],
+ "tokenizer_max_length": 48,
+ "dtype": "float32",
+ # Training-only keys that must be filtered out of the dataclass.
+ "optimizer_lr": 2.5e-5,
+ "scheduler_warmup_steps": 1000,
+ "input_features": {
+ "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+ "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+ "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+ "observation.state": {"type": "STATE", "shape": [32]},
+ },
+}
+
+_EXPECTED_CAMERA_ORDER = [
+ "observation.images.base_0_rgb",
+ "observation.images.left_wrist_0_rgb",
+ "observation.images.right_wrist_0_rgb",
+]
+
+
+def test_config_camera_order_from_input_features():
+ c = Pi0Config.from_model_config(_LEROBOT_CFG)
+ assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
+
+
+def test_config_filters_training_only_keys():
+ c = Pi0Config.from_model_config(_LEROBOT_CFG)
+ assert not hasattr(c, "optimizer_lr")
+ assert not hasattr(c, "scheduler_warmup_steps")
+
+
+def test_config_coerces_image_resolution_to_tuple():
+ c = Pi0Config.from_model_config(_LEROBOT_CFG)
+ assert c.image_resolution == (224, 224)
+ assert isinstance(c.image_resolution, tuple)
+
+
+def test_config_non_square_resolution_raises():
+ with pytest.raises(ValueError):
+ Pi0Config(image_resolution=(224, 256))
+
+
+@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
+def test_config_malformed_resolution_raises(bad):
+ with pytest.raises(ValueError):
+ Pi0Config(image_resolution=bad)
+
+
+def test_config_defaults_and_empty():
+ c = Pi0Config.from_model_config(None)
+ assert c.chunk_size == 50
+ assert c.max_action_dim == 32
+ assert c.max_state_dim == 32
+ assert c.num_inference_steps == 10
+ assert c.image_resolution == (224, 224)
+ assert c.tokenizer_max_length == 48
+ assert c.max_cameras == 3
+ assert c.norm_stats is None
+ assert c.image_key_map == {}
+
+
+def test_config_deploy_yaml_override():
+ """Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
+ c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
+ assert c.dtype == "bfloat16"
+ assert c.num_inference_steps == 4
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: sinusoidal timestep embedding
+# ----------------------------------------------------------------------------
+def test_sinusoidal_embedding_shape():
+ emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
+ assert emb.shape == (3, 64)
+
+
+def test_sinusoidal_embedding_values_at_zero():
+ emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
+ assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
+ assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
+
+
+def test_sinusoidal_embedding_odd_dim_raises():
+ with pytest.raises(ValueError):
+ create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
+
+
+def test_sinusoidal_embedding_non_1d_raises():
+ with pytest.raises(ValueError):
+ create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
+
+
+def test_sinusoidal_embedding_distinct_timesteps_differ():
+ emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
+ for i in range(emb.shape[0]):
+ for j in range(i + 1, emb.shape[0]):
+ assert not torch.allclose(emb[i], emb[j])
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: attention mask construction
+# ----------------------------------------------------------------------------
+def test_make_att_2d_masks_bidirectional():
+ result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
+ assert result.all()
+
+
+def test_make_att_2d_masks_causal():
+ result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
+ expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
+ assert torch.equal(result[0], expected)
+
+
+def test_make_att_2d_masks_prefix_lm():
+ result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
+ assert result[0, 0, 0]
+ assert not result[0, 0, 3]
+ assert result[0, 3, 0]
+ assert not result[0, 3, 4]
+
+
+def test_make_att_2d_masks_respects_padding():
+ pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool) # last token is padding
+ att = torch.zeros(1, 4, dtype=torch.long) # full bidirectional
+ result = make_att_2d_masks(pad, att)
+ assert not result[0, :, 3].any(), "Nothing should attend to padded key"
+ assert not result[0, 3, :].any(), "Padded query should attend to nothing"
+
+
+def test_prepare_attention_masks_4d():
+ mask_2d = torch.tensor([[[True, False], [True, True]]])
+ mask_4d = prepare_attention_masks_4d(mask_2d)
+ assert mask_4d.shape == (1, 1, 2, 2)
+ assert mask_4d[0, 0, 0, 0].item() == 0.0
+ assert mask_4d[0, 0, 1, 1].item() == 0.0
+ expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
+ assert torch.equal(mask_4d[0, 0, 0, 1], expected)
+
+
+def test_attention_mask_value_finite_and_negative():
+ assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
+ assert OPENPI_ATTENTION_MASK_VALUE < -1e30
+
+
+# ----------------------------------------------------------------------------
+# Normalization round-trips
+# ----------------------------------------------------------------------------
+def test_build_norm_buffers_missing():
+ assert _build_norm_buffers(None, "state") is None
+ assert _build_norm_buffers({}, "state") is None
+ assert _build_norm_buffers({"state": None}, "state") is None
+
+
+def test_build_norm_buffers_mean_std():
+ stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
+ out = _build_norm_buffers(stats, "state")
+ assert out["mode"] == "mean_std"
+ assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
+ assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
+
+
+def test_build_norm_buffers_min_max():
+ stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
+ out = _build_norm_buffers(stats, "action")
+ assert out["mode"] == "min_max"
+ assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
+ assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
+
+
+def test_apply_norm_mean_std_forward_inverse():
+ stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
+ x = torch.tensor([[1.0, 0.3]])
+ z = _apply_norm(x, stats, inverse=False)
+ assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
+ x_back = _apply_norm(z, stats, inverse=True)
+ assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_min_max_forward_inverse():
+ stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
+ x = torch.tensor([[0.25, 0.5]])
+ z = _apply_norm(x, stats, inverse=False)
+ assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
+ x_back = _apply_norm(z, stats, inverse=True)
+ assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_preserves_padded_tail():
+ stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
+ x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
+ z = _apply_norm(x, stats, inverse=False)
+ assert torch.allclose(z[:, :2], x[:, :2])
+ assert torch.allclose(z[:, 2:], x[:, 2:])
+
+
+def test_apply_norm_noop_when_stats_none():
+ x = torch.randn(2, 5)
+ assert torch.equal(_apply_norm(x, None, inverse=False), x)
+ assert torch.equal(_apply_norm(x, None, inverse=True), x)
+
+
+# ----------------------------------------------------------------------------
+# Gemma variant configs
+# ----------------------------------------------------------------------------
+def test_gemma_2b_dims():
+ c = get_gemma_config("gemma_2b")
+ assert c.width == 2048
+ assert c.depth == 18
+
+
+def test_gemma_300m_dims():
+ c = get_gemma_config("gemma_300m")
+ assert c.width == 1024
+ assert c.depth == 18
+
+
+def test_gemma_unknown_variant_raises():
+ with pytest.raises(ValueError):
+ get_gemma_config("gemma_7b")
+
+
+# ----------------------------------------------------------------------------
+# load_weights remap rules (tiny model; needs torch + transformers)
+# ----------------------------------------------------------------------------
+def _tiny_pi0_model():
+ """Construct a tiny π0 model for remap tests (full Gemma depth but small
+ action/state dims — the PaliGemma backbone is always full-size)."""
+ cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
+ return Pi0ForActionPrediction(cfg)
+
+
+@pytest.fixture(autouse=True)
+def _silence_pi0_loader(caplog):
+ # The remap tests feed a single synthetic weight; the loader's
+ # "missing params" warning is expected noise here.
+ logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
+ yield
+
+
+@pytest.mark.slow
+def test_remap_paligemma_submodules_to_nested_layout():
+ model = _tiny_pi0_model()
+ params = dict(model.named_parameters())
+ target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
+ flat_key = "model." + target.replace(
+ "paligemma_with_expert.paligemma.model.vision_tower.",
+ "paligemma_with_expert.paligemma.vision_tower.",
+ )
+ payload = torch.full(params[target].shape, 0.12345)
+ before = params[target].detach().clone()
+ model.load_weights([(flat_key, payload)])
+ after = dict(model.named_parameters())[target]
+ assert not torch.equal(before, after), "flat→nested remap broken"
+ assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_lm_head_remap_to_embed_tokens():
+ model = _tiny_pi0_model()
+ params = dict(model.named_parameters())
+ target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+ assert target in params
+ payload = torch.full(params[target].shape, 0.5)
+ before = params[target].detach().clone()
+ model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
+ after = dict(model.named_parameters())[target]
+ assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
+ assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_vision_tower_version_robust_remap():
+ """transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
+ vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
+ load_weights must adapt so vision weights aren't silently dropped (which
+ would leave the vision encoder at random init)."""
+ model = _tiny_pi0_model()
+ params = dict(model.named_parameters())
+ # A representative vision-tower param as it exists on the model.
+ target = next(
+ k
+ for k in params
+ if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
+ )
+ vt = "paligemma_with_expert.paligemma.model.vision_tower."
+ rest = target[len(vt) :]
+ # Build the checkpoint key in whichever layout differs from the model's.
+ if rest.startswith("vision_model."):
+ ckpt_key = vt + rest[len("vision_model.") :]
+ else:
+ ckpt_key = vt + "vision_model." + rest
+ payload = torch.full(params[target].shape, 0.073)
+ before = params[target].detach().clone()
+ model.load_weights([("model." + ckpt_key, payload)])
+ after = dict(model.named_parameters())[target]
+ assert not torch.equal(before, after), "vision_tower version-robust remap broken"
+ assert torch.allclose(after, payload)
+
+
+# ----------------------------------------------------------------------------
+# Image processor
+# ----------------------------------------------------------------------------
+def test_resize_with_pad_square():
+ out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
+ assert out.shape == (1, 3, 224, 224)
+
+
+def test_resize_with_pad_wide():
+ """Wide source (H < W): padding applied to top/bottom (rows)."""
+ out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
+ assert out.shape == (1, 3, 224, 224)
+ assert out[0, 0, 0, 112].item() == -1.0 # top row is padded
+
+
+def test_resize_with_pad_tall():
+ """Tall source (H > W): padding applied to left/right (cols)."""
+ out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
+ assert out.shape == (1, 3, 224, 224)
+ assert out[0, 0, 112, 0].item() == -1.0 # left-most column padded
+
+
+def test_resize_with_pad_rejects_non_4d():
+ with pytest.raises(ValueError):
+ resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
+
+
+def test_pil_to_tensor():
+ import numpy as np
+ from PIL import Image
+
+ img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
+ t = pil_image_to_tensor(img)
+ assert t.shape == (1, 3, 10, 10)
+ expected = 128.0 / 255.0 * 2.0 - 1.0
+ assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
+
+
+def test_pil_to_tensor_rgba_to_rgb():
+ import numpy as np
+ from PIL import Image
+
+ rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
+ img = Image.fromarray(rgba, mode="RGBA")
+ t = pil_image_to_tensor(img)
+ assert t.shape == (1, 3, 8, 8)
+
+
+def test_make_empty_image():
+ proc = Pi0ImageProcessor(image_size=224)
+ empty = proc.make_empty_image()
+ assert empty.shape == (1, 3, 224, 224)
+ assert (empty == -1.0).all()
+
+
+def test_preprocess_single_resizes_ndarray():
+ import numpy as np
+
+ proc = Pi0ImageProcessor(image_size=224)
+ arr = np.full((120, 120, 3), 100, dtype=np.uint8)
+ t = proc.preprocess_single(arr)
+ assert t.shape == (1, 3, 224, 224)
+
+
+# ----------------------------------------------------------------------------
+# build_model_inputs: camera ordering, padding, masks, state
+# ----------------------------------------------------------------------------
+class _FakeTokenizer:
+ """Minimal stand-in for the PaliGemma tokenizer (no download)."""
+
+ def __call__(
+ self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
+ ):
+ ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
+ ids = ids[:max_length]
+ attn = [1] * len(ids)
+ while len(ids) < max_length:
+ ids.append(0)
+ attn.append(0)
+ return {"input_ids": ids, "attention_mask": attn}
+
+
+def test_build_model_inputs_camera_order_and_padding():
+ import numpy as np
+
+ config = Pi0Config.from_model_config(_LEROBOT_CFG) # 3-camera config
+ device = torch.device("cpu")
+ # Provide only 2 of the 3 declared cameras.
+ robot_obs = {
+ "images": {
+ "observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
+ "observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
+ },
+ "state": [0.1] * 14, # raw state shorter than max_state_dim=32
+ "prompt": "pick up the block",
+ }
+ images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
+ robot_obs, config, _FakeTokenizer(), device
+ )
+ assert len(images) == 3
+ assert all(img.shape == (1, 3, 224, 224) for img in images)
+ masks = [bool(m.item()) for m in image_masks]
+ assert masks == [True, True, False]
+ # The padded (3rd) camera must be all -1.
+ assert (images[2] == -1.0).all()
+ assert lang_tokens.shape == (1, 48)
+ assert lang_masks.shape == (1, 48)
+ assert state.shape == (1, 32)
+ # State padded with zeros beyond the 14 real entries.
+ assert torch.allclose(state[0, 14:], torch.zeros(18))
+
+
+def test_build_model_inputs_state_truncated():
+ config = Pi0Config.from_model_config(_LEROBOT_CFG)
+ robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
+ _, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
+ assert state.shape == (1, 32)
+
+
+def test_build_model_inputs_image_key_map():
+ import numpy as np
+
+ cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
+ robot_obs = {
+ "images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
+ "prompt": "go",
+ }
+ _, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
+ masks = [bool(m.item()) for m in image_masks]
+ assert masks == [True, False, False] # only base_0_rgb (mapped) is present
+
+
+# ----------------------------------------------------------------------------
+# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
+#
+# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
+# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
+# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
+# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
+# output to golden values established at 5.3.0 and re-check them here: running on
+# the production transformers proves the port is version-stable and guards the
+# vision_tower / embed_scale version-robust paths against regressions.
+#
+# PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
+# pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
+# ----------------------------------------------------------------------------
+import os # noqa: E402
+
+_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
+
+# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
+_GOLDEN_SUM = -67.808113
+_GOLDEN_STD = 0.252789
+_GOLDEN_MEAN = -0.021190
+_GOLDEN_EMBED_IMAGE_STD = 2.17603
+
+
+def _selfconsist_inputs(device):
+ # NB: draw order is load-bearing — the golden values were established with
+ # state drawn first, then the 3 images, then lang (shared RNG, seed 0).
+ g = torch.Generator(device="cpu").manual_seed(0)
+ b = 2
+ state = torch.randn(b, 32, generator=g)
+ images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
+ masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
+ lang = torch.randint(0, 1000, (b, 48), generator=g)
+ lmask = torch.ones(b, 48, dtype=torch.bool)
+ ng = torch.Generator(device="cpu").manual_seed(42)
+ noise = torch.randn(b, 50, 32, generator=ng)
+ return (
+ [im.to(device) for im in images],
+ [m.to(device) for m in masks],
+ lang.to(device),
+ lmask.to(device),
+ state.to(device),
+ noise.to(device),
+ )
+
+
+def _selfconsist_model(device):
+ import safetensors.torch
+
+ model = Pi0ForActionPrediction(
+ Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
+ )
+ state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
+ filled = model.load_weights(list(state.items()))
+ expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
+ missing = expected - (filled or set())
+ assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
+ return model.to(device).eval()
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_actions():
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ model = _selfconsist_model(device)
+ images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
+ actions = (
+ model.sample_actions(
+ images=images,
+ image_masks=masks,
+ lang_tokens=lang,
+ lang_masks=lmask,
+ state=state,
+ noise=noise,
+ num_steps=10,
+ )
+ .float()
+ .cpu()
+ )
+ s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
+ assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
+ assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
+ assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_embed_image():
+ """Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ model = _selfconsist_model(device)
+ g = torch.Generator(device="cpu").manual_seed(7)
+ img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
+ std = model.paligemma_with_expert.embed_image(img).float().std().item()
+ assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
+ f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
+ "likely a vision_tower weight-load regression under this transformers version."
+ )
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py
index 62d17f705..23f0bd467 100644
--- a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import json
from types import SimpleNamespace
diff --git a/tests/diffusion/models/sana_video/test_pipeline_sana_video.py b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py
new file mode 100644
index 000000000..f6f469aff
--- /dev/null
+++ b/tests/diffusion/models/sana_video/test_pipeline_sana_video.py
@@ -0,0 +1,992 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+
+from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _DummyComponent:
+ pass
+
+
+def _make_request_batch(prompt, request_id="sana-video-test", **sampling_overrides):
+ sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+ request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id)
+ return DiffusionRequestBatch([request])
+
+
+def test_sana_video_pipeline_import_and_registry():
+ from vllm_omni.diffusion.models.sana_video import (
+ SanaImageToVideoPipeline,
+ SanaVideoPipeline,
+ SanaVideoTransformer3DModel,
+ get_sana_video_i2v_post_process_func,
+ get_sana_video_i2v_pre_process_func,
+ get_sana_video_post_process_func,
+ )
+ from vllm_omni.diffusion.registry import (
+ _DIFFUSION_MODELS,
+ _DIFFUSION_POST_PROCESS_FUNCS,
+ _DIFFUSION_PRE_PROCESS_FUNCS,
+ )
+
+ assert SanaImageToVideoPipeline is not None
+ assert SanaVideoPipeline is not None
+ assert SanaVideoTransformer3DModel is not None
+ assert get_sana_video_i2v_post_process_func is not None
+ assert get_sana_video_i2v_pre_process_func is not None
+ assert get_sana_video_post_process_func is not None
+ assert _DIFFUSION_MODELS["SanaVideoPipeline"] == (
+ "sana_video",
+ "pipeline_sana_video",
+ "SanaVideoPipeline",
+ )
+ assert _DIFFUSION_POST_PROCESS_FUNCS["SanaVideoPipeline"] == "get_sana_video_post_process_func"
+ assert _DIFFUSION_MODELS["SanaImageToVideoPipeline"] == (
+ "sana_video",
+ "pipeline_sana_video_i2v",
+ "SanaImageToVideoPipeline",
+ )
+ assert _DIFFUSION_POST_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_post_process_func"
+ assert _DIFFUSION_PRE_PROCESS_FUNCS["SanaImageToVideoPipeline"] == "get_sana_video_i2v_pre_process_func"
+
+
+def test_component_discovery_declarations():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ assert SanaVideoPipeline._dit_modules == ["transformer"]
+ assert SanaVideoPipeline._encoder_modules == ["text_encoder"]
+ assert SanaVideoPipeline._vae_modules == ["vae"]
+ assert SanaVideoPipeline.supports_step_execution is False
+
+
+def test_sana_video_loads_concrete_gemma_tokenizer(monkeypatch):
+ from vllm_omni.diffusion.models.sana_video import pipeline_sana_video
+
+ captured: dict[str, object] = {}
+ expected = object()
+
+ def fake_from_pretrained_with_prefetch(loader, model, **kwargs):
+ captured.update(loader=loader, model=model, kwargs=kwargs)
+ return expected
+
+ monkeypatch.setattr(
+ pipeline_sana_video,
+ "from_pretrained_with_prefetch",
+ fake_from_pretrained_with_prefetch,
+ )
+
+ result = pipeline_sana_video._load_sana_tokenizer(
+ "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+ ["tokenizer", "text_encoder"],
+ local_files_only=False,
+ )
+
+ assert result is expected
+ assert captured["loader"] == pipeline_sana_video.GemmaTokenizer.from_pretrained
+ assert captured["model"] == "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
+ assert captured["kwargs"] == {
+ "subfolder": "tokenizer",
+ "prefetch_list": ["tokenizer", "text_encoder"],
+ "local_files_only": False,
+ }
+
+
+def test_sana_video_declares_extra_body_params():
+ from vllm_omni.model_extras import get_extra_body_params
+
+ assert get_extra_body_params("SanaVideoPipeline") == {
+ "clean_caption",
+ "motion_score",
+ "use_resolution_binning",
+ }
+ assert get_extra_body_params("SanaImageToVideoPipeline") == {
+ "clean_caption",
+ "motion_score",
+ "use_resolution_binning",
+ }
+
+
+def test_sana_video_i2v_preprocesses_image_and_preserves_aspect_ratio():
+ from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func
+
+ request = OmniDiffusionRequest(
+ prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (640, 360))}},
+ sampling_params=OmniDiffusionSamplingParams(),
+ request_id="sana-video-i2v-preprocess",
+ )
+ processed = get_sana_video_i2v_pre_process_func(SimpleNamespace())(request)
+
+ assert processed.sampling_params.height == 448
+ assert processed.sampling_params.width == 832
+ assert processed.prompt["multi_modal_data"]["image"].size == (832, 448)
+
+
+def test_sana_video_i2v_remote_720p_uses_loaded_transformer_config():
+ from vllm_omni.diffusion.models.sana_video import get_sana_video_i2v_pre_process_func
+
+ od_config = SimpleNamespace(
+ model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+ tf_model_config={"sample_size": 22},
+ )
+ request = OmniDiffusionRequest(
+ prompt={"prompt": "a robot walks", "multi_modal_data": {"image": Image.new("RGB", (1280, 704))}},
+ sampling_params=OmniDiffusionSamplingParams(),
+ request_id="sana-video-i2v-remote-720p",
+ )
+
+ processed = get_sana_video_i2v_pre_process_func(od_config)(request)
+
+ assert processed.sampling_params.height == 704
+ assert processed.sampling_params.width == 1280
+ assert processed.prompt["multi_modal_data"]["image"].size == (1280, 704)
+
+
+@pytest.mark.parametrize(
+ "frames",
+ [
+ [[Image.new("RGB", (16, 16)) for _ in range(2)]],
+ np.zeros((1, 2, 16, 16, 3), dtype=np.float32),
+ [[np.zeros((16, 16, 3), dtype=np.float32) for _ in range(2)]],
+ ],
+)
+@pytest.mark.parametrize(
+ "factory_name",
+ ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_diffusers_preserves_postprocessed_frames(frames, factory_name):
+ from vllm_omni.diffusion.models import sana_video
+
+ post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers"))
+
+ result = post_process(frames)
+
+ assert result["payload"]["video"] is frames
+
+
+@pytest.mark.parametrize(
+ "factory_name",
+ ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_diffusers_still_postprocesses_native_tensor(factory_name):
+ from vllm_omni.diffusion.models import sana_video
+
+ decoded_video = torch.zeros(1, 3, 2, 16, 16)
+ post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="diffusers"))
+
+ result = post_process(decoded_video)
+
+ assert isinstance(result["payload"]["video"], np.ndarray)
+ assert result["payload"]["video"].shape == (1, 2, 16, 16, 3)
+
+
+@pytest.mark.parametrize(
+ "factory_name",
+ ["get_sana_video_post_process_func", "get_sana_video_i2v_post_process_func"],
+)
+def test_sana_video_postprocess_preserves_requested_latents(factory_name):
+ from vllm_omni.diffusion.models import sana_video
+
+ latents = torch.randn(1, 4, 2, 2, 2)
+ post_process = getattr(sana_video, factory_name)(SimpleNamespace(diffusion_load_format="native"))
+
+ result = post_process(latents, sampling_params=SimpleNamespace(output_type="latent"))
+
+ assert result is latents
+
+
+def test_sana_video_720p_model_id_fallback(monkeypatch):
+ from vllm_omni.diffusion.models.sana_video import pipeline_sana_video
+
+ monkeypatch.setattr(pipeline_sana_video, "get_hf_file_to_dict", lambda *_args, **_kwargs: None)
+ od_config = SimpleNamespace(
+ model="Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+ tf_model_config={},
+ )
+
+ assert pipeline_sana_video.resolve_sana_video_sample_size(od_config) == 22
+
+
+@pytest.mark.parametrize(
+ ("requested_output_type", "expected_internal_output_type"),
+ [
+ (None, "raw"),
+ ("np", "raw"),
+ ("latent", "latent"),
+ ],
+)
+def test_sana_video_i2v_forward_maps_image_request(requested_output_type, expected_internal_output_type):
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate_i2v(**kwargs):
+ calls.append(kwargs)
+ return torch.zeros(kwargs["num_videos_per_prompt"], 3, 9, 192, 320)
+
+ pipeline._generate_i2v = fake_generate_i2v
+ image = Image.new("RGB", (320, 192))
+ custom_timesteps = torch.tensor([900, 100])
+ req = _make_request_batch(
+ {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}},
+ height=192,
+ width=320,
+ num_frames=9,
+ num_inference_steps=2,
+ timesteps=custom_timesteps,
+ guidance_scale=4.5,
+ num_outputs_per_prompt=2,
+ eta=0.25,
+ seed=42,
+ output_type=requested_output_type,
+ extra_args={"clean_caption": True, "motion_score": 30, "use_resolution_binning": False},
+ )
+
+ output = pipeline.forward(req)
+
+ assert output.output.shape == (2, 3, 9, 192, 320)
+ assert calls[0]["image"] is image
+ assert calls[0]["prompt"] == "a robot walks motion score: 30."
+ assert calls[0]["negative_prompt"] == "blurry"
+ assert calls[0]["frames"] == 9
+ assert calls[0]["timesteps"] is custom_timesteps
+ assert calls[0]["sigmas"] is None
+ assert calls[0]["num_videos_per_prompt"] == 2
+ assert calls[0]["eta"] == 0.25
+ assert calls[0]["clean_caption"] is True
+ assert calls[0]["generator"].initial_seed() == 42
+ assert calls[0]["output_type"] == expected_internal_output_type
+
+
+def test_sana_video_i2v_forward_maps_custom_sigmas():
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate_i2v(**kwargs):
+ calls.append(kwargs)
+ return torch.zeros(1, 3, 9, 192, 320)
+
+ pipeline._generate_i2v = fake_generate_i2v
+ sigmas = [1.0, 0.5]
+ output = pipeline.forward(
+ _make_request_batch(
+ {
+ "prompt": "a robot walks",
+ "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+ },
+ height=192,
+ width=320,
+ num_frames=9,
+ sigmas=sigmas,
+ generator_device="cpu",
+ extra_args={"use_resolution_binning": False},
+ )
+ )
+
+ assert output.output.shape == (1, 3, 9, 192, 320)
+ assert calls[0]["timesteps"] is None
+ assert calls[0]["sigmas"] is sigmas
+
+
+def test_sana_video_i2v_omitted_num_frames_uses_model_default():
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate_i2v(**kwargs):
+ calls.append(kwargs)
+ return torch.zeros(1, 3, 81, 2, 2)
+
+ pipeline._generate_i2v = fake_generate_i2v
+ req = _make_request_batch(
+ {
+ "prompt": "a robot walks",
+ "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+ },
+ height=192,
+ width=320,
+ num_inference_steps=1,
+ generator_device="cpu",
+ extra_args={"use_resolution_binning": False},
+ )
+
+ assert req.sampling_params.num_frames == 1
+ output = pipeline.forward(req)
+
+ assert output.output.shape == (1, 3, 81, 2, 2)
+ assert calls[0]["frames"] == 81
+
+
+def test_sana_video_i2v_clean_caption_request_is_explicitly_rejected():
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ pipeline.text_encoder = SimpleNamespace(dtype=torch.float32)
+ pipeline.tokenizer = None
+ pipeline.check_inputs = lambda **_kwargs: None
+
+ req = _make_request_batch(
+ {
+ "prompt": "a robot walks",
+ "multi_modal_data": {"image": Image.new("RGB", (320, 192))},
+ },
+ height=192,
+ width=320,
+ num_frames=9,
+ generator_device="cpu",
+ extra_args={"clean_caption": True, "use_resolution_binning": False},
+ )
+
+ with pytest.raises(ValueError, match="does not support `clean_caption=True`"):
+ pipeline.forward(req)
+
+
+@pytest.mark.parametrize(
+ ("custom_timesteps", "custom_sigmas"),
+ [
+ (torch.tensor([4.0, 2.0]), None),
+ (None, [1.0, 0.5]),
+ ],
+)
+def test_sana_video_i2v_consumes_batch_schedule_and_step_kwargs(custom_timesteps, custom_sigmas):
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ calls = SimpleNamespace(
+ encode_prompt=[],
+ prepare_latents=[],
+ set_timesteps=[],
+ step=[],
+ transformer=[],
+ )
+
+ class StubTransformer:
+ dtype = torch.float32
+ config = SimpleNamespace(
+ sample_size=30,
+ in_channels=4,
+ out_channels=4,
+ patch_size=(1, 1, 1),
+ )
+
+ def __call__(
+ self,
+ hidden_states,
+ *,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ timestep,
+ return_dict,
+ ):
+ calls.transformer.append(
+ (
+ hidden_states.shape,
+ encoder_hidden_states.shape,
+ encoder_attention_mask.shape,
+ timestep.shape,
+ return_dict,
+ )
+ )
+ return (torch.zeros_like(hidden_states),)
+
+ class StubScheduler:
+ order = 1
+
+ def set_timesteps(
+ self,
+ num_inference_steps=None,
+ *,
+ device=None,
+ timesteps=None,
+ sigmas=None,
+ ):
+ calls.set_timesteps.append(
+ {
+ "num_inference_steps": num_inference_steps,
+ "timesteps": timesteps,
+ "sigmas": sigmas,
+ }
+ )
+ schedule_length = len(timesteps if timesteps is not None else sigmas)
+ self.timesteps = torch.arange(schedule_length, 0, -1, dtype=torch.float32, device=device)
+
+ def step(
+ self,
+ _noise_pred,
+ _timestep,
+ current_latents,
+ *,
+ eta,
+ generator,
+ return_dict,
+ ):
+ calls.step.append({"eta": eta, "generator": generator, "return_dict": return_dict})
+ return (current_latents + 1,)
+
+ class VaeMustNotBeUsed:
+ @property
+ def dtype(self):
+ pytest.fail("latent output must return before accessing the VAE")
+
+ def decode(self, *_args, **_kwargs):
+ pytest.fail("latent output must not call VAE decode")
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = StubTransformer()
+ pipeline.scheduler = StubScheduler()
+ pipeline.vae = VaeMustNotBeUsed()
+ pipeline.check_inputs = lambda **_kwargs: None
+ pipeline.video_processor = SimpleNamespace(
+ preprocess=lambda _image, height, width: torch.zeros(1, 3, height, width),
+ )
+ pipeline.progress_bar = lambda **_kwargs: nullcontext(SimpleNamespace(update=lambda: None))
+
+ def capture_encode_prompt(*_args, **kwargs):
+ calls.encode_prompt.append(kwargs)
+ batch_size = kwargs["num_videos_per_prompt"]
+ return (
+ torch.zeros(batch_size, 1, 1),
+ torch.ones(batch_size, 1, dtype=torch.bool),
+ torch.zeros(batch_size, 1, 1),
+ torch.ones(batch_size, 1, dtype=torch.bool),
+ )
+
+ pipeline.encode_prompt = capture_encode_prompt
+ initial_latents = torch.zeros(2, 4, 2, 2, 2)
+
+ def capture_prepare_latents(*_args, **kwargs):
+ calls.prepare_latents.append(kwargs)
+ return initial_latents
+
+ pipeline._prepare_i2v_latents = capture_prepare_latents
+ generator = torch.Generator(device="cpu").manual_seed(42)
+
+ output = pipeline._generate_i2v(
+ image=Image.new("RGB", (2, 2)),
+ prompt="test",
+ negative_prompt="",
+ height=2,
+ width=2,
+ frames=2,
+ num_inference_steps=50,
+ timesteps=custom_timesteps,
+ sigmas=custom_sigmas,
+ guidance_scale=4.5,
+ num_videos_per_prompt=2,
+ eta=0.25,
+ generator=generator,
+ latents=None,
+ clean_caption=True,
+ use_resolution_binning=False,
+ max_sequence_length=1,
+ output_type="latent",
+ )
+
+ assert calls.encode_prompt[0]["num_videos_per_prompt"] == 2
+ assert calls.encode_prompt[0]["clean_caption"] is True
+ assert calls.prepare_latents[0]["batch_size"] == 2
+ assert len(calls.set_timesteps) == 1
+ assert calls.set_timesteps[0]["num_inference_steps"] is None
+ assert calls.set_timesteps[0]["timesteps"] is custom_timesteps
+ assert calls.set_timesteps[0]["sigmas"] is custom_sigmas
+ assert len(calls.step) == 2
+ assert all(call == {"eta": 0.25, "generator": generator, "return_dict": False} for call in calls.step)
+ assert all(
+ call
+ == (
+ torch.Size([4, 4, 2, 2, 2]),
+ torch.Size([4, 1, 1]),
+ torch.Size([4, 1]),
+ torch.Size([4, 1, 2, 2, 2]),
+ False,
+ )
+ for call in calls.transformer
+ )
+ expected = initial_latents.clone()
+ expected[:, :, 1:] += 2
+ torch.testing.assert_close(output, expected)
+
+
+def test_sana_video_i2v_uses_diffusers_complex_instruction_default():
+ import inspect
+
+ from diffusers import SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline
+
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ reference_default = (
+ inspect.signature(DiffusersSanaImageToVideoPipeline.__call__).parameters["complex_human_instruction"].default
+ )
+ native_default = (
+ inspect.signature(SanaImageToVideoPipeline._generate_i2v).parameters["complex_human_instruction"].default
+ )
+ assert native_default == reference_default
+
+ class StopAfterPromptEncodingError(Exception):
+ pass
+
+ pipeline = object.__new__(SanaImageToVideoPipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.check_inputs = lambda **_kwargs: None
+ captured = {}
+
+ def capture_encode_prompt(*_args, **kwargs):
+ captured.update(kwargs)
+ raise StopAfterPromptEncodingError
+
+ pipeline.encode_prompt = capture_encode_prompt
+ with pytest.raises(StopAfterPromptEncodingError):
+ pipeline._generate_i2v(
+ image=Image.new("RGB", (320, 192)),
+ prompt="a robot walks",
+ negative_prompt="blurry",
+ height=192,
+ width=320,
+ frames=9,
+ num_inference_steps=2,
+ guidance_scale=4.5,
+ generator=torch.Generator(device="cpu").manual_seed(42),
+ latents=None,
+ use_resolution_binning=False,
+ max_sequence_length=300,
+ )
+
+ assert captured["complex_human_instruction"] == reference_default
+
+
+def test_sana_video_t2v_complex_instruction_has_no_mutable_default():
+ import inspect
+
+ from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+ SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION,
+ SanaVideoPipeline,
+ )
+
+ parameter = inspect.signature(SanaVideoPipeline._generate).parameters["complex_human_instruction"]
+ assert parameter.default is None
+ assert isinstance(SANA_VIDEO_COMPLEX_HUMAN_INSTRUCTION, tuple)
+
+
+def test_diffusers_adapter_t2v_omitted_num_frames_uses_sana_default():
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+ adapter = object.__new__(DiffusersAdapterPipeline)
+ adapter._accept_call_kwargs = {"prompt", "negative_prompt", "frames", "generator"}
+ adapter._pipeline_utils = SanaVideoPipelineUtils()
+ adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+
+ req = _make_request_batch(
+ {"prompt": "a robot walks", "negative_prompt": "blurry"},
+ seed=42,
+ generator_device="cpu",
+ )
+ assert req.sampling_params.num_frames == 1
+
+ kwargs = adapter._build_call_kwargs(req)
+
+ assert kwargs["frames"] == 81
+ assert "num_frames" not in kwargs
+ assert kwargs["generator"].initial_seed() == 42
+
+
+def test_diffusers_adapter_i2v_omitted_num_frames_uses_sana_default():
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+ adapter = object.__new__(DiffusersAdapterPipeline)
+ adapter._accept_call_kwargs = {"prompt", "negative_prompt", "image", "frames", "generator"}
+ adapter._pipeline_utils = SanaVideoPipelineUtils()
+ adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+ image = Image.new("RGB", (320, 192))
+
+ req = _make_request_batch(
+ {"prompt": "a robot walks", "negative_prompt": "blurry", "multi_modal_data": {"image": image}},
+ seed=42,
+ generator_device="cpu",
+ )
+ assert req.sampling_params.num_frames == 1
+
+ kwargs = adapter._build_call_kwargs(req)
+
+ assert kwargs["image"] is image
+ assert kwargs["prompt"] == "a robot walks"
+ assert kwargs["negative_prompt"] == "blurry"
+ assert kwargs["frames"] == 81
+ assert kwargs["generator"].initial_seed() == 42
+
+
+@pytest.mark.parametrize(
+ ("sampling_overrides", "expected_frames"),
+ [
+ ({}, 49),
+ ({"num_frames": 9}, 9),
+ ],
+)
+def test_diffusers_adapter_sana_resolves_frame_precedence(sampling_overrides, expected_frames):
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+ adapter = object.__new__(DiffusersAdapterPipeline)
+ adapter._accept_call_kwargs = {"prompt", "frames", "generator"}
+ adapter._pipeline_utils = SanaVideoPipelineUtils()
+ adapter.od_config = SimpleNamespace(diffusers_call_kwargs={"frames": 49}, output_type=None)
+
+ kwargs = adapter._build_call_kwargs(
+ _make_request_batch(
+ "a robot walks",
+ seed=42,
+ generator_device="cpu",
+ **sampling_overrides,
+ )
+ )
+
+ assert kwargs["frames"] == expected_frames
+
+
+@pytest.mark.parametrize("pipeline_class_name", ["SanaVideoPipeline", "SanaImageToVideoPipeline"])
+def test_diffusers_adapter_selects_sana_pipeline_utils(pipeline_class_name):
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import (
+ SanaVideoPipelineUtils,
+ get_pipeline_utils,
+ )
+
+ assert isinstance(get_pipeline_utils(pipeline_class_name), SanaVideoPipelineUtils)
+
+
+def test_diffusers_adapter_resolves_requested_sana_i2v_pipeline():
+ from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline
+
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import (
+ SanaVideoPipelineUtils,
+ get_pipeline_utils_for_config,
+ resolve_diffusers_pipeline_class,
+ )
+
+ od_config = SimpleNamespace(
+ model_class_name="SanaImageToVideoPipeline",
+ diffusers_pipeline_cls=SanaVideoPipeline,
+ )
+
+ assert isinstance(get_pipeline_utils_for_config(od_config), SanaVideoPipelineUtils)
+ assert resolve_diffusers_pipeline_class(od_config) is SanaImageToVideoPipeline
+
+
+def test_diffusers_adapter_pipeline_override_is_sana_specific():
+ from diffusers import DiffusionPipeline
+
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import resolve_diffusers_pipeline_class
+
+ class OtherPipeline(DiffusionPipeline):
+ pass
+
+ od_config = SimpleNamespace(
+ model_class_name="DiffusersAdapterPipeline",
+ diffusers_pipeline_cls=OtherPipeline,
+ )
+
+ assert resolve_diffusers_pipeline_class(od_config) is OtherPipeline
+
+
+def test_diffusers_adapter_loads_resolved_pipeline_class(monkeypatch):
+ from vllm_omni.diffusion.models.diffusers_adapter import pipeline_diffusers_adapter
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+
+ loaded = []
+ expected_pipeline = object()
+
+ class ResolvedPipeline:
+ @classmethod
+ def from_pretrained(cls, model_id, **kwargs):
+ loaded.append((model_id, kwargs))
+ return expected_pipeline
+
+ monkeypatch.setattr(
+ pipeline_diffusers_adapter,
+ "resolve_diffusers_pipeline_class",
+ lambda _od_config: ResolvedPipeline,
+ )
+ adapter = object.__new__(DiffusersAdapterPipeline)
+ adapter.od_config = SimpleNamespace(diffusers_pipeline_cls=object)
+
+ result = adapter._load_pipeline_from_pretrained("sana-checkpoint", {"torch_dtype": torch.bfloat16})
+
+ assert result is expected_pipeline
+ assert loaded == [("sana-checkpoint", {"torch_dtype": torch.bfloat16})]
+
+
+def test_diffusers_adapter_sana_i2v_override_reports_image_capability():
+ from diffusers import SanaVideoPipeline
+
+ from vllm_omni.diffusion.io_support import supports_multimodal_input
+
+ od_config = SimpleNamespace(
+ diffusion_load_format="diffusers",
+ model_class_name="SanaImageToVideoPipeline",
+ diffusers_pipeline_cls=SanaVideoPipeline,
+ )
+
+ assert supports_multimodal_input(od_config) == (True, False)
+
+
+def test_diffusers_adapter_sana_i2v_warmup_supplies_image():
+ from diffusers import SanaVideoPipeline
+
+ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+ engine = object.__new__(DiffusionEngine)
+ engine.od_config = SimpleNamespace(
+ diffusion_load_format="diffusers",
+ model_class_name="SanaImageToVideoPipeline",
+ diffusers_pipeline_cls=SanaVideoPipeline,
+ )
+ captured_requests = []
+ engine.pre_process_func = lambda request: request
+
+ def capture_request(request):
+ captured_requests.append(request)
+ return SimpleNamespace(error=None)
+
+ engine.add_req_and_wait_for_response = capture_request
+
+ engine._dummy_run()
+
+ assert len(captured_requests) == 1
+ prompt = captured_requests[0].prompt
+ assert isinstance(prompt, dict)
+ assert isinstance(prompt["multi_modal_data"]["image"], Image.Image)
+
+
+def test_diffusers_adapter_disables_resolution_binning_for_warmup():
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
+ from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import SanaVideoPipelineUtils
+
+ adapter = object.__new__(DiffusersAdapterPipeline)
+ adapter._accept_call_kwargs = {"prompt", "frames", "use_resolution_binning", "generator"}
+ adapter._pipeline_utils = SanaVideoPipelineUtils()
+ adapter.od_config = SimpleNamespace(diffusers_call_kwargs={}, output_type=None)
+
+ kwargs = adapter._build_call_kwargs(
+ _make_request_batch(
+ "dummy run",
+ request_id=DUMMY_DIFFUSION_REQUEST_ID,
+ generator_device="cpu",
+ )
+ )
+
+ assert kwargs["use_resolution_binning"] is False
+ assert kwargs["frames"] == 1
+
+
+def test_pipeline_is_torch_module_and_supports_eval():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ nn.Module.__init__(pipeline)
+
+ assert isinstance(pipeline, nn.Module)
+ assert pipeline.eval() is pipeline
+ assert pipeline.training is False
+
+
+@pytest.mark.parametrize(
+ ("vae_scale", "patch_size", "valid_size", "invalid_size", "alignment"),
+ [
+ (8, (1, 2, 2), 624, 632, 16),
+ (32, (1, 1, 1), 704, 712, 32),
+ ],
+)
+def test_check_inputs_uses_variant_spatial_alignment(vae_scale, patch_size, valid_size, invalid_size, alignment):
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ pipeline.vae_scale_factor_spatial = vae_scale
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(patch_size=patch_size))
+
+ pipeline.check_inputs(prompt="test", height=valid_size, width=valid_size)
+ with pytest.raises(ValueError, match=f"divisible by {alignment}"):
+ pipeline.check_inputs(prompt="test", height=invalid_size, width=valid_size)
+
+
+@pytest.mark.parametrize(
+ ("requested_output_type", "expected_internal_output_type"),
+ [
+ (None, "raw"),
+ ("np", "raw"),
+ ("latent", "latent"),
+ ],
+)
+def test_forward_maps_omni_request_to_sana_generation_args(requested_output_type, expected_internal_output_type):
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate(**kwargs):
+ calls.append(kwargs)
+ return SimpleNamespace(frames=torch.zeros(1, 3, 9, 192, 320))
+
+ pipeline._generate = fake_generate
+ req = _make_request_batch(
+ {"prompt": "a robot walks", "negative_prompt": "blurry"},
+ height=192,
+ width=320,
+ num_frames=9,
+ num_inference_steps=2,
+ guidance_scale=4.5,
+ seed=42,
+ output_type=requested_output_type,
+ extra_args={"motion_score": 30, "use_resolution_binning": False},
+ )
+
+ output = pipeline.forward(req)
+
+ assert output.output.shape == (1, 3, 9, 192, 320)
+ assert calls[0]["prompt"] == "a robot walks motion score: 30."
+ assert calls[0]["negative_prompt"] == "blurry"
+ assert calls[0]["height"] == 192
+ assert calls[0]["width"] == 320
+ assert calls[0]["frames"] == 9
+ assert calls[0]["num_inference_steps"] == 2
+ assert calls[0]["guidance_scale"] == 4.5
+ assert calls[0]["use_resolution_binning"] is False
+ assert calls[0]["generator"].initial_seed() == 42
+ assert calls[0]["output_type"] == expected_internal_output_type
+
+
+def test_sana_video_t2v_omitted_num_frames_uses_model_default():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate(**kwargs):
+ calls.append(kwargs)
+ return SimpleNamespace(frames=torch.zeros(1, 3, 81, 2, 2))
+
+ pipeline._generate = fake_generate
+ req = _make_request_batch(
+ "a robot walks",
+ height=192,
+ width=320,
+ num_inference_steps=1,
+ generator_device="cpu",
+ )
+
+ assert req.sampling_params.num_frames == 1
+ output = pipeline.forward(req)
+
+ assert output.output.shape == (1, 3, 81, 2, 2)
+ assert calls[0]["frames"] == 81
+
+
+def test_sana_video_t2v_dummy_run_keeps_single_frame():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=30))
+ calls = []
+
+ def fake_generate(**kwargs):
+ calls.append(kwargs)
+ return SimpleNamespace(frames=torch.zeros(1, 3, 1, 192, 320))
+
+ pipeline._generate = fake_generate
+
+ output = pipeline.forward(
+ _make_request_batch(
+ "dummy run",
+ request_id=DUMMY_DIFFUSION_REQUEST_ID,
+ height=192,
+ width=320,
+ num_frames=1,
+ num_inference_steps=1,
+ generator_device="cpu",
+ )
+ )
+
+ assert output.output.shape == (1, 3, 1, 192, 320)
+ assert calls[0]["frames"] == 1
+
+
+@pytest.mark.parametrize(
+ ("sample_size", "expected_height", "expected_width"),
+ [
+ (30, 480, 832),
+ (22, 704, 1280),
+ ],
+)
+def test_t2v_uses_variant_default_resolution(sample_size, expected_height, expected_width):
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+ pipeline.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=sample_size))
+ calls = []
+
+ def fake_generate(**kwargs):
+ calls.append(kwargs)
+ return SimpleNamespace(frames=torch.zeros(1, 3, 9, expected_height, expected_width))
+
+ pipeline._generate = fake_generate
+ output = pipeline.forward(
+ _make_request_batch(
+ "a robot walks",
+ num_frames=9,
+ num_inference_steps=1,
+ guidance_scale=1.0,
+ seed=42,
+ generator_device="cpu",
+ )
+ )
+
+ assert output.output.shape == (1, 3, 9, expected_height, expected_width)
+ assert calls[0]["height"] == expected_height
+ assert calls[0]["width"] == expected_width
+
+
+def test_forward_requires_exactly_one_nonempty_prompt():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ pipeline = object.__new__(SanaVideoPipeline)
+
+ with pytest.raises(ValueError, match="Prompt is required"):
+ pipeline.forward(_make_request_batch(""))
+
+ first = OmniDiffusionRequest(
+ prompt="first",
+ sampling_params=OmniDiffusionSamplingParams(),
+ request_id="first",
+ )
+ second = OmniDiffusionRequest(
+ prompt="second",
+ sampling_params=OmniDiffusionSamplingParams(),
+ request_id="second",
+ )
+ with pytest.raises(ValueError, match="exactly one prompt"):
+ pipeline.forward(DiffusionRequestBatch([first, second]))
diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py
new file mode 100644
index 000000000..804692d7a
--- /dev/null
+++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.py
@@ -0,0 +1,602 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_TINY_CONFIG = {
+ "in_channels": 4,
+ "out_channels": 4,
+ "num_attention_heads": 2,
+ "attention_head_dim": 12,
+ "num_layers": 1,
+ "num_cross_attention_heads": 2,
+ "cross_attention_head_dim": 12,
+ "cross_attention_dim": 24,
+ "caption_channels": 8,
+ "mlp_ratio": 2.0,
+ "patch_size": (1, 2, 2),
+ "sample_size": 4,
+ "rope_max_seq_len": 32,
+}
+
+_GOLDEN_PREFIX = torch.tensor(
+ [
+ -0.1159307063,
+ 0.2003862262,
+ -0.3754127920,
+ -0.2403523028,
+ -1.2376221418,
+ -0.6513092518,
+ -0.3358457983,
+ -0.9894289970,
+ 0.7858567238,
+ 0.5089095831,
+ -0.5362522602,
+ -0.3732060790,
+ ]
+)
+
+
+@pytest.mark.parametrize(
+ ("elementwise_affine", "bias", "expected_state_keys"),
+ [
+ (False, False, set()),
+ (False, True, set()),
+ (True, False, {"weight"}),
+ (True, True, {"weight", "bias"}),
+ ],
+)
+@pytest.mark.parametrize(
+ ("input_dtype", "parameter_dtype"),
+ [
+ (torch.float16, torch.float16),
+ (torch.bfloat16, torch.bfloat16),
+ (torch.float32, torch.float32),
+ (torch.float64, torch.float64),
+ (torch.float16, torch.float32),
+ (torch.bfloat16, torch.float32),
+ (torch.float32, torch.float16),
+ (torch.float32, torch.bfloat16),
+ ],
+)
+def test_sana_rms_norm_matches_diffusers(
+ elementwise_affine,
+ bias,
+ expected_state_keys,
+ input_dtype,
+ parameter_dtype,
+):
+ from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaRMSNorm
+
+ reference = DiffusersRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(
+ dtype=parameter_dtype
+ )
+ actual = SanaRMSNorm(8, eps=1e-5, elementwise_affine=elementwise_affine, bias=bias).to(dtype=parameter_dtype)
+
+ assert actual.eps == reference.eps
+ assert actual.elementwise_affine == reference.elementwise_affine
+ assert actual.dim == reference.dim
+ assert set(actual.state_dict()) == expected_state_keys
+ assert set(actual.state_dict()) == set(reference.state_dict())
+
+ if elementwise_affine:
+ weight = torch.linspace(0.5, 1.5, 8, dtype=parameter_dtype)
+ actual.weight.data.copy_(weight)
+ reference.weight.data.copy_(weight)
+ if bias:
+ norm_bias = torch.linspace(-0.25, 0.25, 8, dtype=parameter_dtype)
+ actual.bias.data.copy_(norm_bias)
+ reference.bias.data.copy_(norm_bias)
+ else:
+ assert actual.weight is None
+ assert actual.bias is None
+
+ hidden_states = torch.linspace(-2.0, 2.0, 48, dtype=input_dtype).reshape(2, 3, 8)
+ expected = reference(hidden_states)
+ result = actual(hidden_states)
+
+ assert result.dtype == expected.dtype
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+ ("num_channels", "flip_sin_to_cos", "downscale_freq_shift", "scale"),
+ [
+ (256, True, 0, 1),
+ (7, False, 1, 0.5),
+ (12, True, 0.5, 2),
+ ],
+)
+@pytest.mark.parametrize("timestep_dtype", [torch.int64, torch.float32, torch.float64])
+def test_sana_timesteps_matches_diffusers(
+ num_channels,
+ flip_sin_to_cos,
+ downscale_freq_shift,
+ scale,
+ timestep_dtype,
+):
+ from diffusers.models.embeddings import Timesteps
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaTimesteps
+
+ reference = Timesteps(
+ num_channels=num_channels,
+ flip_sin_to_cos=flip_sin_to_cos,
+ downscale_freq_shift=downscale_freq_shift,
+ scale=scale,
+ )
+ actual = SanaTimesteps(
+ num_channels=num_channels,
+ flip_sin_to_cos=flip_sin_to_cos,
+ downscale_freq_shift=downscale_freq_shift,
+ scale=scale,
+ )
+ timesteps = torch.tensor([0, 1, 500, 999], dtype=timestep_dtype)
+
+ assert actual.state_dict() == reference.state_dict() == {}
+ assert actual.num_channels == reference.num_channels
+ assert actual.flip_sin_to_cos == reference.flip_sin_to_cos
+ assert actual.downscale_freq_shift == reference.downscale_freq_shift
+ assert actual.scale == reference.scale
+ torch.testing.assert_close(actual(timesteps), reference(timesteps), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "sample_shape", "condition_shape"),
+ [
+ ({}, (2, 256), None),
+ (
+ {
+ "out_dim": 12,
+ "post_act_fn": "silu",
+ "cond_proj_dim": 6,
+ "sample_proj_bias": False,
+ },
+ (2, 256),
+ (2, 6),
+ ),
+ ],
+)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_timestep_embedding_matches_diffusers(
+ kwargs,
+ sample_shape,
+ condition_shape,
+ dtype,
+):
+ from diffusers.models.embeddings import TimestepEmbedding
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaTimestepEmbedding,
+ )
+
+ torch.manual_seed(17)
+ reference = TimestepEmbedding(256, 24, **kwargs).to(dtype=dtype)
+ actual = SanaTimestepEmbedding(256, 24, **kwargs).to(dtype=dtype)
+ assert set(actual.state_dict()) == set(reference.state_dict())
+ actual.load_state_dict(reference.state_dict())
+
+ sample = torch.randn(sample_shape, dtype=dtype)
+ condition = torch.randn(condition_shape, dtype=dtype) if condition_shape else None
+ expected = reference(sample, condition)
+ result = actual(sample, condition)
+
+ assert result.shape == expected.shape
+ assert result.dtype == expected.dtype
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_sana_timestep_embedding_rejects_unsupported_activation():
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaTimestepEmbedding,
+ )
+
+ with pytest.raises(ValueError, match="act_fn='silu'"):
+ SanaTimestepEmbedding(8, 16, act_fn="gelu")
+ with pytest.raises(ValueError, match="post_act_fn"):
+ SanaTimestepEmbedding(8, 16, post_act_fn="gelu")
+
+ module = SanaTimestepEmbedding(8, 16)
+ with pytest.raises(ValueError, match="cond_proj_dim"):
+ module(torch.randn(1, 8), torch.randn(1, 4))
+
+
+@pytest.mark.parametrize("use_additional_conditions", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_ada_layer_norm_single_matches_diffusers(
+ use_additional_conditions,
+ dtype,
+):
+ from diffusers.models.normalization import AdaLayerNormSingle
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaAdaLayerNormSingle,
+ )
+
+ torch.manual_seed(19)
+ reference = AdaLayerNormSingle(
+ 24,
+ use_additional_conditions=use_additional_conditions,
+ ).to(dtype=dtype)
+ actual = SanaAdaLayerNormSingle(
+ 24,
+ use_additional_conditions=use_additional_conditions,
+ ).to(dtype=dtype)
+ assert set(actual.state_dict()) == set(reference.state_dict())
+ actual.load_state_dict(reference.state_dict())
+
+ timestep = torch.tensor([10, 500])
+ added_cond_kwargs = None
+ if use_additional_conditions:
+ added_cond_kwargs = {
+ "resolution": torch.tensor([[480, 832], [704, 1280]]),
+ "aspect_ratio": torch.tensor([[832 / 480], [1280 / 704]]),
+ }
+
+ expected = reference(
+ timestep,
+ added_cond_kwargs=added_cond_kwargs,
+ batch_size=2,
+ hidden_dtype=dtype,
+ )
+ result = actual(
+ timestep,
+ added_cond_kwargs=added_cond_kwargs,
+ batch_size=2,
+ hidden_dtype=dtype,
+ )
+
+ assert len(result) == len(expected) == 2
+ for actual_tensor, expected_tensor in zip(result, expected):
+ assert actual_tensor.shape == expected_tensor.shape
+ assert actual_tensor.dtype == expected_tensor.dtype
+ torch.testing.assert_close(actual_tensor, expected_tensor, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("act_fn", ["gelu_tanh", "silu", "silu_fp32"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_sana_pixart_text_projection_matches_diffusers(act_fn, dtype):
+ from diffusers.models.embeddings import PixArtAlphaTextProjection
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaPixArtAlphaTextProjection,
+ )
+
+ torch.manual_seed(23)
+ reference = PixArtAlphaTextProjection(
+ in_features=8,
+ hidden_size=24,
+ out_features=12,
+ act_fn=act_fn,
+ ).to(dtype=dtype)
+ actual = SanaPixArtAlphaTextProjection(
+ in_features=8,
+ hidden_size=24,
+ out_features=12,
+ act_fn=act_fn,
+ ).to(dtype=dtype)
+ assert set(actual.state_dict()) == set(reference.state_dict())
+ actual.load_state_dict(reference.state_dict())
+
+ caption = torch.randn(2, 5, 8, dtype=dtype)
+ expected = reference(caption)
+ result = actual(caption)
+
+ assert result.shape == expected.shape
+ assert result.dtype == expected.dtype
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("freqs_dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize(("dim", "max_seq_len", "theta"), [(4, 8, 10000.0), (12, 32, 256.0)])
+def test_native_rope_matches_diffusers(dim, max_seq_len, theta, freqs_dtype):
+ from diffusers.models.embeddings import get_1d_rotary_pos_embed
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ _get_1d_rotary_pos_embed,
+ )
+
+ expected_cos, expected_sin = get_1d_rotary_pos_embed(
+ dim,
+ max_seq_len,
+ theta,
+ use_real=True,
+ repeat_interleave_real=True,
+ freqs_dtype=freqs_dtype,
+ )
+ actual_cos, actual_sin = _get_1d_rotary_pos_embed(dim, max_seq_len, theta, freqs_dtype)
+
+ torch.testing.assert_close(actual_cos, expected_cos, rtol=0, atol=0)
+ torch.testing.assert_close(actual_sin, expected_sin, rtol=0, atol=0)
+
+
+def test_tiny_transformer_matches_diffusers_and_frozen_output():
+ from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+ from diffusers.configuration_utils import ConfigMixin
+ from diffusers.models.modeling_utils import ModelMixin
+
+ from vllm_omni.diffusion.attention.layer import Attention
+ from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaLinearAttention,
+ SanaRMSNorm,
+ SanaVideoTransformerConfig,
+ SanaVideoTransformerOutput,
+ )
+
+ torch.manual_seed(7)
+ reference = DiffusersTransformer(**_TINY_CONFIG).eval()
+ model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+ assert not isinstance(model, ModelMixin)
+ assert not isinstance(model, ConfigMixin)
+ assert isinstance(model.config, SanaVideoTransformerConfig)
+ assert set(model.state_dict()) == set(reference.state_dict())
+ model.load_state_dict(reference.state_dict())
+
+ block = model.transformer_blocks[0]
+ assert isinstance(block.attn1, SanaLinearAttention)
+ assert isinstance(block.attn2.attn, Attention)
+ assert block.attn2.attn.role == "cross"
+ assert isinstance(block.attn1.norm_q, SanaRMSNorm)
+ assert isinstance(block.attn1.norm_k, SanaRMSNorm)
+ assert isinstance(block.attn2.norm_q, SanaRMSNorm)
+ assert isinstance(block.attn2.norm_k, SanaRMSNorm)
+ assert isinstance(model.caption_norm, SanaRMSNorm)
+
+ torch.manual_seed(11)
+ hidden_states = torch.randn(1, 4, 3, 4, 4)
+ encoder_hidden_states = torch.randn(1, 5, 8)
+ encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]])
+ timestep = torch.tensor([500.0])
+
+ with torch.no_grad():
+ expected = reference(
+ hidden_states,
+ encoder_hidden_states,
+ timestep,
+ encoder_attention_mask=encoder_attention_mask,
+ ).sample
+ actual_output = model(
+ hidden_states,
+ encoder_hidden_states,
+ timestep,
+ encoder_attention_mask=encoder_attention_mask,
+ )
+ actual = actual_output.sample
+
+ assert isinstance(actual_output, SanaVideoTransformerOutput)
+ torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ torch.testing.assert_close(actual.flatten()[: len(_GOLDEN_PREFIX)], _GOLDEN_PREFIX, rtol=1e-5, atol=1e-5)
+
+
+def test_encoder_attention_mask_uses_bool_sdpa_contract(monkeypatch):
+ from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+ torch.manual_seed(37)
+ model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+ captured = {}
+ omni_attention = model.transformer_blocks[0].attn2.attn
+
+ def fake_sdpa_forward(query, key, value, attn_metadata):
+ captured["query_shape"] = query.shape
+ captured["key_shape"] = key.shape
+ captured["mask"] = attn_metadata.attn_mask.detach().clone()
+ return torch.zeros_like(query)
+
+ def unexpected_backend_forward(*args, **kwargs):
+ raise AssertionError("masked SANA cross-attention must use SDPA")
+
+ monkeypatch.setattr(omni_attention.sdpa_fallback, "forward", fake_sdpa_forward)
+ monkeypatch.setattr(omni_attention, "forward", unexpected_backend_forward)
+
+ hidden_states = torch.randn(2, 4, 3, 4, 4)
+ encoder_hidden_states = torch.randn(2, 5, 8)
+ input_mask = torch.tensor(
+ [
+ [1, 1, 1, 0, 0],
+ [1, 1, 1, 1, 0],
+ ],
+ dtype=torch.int64,
+ )
+
+ with torch.no_grad():
+ output = model(
+ hidden_states,
+ encoder_hidden_states,
+ torch.tensor([100.0, 500.0]),
+ encoder_attention_mask=input_mask,
+ ).sample
+
+ assert output.shape == hidden_states.shape
+ assert captured["query_shape"][:2] == (2, 12)
+ assert captured["key_shape"][:2] == (2, 5)
+ assert captured["query_shape"][1] != captured["key_shape"][1]
+ assert captured["mask"].shape == (2, 5)
+ assert captured["mask"].dtype == torch.bool
+ assert torch.equal(captured["mask"], input_mask.bool())
+
+
+def test_encoder_attention_mask_rejects_additive_bias():
+ from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+ model = SanaVideoTransformer3DModel(**_TINY_CONFIG).eval()
+ additive_mask = torch.tensor([[0.0, 0.0, 0.0, -10000.0, -10000.0]])
+
+ with pytest.raises(TypeError, match="not a floating-point additive attention bias"):
+ model(
+ torch.randn(1, 4, 3, 4, 4),
+ torch.randn(1, 5, 8),
+ torch.tensor([500.0]),
+ encoder_attention_mask=additive_mask,
+ )
+
+
+def test_linear_attention_requires_rotary_embeddings():
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention
+
+ attention = SanaLinearAttention(
+ dim=24,
+ num_heads=2,
+ head_dim=12,
+ dropout=0.0,
+ bias=True,
+ qk_norm="rms_norm_across_heads",
+ )
+
+ with pytest.raises(ValueError, match="requires rotary_emb"):
+ attention(torch.randn(1, 4, 24))
+
+
+def test_linear_attention_bfloat16_matches_diffusers_mixed_precision():
+ from diffusers.models.attention_processor import Attention as DiffusersAttention
+ from diffusers.models.transformers.transformer_sana_video import SanaLinearAttnProcessor3_0
+
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaLinearAttention
+
+ torch.manual_seed(17)
+ reference = DiffusersAttention(
+ query_dim=24,
+ heads=2,
+ dim_head=12,
+ kv_heads=2,
+ qk_norm="rms_norm_across_heads",
+ dropout=0.0,
+ bias=True,
+ cross_attention_dim=None,
+ processor=SanaLinearAttnProcessor3_0(),
+ ).to(dtype=torch.bfloat16)
+ actual = SanaLinearAttention(
+ dim=24,
+ num_heads=2,
+ head_dim=12,
+ dropout=0.0,
+ bias=True,
+ qk_norm="rms_norm_across_heads",
+ ).to(dtype=torch.bfloat16)
+ actual.load_state_dict(reference.state_dict())
+
+ hidden_states = torch.randn(1, 5, 24, dtype=torch.bfloat16)
+ rotary_emb = (
+ torch.randn(1, 5, 1, 12, dtype=torch.bfloat16),
+ torch.randn(1, 5, 1, 12, dtype=torch.bfloat16),
+ )
+
+ with torch.no_grad():
+ expected = reference(hidden_states, rotary_emb=rotary_emb)
+ result = actual(hidden_states, rotary_emb=rotary_emb)
+
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
+
+
+def test_transformer_block_without_cross_attention_keeps_feed_forward_norm():
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerBlock
+
+ block = SanaVideoTransformerBlock(
+ dim=24,
+ num_attention_heads=2,
+ attention_head_dim=12,
+ num_cross_attention_heads=2,
+ cross_attention_head_dim=12,
+ cross_attention_dim=None,
+ mlp_ratio=2.0,
+ )
+
+ assert block.attn2 is None
+ assert isinstance(block.norm2, torch.nn.LayerNorm)
+
+
+def test_guidance_transformer_matches_diffusers_and_tuple_output():
+ from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+
+ from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+ SanaTimestepEmbedding,
+ SanaTimesteps,
+ )
+
+ config = _TINY_CONFIG | {"guidance_embeds": True}
+ torch.manual_seed(29)
+ reference = DiffusersTransformer(**config).eval()
+ model = SanaVideoTransformer3DModel(**config).eval()
+ assert set(model.state_dict()) == set(reference.state_dict())
+ model.load_state_dict(reference.state_dict())
+ assert isinstance(model.time_embed.time_proj, SanaTimesteps)
+ assert isinstance(model.time_embed.timestep_embedder, SanaTimestepEmbedding)
+ assert isinstance(model.time_embed.guidance_condition_proj, SanaTimesteps)
+ assert isinstance(model.time_embed.guidance_embedder, SanaTimestepEmbedding)
+
+ torch.manual_seed(31)
+ hidden_states = torch.randn(1, 4, 3, 4, 4)
+ encoder_hidden_states = torch.randn(1, 5, 8)
+ encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 0]])
+ timestep = torch.tensor([500.0])
+ guidance = torch.tensor([6.0])
+
+ with torch.no_grad():
+ expected = reference(
+ hidden_states,
+ encoder_hidden_states,
+ timestep,
+ guidance=guidance,
+ encoder_attention_mask=encoder_attention_mask,
+ return_dict=False,
+ )
+ actual = model(
+ hidden_states,
+ encoder_hidden_states,
+ timestep,
+ guidance=guidance,
+ encoder_attention_mask=encoder_attention_mask,
+ return_dict=False,
+ )
+
+ assert isinstance(actual, tuple)
+ assert len(actual) == len(expected) == 1
+ torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0)
+
+
+def test_native_transformer_config_filters_diffusers_metadata():
+ from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformerConfig
+
+ config = SanaVideoTransformerConfig.from_dict(
+ _TINY_CONFIG | {"_class_name": "SanaVideoTransformer3DModel", "_diffusers_version": "0.38.0"}
+ )
+ assert config.patch_size == (1, 2, 2)
+
+ with pytest.raises(ValueError, match="unsupported_field"):
+ SanaVideoTransformerConfig.from_dict(_TINY_CONFIG | {"unsupported_field": True})
+
+
+def test_dual_vae_variant_resolution():
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+
+ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+ DistributedAutoencoderKLLTX2Video,
+ )
+ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+ DistributedAutoencoderKLWan,
+ )
+ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+ from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+ SanaVideoPipeline,
+ _resolve_vae_class_and_dtype,
+ )
+
+ assert not issubclass(SanaVideoPipeline, DiffusionPipeline)
+ assert issubclass(SanaVideoPipeline, ProgressBarMixin)
+
+ vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLWan", torch.bfloat16)
+ assert vae_class is DistributedAutoencoderKLWan
+ assert dtype is torch.float32
+
+ vae_class, dtype = _resolve_vae_class_and_dtype("AutoencoderKLLTX2Video", torch.bfloat16)
+ assert vae_class is DistributedAutoencoderKLLTX2Video
+ assert dtype is torch.bfloat16
+
+ with pytest.raises(ValueError, match="Unsupported SANA-Video VAE"):
+ _resolve_vae_class_and_dtype("UnknownVAE", torch.bfloat16)
diff --git a/tests/diffusion/models/sana_wm/__init__.py b/tests/diffusion/models/sana_wm/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py
new file mode 100644
index 000000000..cea8aeb76
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_cam_reduce_scatter.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's camera-branch reduce-scatter.
+
+``out_proj_cam`` is row-parallel, so every rank produces a partial sum over the
+full hidden width, while the GDN stream it is added to is TP-local. Instead of
+all-reducing and then discarding all but one slice, the branch reduce-scatters
+straight to the local shard. These tests pin the shard each rank ends up with,
+including the bias slice, on both code paths:
+
+* the process-group path, which drives ``torch.distributed.reduce_scatter_tensor``
+ directly because vLLM-Omni's diffusion ``GroupCoordinator`` does not implement
+ ``reduce_scatter``;
+* the all-reduce-and-slice fallback taken when the coordinator exposes no
+ ``device_group``.
+
+CPU-only: the collectives and the rank/world-size lookups are mocked, so a
+single process can simulate an arbitrary TP degree.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import SanaWmSelfAttention
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+HIDDEN = 24
+
+
+def _host(bias: torch.Tensor | None) -> SimpleNamespace:
+ """A stand-in ``self`` exposing only what the method reads."""
+ return SimpleNamespace(out_proj_cam=SimpleNamespace(bias=bias))
+
+
+def _fake_reduce_scatter_tensor(rank: int, tp_size: int):
+ """Model ``reduce_scatter_tensor``: split dim 0 and keep chunk ``rank``.
+
+ The test drives a single process, so the other ranks' partials are already
+ folded into the input; scattering is all this has to reproduce.
+ """
+
+ def _impl(output: torch.Tensor, input_: torch.Tensor, group=None) -> None:
+ output.copy_(input_.chunk(tp_size, dim=0)[rank])
+
+ return _impl
+
+
+def _tp_group(with_device_group: bool):
+ return SimpleNamespace(device_group=object() if with_device_group else None)
+
+
+@pytest.mark.parametrize("tp_size", [2, 4])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_process_group_path_matches_all_reduce_then_slice(tp_size: int, rank: int) -> None:
+ """The scattered shard equals the slice the all-reduce path would take."""
+ local_width = HIDDEN // tp_size
+ torch.manual_seed(0)
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+ reference = torch.zeros(2, 3, local_width)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+ patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+ patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(True)),
+ patch("torch.distributed.reduce_scatter_tensor", side_effect=_fake_reduce_scatter_tensor(rank, tp_size)),
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+ start = rank * local_width
+ expected = (contrib + bias)[..., start : start + local_width]
+ assert out.shape == (2, 3, local_width)
+ torch.testing.assert_close(out, expected)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_fallback_path_matches_the_process_group_path(rank: int) -> None:
+ """Without a device_group the method all-reduces and slices, same result."""
+ tp_size = 2
+ local_width = HIDDEN // tp_size
+ torch.manual_seed(1)
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+ reference = torch.zeros(2, 3, local_width)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size),
+ patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank),
+ patch(f"{_MODULE}.get_tp_group", return_value=_tp_group(False)),
+ patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t),
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, reference)
+
+ start = rank * local_width
+ torch.testing.assert_close(out, (contrib + bias)[..., start : start + local_width])
+
+
+def test_tp1_adds_bias_without_a_collective() -> None:
+ contrib = torch.randn(2, 3, HIDDEN)
+ bias = torch.randn(HIDDEN)
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1),
+ patch("torch.distributed.reduce_scatter_tensor") as collective,
+ ):
+ out = SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(bias), contrib, contrib)
+
+ assert collective.call_count == 0
+ torch.testing.assert_close(out, contrib + bias)
+
+
+def test_rejects_a_width_that_does_not_shard() -> None:
+ """A local width that is not 1/tp of the contribution is a wiring bug."""
+ contrib = torch.randn(1, 2, HIDDEN)
+ reference = torch.zeros(1, 2, HIDDEN // 3) # tp=2 would need HIDDEN // 2
+
+ with (
+ patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2),
+ patch("torch.distributed.reduce_scatter_tensor"),
+ pytest.raises(ValueError, match="camera contribution width mismatch"),
+ ):
+ SanaWmSelfAttention._reduce_scatter_cam_contrib(_host(None), contrib, reference)
diff --git a/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py
new file mode 100644
index 000000000..b01ef02ce
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_fused_qk_rms_norm.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Equivalence tests for SANA-WM's TensorParallelRMSNorm q/k fusion.
+
+`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor
+and issues a single tensor-parallel all-reduce instead of two. These tests lock
+in that it is numerically identical to applying `norm_q`/`norm_k` separately,
+and that it really issues exactly one (fused) all-reduce.
+
+CPU-only: the collective is mocked, so a single process simulates one rank
+(all-reduce is identity). This exercises the cat/split logic a real multi-rank
+run depends on.
+"""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.sana_wm.sana_wm_transformer import (
+ TensorParallelRMSNorm,
+ fused_qk_rms_norm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODULE = "vllm_omni.diffusion.models.sana_wm.sana_wm_transformer"
+
+
+def _make_norms(dim: int, tp_size: int) -> tuple[TensorParallelRMSNorm, TensorParallelRMSNorm]:
+ """Two norms with distinct random weights so a swapped q/k slice is caught."""
+ torch.manual_seed(0)
+ norm_q = TensorParallelRMSNorm(dim, tp_size=tp_size)
+ norm_k = TensorParallelRMSNorm(dim, tp_size=tp_size)
+ norm_q.weight.data.normal_()
+ norm_k.weight.data.normal_()
+ return norm_q, norm_k
+
+
+def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+ """Single-rank all-reduce: the reduction over one rank is the identity."""
+ return tensor
+
+
+def test_fused_matches_separate() -> None:
+ dim = 64
+ norm_q, norm_k = _make_norms(dim, tp_size=2)
+ q = torch.randn(2, 5, dim)
+ k = torch.randn(2, 5, dim)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce):
+ ref_q, ref_k = norm_q(q), norm_k(k)
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+
+ torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0)
+
+
+def test_fused_issues_one_all_reduce() -> None:
+ dim = 32
+ norm_q, norm_k = _make_norms(dim, tp_size=4)
+ q = torch.randn(1, 3, dim)
+ k = torch.randn(1, 3, dim)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 1
+ # Both sums travel in one tensor: last dim is 2, not 1.
+ assert collective.call_args.args[0].shape[-1] == 2
+
+ collective.reset_mock()
+ norm_q(q), norm_k(k)
+ assert collective.call_count == 2
+
+
+@pytest.mark.parametrize(
+ "norms",
+ [
+ (RMSNorm(16), RMSNorm(16)),
+ (nn.Identity(), nn.Identity()),
+ ],
+ ids=["plain_rms_norm", "identity"],
+)
+def test_falls_back_for_non_tp_norms(norms: tuple[nn.Module, nn.Module]) -> None:
+ """qk_norm=False and TP=1 both hand in norms the fused path cannot use."""
+ norm_q, norm_k = norms
+ q = torch.randn(2, 4, 16)
+ k = torch.randn(2, 4, 16)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 0
+
+ torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
+
+
+def test_tp1_skips_the_collective() -> None:
+ norm_q, norm_k = _make_norms(16, tp_size=1)
+ q = torch.randn(2, 4, 16)
+ k = torch.randn(2, 4, 16)
+
+ with patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce) as collective:
+ fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k)
+ assert collective.call_count == 0
+
+ torch.testing.assert_close(fused_q, norm_q(q), rtol=0, atol=0)
+ torch.testing.assert_close(fused_k, norm_k(k), rtol=0, atol=0)
diff --git a/tests/diffusion/models/sana_wm/test_request.py b/tests/diffusion/models/sana_wm/test_request.py
new file mode 100644
index 000000000..7e40eedaf
--- /dev/null
+++ b/tests/diffusion/models/sana_wm/test_request.py
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SANA-WM request normalization."""
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_wm.request import (
+ SANA_WM_DEFAULT_HEIGHT,
+ SANA_WM_DEFAULT_NUM_FRAMES,
+ SANA_WM_DEFAULT_WIDTH,
+ normalize_sana_wm_payload,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _prompt(**sana_wm):
+ """A valid request, with the action rollout kept in step with num_frames."""
+ block = {"num_frames": SANA_WM_DEFAULT_NUM_FRAMES, **sana_wm}
+ block.setdefault("action", f"w-{block['num_frames'] - 1}")
+ return {
+ "prompt": "a quiet city street",
+ "multi_modal_data": {"image": object()},
+ "sana_wm": block,
+ }
+
+
+def _identity_poses(count):
+ return [[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]] for _ in range(count)]
+
+
+def _camera_prompt(**camera):
+ prompt = _prompt(num_frames=9)
+ prompt["sana_wm"].pop("action")
+ prompt["sana_wm"]["camera"] = {"poses": _identity_poses(9), **camera}
+ return prompt
+
+
+def test_defaults_are_latent_aligned():
+ payload = normalize_sana_wm_payload(_prompt())["additional_information"]["sana_wm"]
+ assert payload["num_frames"] == SANA_WM_DEFAULT_NUM_FRAMES
+ assert payload["height"] == SANA_WM_DEFAULT_HEIGHT
+ assert payload["width"] == SANA_WM_DEFAULT_WIDTH
+
+
+@pytest.mark.parametrize("num_frames", [9, 33, 161])
+def test_accepts_aligned_geometry(num_frames):
+ payload = normalize_sana_wm_payload(_prompt(num_frames=num_frames))["additional_information"]["sana_wm"]
+ assert payload["num_frames"] == num_frames
+
+
+@pytest.mark.parametrize("num_frames", [8, 10, 160])
+def test_rejects_unaligned_num_frames(num_frames):
+ # The VAE compresses 8x temporally; an unaligned count would be floored.
+ with pytest.raises(ValueError, match="num_frames"):
+ normalize_sana_wm_payload(_prompt(num_frames=num_frames))
+
+
+@pytest.mark.parametrize(("height", "width"), [(700, 1280), (704, 1270), (703, 1279)])
+def test_rejects_unaligned_resolution(height, width):
+ with pytest.raises(ValueError, match="divisible by 32"):
+ normalize_sana_wm_payload(_prompt(height=height, width=width))
+
+
+def test_normalization_is_idempotent():
+ once = normalize_sana_wm_payload(_prompt(num_frames=33))
+ twice = normalize_sana_wm_payload(once)
+ assert once["additional_information"]["sana_wm"] == twice["additional_information"]["sana_wm"]
+
+
+def test_rejects_action_shorter_than_num_frames():
+ # Used to shrink the camera tensor silently, then fail deep in the
+ # transformer on a camera/latent frame mismatch.
+ with pytest.raises(ValueError, match="rolls out 33 frames but num_frames is 161"):
+ normalize_sana_wm_payload(_prompt(action="w-32"))
+
+
+def test_rejects_action_longer_than_num_frames():
+ with pytest.raises(ValueError, match="rolls out 161 frames but num_frames is 9"):
+ normalize_sana_wm_payload(_prompt(num_frames=9, action="w-160"))
+
+
+def test_rejects_oversized_action_without_expanding_it():
+ # A 12-byte request body: expanding this rollout before checking the frame
+ # count allocated ~16 GB. Rejection has to be arithmetic, so this test also
+ # fails (by hanging/OOM) if the duration is ever expanded again.
+ with pytest.raises(ValueError, match="rolls out 1000000001 frames"):
+ normalize_sana_wm_payload(_prompt(action="w-1000000000"))
+
+
+@pytest.mark.parametrize("action", ["", " ", "w-0", "w-abc", "w", "x-4"])
+def test_rejects_malformed_action(action):
+ with pytest.raises(ValueError):
+ normalize_sana_wm_payload(_prompt(action=action))
+
+
+@pytest.mark.parametrize("value", [0, -1, float("nan"), float("inf")])
+def test_rejects_unusable_focal_length(value):
+ # compute_raymap divides by fx/fy; a zero or non-finite focal length turns
+ # the whole Plucker map into NaN and the response into a NaN video.
+ intrinsics = {"fx": value, "fy": 640, "cx": 640, "cy": 352}
+ with pytest.raises(ValueError, match="fx"):
+ normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_principal_point(value):
+ intrinsics = {"fx": 640, "fy": 640, "cx": value, "cy": 352}
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(_prompt(intrinsics=intrinsics))
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf")])
+def test_rejects_non_finite_speeds(value):
+ # `parsed <= 0` is False for NaN, so the positivity check alone let it pass.
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(_prompt(translation_speed=value))
+
+
+def test_rejects_non_finite_pose_entry():
+ prompt = _camera_prompt()
+ prompt["sana_wm"]["camera"]["poses"][3][2][1] = float("nan")
+ with pytest.raises(ValueError, match="must be finite"):
+ normalize_sana_wm_payload(prompt)
+
+
+def test_accepts_default_camera_metadata():
+ payload = normalize_sana_wm_payload(_camera_prompt())["additional_information"]["sana_wm"]
+ assert payload["camera"]["poses"]
+ # The unread fields are not carried into the canonical payload.
+ assert set(payload["camera"]) == {"poses"}
+
+
+@pytest.mark.parametrize(
+ ("field", "value"),
+ [("format", "w2c_4x4"), ("coordinate_system", "opengl")],
+)
+def test_rejects_unsupported_camera_convention(field, value):
+ # Nothing downstream reads these, so accepting one would mean interpreting
+ # the poses under the wrong convention.
+ with pytest.raises(ValueError, match=field.split("_")[0]):
+ normalize_sana_wm_payload(_camera_prompt(**{field: value}))
diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py
new file mode 100644
index 000000000..26af9a259
--- /dev/null
+++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_logits_processor.py
@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression test: SenseNovaU1ForCausalLM instantiation under the diffusion
+config shim must not crash on missing ``head_dtype`` and the resulting
+LogitsProcessor must have ``head_dtype is None`` (= use model dtype)."""
+
+import os
+
+import pytest
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+from vllm.distributed.parallel_state import (
+ cleanup_dist_env_and_memory,
+ init_distributed_environment,
+ initialize_model_parallel,
+)
+
+from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+ SenseNovaU1ForCausalLM,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+from vllm_omni.transformers_utils.configs.sensenova_u1 import (
+ SenseNovaU1Config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _build_tiny_sensenova_u1(quant_config=None, prefix="model"):
+ """Build a minimal SenseNovaU1ForCausalLM for unit tests."""
+ config = SenseNovaU1Config(
+ llm_config={
+ "hidden_size": 64,
+ "num_attention_heads": 2,
+ "num_hidden_layers": 1,
+ "intermediate_size": 128,
+ "vocab_size": 32,
+ "max_position_embeddings": 128,
+ "max_position_embeddings_hw": 128,
+ },
+ )
+ return SenseNovaU1ForCausalLM(config.llm_config, quant_config=quant_config, prefix=prefix)
+
+
+def test_logits_processor_head_dtype_under_diffusion_shim():
+ # Initialize the distributed environment
+ os.environ.setdefault("MASTER_ADDR", "localhost")
+ os.environ.setdefault("MASTER_PORT", "29543")
+ init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+ initialize_model_parallel()
+
+ # Create a vLLM config with a wrapped diffusion config
+ fake_diff_config = _DiffusionVllmModelConfig(
+ model="sensenova-u1-test",
+ dtype=torch.bfloat16,
+ max_model_len=8192,
+ original_max_model_len=8192,
+ )
+ vllm_config = VllmConfig(
+ device_config=DeviceConfig(device="cpu"),
+ )
+ vllm_config.modelconfig = fake_diff_config # type: ignore[assignment]
+
+ # Build a tiny version of the Causal LM component
+ with set_current_vllm_config(vllm_config):
+ model = _build_tiny_sensenova_u1()
+
+ # Ensure that we have a set head_dtype attribute. Currently, the head_dtype
+ # is set to None
+ head_dtype = model.logits_processor.head_dtype
+ assert head_dtype is None
+ cleanup_dist_env_and_memory()
diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py
index 9d5b11d3b..12c0ebf45 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py
@@ -8,15 +8,28 @@ import torch
from PIL import Image
from torch import nn
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
Wan22I2VPipeline,
+ get_wan22_i2v_post_process_func,
get_wan22_i2v_pre_process_func,
)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
+ video = torch.zeros(1, 4, 1, 2, 2)
+
+ output = get_wan22_i2v_post_process_func(SimpleNamespace())(
+ video,
+ sampling_params=SimpleNamespace(output_type="latent"),
+ )
+
+ assert output is video
+
+
def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
pipeline = object.__new__(Wan22I2VPipeline)
nn.Module.__init__(pipeline)
@@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
return pipeline
+def _make_i2v_sampling(**overrides):
+ values = {
+ "height": 16,
+ "width": 16,
+ "num_frames": 5,
+ "num_inference_steps": 1,
+ "guidance_scale_provided": True,
+ "guidance_scale": 1.0,
+ "guidance_scale_2": None,
+ "guidance_scale_2_provided": False,
+ "boundary_ratio": None,
+ "generator": None,
+ "seed": None,
+ "num_outputs_per_prompt": 1,
+ "max_sequence_length": 8,
+ "latents": None,
+ "output_type": "latent",
+ "extra_args": {},
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
request = SimpleNamespace(
@@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess(missing_image)
+def _make_i2v_preprocess_request(last_image):
+ return SimpleNamespace(
+ prompt={
+ "prompt": "p",
+ "multi_modal_data": {
+ "image": Image.new("RGB", (320, 160), "red"),
+ "last_image": last_image,
+ },
+ },
+ sampling_params=SimpleNamespace(height=16, width=16),
+ )
+
+
+def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+ result = preprocess(_make_i2v_preprocess_request([]))
+
+ assert result.prompt["multi_modal_data"]["last_image"] is None
+ assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
+
+
+def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+ last_image = Image.new("RGB", (16, 16), "blue")
+
+ result = preprocess(_make_i2v_preprocess_request([last_image]))
+
+ assert result.prompt["multi_modal_data"]["last_image"] is last_image
+ assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+@pytest.mark.parametrize(
+ "last_image",
+ [
+ Image.new("RGB", (16, 16), "blue"),
+ torch.zeros(3, 16, 16),
+ ],
+)
+def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+ result = preprocess(_make_i2v_preprocess_request(last_image))
+
+ assert result.prompt["multi_modal_data"]["last_image"] is last_image
+ assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+ path = tmp_path / "last.png"
+ Image.new("RGB", (16, 16), "blue").save(path)
+
+ result = preprocess(_make_i2v_preprocess_request(str(path)))
+
+ last_image = result.prompt["multi_modal_data"]["last_image"]
+ assert isinstance(last_image, Image.Image)
+ assert last_image.mode == "RGB"
+ assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
+
+
+def test_i2v_preprocess_rejects_multiple_last_images() -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+ with pytest.raises(ValueError, match="at most one last_image"):
+ preprocess(
+ _make_i2v_preprocess_request(
+ [
+ Image.new("RGB", (16, 16), "blue"),
+ Image.new("RGB", (16, 16), "green"),
+ ]
+ )
+ )
+
+
+def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
+ preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
+
+ with pytest.raises(TypeError, match="Unsupported last_image format"):
+ preprocess(_make_i2v_preprocess_request({"not": "an image"}))
+
+
def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
latents = torch.zeros(1, 4, 2, 4, 4)
@@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
assert first_frame_mask.shape == (1, 1, 2, 2, 2)
assert first_frame_mask[:, :, 0].sum() == 0
assert first_frame_mask[:, :, 1].sum() == 4
+
+
+def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
+ pipeline = _make_i2v_pipeline(expand_timesteps=True)
+ generators = [
+ torch.Generator(device="cpu").manual_seed(1),
+ torch.Generator(device="cpu").manual_seed(2),
+ ]
+
+ latents, condition, first_frame_mask = pipeline.prepare_latents(
+ image=torch.zeros(2, 3, 16, 16),
+ batch_size=2,
+ num_channels_latents=4,
+ height=16,
+ width=16,
+ num_frames=5,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ generator=generators,
+ )
+
+ assert latents.shape == (2, 4, 2, 2, 2)
+ assert condition.shape == (2, 4, 1, 2, 2)
+ assert first_frame_mask.shape == (2, 1, 2, 2, 2)
+
+
+def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+ SimpleNamespace(is_available=lambda: False),
+ )
+ pipeline = _make_i2v_pipeline(expand_timesteps=True)
+ pipeline.scheduler = StubScheduler([9])
+ pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+ pipeline._sample_solver = "unipc"
+ pipeline._flow_shift = 5.0
+ pipeline.boundary_ratio = 0.875
+ pipeline.has_image_encoder = False
+ pipeline._guidance_scale = None
+ pipeline._guidance_scale_2 = None
+ pipeline._num_timesteps = None
+ pipeline._current_timestep = None
+ pipeline.check_inputs = lambda **kwargs: None
+ prepare_call = {}
+
+ def _fake_encode_prompt(**kwargs):
+ batch_size = len(kwargs["prompt"])
+ n = batch_size * kwargs["num_videos_per_prompt"]
+ return torch.zeros(n, 2, 3), None
+
+ def _fake_prepare_latents(**kwargs):
+ prepare_call.update(kwargs)
+ batch_size = kwargs["batch_size"]
+ return (
+ kwargs["latents"],
+ torch.zeros(batch_size, 4, 1, 2, 2),
+ torch.ones(batch_size, 1, 2, 2, 2),
+ )
+
+ pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
+ pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
+ pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
+
+ gen_a = torch.Generator(device="cpu").manual_seed(1)
+ gen_b = torch.Generator(device="cpu").manual_seed(2)
+ latents_a = torch.zeros(2, 4, 2, 2, 2)
+ latents_b = torch.ones(2, 4, 2, 2, 2)
+ image_a = torch.zeros(1, 3, 16, 16)
+ image_b = torch.ones(1, 3, 16, 16)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
+ sampling_params=_make_i2v_sampling(
+ generator=gen_a,
+ latents=latents_a,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
+ sampling_params=_make_i2v_sampling(
+ generator=gen_b,
+ latents=latents_b,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ ]
+ )
+
+ outputs = pipeline.forward(batch)
+
+ assert prepare_call["batch_size"] == 4
+ assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+ torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+ torch.testing.assert_close(
+ prepare_call["image"],
+ torch.cat([image_a, image_a, image_b, image_b]),
+ )
+ assert len(outputs) == 2
+ torch.testing.assert_close(outputs[0].output, latents_a)
+ torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
+ SimpleNamespace(is_available=lambda: False),
+ )
+ pipeline = _make_i2v_pipeline(expand_timesteps=True)
+ pipeline.scheduler = StubScheduler([9])
+ pipeline.od_config = SimpleNamespace(flow_shift=5.0)
+ pipeline._sample_solver = "unipc"
+ pipeline._flow_shift = 5.0
+ pipeline.boundary_ratio = 0.875
+ pipeline.has_image_encoder = False
+ pipeline._guidance_scale = None
+ pipeline._guidance_scale_2 = None
+ pipeline._num_timesteps = None
+ pipeline._current_timestep = None
+ pipeline.check_inputs = lambda **kwargs: None
+ pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
+ sampling_params=_make_i2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
+ sampling_params=_make_i2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="image condition"):
+ pipeline.forward(batch)
+
+
+def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
+ pipeline = _make_i2v_pipeline(expand_timesteps=True)
+ image = torch.zeros(1, 3, 16, 16)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "multi_modal_data": {"image": image, "last_image": image},
+ },
+ sampling_params=_make_i2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt": "second", "multi_modal_data": {"image": image}},
+ sampling_params=_make_i2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="mix of provided and missing last_image"):
+ pipeline.forward(batch)
diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
index 8c06df034..de5abeab0 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import importlib
from contextlib import contextmanager
from types import SimpleNamespace
@@ -9,6 +10,7 @@ import torch
from torch import nn
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -82,6 +84,7 @@ def _make_pipeline() -> Wan22Pipeline:
pipeline.vae_scale_factor_spatial = 8
pipeline.boundary_ratio = 0.875
pipeline.expand_timesteps = False
+ pipeline.is_dmd = False
pipeline._guidance_scale = None
pipeline._guidance_scale_2 = None
pipeline._num_timesteps = None
@@ -93,6 +96,29 @@ def _make_pipeline() -> Wan22Pipeline:
return pipeline
+def _make_sampling(**overrides):
+ values: dict[str, object] = {
+ "height": None,
+ "width": None,
+ "num_frames": 1,
+ "num_inference_steps": 2,
+ "guidance_scale_provided": True,
+ "guidance_scale": 1.0,
+ "guidance_scale_2": None,
+ "guidance_scale_2_provided": False,
+ "boundary_ratio": None,
+ "generator": None,
+ "seed": None,
+ "num_outputs_per_prompt": 1,
+ "max_sequence_length": 32,
+ "latents": None,
+ "output_type": "latent",
+ "extra_args": {},
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
@pytest.mark.parametrize(
("sampling_params_kwargs", "expected_low", "expected_high"),
[
@@ -129,9 +155,10 @@ def test_forward_delegates_denoising_to_diffuse(
)
batch = DiffusionRequestBatch(requests=[mock_req])
- output = pipeline.forward(batch)
+ outputs = pipeline.forward(batch)
- assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+ assert len(outputs) == 1
+ assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
assert captured["guidance_low"] == expected_low
@@ -142,6 +169,134 @@ def test_forward_delegates_denoising_to_diffuse(
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
+def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
+ pipeline = _make_pipeline()
+ encode_call = {}
+ prepare_call = {}
+
+ def _fake_encode_prompt(**kwargs):
+ encode_call.update(kwargs)
+ batch_size = len(kwargs["prompt"])
+ n = batch_size * kwargs["num_videos_per_prompt"]
+ return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
+
+ def _fake_prepare_latents(**kwargs):
+ prepare_call.update(kwargs)
+ return kwargs["latents"]
+
+ pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
+ pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
+ pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
+
+ gen_a = torch.Generator(device="cpu").manual_seed(1)
+ gen_b = torch.Generator(device="cpu").manual_seed(2)
+ latents_a = torch.zeros(2, 4, 1, 2, 2)
+ latents_b = torch.ones(2, 4, 1, 2, 2)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={"prompt": "first", "negative_prompt": "bad first"},
+ sampling_params=_make_sampling(
+ generator=gen_a,
+ latents=latents_a,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt": "second", "negative_prompt": "bad second"},
+ sampling_params=_make_sampling(
+ generator=gen_b,
+ latents=latents_b,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ ]
+ )
+
+ outputs = pipeline.forward(batch)
+
+ assert encode_call["prompt"] == ["first", "second"]
+ assert encode_call["negative_prompt"] == ["bad first", "bad second"]
+ assert prepare_call["batch_size"] == 4
+ assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+ torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+ assert len(outputs) == 2
+ torch.testing.assert_close(outputs[0].output, latents_a)
+ torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_forward_batches_precomputed_prompt_embeddings() -> None:
+ pipeline = _make_pipeline()
+ diffuse_call = {}
+ pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run") # type: ignore[method-assign]
+ pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2) # type: ignore[method-assign]
+
+ def _fake_diffuse(**kwargs):
+ diffuse_call.update(kwargs)
+ return kwargs["latents"]
+
+ pipeline.diffuse = _fake_diffuse # type: ignore[method-assign]
+ embeds_a = torch.zeros(3, 4)
+ embeds_b = torch.ones(3, 4)
+ negative_a = torch.full((3, 4), 2.0)
+ negative_b = torch.full((3, 4), 3.0)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
+ sampling_params=_make_sampling(guidance_scale=4.0),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
+ sampling_params=_make_sampling(guidance_scale=4.0),
+ ),
+ ]
+ )
+
+ outputs = pipeline.forward(batch)
+
+ torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
+ torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
+ assert len(outputs) == 2
+
+
+def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
+ pipeline = _make_pipeline()
+ kwargs = {
+ "num_channels_latents": 4,
+ "height": 16,
+ "width": 16,
+ "num_frames": 5,
+ "dtype": torch.float32,
+ "device": torch.device("cpu"),
+ }
+
+ batched = Wan22Pipeline.prepare_latents(
+ pipeline,
+ batch_size=2,
+ generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
+ **kwargs,
+ )
+ singles = torch.cat(
+ [
+ Wan22Pipeline.prepare_latents(
+ pipeline,
+ batch_size=1,
+ generator=torch.Generator().manual_seed(seed),
+ **kwargs,
+ )
+ for seed in (1, 2)
+ ]
+ )
+
+ torch.testing.assert_close(batched, singles)
+ assert not torch.equal(batched[0], batched[1])
+
+
def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
pipeline = _make_pipeline()
latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
@@ -186,3 +341,125 @@ def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
(3.0, 3, 28.0, False),
]
assert torch.equal(result, torch.full_like(latents, 10.0))
+
+
+class _StubDMDScheduler:
+ def __init__(self) -> None:
+ self.predict_clean_calls: list[tuple[float, float, float]] = []
+ self.add_noise_calls: list[tuple[float, float, float]] = []
+
+ def predict_clean(self, model_output, sample, timestep):
+ self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep)))
+ return sample - model_output
+
+ def add_noise(self, clean_sample, noise, timestep):
+ self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep)))
+ return clean_sample + 10.0
+
+
+def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None:
+ pipeline = _make_pipeline()
+ pipeline.is_dmd = True
+ pipeline.scheduler = _StubDMDScheduler()
+ latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32)
+ timesteps = torch.tensor([1000.0, 757.0, 522.0])
+
+ pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents) # type: ignore[method-assign]
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor",
+ lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]),
+ )
+
+ result = pipeline.diffuse(
+ latents=latents,
+ timesteps=timesteps,
+ prompt_embeds=torch.zeros(1, 8),
+ negative_prompt_embeds=None,
+ guidance_low=1.0,
+ guidance_high=1.0,
+ boundary_timestep=None,
+ dtype=torch.float32,
+ attention_kwargs={},
+ generator=torch.Generator(device="cpu").manual_seed(1),
+ )
+
+ assert pipeline.scheduler.predict_clean_calls == [
+ (1.0, 0.0, 1000.0),
+ (1.0, 9.0, 757.0),
+ (1.0, 18.0, 522.0),
+ ]
+ assert pipeline.scheduler.add_noise_calls == [
+ (-1.0, 2.0, 757.0),
+ (8.0, 2.0, 522.0),
+ ]
+ torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]]))
+
+
+def _make_gate_loading_pipeline():
+ pipeline = Wan22Pipeline.__new__(Wan22Pipeline)
+ nn.Module.__init__(pipeline)
+ gate = WanSelfAttention.__new__(WanSelfAttention)
+ nn.Module.__init__(gate)
+ gate.to_gate_compress = nn.Linear(1, 1)
+ pipeline.gate_holder = gate
+ return pipeline, gate
+
+
+@pytest.mark.parametrize(
+ ("module_name", "class_name"),
+ [
+ ("pipeline_wan2_2", "Wan22Pipeline"),
+ ("pipeline_wan2_2_i2v", "Wan22I2VPipeline"),
+ ("pipeline_wan2_2_s2v", "Wan22S2VPipeline"),
+ ("pipeline_wan2_2_vace", "Wan22VACEPipeline"),
+ ],
+)
+def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None:
+ module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}")
+ pipeline_cls = getattr(module, class_name)
+ pipeline = pipeline_cls.__new__(pipeline_cls)
+ expected = {"loaded"}
+
+ def fake_loader(model, weights):
+ assert model is pipeline
+ assert list(weights) == [("weight", torch.ones(1))]
+ return expected
+
+ monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader)
+
+ assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected
+
+
+def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None:
+ pipeline, gate = _make_gate_loading_pipeline()
+
+ class _Loader:
+ def __init__(self, model):
+ del model
+
+ def load_weights(self, weights):
+ return {name for name, _ in weights}
+
+ monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+ pipeline.load_weights(iter((("other.weight", torch.ones(1)),)))
+
+ assert pipeline.has_gate_compress_weights is False
+ assert gate.to_gate_compress is None
+
+
+def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None:
+ pipeline, gate = _make_gate_loading_pipeline()
+ original_gate = gate.to_gate_compress
+
+ class _Loader:
+ def __init__(self, model):
+ del model
+
+ def load_weights(self, weights):
+ return {name for name, _ in weights}
+
+ monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
+ pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),)))
+
+ assert pipeline.has_gate_compress_weights is True
+ assert gate.to_gate_compress is original_gate
diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py
index 314717869..5a37e14f2 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py
@@ -10,6 +10,7 @@ import torch
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
create_transformer_from_config,
+ get_wan22_post_process_func,
load_transformer_config,
retrieve_latents,
)
@@ -26,6 +27,17 @@ class _LatentDist:
return torch.tensor([2.0])
+def test_wan22_postprocess_honors_request_output_type() -> None:
+ video = torch.zeros(1, 4, 1, 2, 2)
+
+ output = get_wan22_post_process_func(SimpleNamespace())(
+ video,
+ sampling_params=SimpleNamespace(output_type="latent"),
+ )
+
+ assert output is video
+
+
def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
generator = torch.Generator(device="cpu")
diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py
index 4440a8f35..e5bae65df 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py
@@ -1,14 +1,307 @@
+from types import SimpleNamespace
from unittest.mock import MagicMock, patch
+import numpy as np
+import PIL.Image
import pytest
import torch
import torch.nn as nn
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
+ Wan22S2VPipeline,
+ _make_clip_generators,
+)
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def _make_s2v_sampling(**overrides):
+ values = {
+ "height": 16,
+ "width": 16,
+ "num_frames": 8,
+ "num_inference_steps": 1,
+ "guidance_scale_provided": True,
+ "guidance_scale": 1.0,
+ "generator": None,
+ "seed": None,
+ "num_outputs_per_prompt": 1,
+ "max_sequence_length": 8,
+ "latents": None,
+ "output_type": "latent",
+ "extra_args": {},
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
+def _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
+ pipeline = object.__new__(Wan22S2VPipeline)
+ nn.Module.__init__(pipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = SimpleNamespace(dtype=torch.float32)
+ pipeline.vae_scale_factor_spatial = 8
+ pipeline.resolution_divisor = 16
+ pipeline.motion_frames = 7
+ pipeline.drop_first_motion = True
+ pipeline._DEFAULT_INFER_FRAMES = 8
+ pipeline._guidance_scale = None
+ pipeline.check_inputs = lambda *args, **kwargs: None
+ pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
+ return pipeline
+
+
+def test_s2v_predict_noise_keeps_batch_dimension() -> None:
+ pipeline = object.__new__(Wan22S2VPipeline)
+ nn.Module.__init__(pipeline)
+ pipeline.device = torch.device("cpu")
+ transformer = MagicMock()
+ transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
+ transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
+ pipeline.transformer = transformer
+
+ result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))
+
+ assert result.shape == (1, 16, 2, 2, 2)
+
+
+def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
+ explicit_generator = torch.Generator(device="cpu").manual_seed(999)
+
+ generators = _make_clip_generators(
+ seeds=[1234, 5678, None, None],
+ request_generators=[
+ torch.Generator(device="cpu").manual_seed(1234),
+ torch.Generator(device="cpu").manual_seed(5678),
+ explicit_generator,
+ None,
+ ],
+ clip_index=2,
+ device=torch.device("cpu"),
+ )
+
+ assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
+ assert generators[2] is explicit_generator
+
+
+def test_s2v_forward_rejects_different_num_repeat() -> None:
+ pipeline = _make_s2v_validation_pipeline()
+ pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign]
+ image = PIL.Image.new("RGB", (16, 16))
+ audio = np.zeros(16000, dtype=np.float32)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "multi_modal_data": {"image": image, "audio": audio},
+ "additional_information": {"num_repeat": 1},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "multi_modal_data": {"image": image, "audio": audio},
+ "additional_information": {"num_repeat": 2},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="num_repeat"):
+ pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
+ pipeline = _make_s2v_validation_pipeline()
+ pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign]
+ torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
+ 1,
+ 16 if audio[0] == 0 else 24,
+ )
+ image = PIL.Image.new("RGB", (16, 16))
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="audio embedding time lengths"):
+ pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_init_first_frame() -> None:
+ pipeline = _make_s2v_validation_pipeline()
+ image = PIL.Image.new("RGB", (16, 16))
+ audio = np.zeros(16000, dtype=np.float32)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "multi_modal_data": {"image": image, "audio": audio},
+ "additional_information": {"init_first_frame": True},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "multi_modal_data": {"image": image, "audio": audio},
+ "additional_information": {"init_first_frame": False},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="init_first_frame"):
+ pipeline.forward(batch)
+
+
+def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
+ pipeline = _make_s2v_validation_pipeline()
+ pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign]
+ image = PIL.Image.new("RGB", (16, 16))
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
+ },
+ sampling_params=_make_s2v_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
+ pipeline.forward(batch)
+
+
+def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
+ pipeline = object.__new__(Wan22S2VPipeline)
+ nn.Module.__init__(pipeline)
+ pipeline.device = torch.device("cpu")
+ pipeline.transformer = MagicMock()
+ pipeline.transformer.dtype = torch.float32
+ pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
+ pipeline.transformer.casual_audio_encoder = None
+ pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
+ pipeline.vae = MagicMock()
+ pipeline.vae.dtype = torch.float32
+ pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
+ pipeline.od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ parallel_config=SimpleNamespace(use_hsdp=False),
+ )
+ pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
+ pipeline.vae_scale_factor_spatial = 8
+ pipeline.resolution_divisor = 16
+ pipeline.motion_frames = 7
+ pipeline.drop_first_motion = True
+ pipeline._DEFAULT_INFER_FRAMES = 8
+ pipeline._guidance_scale = None
+ pipeline._num_timesteps = None
+ pipeline.check_inputs = lambda *args, **kwargs: None
+ pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
+ pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
+ pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
+ pipeline.prepare_motion_latents = MagicMock(
+ side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
+ )
+ pipeline._denormalize_latents = lambda latents: latents
+ pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
+
+ image = PIL.Image.new("RGB", (16, 16))
+ audio_a = np.zeros(16000, dtype=np.float32)
+ audio_b = np.ones(16000, dtype=np.float32)
+ latents_a = torch.zeros(2, 16, 2, 2, 2)
+ latents_b = torch.ones(2, 16, 2, 2, 2)
+ generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
+ generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "negative_prompt": "negative first",
+ "multi_modal_data": {"image": image, "audio": audio_a},
+ },
+ sampling_params=_make_s2v_sampling(
+ num_outputs_per_prompt=2,
+ generator=generators_a,
+ latents=latents_a,
+ ),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "negative_prompt": "negative second",
+ "multi_modal_data": {"image": image, "audio": audio_b},
+ },
+ sampling_params=_make_s2v_sampling(
+ num_outputs_per_prompt=2,
+ generator=generators_b,
+ latents=latents_b,
+ ),
+ ),
+ ]
+ )
+
+ with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
+ platform.is_available.return_value = False
+ outputs = pipeline.forward(batch)
+
+ assert len(outputs) == 2
+ assert outputs[0].output[0].shape[0] == 2
+ assert outputs[1].output[0].shape[0] == 2
+ assert outputs[0].output[1].shape == (16000,)
+ assert outputs[1].output[1].shape == (16000,)
+ np.testing.assert_array_equal(outputs[0].output[1], audio_a)
+ np.testing.assert_array_equal(outputs[1].output[1], audio_b)
+ assert np.shares_memory(outputs[0].output[1], audio_a)
+ assert np.shares_memory(outputs[1].output[1], audio_b)
+ torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
+ assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
+ assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
+ assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]
+
+
def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
model = object.__new__(WanS2VTransformer3DModel)
nn.Module.__init__(model)
@@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean
# -- Build request --
- from vllm_omni.diffusion.request import OmniDiffusionRequest
-
- sampling_params = MagicMock()
- sampling_params.height = 704
- sampling_params.width = 1024
- sampling_params.num_frames = 80
- sampling_params.num_inference_steps = 5
- sampling_params.guidance_scale = 4.5
- sampling_params.guidance_scale_provided = True
- sampling_params.generator = None
- sampling_params.seed = 42
+ sampling_params = _make_s2v_sampling(
+ height=704,
+ width=1024,
+ num_frames=80,
+ num_inference_steps=5,
+ guidance_scale=4.5,
+ generator=torch.Generator(device="cpu").manual_seed(42),
+ max_sequence_length=512,
+ )
ref_image = PIL.Image.new("RGB", (1024, 704))
audio_data = np.zeros(16000, dtype=np.float32)
- req = MagicMock(spec=OmniDiffusionRequest)
- req.prompts = [
- {
- "prompt": "test prompt",
- "negative_prompt": "bad quality",
- "multi_modal_data": {"image": ref_image, "audio": audio_data},
- "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
- }
- ]
- req.sampling_params = sampling_params
+ req = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="test",
+ prompt={
+ "prompt": "test prompt",
+ "negative_prompt": "bad quality",
+ "multi_modal_data": {"image": ref_image, "audio": audio_data},
+ "additional_information": {
+ "audio_path": audio_data,
+ "pose_video": None,
+ "init_first_frame": False,
+ },
+ },
+ sampling_params=sampling_params,
+ )
+ ]
+ )
# -- Mock methods that use complex internal state --
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
@@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)
# Mock predict_noise_maybe_with_cfg from CFGParallelMixin
- pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
+ pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)
# Mock platform methods
@@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
with patch(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
):
- result = Wan22S2VPipeline.forward(pipeline, req=req)
+ result = Wan22S2VPipeline.forward(pipeline, req=req)[0]
# -- Assertions --
# Text encoder was called
diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py
index 714b79d50..6534b7cff 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py
@@ -8,13 +8,14 @@ import torch
from PIL import Image
from torch import nn
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
Wan22VACEPipeline,
create_vace_transformer_from_config,
get_wan22_vace_pre_process_func,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
return pipeline
+def _make_vace_sampling(**overrides):
+ values = {
+ "height": 16,
+ "width": 16,
+ "num_frames": 5,
+ "num_inference_steps": 1,
+ "guidance_scale_provided": True,
+ "guidance_scale": 1.0,
+ "boundary_ratio": None,
+ "generator": None,
+ "seed": None,
+ "num_outputs_per_prompt": 1,
+ "max_sequence_length": 8,
+ "latents": None,
+ "output_type": "latent",
+ "extra_args": {},
+ }
+ values.update(overrides)
+ return SimpleNamespace(**values)
+
+
def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
ref = Image.new("RGB", (320, 160), "green")
@@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
torch.testing.assert_close(result, torch.ones_like(latents))
+
+
+def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+ SimpleNamespace(is_available=lambda: False),
+ )
+ pipeline = _make_vace_pipeline()
+ pipeline.transformer.vace_patch_embedding = None
+ pipeline.transformer_2 = None
+ pipeline.scheduler = StubScheduler([9])
+ pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+ pipeline._sample_solver = "unipc"
+ pipeline._flow_shift = 3.0
+ pipeline.boundary_ratio = None
+ pipeline._guidance_scale = None
+ pipeline._num_timesteps = None
+ pipeline._current_timestep = None
+ pipeline.check_inputs = lambda **kwargs: None
+ prepare_call = {}
+
+ def _fake_encode_prompt(**kwargs):
+ batch_size = len(kwargs["prompt"])
+ n = batch_size * kwargs["num_videos_per_prompt"]
+ return torch.zeros(n, 2, 3), None
+
+ def _fake_prepare_latents(**kwargs):
+ prepare_call.update(kwargs)
+ return kwargs["latents"]
+
+ pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
+ pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
+ pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
+ gen_a = torch.Generator(device="cpu").manual_seed(1)
+ gen_b = torch.Generator(device="cpu").manual_seed(2)
+ latents_a = torch.zeros(2, 4, 2, 2, 2)
+ latents_b = torch.ones(2, 4, 2, 2, 2)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={"prompt": "first"},
+ sampling_params=_make_vace_sampling(
+ generator=gen_a,
+ latents=latents_a,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={"prompt": "second"},
+ sampling_params=_make_vace_sampling(
+ generator=gen_b,
+ latents=latents_b,
+ num_outputs_per_prompt=2,
+ ),
+ ),
+ ]
+ )
+
+ outputs = pipeline.forward(batch)
+
+ assert prepare_call["batch_size"] == 4
+ assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
+ torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
+ assert len(outputs) == 2
+ torch.testing.assert_close(outputs[0].output, latents_a)
+ torch.testing.assert_close(outputs[1].output, latents_b)
+
+
+def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+ SimpleNamespace(is_available=lambda: False),
+ )
+ pipeline = _make_vace_pipeline()
+ pipeline.transformer.vace_patch_embedding = object()
+ pipeline.transformer_2 = None
+ pipeline.scheduler = StubScheduler([9])
+ pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+ pipeline._sample_solver = "unipc"
+ pipeline._flow_shift = 3.0
+ pipeline.boundary_ratio = None
+ pipeline._guidance_scale = None
+ pipeline._num_timesteps = None
+ pipeline._current_timestep = None
+ pipeline.check_inputs = lambda **kwargs: None
+ pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
+
+ def _fake_preprocess(video, **kwargs):
+ del kwargs
+ mask = torch.ones_like(video)
+ return video, mask, [[]]
+
+ pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
+ pipeline.prepare_video_latents = ( # type: ignore[method-assign]
+ lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
+ )
+ pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2) # type: ignore[method-assign]
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
+ },
+ sampling_params=_make_vace_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
+ },
+ sampling_params=_make_vace_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="VACE condition"):
+ pipeline.forward(batch)
+
+
+def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
+ SimpleNamespace(is_available=lambda: False),
+ )
+ pipeline = _make_vace_pipeline()
+ pipeline.transformer.vace_patch_embedding = object()
+ pipeline.transformer_2 = None
+ pipeline.scheduler = StubScheduler([9])
+ pipeline.od_config = SimpleNamespace(flow_shift=3.0)
+ pipeline._sample_solver = "unipc"
+ pipeline._flow_shift = 3.0
+ pipeline.boundary_ratio = None
+ pipeline._guidance_scale = None
+ pipeline._num_timesteps = None
+ pipeline._current_timestep = None
+ pipeline.check_inputs = lambda **kwargs: None
+ pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
+
+ def _fake_preprocess(video, mask, reference_images, **kwargs):
+ del kwargs
+ return video, mask, [reference_images]
+
+ pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
+ pipeline.prepare_video_latents = ( # type: ignore[method-assign]
+ lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
+ )
+ pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2) # type: ignore[method-assign]
+ video = torch.zeros(1, 3, 5, 16, 16)
+ mask = torch.ones_like(video)
+ reference = torch.zeros(3, 16, 16)
+ batch = DiffusionRequestBatch(
+ requests=[
+ SimpleNamespace(
+ request_id="a",
+ prompt={
+ "prompt": "first",
+ "additional_information": {
+ "source_video": video,
+ "mask": mask,
+ "reference_images": [reference],
+ },
+ },
+ sampling_params=_make_vace_sampling(),
+ ),
+ SimpleNamespace(
+ request_id="b",
+ prompt={
+ "prompt": "second",
+ "additional_information": {
+ "source_video": video,
+ "mask": mask,
+ "reference_images": [reference, reference],
+ },
+ },
+ sampling_params=_make_vace_sampling(),
+ ),
+ ]
+ )
+
+ with pytest.raises(ValueError, match="same number of reference images"):
+ pipeline.forward(batch)
diff --git a/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py
new file mode 100644
index 000000000..8c9208014
--- /dev/null
+++ b/tests/diffusion/models/wan2_2/test_wan_euler_scheduler_dmd.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_dmd_predict_clean_and_add_noise_match_flow_formulas() -> None:
+ scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+ timestep = torch.tensor(757.0)
+ sigma = scheduler.sigma_for_timestep(timestep)
+ sample = torch.tensor([2.0, -1.0])
+ model_output = torch.tensor([0.5, 0.25])
+ noise = torch.tensor([-0.5, 1.5])
+
+ clean = scheduler.predict_clean(model_output, sample, timestep)
+ expected_clean = sample - sigma * model_output
+ torch.testing.assert_close(clean, expected_clean)
+
+ renoised = scheduler.add_noise(clean, noise, timestep)
+ expected_renoised = (1.0 - sigma) * expected_clean + sigma * noise
+ torch.testing.assert_close(renoised, expected_renoised)
+
+
+def test_dmd_sigma_requires_scalar_timestep() -> None:
+ scheduler = WanEulerScheduler(num_train_timesteps=1000, shift=8.0)
+ with pytest.raises(ValueError, match="scalar timestep"):
+ scheduler.sigma_for_timestep(torch.tensor([757.0, 522.0]))
diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py
new file mode 100644
index 000000000..a3efbb40a
--- /dev/null
+++ b/tests/diffusion/offloader/test_cuda_host_registration.py
@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import ctypes
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+ CudaHostRegistration,
+ _coalesce_ranges,
+ _mapped_regions,
+)
+from vllm_omni.diffusion.offloader.host_registration import (
+ HostRegistrationBudgetError,
+ HostRegistrationCleanupError,
+ HostRegistrationError,
+ register_host_mappings,
+)
+from vllm_omni.host_weight_runtime import MappedHostRegion
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeRuntime:
+ def __init__(
+ self,
+ register_results: list[int | Exception],
+ unregister_results: list[int | Exception] | None = None,
+ *,
+ read_only_supported: bool = True,
+ attribute_result: int = 0,
+ ) -> None:
+ self._register_results = iter(register_results)
+ self._unregister_results = iter(unregister_results or [])
+ self._read_only_supported = read_only_supported
+ self._attribute_result = attribute_result
+ self.attribute_queries: list[tuple[int, int]] = []
+ self.registered: list[tuple[int, int, int]] = []
+ self.unregistered: list[int] = []
+
+ def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int:
+ self.attribute_queries.append((attribute, device))
+ ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported)
+ return self._attribute_result
+
+ def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+ self.registered.append((pointer, size, flags))
+ result = next(self._register_results)
+ if isinstance(result, Exception):
+ raise result
+ return result
+
+ def cudaHostUnregister(self, pointer: int) -> int:
+ self.unregistered.append(pointer)
+ result = next(self._unregister_results, 0)
+ if isinstance(result, Exception):
+ raise result
+ return result
+
+ @staticmethod
+ def cudaGetErrorString(error: int) -> str:
+ return f"error-{error}"
+
+
+def _region(
+ file_name: str,
+ address: int,
+ length: int,
+ *,
+ read_only: bool = True,
+) -> MappedHostRegion:
+ return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only)
+
+
+@pytest.fixture(autouse=True)
+def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0)
+
+
+def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None:
+ sentinel = object()
+ regions = (_region("weights", 0x1000, 4096),)
+
+ def create(actual_regions, *, max_bytes):
+ assert actual_regions is regions
+ assert max_bytes == 4096
+ return sentinel
+
+ monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+ assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel
+ with pytest.raises(HostRegistrationError, match="not supported on cpu"):
+ register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None)
+
+
+def test_ranges_align_per_file_without_merging_adjacent_files() -> None:
+ assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == (
+ registration_module._AddressRange(0x1000, 0x3000),
+ )
+ assert _mapped_regions(
+ (
+ _region("first.safetensors", 0x1000, 4096),
+ _region("second.safetensors", 0x2000, 4096),
+ )
+ ) == (
+ registration_module._AddressRange(0x1000, 0x2000),
+ registration_module._AddressRange(0x2000, 0x3000),
+ )
+
+
+def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0])
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+ with pytest.raises(HostRegistrationError, match="is writable"):
+ CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None)
+ with pytest.raises(HostRegistrationBudgetError, match="exceeding"):
+ CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096)
+ assert runtime.registered == []
+
+
+def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0], read_only_supported=False)
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+ with pytest.raises(HostRegistrationError, match="does not support read-only"):
+ CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+ assert runtime.attribute_queries == [
+ (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0)
+ ]
+ assert runtime.registered == []
+
+
+def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0], attribute_result=7)
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+ consumed: list[tuple[object, int]] = []
+ monkeypatch.setattr(
+ registration_module,
+ "_consume_last_cuda_error",
+ lambda actual_runtime, error: consumed.append((actual_runtime, error)),
+ )
+
+ with pytest.raises(HostRegistrationError, match="error-7"):
+ CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None)
+ assert consumed == [(runtime, 7)]
+
+
+def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0, 7])
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+ monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+ with pytest.raises(HostRegistrationError, match="error-7"):
+ CudaHostRegistration.create(
+ (
+ _region("first", 0x1000, 4096),
+ _region("second", 0x9000, 4096),
+ ),
+ max_bytes=None,
+ )
+ assert runtime.unregistered == [0x1000]
+
+
+def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0, 7], unregister_results=[9, 0])
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+ monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+
+ with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error:
+ CudaHostRegistration.create(
+ (
+ _region("first", 0x1000, 4096),
+ _region("second", 0x9000, 4096),
+ ),
+ max_bytes=None,
+ )
+ registration = error.value.active_registration
+ assert registration is not None
+ assert registration.close() == ()
+ assert runtime.unregistered == [0x1000, 0x1000]
+
+
+def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None:
+ runtime = _FakeRuntime([0], unregister_results=[9, 0])
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+ monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None)
+ registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None)
+
+ assert registration.total_bytes == 4096
+ assert registration.region_count == 1
+ assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",)
+ assert registration.close() == ()
+ assert runtime.unregistered == [0x1000, 0x1000]
+
+
+@pytest.mark.parametrize("pending_error", [0, 801])
+def test_consume_last_cuda_error_accepts_cleared_or_matching_state(
+ monkeypatch: pytest.MonkeyPatch,
+ pending_error: int,
+) -> None:
+ runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error)
+ monkeypatch.setattr(
+ registration_module.ctypes,
+ "CDLL",
+ lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"),
+ )
+
+ registration_module._consume_last_cuda_error(runtime, 801)
+
+
+def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None:
+ global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801)
+ monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime)
+
+ registration_module._consume_last_cuda_error(SimpleNamespace(), 801)
+
+
+def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ runtime = _FakeRuntime([0, 7])
+ monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False)
+ monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+ def missing_global_symbol(_name: object) -> None:
+ raise OSError("cudart symbols are local")
+
+ monkeypatch.setattr(
+ registration_module.ctypes,
+ "CDLL",
+ missing_global_symbol,
+ )
+
+ with pytest.raises(HostRegistrationError, match="cannot clear") as error:
+ CudaHostRegistration.create(
+ (
+ _region("first", 0x1000, 4096),
+ _region("second", 0x9000, 4096),
+ ),
+ max_bytes=None,
+ )
+
+ assert not isinstance(error.value, HostRegistrationCleanupError)
+ assert runtime.unregistered == [0x1000]
diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py
index 810f0dd52..80a03180f 100644
--- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py
+++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py
@@ -1,25 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import gc
+from typing import Any
import numpy as np
import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
-AUDIO_MODEL = {
+AUDIO_MODEL: dict[str, dict[str, int | None]] = {
"stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
}
-IMAGE_MODELS = {
+IMAGE_MODELS: dict[str, dict[str, int | None]] = {
"riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
}
-MODELS = {**AUDIO_MODEL, **IMAGE_MODELS}
+MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_MODELS}
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
@@ -29,12 +33,12 @@ MODEL_MARKS = {
_GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
-AUDIO_MODEL_PARAMS = {
+AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
"runner_params": {},
"sampler_params": {},
}
-IMAGE_MODEL_PARAMS = {
+IMAGE_MODEL_PARAMS: dict[str, dict[str, Any]] = {
"runner_params": {},
"sampler_params": {
"height": 256,
@@ -43,7 +47,7 @@ IMAGE_MODEL_PARAMS = {
}
-def inference(model_name: str, offload: bool = True):
+def inference(model_name: str, offload: bool = True) -> tuple[float, Any]:
gc.collect()
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
@@ -133,6 +137,7 @@ def test_cpu_offload_diffusion_model(model_name: str):
threshold = MODELS[model_name][platform]
if threshold is None:
pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+ assert threshold is not None
assert offload_peak_memory + threshold < no_offload_peak_memory, (
f"Offload peak memory {offload_peak_memory} MB should be less than "
diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
index 26944b202..8e083605f 100644
--- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py
@@ -1,36 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import gc
+from typing import Any
+
import numpy as np
import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
-AUDIO_MODEL = {
- "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 1500},
+AUDIO_MODEL: dict[str, dict[str, int | None]] = {
+ # The inference peak includes backend workspaces as well as model weights.
+ # On ROCm, large MIOpen workspaces and allocator fragmentation can mask the
+ # resident-weight reduction, so use a conservative floor that still catches
+ # a disabled/no-op layerwise offloader.
+ "stabilityai/stable-audio-open-1.0": {"cuda": 1500, "rocm": 512},
}
-IMAGE_VIDEO_MODELS = {
+IMAGE_VIDEO_MODELS: dict[str, dict[str, int | None]] = {
"riverclouds/qwen_image_random": {"cuda": 4500, "rocm": None},
# "Wan-AI/Wan2.2-T2V-A14B-Diffusers": {"cuda": 45000, "rocm": None},
}
-MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+MODELS: dict[str, dict[str, int | None]] = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
"stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
}
-AUDIO_MODEL_PARAMS = {
+AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = {
"runner_params": {},
"sampler_params": {},
}
-IMAGE_VIDEO_MODELS_PARAMS = {
+IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = {
"runner_params": {"boundary_ratio": 0.875, "flow_shift": 5.0},
"sampler_params": {"height": 480, "width": 640, "num_frames": 5},
}
@@ -55,11 +65,12 @@ def run_inference(
model_name: str,
layerwise_offload: bool = False,
num_inference_steps: int = 3,
-) -> float:
+) -> tuple[float, Any]:
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
- monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
- monitor.start()
+ with current_omni_platform.device(device_index):
+ free_bytes, total_bytes = current_omni_platform.mem_get_info()
+ initial_used_mb = (total_bytes - free_bytes) / (1024**2)
if model_name in AUDIO_MODEL:
params = AUDIO_MODEL_PARAMS
@@ -73,22 +84,37 @@ def run_inference(
# cache_backend="cache_dit",
**params["runner_params"],
) as runner:
+ # Measure steady-state inference memory, not model construction. Enabling
+ # layerwise offload first loads the model and then replaces each block's
+ # device storage with CPU-backed weights. Monitoring that transition
+ # captures both the original model and temporary staging allocations,
+ # which is not representative of layerwise-offloaded inference.
+ monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
current_omni_platform.reset_peak_memory_stats()
+ monitor.start()
- # Refer to tests/e2e/offline_inference/test_wan22.py
- # Use minimal settings for testing
- output = runner.omni.generate(
- "A cat sitting on a table",
- OmniDiffusionSamplingParams(
- generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
- guidance_scale=1.0,
- num_inference_steps=num_inference_steps,
- **params["sampler_params"],
- ),
- )
+ try:
+ # Refer to tests/e2e/offline_inference/test_wan22.py
+ # Use minimal settings for testing
+ output = runner.omni.generate(
+ "A cat sitting on a table",
+ OmniDiffusionSamplingParams(
+ generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+ guidance_scale=1.0,
+ num_inference_steps=num_inference_steps,
+ **params["sampler_params"],
+ ),
+ )
+ finally:
+ monitor.stop()
- peak = monitor.peak_used_mb
- monitor.stop()
+ # DeviceMemoryMonitor reports absolute device usage. Subtract this run's
+ # starting usage so process-wide compiler/workspace caches retained from a
+ # previous run do not make the second measurement order-dependent.
+ peak = max(0.0, monitor.peak_used_mb - initial_used_mb)
+
+ gc.collect()
+ current_omni_platform.empty_cache()
return peak, output
@@ -140,10 +166,13 @@ def test_layerwise_offload_diffusion_model(model_name: str):
if expected_saved_memory is None:
pytest.skip(f"Threshold not defined for {platform} on {model_name}")
+ assert expected_saved_memory is not None
# Verify that layerwise offloading significantly reduces memory usage
- # Passes only if the actual savings exceeds the expected savings
- assert layerwise_offload_peak_memory + expected_saved_memory < no_offload_peak_memory, (
+ # Passes only if the actual savings meets the expected savings
+ actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory
+ assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, (
f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
- f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
+ f"should be at least {expected_saved_memory} MB less than no offload peak memory "
+ f"{no_offload_peak_memory} MB (actual savings: {actual_saved_memory} MB)"
)
diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
index fd004d056..934fe3299 100644
--- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py
+++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
@@ -1,12 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for DistributedLayerwiseOffloadHook and backend utilities."""
import gc
import json
-import os
-import socket
+import weakref
from contextlib import contextmanager
from types import SimpleNamespace
@@ -18,6 +17,13 @@ from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.data import validate_dlo_host_registration_options
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+ HostWeightPlan,
+ TensorBinding,
+ build_checkpoint_mmap_plan,
+)
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.block_discovery import (
get_blocks_attr_names,
@@ -29,8 +35,17 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
DistributedLayerwiseOffloadHook,
PinnedResidentLayerGroup,
)
+from vllm_omni.diffusion.offloader.host_registration import (
+ HostRegistrationCleanupError,
+ HostRegistrationError,
+)
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
-from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
+from vllm_omni.diffusion.offloader.offload_plan import (
+ OffloadPlan,
+ get_offload_plan,
+)
+from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state
+from vllm_omni.host_weight_runtime import MappedHostRegion
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -48,33 +63,19 @@ class DummyEvent:
def record(self, _stream) -> None:
return None
+ def synchronize(self) -> None:
+ return None
+
@contextmanager
def dummy_stream(_stream):
yield None
-def _find_free_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
- os.environ["RANK"] = str(rank)
- os.environ["LOCAL_RANK"] = str(rank)
- os.environ["WORLD_SIZE"] = str(world_size)
- os.environ["MASTER_ADDR"] = "127.0.0.1"
- os.environ["MASTER_PORT"] = str(master_port)
-
-
def _cleanup_distributed() -> None:
if dist.is_initialized():
dist.destroy_process_group()
- for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
- os.environ.pop(key, None)
-
gc.collect()
if current_omni_platform.is_available():
current_omni_platform.empty_cache()
@@ -83,10 +84,7 @@ def _cleanup_distributed() -> None:
@pytest.fixture(scope="module")
def dist_group():
- master_port = _find_free_port()
- _set_dist_env(rank=0, world_size=1, master_port=master_port)
-
- dist.init_process_group("gloo", rank=0, world_size=1)
+ dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
try:
yield
finally:
@@ -139,6 +137,268 @@ class TestDistributedLayerwiseOffloadHook:
assert shard.numel() == 4
assert torch.equal(shard, _make_values(10.0))
+ def test_prefetch_preserves_transposed_weight_stride(
+ self,
+ dist_group,
+ patched_offload_runtime,
+ ):
+ """Online-FP8 Cutlass weights must retain their transposed layout."""
+
+ class StridedBlock(nn.Module):
+ def __init__(self, start: float):
+ super().__init__()
+ base = torch.arange(start, start + 12).reshape(3, 4)
+ self.weight = nn.Parameter(base.t(), requires_grad=False)
+
+ current_block = StridedBlock(1.0)
+ next_block = StridedBlock(20.0)
+ expected = next_block.weight.detach().clone()
+ expected_stride = next_block.weight.stride()
+ assert expected_stride == (1, 4)
+
+ hook = DistributedLayerwiseOffloadHook(
+ next_block=next_block,
+ device=torch.device("cpu"),
+ dp_group=None,
+ dp_size=1,
+ rank=0,
+ copy_stream=DummyStream(),
+ comm_stream=DummyStream(),
+ pin_memory=False,
+ )
+ hook.initialize_hook(current_block)
+ hook.prefetch_layer(slot=0, non_blocking=False)
+
+ assert next_block.weight.stride() == expected_stride
+ assert torch.equal(next_block.weight, expected)
+
+ def test_allgather_reconstructs_online_fp8_weight_and_scale(
+ self,
+ patched_offload_runtime,
+ monkeypatch,
+ ):
+ """DLO must reconstruct finalized online-FP8 tensors byte-for-byte."""
+
+ class OnlineFp8Block(nn.Module):
+ def __init__(self):
+ super().__init__()
+ logical_weight = torch.tensor(
+ [
+ [1.0, 2.0, 3.0, 4.0],
+ [5.0, 6.0, 7.0, 8.0],
+ [9.0, 10.0, 11.0, 12.0],
+ ],
+ dtype=torch.float32,
+ ).to(torch.float8_e4m3fn)
+ self.weight = nn.Parameter(logical_weight.t(), requires_grad=False)
+ self.weight_scale = nn.Parameter(torch.tensor([0.25]), requires_grad=False)
+
+ current_block = OnlineFp8Block()
+ rank0_block = OnlineFp8Block()
+ rank1_block = OnlineFp8Block()
+ expected_weight = rank0_block.weight.detach().clone()
+ expected_scale = rank0_block.weight_scale.detach().clone()
+ expected_stride = rank0_block.weight.stride()
+
+ rank0_hook = DistributedLayerwiseOffloadHook(
+ next_block=rank0_block,
+ device=torch.device("cpu"),
+ dp_group=object(),
+ dp_size=2,
+ rank=0,
+ copy_stream=DummyStream(),
+ comm_stream=DummyStream(),
+ pin_memory=False,
+ )
+ rank1_hook = DistributedLayerwiseOffloadHook(
+ next_block=rank1_block,
+ device=torch.device("cpu"),
+ dp_group=object(),
+ dp_size=2,
+ rank=1,
+ copy_stream=DummyStream(),
+ comm_stream=DummyStream(),
+ pin_memory=False,
+ )
+ rank0_hook.initialize_hook(current_block)
+ rank1_hook.initialize_hook(OnlineFp8Block())
+ rank0_hook.gpu_shard_buffers = [
+ {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+ ]
+
+ def fake_allgather(output, local_shard, *, group):
+ del group
+ remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+ shard_size = local_shard.numel()
+ output[:shard_size].copy_(local_shard)
+ output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+ monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+ rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+ assert rank0_block.weight.stride() == expected_stride
+ torch.testing.assert_close(
+ rank0_block.weight.float(),
+ expected_weight.float(),
+ rtol=0,
+ atol=0,
+ )
+ torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
+
+ def test_allgather_reconstructs_online_int8_weight_and_scale(
+ self,
+ patched_offload_runtime,
+ monkeypatch,
+ ):
+ """DLO must reconstruct finalized online-INT8 tensors byte-for-byte.
+
+ One block carries both finalized online-INT8 layouts: the NPU method
+ leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
+ per-channel scale, while the CUDA method leaves a transposed-view
+ weight (stride (1, K)) with a scalar fp32 scale. The int8 and fp32
+ dtype groups must survive the flat-shard AllGather independently.
+ """
+
+ class OnlineInt8Block(nn.Module):
+ def __init__(self):
+ super().__init__()
+ logical = torch.arange(1.0, 13.0).reshape(3, 4)
+ # NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
+ self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
+ self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
+ # CUDA online int8: Parameter(qweight.t().data) -> transposed
+ # view (stride (1, K)) over the (N, K) storage, scalar scale
+ self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
+ self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+
+ current_block = OnlineInt8Block()
+ rank0_block = OnlineInt8Block()
+ rank1_block = OnlineInt8Block()
+ expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+ expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+ def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+ return DistributedLayerwiseOffloadHook(
+ next_block=block,
+ device=torch.device("cpu"),
+ dp_group=object(),
+ dp_size=2,
+ rank=rank,
+ copy_stream=DummyStream(),
+ comm_stream=DummyStream(),
+ pin_memory=False,
+ )
+
+ rank0_hook = make_hook(rank0_block, 0)
+ rank1_hook = make_hook(rank1_block, 1)
+ rank0_hook.initialize_hook(current_block)
+ rank1_hook.initialize_hook(OnlineInt8Block())
+ rank0_hook.gpu_shard_buffers = [
+ {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+ ]
+
+ def fake_allgather(output, local_shard, *, group):
+ del group
+ remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+ shard_size = local_shard.numel()
+ output[:shard_size].copy_(local_shard)
+ output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+ monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+ rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+ for name, tensor in rank0_block.named_parameters():
+ assert tensor.stride() == expected_strides[name]
+ torch.testing.assert_close(
+ tensor.float(),
+ expected[name].float(),
+ rtol=0,
+ atol=0,
+ )
+
+ def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
+ self,
+ patched_offload_runtime,
+ monkeypatch,
+ ):
+ """DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.
+
+ One block carries both finalized online-MXFP8 layouts: the NPU method
+ leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
+ (K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
+ (XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
+ view over a contiguous (K/32, N) buffer. The fp8 and e8m0 dtype
+ groups must survive the flat-shard AllGather independently.
+ """
+
+ class OnlineMxfp8Block(nn.Module):
+ def __init__(self):
+ super().__init__()
+ logical = torch.arange(1.0, 13.0).reshape(3, 4)
+ # Build scales from raw e8m0 bit patterns (1 byte/elem).
+ npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
+ kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
+ # NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
+ # weight, contiguous (K_groups/2, N, 2) e8m0 scale
+ self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
+ self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
+ # vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
+ # scale kept as a transposed view (stride (1, K/32)) over the
+ # contiguous (K/32, N) storage
+ self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
+ self.kernel_weight_scale = nn.Parameter(
+ kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
+ )
+
+ current_block = OnlineMxfp8Block()
+ rank0_block = OnlineMxfp8Block()
+ rank1_block = OnlineMxfp8Block()
+ expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+ expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+ def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+ return DistributedLayerwiseOffloadHook(
+ next_block=block,
+ device=torch.device("cpu"),
+ dp_group=object(),
+ dp_size=2,
+ rank=rank,
+ copy_stream=DummyStream(),
+ comm_stream=DummyStream(),
+ pin_memory=False,
+ )
+
+ rank0_hook = make_hook(rank0_block, 0)
+ rank1_hook = make_hook(rank1_block, 1)
+ rank0_hook.initialize_hook(current_block)
+ rank1_hook.initialize_hook(OnlineMxfp8Block())
+ rank0_hook.gpu_shard_buffers = [
+ {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+ ]
+
+ def fake_allgather(output, local_shard, *, group):
+ del group
+ remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+ shard_size = local_shard.numel()
+ output[:shard_size].copy_(local_shard)
+ output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+ monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+ rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+ for name, tensor in rank0_block.named_parameters():
+ assert tensor.stride() == expected_strides[name]
+ # Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
+ # version-sensitive, while contiguous().view(uint8) is an exact
+ # byte-level comparison for both contiguous and transposed-view
+ # parameters.
+ torch.testing.assert_close(
+ tensor.detach().contiguous().view(torch.uint8),
+ expected[name].contiguous().view(torch.uint8),
+ rtol=0,
+ atol=0,
+ )
+
def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
"""DTensor wrapper should be preserved through prefetch/offload cycle."""
current_block = TinyBlock(_make_values(1.0))
@@ -316,8 +576,9 @@ class TestDistributedLayerwiseOffloadHook:
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
- next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
- next_block.weight.mmap_weight_transform_pending = True
+
+ def transform(tensor):
+ return tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -326,17 +587,17 @@ class TestDistributedLayerwiseOffloadHook:
dp_size=2,
rank=0,
pin_memory=False,
+ tensor_transforms={id(next_block.weight): transform},
)
hook.initialize_hook(block)
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
- def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
+ def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
block = nn.Module()
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
next_block = nn.Module()
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
- next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
hook = DistributedLayerwiseOffloadHook(
next_block=next_block,
@@ -350,6 +611,82 @@ class TestDistributedLayerwiseOffloadHook:
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
+ def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime):
+ current_block = nn.Linear(2, 2, bias=False)
+ next_block = nn.Linear(2, 2, bias=False)
+ next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
+ source_storage = next_block.weight.untyped_storage().data_ptr()
+
+ def transform(tensor):
+ return tensor.t()
+
+ hook = DistributedLayerwiseOffloadHook(
+ next_block=next_block,
+ device=torch.device("cpu"),
+ dp_group=None,
+ dp_size=1,
+ rank=0,
+ pin_memory=False,
+ rank_local_mmap=True,
+ tensor_transforms={id(next_block.weight): transform},
+ )
+ hook.initialize_hook(current_block)
+
+ source = hook.cpu_sources[torch.float32][0]["tensor"]
+ assert source.untyped_storage().data_ptr() == source_storage
+ assert hook.cpu_shards == {}
+ assert next_block.weight.numel() == 0
+
+ hook.cpu_staging_buffers = [
+ {torch.float32: torch.empty(4)},
+ {torch.float32: torch.empty(4)},
+ ]
+ hook.prefetch_layer(slot=0, non_blocking=False)
+
+ assert torch.equal(
+ next_block.weight,
+ torch.tensor([[0.0, 2.0], [1.0, 3.0]]),
+ )
+ assert next_block.weight.stride() == (1, 2)
+ assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))
+
+ def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
+ hooks = []
+ for size in (4, 9):
+ current_block = nn.Linear(1, 1, bias=False)
+ next_block = nn.Module()
+ next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32))
+ hook = DistributedLayerwiseOffloadHook(
+ next_block=next_block,
+ device=torch.device("cpu"),
+ dp_group=None,
+ dp_size=1,
+ rank=0,
+ pin_memory=False,
+ shared_buffers=[None, None],
+ rank_local_mmap=True,
+ )
+ hook.initialize_hook(current_block)
+ hooks.append(hook)
+
+ resident_block = nn.Module()
+ resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32))
+ resident_group = PinnedResidentLayerGroup(
+ [resident_block],
+ device=torch.device("cpu"),
+ copy_stream=DummyStream(),
+ pin_memory=False,
+ rank_local_mmap=True,
+ )
+ staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers(
+ hooks,
+ resident_group,
+ )
+
+ assert len(staging) == 2
+ assert staging[0][torch.float32].numel() == 12
+ assert staging[1][torch.float32].numel() == 12
+
class TestPinnedResidentLayerGroup:
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -386,23 +723,59 @@ class TestPinnedResidentLayerGroup:
group.offload()
- def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
- module = nn.Linear(2, 2, bias=False)
+ def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime):
+ block = nn.Module()
+ weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
+ expected = weight.detach().clone()
+ expected_stride = weight.stride()
+ block.weight = nn.Parameter(weight)
+ group = PinnedResidentLayerGroup(
+ [block],
+ device=torch.device("cpu"),
+ copy_stream=DummyStream(),
+ pin_memory=False,
+ )
- def transform(tensor):
- return tensor.flip(0)
+ group.load()
+ assert torch.equal(block.weight, expected)
+ assert block.weight.stride() == expected_stride
+ group.offload()
- module.weight.mmap_weight_transform = transform
+ def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
+ blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
+ expected = []
+ source_storage = []
+ for index, block in enumerate(blocks):
+ block.weight.data.fill_(index + 1)
+ expected.append(block.weight.detach().clone())
+ source_storage.append(block.weight.untyped_storage().data_ptr())
- backend = object.__new__(DistributedLayerwiseOffloadBackend)
- backend._remember_mmap_param_attrs(module)
- module.to_empty(device="meta")
+ group = PinnedResidentLayerGroup(
+ blocks,
+ device=torch.device("cpu"),
+ copy_stream=DummyStream(),
+ pin_memory=False,
+ rank_local_mmap=True,
+ )
- replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
- backend._attach_mmap_param_attrs("weight", replacement, module.weight)
+ for index, state in enumerate(group._states):
+ source = state["cpu_sources"][torch.float32][0]["tensor"]
+ assert source.untyped_storage().data_ptr() == source_storage[index]
+ assert state["cpu_shards"] == {}
+ assert len(group._cpu_staging_buffers) == 2
- assert replacement.mmap_weight_transform is transform
- assert replacement.mmap_weight_transform_pending is True
+ group.load()
+ for block, weight in zip(blocks, expected):
+ assert torch.equal(block.weight, weight)
+
+ blocks[0].weight.data.zero_()
+ group.offload()
+ group.load()
+ for block, weight in zip(blocks, expected):
+ assert torch.equal(block.weight, weight)
+
+ group.offload()
+ assert all(block.weight.numel() == 0 for block in blocks)
def test_shared_allgather_output_is_narrowed_to_current_block(
self, dist_group, patched_offload_runtime, monkeypatch
@@ -484,6 +857,271 @@ class _SingleBlockModel(nn.Module):
self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+class _HWRPipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = _SingleBlockModel(num_blocks=3)
+
+
+class _FakeHostWeightLease:
+ def __init__(self, resolution_id: str, events: list[str] | None = None):
+ self.closed = False
+ self.provenance = SimpleNamespace(resolution_id=resolution_id)
+ self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),)
+ self._events = events
+
+ def close(self):
+ if self._events is not None:
+ self._events.append("lease")
+ self.closed = True
+
+
+class _FakeLeaseCarrier:
+ def __init__(self, lease: _FakeHostWeightLease):
+ self.lease = lease
+ self.taken = False
+
+ @property
+ def closed(self):
+ return self.taken or self.lease.closed
+
+ def take(self):
+ if self.taken:
+ raise RuntimeError("already taken")
+ self.taken = True
+ return self.lease
+
+ def close(self):
+ if not self.taken:
+ self.lease.close()
+
+
+def _fake_hwr_plan(resolution_id: str = "hwr-test"):
+ lease = _FakeHostWeightLease(resolution_id)
+ carrier = _FakeLeaseCarrier(lease)
+ plan = HostWeightPlan(
+ backing_kind="host_weight_runtime",
+ bindings={},
+ lease_carrier=carrier,
+ )
+ return plan, carrier, lease
+
+
+def _hwr_backend():
+ plan, carrier, lease = _fake_hwr_plan()
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=False,
+ dp_size=1,
+ dlo_use_allgather=False,
+ ),
+ torch.device("cpu"),
+ host_weight_plan=plan,
+ )
+ return _HWRPipeline(), backend, carrier, lease
+
+
+def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime):
+ current_block = nn.Linear(2, 2, bias=False)
+ next_block = nn.Linear(2, 2, bias=False)
+ expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+ next_block.weight.data.copy_(expected)
+ hook = DistributedLayerwiseOffloadHook(
+ next_block=next_block,
+ device=torch.device("cpu"),
+ dp_group=None,
+ dp_size=1,
+ rank=0,
+ pin_memory=False,
+ rank_local_mmap=True,
+ )
+ hook.initialize_hook(current_block)
+ hook.registered_mmap = True
+ hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging") # type: ignore[method-assign]
+
+ hook.prefetch_layer(slot=0, non_blocking=True)
+
+ assert torch.equal(next_block.weight, expected)
+
+
+def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime):
+ block = nn.Linear(2, 2, bias=False)
+ expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+ block.weight.data.copy_(expected)
+ group = PinnedResidentLayerGroup(
+ [block],
+ device=torch.device("cpu"),
+ copy_stream=DummyStream(),
+ pin_memory=False,
+ rank_local_mmap=True,
+ )
+ group.registered_mmap = True
+ group._cpu_staging_buffers.clear()
+
+ group.load()
+
+ assert torch.equal(block.weight, expected)
+
+
+def test_hwr_registration_uses_transport_budget(
+ monkeypatch: pytest.MonkeyPatch,
+ patched_offload_runtime,
+):
+ dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+ plan, _, lease = _fake_hwr_plan()
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=True,
+ dp_size=1,
+ dlo_use_allgather=False,
+ dlo_host_registration_limit_gib=1.5,
+ ),
+ torch.device("cuda"),
+ host_weight_plan=plan,
+ )
+ backend._host_weight_lease = lease # type: ignore[assignment]
+ calls: list[tuple[object, int | None]] = []
+
+ class Registration:
+ total_bytes = 4096
+ region_count = 1
+
+ @staticmethod
+ def close() -> tuple[str, ...]:
+ return ()
+
+ def register(regions, *, device, max_bytes):
+ assert device == torch.device("cuda")
+ calls.append((regions, max_bytes))
+ return Registration()
+
+ monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+ pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True)
+
+ assert backend._try_register_hwr_mmap((pinned_source,)) # type: ignore[arg-type]
+ assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))]
+ assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+ backend._release_registered_mmap()
+
+ assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+
+
+def test_hwr_registration_budget_validation_is_transport_scoped():
+ assert (
+ validate_dlo_host_registration_options(
+ limit_gib=1.5,
+ enable_dlo=True,
+ use_allgather=False,
+ hwr_mode="preferred",
+ )
+ == 1.5
+ )
+ invalid = (
+ dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+ dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"),
+ dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"),
+ dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"),
+ dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"),
+ )
+ for options in invalid:
+ with pytest.raises(ValueError):
+ validate_dlo_host_registration_options(**options)
+
+
+def test_unregistration_precedes_lease_close_and_retries_failures(
+ monkeypatch: pytest.MonkeyPatch,
+ patched_offload_runtime,
+):
+ dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+ events: list[str] = []
+ responses = iter([("busy",), ()])
+
+ class Registration:
+ @staticmethod
+ def close() -> tuple[str, ...]:
+ events.append("unregister")
+ return next(responses)
+
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=True,
+ dp_size=1,
+ dlo_use_allgather=False,
+ ),
+ torch.device("cpu"),
+ )
+ lease = _FakeHostWeightLease("retry", events)
+ backend._host_weight_lease = lease # type: ignore[assignment]
+ backend._host_registration = Registration()
+ backend._using_rank_local_mmap = True
+ backend._using_registered_mmap = True
+ backend.enabled = True
+ monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+ with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+ backend.disable()
+ assert not lease.closed
+ assert backend._host_registration is not None
+ assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)]
+
+ backend.disable()
+
+ assert lease.closed
+ assert backend._host_registration is None
+ assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == []
+ assert events == ["unregister", "unregister", "lease"]
+
+
+def test_failed_unregistration_retains_lease_after_backend_is_dropped(
+ monkeypatch: pytest.MonkeyPatch,
+ patched_offload_runtime,
+):
+ events: list[str] = []
+
+ class Registration:
+ @staticmethod
+ def close() -> tuple[str, ...]:
+ events.append("unregister")
+ return ("busy",)
+
+ dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=True,
+ dp_size=1,
+ dlo_use_allgather=False,
+ ),
+ torch.device("cpu"),
+ )
+ lease = _FakeHostWeightLease("retained", events)
+ registration = Registration()
+ backend._host_weight_lease = lease # type: ignore[assignment]
+ backend._host_registration = registration
+ backend._using_rank_local_mmap = True
+ backend.enabled = True
+ monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+ lease_ref = weakref.ref(lease)
+
+ with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"):
+ backend.disable()
+ del backend
+ del lease
+ gc.collect()
+
+ retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1]
+ assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)]
+ assert lease_ref() is retained_lease
+ assert not retained_lease.closed
+
+ dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear()
+ retained_lease.close()
+
+
class _MultiBlockModel(nn.Module):
_layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
@@ -545,7 +1183,6 @@ class TestMmapWeightLoading:
OffloadConfig(
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
pin_cpu_memory=False,
- model_path=str(tmp_path),
),
torch.device("cpu"),
)
@@ -553,13 +1190,167 @@ class TestMmapWeightLoading:
dits=[pipeline.transformer],
dit_names=["transformer"],
)
+ plan = HostWeightPlan(
+ backing_kind="checkpoint_mmap",
+ bindings={
+ name: TensorBinding(
+ checkpoint_key=name,
+ file_path=str(weight_file),
+ )
+ for name in weights
+ },
+ )
- backend._load_weights_via_mmap(pipeline, modules)
+ backend._load_weights_via_mmap(pipeline, modules, plan)
assert pipeline.transformer.post_load_calls == 1
assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16
+ def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
+ self,
+ tmp_path,
+ patched_offload_runtime,
+ ):
+ class Transformer(nn.Module):
+ _layerwise_offload_blocks_attrs = ["blocks"]
+
+ def __init__(self):
+ super().__init__()
+ self.blocks = nn.ModuleList(
+ [
+ nn.Linear(2, 2, bias=False),
+ nn.Linear(2, 2, bias=False),
+ ]
+ )
+
+ class Pipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = Transformer()
+
+ @staticmethod
+ def _remap_ckpt_key(key):
+ return key
+
+ pipeline = Pipeline()
+ weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
+ checkpoint_file = tmp_path / "model.safetensors"
+ save_file(weights, str(checkpoint_file))
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path=str(tmp_path),
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+ assert result.plan is not None
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=False,
+ dp_size=1,
+ dlo_use_allgather=True,
+ ),
+ torch.device("cpu"),
+ host_weight_plan=result.plan,
+ )
+
+ backend.enable(pipeline)
+
+ assert backend._using_rank_local_mmap
+ assert backend.dp_group is None
+ assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
+ backend.disable()
+
+ def test_hwr_carrier_is_taken_and_released_after_bounded_staging(self, patched_offload_runtime):
+ """A warm HWR plan uses the existing two-slot rank-local transport."""
+ pipeline, backend, carrier, lease = _hwr_backend()
+ model_bytes = sum(param.numel() * param.element_size() for param in pipeline.transformer.parameters())
+
+ backend.enable(pipeline)
+
+ assert carrier.taken
+ assert backend._using_rank_local_mmap
+ assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+ staging_bytes = sum(
+ buffer.numel() * buffer.element_size()
+ for buffer in backend._all_hook_groups[0][0].cpu_staging_buffers[0].values()
+ )
+ assert staging_bytes * 2 < model_bytes
+ assert not lease.closed
+
+ backend.disable()
+ assert lease.closed
+
+ def test_hwr_registration_failure_falls_back_to_bounded_staging(
+ self,
+ monkeypatch,
+ patched_offload_runtime,
+ ):
+ plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback")
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=True,
+ dp_size=1,
+ dlo_use_allgather=False,
+ ),
+ torch.device("cpu"),
+ host_weight_plan=plan,
+ )
+ pipeline = _HWRPipeline()
+ original_staging_allocator = backend._allocate_shared_cpu_staging_buffers
+
+ def fail_registration(*args, **kwargs):
+ del args, kwargs
+ raise HostRegistrationError("registration unavailable in CPU test")
+
+ def allocate_unpinned_staging(hooks, resident_group=None):
+ # Registration selection still observes pin_cpu_memory=True. The
+ # CPU-only test avoids asking the host for CUDA-pinned allocations.
+ for hook in hooks:
+ hook.pin_memory = False
+ return original_staging_allocator(hooks, resident_group)
+
+ monkeypatch.setattr(dist_backend_module, "register_host_mappings", fail_registration)
+ monkeypatch.setattr(backend, "_allocate_shared_cpu_staging_buffers", allocate_unpinned_staging)
+
+ backend.enable(pipeline)
+
+ assert carrier.taken
+ assert backend._using_rank_local_mmap
+ assert not backend._using_registered_mmap
+ assert backend._host_registration is None
+ assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group)
+ assert not lease.closed
+
+ backend.disable()
+ assert lease.closed
+
+ def test_hwr_backend_failure_drains_partial_setup_before_lease_close(
+ self,
+ monkeypatch,
+ patched_offload_runtime,
+ ):
+ pipeline, backend, carrier, lease = _hwr_backend()
+ monkeypatch.setattr(
+ backend,
+ "_allocate_shared_buffers",
+ lambda hooks: (_ for _ in ()).throw(RuntimeError("device buffer allocation failed")),
+ )
+
+ with pytest.raises(RuntimeError, match="device buffer allocation failed"):
+ backend.enable(pipeline)
+
+ assert carrier.taken
+ assert lease.closed
+ assert not backend.enabled
+ assert not backend._blocks
+ assert not backend._all_hook_groups
+
class TestGetBlocksFromDit:
def test_get_blocks_from_dit_single_block_attr(self):
@@ -825,23 +1616,123 @@ class TestOffloadPlan:
class TestMmapValidation:
- """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
+ """Tests for loader preflight and backend plan realization."""
- def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
- """DLO+AllGather should reject when TP world size > 1."""
- import json
- from types import SimpleNamespace
-
- from safetensors.torch import save_file
-
- from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-
- # Mock TP world size = 2
- monkeypatch.setattr(
- "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
- lambda: 2,
+ def test_component_source_prefix_builds_model_namespace(self, tmp_path):
+ source_root = tmp_path / "partition"
+ checkpoint_dir = source_root / "transformer"
+ checkpoint_dir.mkdir(parents=True)
+ checkpoint_file = checkpoint_dir / "model.safetensors"
+ save_file(
+ {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
+ str(checkpoint_file),
)
+ class Pipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = nn.Module()
+ self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
+ self.weights_sources = [
+ SimpleNamespace(
+ model_or_path=str(source_root),
+ subfolder="transformer",
+ revision=None,
+ prefix="transformer.",
+ )
+ ]
+
+ pipeline = Pipeline()
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=tuple(pipeline.weights_sources),
+ model_path=None,
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.fallback_reason is None
+ assert result.plan is not None
+ assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
+ checkpoint_key="blocks.0.weight",
+ file_path=str(checkpoint_file),
+ )
+ assert result.plan.planned_source_prefixes == frozenset({"transformer."})
+
+ def test_non_dedicated_component_source_falls_back_before_discovery(self):
+ pipeline = nn.Module()
+ pipeline.transformer = nn.Linear(2, 2, bias=False)
+ sources = (
+ SimpleNamespace(
+ model_or_path="unused",
+ subfolder=None,
+ revision=None,
+ prefix="",
+ ),
+ )
+
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=sources,
+ model_path=None,
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.plan is None
+ assert "dedicated component weight source" in result.fallback_reason
+
+ def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
+ source_root = tmp_path / "partition"
+ checkpoint_dir = source_root / "transformer"
+ checkpoint_dir.mkdir(parents=True)
+ save_file(
+ {"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)},
+ str(checkpoint_dir / "model.safetensors"),
+ )
+
+ class Transformer(nn.Module):
+ _layerwise_offload_blocks_attrs = ["blocks"]
+
+ def __init__(self):
+ super().__init__()
+ self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])
+
+ class Pipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = Transformer()
+ self.weights_sources = [
+ SimpleNamespace(
+ model_or_path=str(source_root),
+ subfolder="transformer",
+ revision=None,
+ prefix="transformer.",
+ )
+ ]
+
+ pipeline = Pipeline()
+ original_weight = pipeline.transformer.blocks[0].weight
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=tuple(pipeline.weights_sources),
+ model_path=None,
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.plan is None
+ assert "1 required DiT tensors" in result.fallback_reason
+ assert pipeline.transformer.blocks[0].weight is original_weight
+ assert not original_weight.is_meta
+
+ def test_tp_falls_back_before_checkpoint_discovery(self):
class SimpleModel(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -859,41 +1750,118 @@ class TestMmapValidation:
return key
pipeline = SimplePipeline()
- weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
- save_file(weights, str(tmp_path / "model.safetensors"))
- (tmp_path / "model.safetensors.index.json").write_text(
- json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path="unused",
+ tensor_parallel_size=2,
+ use_hsdp=False,
+ online_quantization=False,
)
- backend = DistributedLayerwiseOffloadBackend(
- OffloadConfig(
- strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
- pin_cpu_memory=False,
- model_path=str(tmp_path),
- ),
- torch.device("cpu"),
- )
- modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+ assert result.plan is None
+ assert result.fallback_reason == "TP=2 requires the ordinary loader"
- with pytest.raises(ValueError, match="Tensor Parallel"):
- backend._load_weights_via_mmap(pipeline, modules)
+ @pytest.mark.parametrize(
+ ("use_hsdp", "online_quantization", "expected_reason"),
+ [
+ (True, False, "HSDP requires the ordinary loader"),
+ (False, True, "online quantization requires the ordinary loader"),
+ ],
+ )
+ def test_runtime_layouts_that_need_loader_callbacks_fall_back(
+ self,
+ use_hsdp,
+ online_quantization,
+ expected_reason,
+ ):
+ pipeline = nn.Module()
+ pipeline.transformer = nn.Linear(2, 2, bias=False)
- def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch):
- """DLO+AllGather should work when TP world size = 1, even if params
- have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1)."""
- import json
- from types import SimpleNamespace
-
- from safetensors.torch import save_file
-
- from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-
- # Mock TP world size = 1 (default)
- monkeypatch.setattr(
- "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
- lambda: 1,
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path="unused",
+ tensor_parallel_size=1,
+ use_hsdp=use_hsdp,
+ online_quantization=online_quantization,
)
+ assert result.plan is None
+ assert result.fallback_reason == expected_reason
+
+ @pytest.mark.parametrize(
+ ("checkpoint_tensor", "expected_reason"),
+ [
+ (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
+ (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
+ ],
+ )
+ def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
+ checkpoint_file = tmp_path / "model.safetensors"
+ save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))
+
+ class Pipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = nn.Linear(2, 2, bias=False)
+
+ @staticmethod
+ def _remap_ckpt_key(key):
+ return key
+
+ pipeline = Pipeline()
+
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path=str(tmp_path),
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.plan is None
+ assert expected_reason in result.fallback_reason
+
+ def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
+ class Transformer(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(2, 2))
+ self.register_buffer("derived", torch.ones(2), persistent=False)
+
+ checkpoint_file = tmp_path / "model.safetensors"
+ save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))
+
+ class Pipeline(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.transformer = Transformer()
+
+ @staticmethod
+ def _remap_ckpt_key(key):
+ return key
+
+ pipeline = Pipeline()
+
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path=str(tmp_path),
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.plan is not None
+ assert set(result.plan.bindings) == {"transformer.weight"}
+
+ def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
def custom_weight_loader(param, weight):
param.data.copy_(weight)
@@ -923,27 +1891,67 @@ class TestMmapValidation:
json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
)
- backend = DistributedLayerwiseOffloadBackend(
- OffloadConfig(
- strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
- pin_cpu_memory=False,
- model_path=str(tmp_path),
- ),
- torch.device("cpu"),
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path=str(tmp_path),
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
)
- modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
- # Should NOT raise — TP=1 is allowed even with custom weight_loader
- backend._load_weights_via_mmap(pipeline, modules)
+ assert result.plan is None
+ assert "custom weight_loader" in result.fallback_reason
+
+ def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
+ from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+ Cosmos3OmniDiffusersPipeline,
+ )
+ from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+ Cosmos3VFMTransformer,
+ )
+
+ transformer = object.__new__(Cosmos3VFMTransformer)
+ nn.Module.__init__(transformer)
+ transformer.proj_in = nn.Linear(2, 2, bias=False)
+ transformer.proj_in.weight.weight_loader = lambda *_args: None
+ pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+ nn.Module.__init__(pipeline)
+ pipeline.transformer = transformer
+
+ source_root = tmp_path / "partition"
+ checkpoint_dir = source_root / "transformer"
+ checkpoint_dir.mkdir(parents=True)
+ checkpoint_file = checkpoint_dir / "model.safetensors"
+ save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
+ sources = (
+ SimpleNamespace(
+ model_or_path=str(source_root),
+ subfolder="transformer",
+ revision=None,
+ prefix="transformer.",
+ ),
+ )
+
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", transformer),),
+ sources=sources,
+ model_path=None,
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
+
+ assert result.plan is not None
+ assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
+ checkpoint_key="proj_in.weight",
+ file_path=str(checkpoint_file),
+ )
def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
"""validate_loaded_weights should be called after mmap loading."""
- import json
- from types import SimpleNamespace
-
- from safetensors.torch import save_file
-
- from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
class ModelWithValidate(nn.Module):
_layerwise_offload_blocks_attrs = ["blocks"]
@@ -985,19 +1993,153 @@ class TestMmapValidation:
backend = DistributedLayerwiseOffloadBackend(
OffloadConfig(
- strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path)
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=False,
),
torch.device("cpu"),
)
modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+ result = build_checkpoint_mmap_plan(
+ pipeline,
+ dit_modules=(("transformer", pipeline.transformer),),
+ sources=(),
+ model_path=str(tmp_path),
+ tensor_parallel_size=1,
+ use_hsdp=False,
+ online_quantization=False,
+ )
- backend._load_weights_via_mmap(pipeline, modules)
+ assert result.plan is not None
+ backend._load_weights_via_mmap(pipeline, modules, result.plan)
assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"
class TestConfigValidation:
"""Tests for configuration validation in OffloadConfig / execute_request."""
+ @pytest.mark.parametrize(
+ ("dp", "sp", "tp", "use_allgather", "expected_dlo_group"),
+ [
+ (1, 1, 1, True, 1),
+ (2, 1, 1, True, 2),
+ (1, 4, 1, True, 4),
+ (2, 2, 2, True, 2),
+ (2, 2, 2, False, 1),
+ (1, 4, 2, False, 1),
+ ],
+ )
+ def test_dlo_group_selection_covers_dp_sp_tp_matrix(
+ self,
+ dp,
+ sp,
+ tp,
+ use_allgather,
+ expected_dlo_group,
+ ):
+ """DLO sharding follows DP, falls back to SP, and never uses TP."""
+
+ config = OffloadConfig.from_od_config(
+ SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=use_allgather,
+ dlo_resident_layers=0,
+ pin_cpu_memory=False,
+ parallel_config=SimpleNamespace(
+ data_parallel_size=dp,
+ use_hsdp=False,
+ hsdp_shard_size=-1,
+ hsdp_replicate_size=1,
+ sequence_parallel_size=sp,
+ tensor_parallel_size=tp,
+ ),
+ )
+ )
+ assert config.dp_size == expected_dlo_group
+ assert config.dlo_use_allgather is use_allgather
+
+ def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
+ import vllm_omni.diffusion.offloader as offloader_module
+
+ od_config = SimpleNamespace(
+ enable_cpu_offload=False,
+ enable_layerwise_offload=False,
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ dlo_resident_layers=0,
+ pin_cpu_memory=False,
+ model="unused",
+ parallel_config=SimpleNamespace(
+ use_hsdp=False,
+ data_parallel_size=1,
+ sequence_parallel_size=1,
+ ),
+ )
+ plan = HostWeightPlan(
+ backing_kind="checkpoint_mmap",
+ bindings={},
+ )
+ monkeypatch.setattr(
+ offloader_module.current_omni_platform,
+ "supports_cpu_offload",
+ lambda: False,
+ )
+
+ with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
+ offloader_module.get_offload_backend(
+ od_config,
+ device=torch.device("cpu"),
+ host_weight_plan=plan,
+ )
+
+ def test_startup_recovery_stays_inside_offloader_boundary(self, monkeypatch):
+ import vllm_omni.diffusion.offloader as offloader_module
+
+ class FailingBackend:
+ def __init__(self):
+ self.disabled = False
+
+ def enable(self, pipeline):
+ del pipeline
+ raise RuntimeError("initial prefetch failed")
+
+ def disable(self):
+ self.disabled = True
+
+ warm_pipeline = nn.Module()
+ canonical_pipeline = nn.Module()
+ plan, carrier, _lease = _fake_hwr_plan("hwr-recovery")
+ failing_backend = FailingBackend()
+ calls = []
+
+ def fake_backend(od_config, device, host_weight_plan):
+ del od_config, device
+ calls.append(host_weight_plan)
+ return failing_backend if host_weight_plan is plan else None
+
+ attach_offload_startup_state(
+ warm_pipeline,
+ OffloadStartupState(
+ host_weight_plan=plan,
+ fresh_model_loader=lambda: canonical_pipeline,
+ allow_fresh_retry=True,
+ ),
+ )
+ monkeypatch.setattr(offloader_module, "get_offload_backend", fake_backend)
+
+ recovered, backend = offloader_module.enable_offload_backend(
+ SimpleNamespace(),
+ warm_pipeline,
+ device=torch.device("cpu"),
+ )
+
+ assert recovered is canonical_pipeline
+ assert backend is None
+ assert failing_backend.disabled
+ assert carrier.closed
+ assert calls == [plan, None]
+
def test_hsdp_with_allgather_rejected(self):
"""HSDP + DLO + AllGather should raise ValueError (double sharding)."""
from vllm_omni.diffusion.offloader.base import OffloadConfig
@@ -1195,8 +2337,9 @@ class TestDynamicSlotTracking:
prev_idx = (i - 1) % num_blocks
# Dynamic slot tracking: read from prev's prefetched slot
- if prefetched_slots[prev_idx] is not None:
- current_slots[i] = prefetched_slots[prev_idx]
+ prefetched_slot = prefetched_slots[prev_idx]
+ if prefetched_slot is not None:
+ current_slots[i] = prefetched_slot
read_slot = current_slots[i]
diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py
index 5fd80e75c..9253f6d03 100644
--- a/tests/diffusion/offloader/test_layerwise_backend.py
+++ b/tests/diffusion/offloader/test_layerwise_backend.py
@@ -1,11 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
import gc
-import os
-import socket
from contextlib import contextmanager
import pytest
@@ -15,6 +13,7 @@ from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
+from tests.helpers.runtime import get_distributed_init_method
from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
from vllm_omni.platforms import current_omni_platform
@@ -39,27 +38,10 @@ def dummy_stream(_stream):
yield None
-def _find_free_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- return int(s.getsockname()[1])
-
-
-def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
- os.environ["RANK"] = str(rank)
- os.environ["LOCAL_RANK"] = str(rank)
- os.environ["WORLD_SIZE"] = str(world_size)
- os.environ["MASTER_ADDR"] = "127.0.0.1"
- os.environ["MASTER_PORT"] = str(master_port)
-
-
def _cleanup_distributed() -> None:
if dist.is_initialized():
dist.destroy_process_group()
- for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
- os.environ.pop(key, None)
-
gc.collect()
if current_omni_platform.is_available():
current_omni_platform.empty_cache()
@@ -68,10 +50,7 @@ def _cleanup_distributed() -> None:
@pytest.fixture(scope="module")
def dist_group():
- master_port = _find_free_port()
- _set_dist_env(rank=0, world_size=1, master_port=master_port)
-
- dist.init_process_group("gloo", rank=0, world_size=1)
+ dist.init_process_group("gloo", rank=0, world_size=1, init_method=get_distributed_init_method())
try:
yield
finally:
@@ -79,11 +58,11 @@ def dist_group():
@pytest.fixture
-def patched_offload_runtime(mocker):
- mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
- mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
- mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
- mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+def patched_offload_runtime(monkeypatch):
+ monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+ monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+ monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+ monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
class TinyBlock(nn.Module):
@@ -128,6 +107,33 @@ class TestLayerwiseOffloadHook:
assert current_block.weight.to_local().shape == torch.Size([4])
assert not hook.is_materialized
+ def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
+ """Online-FP8 Cutlass weights must retain their transposed layout."""
+
+ class StridedBlock(nn.Module):
+ def __init__(self, start: float):
+ super().__init__()
+ base = torch.arange(start, start + 12).reshape(3, 4)
+ self.weight = nn.Parameter(base.t(), requires_grad=False)
+
+ current_block = StridedBlock(1.0)
+ next_block = StridedBlock(20.0)
+ expected = next_block.weight.detach().clone()
+ expected_stride = next_block.weight.stride()
+ assert expected_stride == (1, 4)
+
+ hook = LayerwiseOffloadHook(
+ next_block=next_block,
+ device=torch.device("cpu"),
+ stream=DummyStream(),
+ pin_memory=False,
+ )
+ hook.initialize_hook(current_block)
+ hook.prefetch_layer(non_blocking=False)
+
+ assert next_block.weight.stride() == expected_stride
+ assert torch.equal(next_block.weight, expected)
+
class _DummyBlock(nn.Module):
def __init__(self):
diff --git a/tests/diffusion/offloader/test_module_residency.py b/tests/diffusion/offloader/test_module_residency.py
new file mode 100644
index 000000000..adf3c3afc
--- /dev/null
+++ b/tests/diffusion/offloader/test_module_residency.py
@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+from torch import nn
+
+import vllm_omni.diffusion.offloader.module_residency as residency_module
+from vllm_omni.diffusion.offloader.module_residency import (
+ BoundedAllocatorCache,
+ PinnedModuleStager,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyStream:
+ def wait_stream(self, _stream) -> None:
+ return None
+
+ def wait_event(self, _event) -> None:
+ return None
+
+
+class _DummyEvent:
+ def __init__(self) -> None:
+ self.record_count = 0
+
+ def record(self, _stream) -> None:
+ self.record_count += 1
+
+
+@contextmanager
+def _dummy_stream(_stream):
+ yield
+
+
+@pytest.fixture
+def patched_runtime(monkeypatch, mocker):
+ empty_cache = mocker.Mock()
+ synchronize = mocker.Mock()
+ monkeypatch.setattr(residency_module.current_omni_platform, "Stream", _DummyStream)
+ monkeypatch.setattr(residency_module.current_omni_platform, "Event", _DummyEvent)
+ monkeypatch.setattr(residency_module.current_omni_platform, "current_stream", _DummyStream)
+ monkeypatch.setattr(residency_module.current_omni_platform, "stream", _dummy_stream)
+ monkeypatch.setattr(residency_module.current_omni_platform, "empty_cache", empty_cache)
+ monkeypatch.setattr(residency_module.current_omni_platform, "synchronize", synchronize)
+ return empty_cache, synchronize
+
+
+def _aliased_modules() -> tuple[nn.Module, nn.Module]:
+ parameter_storage = torch.arange(12, dtype=torch.float32)
+ buffer_storage = torch.arange(8, dtype=torch.int64)
+ left = nn.Module()
+ right = nn.Module()
+ left.weight = nn.Parameter(parameter_storage[:8].view(2, 4))
+ right.weight = nn.Parameter(parameter_storage[4:].view(2, 4))
+ left.register_buffer("state", buffer_storage[:6])
+ right.register_buffer("state", buffer_storage[2:])
+ return left, right
+
+
+def test_module_group_preserves_parameter_and_buffer_storage_aliases(patched_runtime, mocker):
+ left, right = _aliased_modules()
+ expected_left_weight = left.weight.detach().clone()
+ expected_right_weight = right.weight.detach().clone()
+ expected_left_state = left.state.clone()
+ expected_right_state = right.state.clone()
+ retention = mocker.Mock(spec=BoundedAllocatorCache)
+ stager = PinnedModuleStager(
+ [left, right],
+ torch.device("cpu"),
+ pin_memory=False,
+ cache_retention=retention,
+ )
+
+ master_weight_ptr = left.weight.untyped_storage().data_ptr()
+ master_state_ptr = left.state.untyped_storage().data_ptr()
+ assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+ assert right.state.untyped_storage().data_ptr() == master_state_ptr
+ assert (left.weight.storage_offset(), right.weight.storage_offset()) == (0, 4)
+ assert (left.state.storage_offset(), right.state.storage_offset()) == (0, 2)
+
+ stager.load()
+
+ assert left.weight.untyped_storage().data_ptr() != master_weight_ptr
+ assert right.weight.untyped_storage().data_ptr() == left.weight.untyped_storage().data_ptr()
+ assert right.state.untyped_storage().data_ptr() == left.state.untyped_storage().data_ptr()
+ left.weight.data[1, 0] = -1
+ left.state[2] = -1
+ assert right.weight[0, 0].item() == -1
+ assert right.state[0].item() == -1
+
+ stager.offload()
+
+ assert left.weight.untyped_storage().data_ptr() == master_weight_ptr
+ assert right.weight.untyped_storage().data_ptr() == master_weight_ptr
+ assert left.state.untyped_storage().data_ptr() == master_state_ptr
+ assert right.state.untyped_storage().data_ptr() == master_state_ptr
+ assert torch.equal(left.weight, expected_left_weight)
+ assert torch.equal(right.weight, expected_right_weight)
+ assert torch.equal(left.state, expected_left_state)
+ assert torch.equal(right.state, expected_right_state)
+ retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_staging_transitions_are_idempotent_and_reuse_one_event(patched_runtime, mocker):
+ module = nn.Linear(2, 2)
+ retention = mocker.Mock(spec=BoundedAllocatorCache)
+ stager = PinnedModuleStager(
+ module,
+ torch.device("cpu"),
+ pin_memory=False,
+ cache_retention=retention,
+ )
+ ready_event = stager._ready_event
+
+ stager.load()
+ first_storage = module.weight.untyped_storage().data_ptr()
+ stager.load()
+ assert module.weight.untyped_storage().data_ptr() == first_storage
+ assert stager._ready_event is ready_event
+ assert ready_event.record_count == 1
+
+ stager.offload()
+ stager.offload()
+ assert not stager.loaded
+ retention.release_if_needed.assert_called_once_with(force=False)
+
+
+def test_failed_load_restores_master_and_forces_cache_release(monkeypatch, patched_runtime, mocker):
+ module = nn.Linear(2, 2)
+ expected = module.weight.detach().clone()
+ retention = mocker.Mock(spec=BoundedAllocatorCache)
+ stager = PinnedModuleStager(
+ module,
+ torch.device("cpu"),
+ pin_memory=False,
+ cache_retention=retention,
+ )
+ master_ptr = module.weight.untyped_storage().data_ptr()
+
+ def fail_after_rebind() -> None:
+ device_storages = [torch.empty_like(group.master) for group in stager._groups]
+ stager._bind(device_storages)
+ module.weight.data.zero_()
+ raise RuntimeError("copy failed")
+
+ monkeypatch.setattr(stager, "_load_once", fail_after_rebind)
+ with pytest.raises(RuntimeError, match="copy failed"):
+ stager.load()
+
+ assert not stager.loaded
+ assert module.weight.untyped_storage().data_ptr() == master_ptr
+ assert torch.equal(module.weight, expected)
+ retention.release_if_needed.assert_called_once_with(force=True)
+
+
+def test_out_of_memory_releases_cache_and_retries_once(monkeypatch, patched_runtime, mocker):
+ module = nn.Linear(2, 2)
+ retention = mocker.Mock(spec=BoundedAllocatorCache)
+ stager = PinnedModuleStager(
+ module,
+ torch.device("cpu"),
+ pin_memory=False,
+ cache_retention=retention,
+ )
+ load_once = stager._load_once
+ attempts = 0
+
+ def fail_once() -> None:
+ nonlocal attempts
+ attempts += 1
+ if attempts == 1:
+ raise torch.OutOfMemoryError("memory pressure")
+ load_once()
+
+ monkeypatch.setattr(stager, "_load_once", fail_once)
+ stager.load()
+
+ assert stager.loaded
+ assert attempts == 2
+ retention.release_if_needed.assert_called_once_with(force=True)
+
+
+@pytest.mark.parametrize(
+ ("reserved", "allocated", "free", "expected_release"),
+ [
+ (20, 10, 80, False),
+ (40, 10, 80, True),
+ (20, 10, 4, True),
+ ],
+)
+def test_allocator_cache_policy_enforces_cache_and_free_memory_bounds(
+ monkeypatch,
+ patched_runtime,
+ reserved,
+ allocated,
+ free,
+ expected_release,
+):
+ empty_cache, _ = patched_runtime
+ monkeypatch.setattr(torch.accelerator, "memory_reserved", lambda _device: reserved)
+ monkeypatch.setattr(torch.accelerator, "memory_allocated", lambda _device: allocated)
+ monkeypatch.setattr(
+ residency_module.current_omni_platform,
+ "get_device_memory",
+ lambda _device: (free, 100),
+ )
+ retention = BoundedAllocatorCache(torch.device("cpu"))
+
+ released = retention.release_if_needed()
+
+ assert released is expected_release
+ assert empty_cache.call_count == int(expected_release)
+
+
+def test_allocator_cache_releases_when_telemetry_is_unavailable(monkeypatch, patched_runtime):
+ empty_cache, _ = patched_runtime
+
+ def unavailable(_device):
+ raise NotImplementedError
+
+ monkeypatch.setattr(torch.accelerator, "memory_reserved", unavailable)
+ retention = BoundedAllocatorCache(torch.device("cpu"))
+
+ assert retention.release_if_needed()
+ empty_cache.assert_called_once_with()
diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py
index dc74248e5..ac3d3181d 100644
--- a/tests/diffusion/offloader/test_sequential_backend.py
+++ b/tests/diffusion/offloader/test_sequential_backend.py
@@ -8,7 +8,13 @@ import torch
from torch import nn
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
-from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
+from vllm_omni.diffusion.offloader.sequential_backend import (
+ ModelLevelOffloadBackend,
+ SequentialOffloadHook,
+ apply_sequential_offload,
+ remove_sequential_offload,
+ sequential_offload_component,
+)
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -76,6 +82,52 @@ def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
assert pipeline.disable_called is True
+def test_sequential_offload_can_begin_with_dit_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+ dit = _create_simple_module()
+ encoder = _create_simple_module()
+ offloaded: list[nn.Module] = []
+ monkeypatch.setattr(
+ SequentialOffloadHook,
+ "_to_cpu",
+ lambda self, module: offloaded.append(module),
+ )
+
+ apply_sequential_offload(
+ dit_modules=[dit],
+ encoder_modules=[encoder],
+ device=torch.device("cpu"),
+ offload_initial_dits=True,
+ )
+
+ assert offloaded == [dit]
+ remove_sequential_offload([dit, encoder])
+
+
+def test_direct_component_activation_failure_still_offloads(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ component = _create_simple_module()
+ apply_sequential_offload(
+ dit_modules=[_create_simple_module()],
+ encoder_modules=[component],
+ device=torch.device("cpu"),
+ )
+ hook = component._hook_registry.get_hook(SequentialOffloadHook._HOOK_NAME)
+ offloaded: list[nn.Module] = []
+
+ def fail_activation(module: nn.Module) -> None:
+ raise RuntimeError("activation failed")
+
+ monkeypatch.setattr(hook, "pre_forward", fail_activation)
+ monkeypatch.setattr(hook, "_to_cpu", lambda module: offloaded.append(module))
+
+ with pytest.raises(RuntimeError, match="activation failed"):
+ with sequential_offload_component(component):
+ pass
+
+ assert offloaded == [component]
+
+
class TestMoveParamsPinMemory:
def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
"""DTensor should skip pin_memory to avoid RuntimeError."""
@@ -85,7 +137,12 @@ class TestMoveParamsPinMemory:
original_isinstance = isinstance
def fake_isinstance(obj, cls):
- if cls.__name__ == "DTensor":
+ # torchada's Tensor.to() checks tuples such as
+ # ``(str, torch.device)``. Keep the DTensor probe compatible with
+ # the normal isinstance contract instead of breaking those calls.
+ if type(cls) is tuple:
+ return any(fake_isinstance(obj, candidate) for candidate in cls)
+ if getattr(cls, "__name__", None) == "DTensor":
return True
return original_isinstance(obj, cls)
diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py
index c8b3837e2..943411779 100644
--- a/tests/diffusion/quantization/test_component_routing.py
+++ b/tests/diffusion/quantization/test_component_routing.py
@@ -9,6 +9,9 @@ import torch
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
+ CompressedTensorsConfig,
+)
from vllm.model_executor.models.utils import WeightsMapper
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
@@ -16,10 +19,11 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
)
from vllm_omni.quantization.component_config import (
ComponentQuantizationConfig,
+ resolve_component_quant_config,
)
from vllm_omni.quantization.inc_config import OmniINCConfig
-pytestmark = [pytest.mark.core_model]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ---------------------------------------------------------------------------
@@ -343,10 +347,62 @@ class TestTalkerVisualRouting:
# ===================================================================
-# 4. ComponentQuantizationConfig.resolve
+# 4. ComponentQuantizationConfig.apply_vllm_mapper
# ===================================================================
+class TestComponentApplyVllmMapper:
+ def test_maps_compressed_tensors_component_and_default(self):
+ thinker = CompressedTensorsConfig(
+ target_scheme_map={},
+ ignore=["thinker.model.layers.0.mlp.gate"],
+ quant_format="pack-quantized",
+ )
+ default = CompressedTensorsConfig(
+ target_scheme_map={},
+ ignore=["talker.model.layers.0.mlp.gate"],
+ quant_format="pack-quantized",
+ )
+ config = ComponentQuantizationConfig(
+ component_configs={"thinker": thinker, "code2wav": None},
+ default_config=default,
+ )
+ mapper = WeightsMapper(
+ orig_to_new_prefix={
+ "thinker.model": "thinker.language_model.model",
+ "talker.model": "talker.language_model.model",
+ }
+ ).get_unstacked_mapper()
+
+ config.apply_vllm_mapper(mapper)
+
+ assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"]
+ assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"]
+
+
+# ===================================================================
+# 5. ComponentQuantizationConfig.resolve
+# ===================================================================
+
+
+class TestResolveComponentQuantConfig:
+ def test_global_config_applies_to_each_quantization_aware_component(self):
+ config = _MockQuantConfig("fp8")
+
+ assert resolve_component_quant_config(config, "transformer") is config
+ assert resolve_component_quant_config(config, "text_encoder") is config
+
+ def test_component_map_is_the_only_scope_override(self):
+ transformer = _MockQuantConfig("fp8")
+ config = ComponentQuantizationConfig(
+ component_configs={"transformer": transformer},
+ default_config=None,
+ )
+
+ assert resolve_component_quant_config(config, "transformer") is transformer
+ assert resolve_component_quant_config(config, "text_encoder") is None
+
+
class TestComponentResolve:
def test_longest_prefix_match(self):
a = _MockQuantConfig("a")
diff --git a/tests/diffusion/quantization/test_gguf_diffusion.py b/tests/diffusion/quantization/test_gguf_diffusion.py
index 243282b0c..ad2e72edc 100644
--- a/tests/diffusion/quantization/test_gguf_diffusion.py
+++ b/tests/diffusion/quantization/test_gguf_diffusion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""End-to-end tests for GGUF quantization on diffusion models.
Validates that GGUF-quantized diffusion models generate valid images and
@@ -19,8 +19,8 @@ import numpy as np
import pytest
import torch
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py
index 482235ff2..ce939dbf3 100644
--- a/tests/diffusion/quantization/test_hsdp_fp8.py
+++ b/tests/diffusion/quantization/test_hsdp_fp8.py
@@ -36,15 +36,23 @@ class _ToyKernel:
# --- helper to build a toy module simulating online-FP8 post-load state ---
-def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
+def _make_fp8_toy_module(
+ out_features: int = 16,
+ in_features: int = 32,
+ *,
+ quant_method_cls=None,
+):
"""Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``.
"""
from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+ if quant_method_cls is None:
+ quant_method_cls = Fp8LinearMethod
+
# Create the quant-method instance without calling __init__ (avoids
# heavyweight upstream config dependencies). isinstance() still works.
- qm = object.__new__(Fp8LinearMethod)
+ qm = object.__new__(quant_method_cls)
qm.fp8_linear = _ToyKernel()
# randn does not support float8 on CPU; zeros does and the test only
@@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight():
assert module.weight.shape == (16, 32)
+def test_rewrites_per_tensor_online_fp8_weight():
+ """Regression test for vLLM's new online FP8 frontend (#45463)."""
+ from vllm.model_executor.layers.quantization.online.fp8 import (
+ Fp8PerTensorOnlineLinearMethod,
+ )
+
+ module = _make_fp8_toy_module(
+ 16,
+ 32,
+ quant_method_cls=Fp8PerTensorOnlineLinearMethod,
+ )
+
+ n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+
+ assert n_rewritten_layers == 1
+ assert module.weight.is_contiguous()
+ assert module.weight.shape == (16, 32)
+ w, *_ = module.quant_method.fp8_linear._get_layer_params(module)
+ assert w.shape == (32, 16)
+ assert w.stride() == (1, 32)
+
+
def test_patched_get_layer_params_returns_transposed_view():
module = _make_fp8_toy_module(16, 32)
prepare_fp8_layers_for_fsdp(module)
diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py
index 617083398..0da5b6dd3 100644
--- a/tests/diffusion/quantization/test_quantization_fp8.py
+++ b/tests/diffusion/quantization/test_quantization_fp8.py
@@ -199,7 +199,7 @@ def _generate_bagel_image(
config_path = get_deploy_config_path("ci/bagel.yaml")
omni_kwargs: dict[str, Any] = {
"model": "ByteDance-Seed/BAGEL-7B-MoT",
- "stage_configs_path": config_path,
+ "deploy_config": config_path,
"stage_init_timeout": 300,
}
if diffusion_quantization_config:
diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py
index fdcef37b5..47908daa0 100644
--- a/tests/diffusion/quantization/test_quantization_quality.py
+++ b/tests/diffusion/quantization/test_quantization_quality.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Quantization quality gate for diffusion models.
@@ -53,6 +53,7 @@ _BENCH_MODULE_NAME = "benchmarks.diffusion.quantization_quality"
if _BENCH_MODULE_NAME not in sys.modules:
_spec = importlib.util.spec_from_file_location(_BENCH_MODULE_NAME, _BENCH_MODULE_PATH)
+ assert _spec is not None and _spec.loader is not None
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_BENCH_MODULE_NAME] = _mod
_spec.loader.exec_module(_mod)
@@ -80,7 +81,7 @@ class QualityTestConfig:
num_frames: int = 5 # only for t2v
seed: int = 42
gpu: str = "H100" # minimum GPU requirement
- negative_prompt: str = ""
+ negative_prompt: str | None = ""
guidance_scale: float | None = None
sigmas: list[float] | None = None
enable_cpu_offload: bool = False
@@ -93,7 +94,7 @@ class QualityTestConfig:
return self.quantized_model
return self.model or ""
- def quantization_ref(self) -> str | None:
+ def quantization_ref(self) -> str | dict[str, object] | None:
if self.quantized_model is not None:
return None
return self.quantization
@@ -211,24 +212,14 @@ QUALITY_CONFIGS = [
quantization="fp8",
task="t2v",
prompt="A serene lakeside sunrise with mist over the water",
- max_lpips=0.10,
- height=256,
- width=256,
- num_frames=25,
- num_inference_steps=8,
- # Preserve the final scheduler trajectory used when this FP8 quality
- # gate was established. Explicit LTX sigmas bypass dynamic shifting.
- sigmas=[
- 1.0,
- 0.92185378074646,
- 0.8327768445014954,
- 0.7303033471107483,
- 0.6111654043197632,
- 0.4709382653236389,
- 0.30347853899002075,
- 0.10000002384185791,
- 0.0,
- ],
+ max_lpips=0.20,
+ height=384,
+ width=512,
+ num_frames=73,
+ num_inference_steps=10,
+ # Do not override the recipe's negative conditioning, guidance, or
+ # scheduler trajectory: this gate follows the supported LTX default.
+ negative_prompt=None,
),
]
@@ -295,8 +286,11 @@ def _generate_image(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
+ prompt = {"prompt": config.prompt}
+ if config.negative_prompt is not None:
+ prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
- {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
+ prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -327,8 +321,11 @@ def _generate_video(omni, config: QualityTestConfig):
).manual_seed(config.seed)
torch.accelerator.reset_peak_memory_stats()
+ prompt = {"prompt": config.prompt}
+ if config.negative_prompt is not None:
+ prompt["negative_prompt"] = config.negative_prompt
outputs = omni.generate(
- {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
+ prompt,
OmniDiffusionSamplingParams(
height=config.height,
width=config.width,
@@ -464,6 +461,7 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
def test_generate_video_forwards_sigmas(monkeypatch):
+ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
@@ -474,7 +472,12 @@ def test_generate_video_forwards_sigmas(monkeypatch):
class DummyOmni:
def generate(self, _prompt, sampling):
captured.sampling = sampling
- return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])]
+ return [
+ OmniRequestOutput.from_diffusion(
+ request_id="req",
+ images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
+ )
+ ]
config = QualityTestConfig(
id="ltx-sigmas",
@@ -490,6 +493,41 @@ def test_generate_video_forwards_sigmas(monkeypatch):
assert captured.sampling.sigmas == [1.0, 0.5]
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ltx_quality_gate_uses_official_eager_defaults(monkeypatch):
+ from vllm_omni.outputs import OmniRequestOutput
+ from vllm_omni.platforms import current_omni_platform
+
+ monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
+ monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
+ monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
+ captured = SimpleNamespace(prompt=None, sampling=None)
+
+ class DummyOmni:
+ def generate(self, prompt, sampling):
+ captured.prompt = prompt
+ captured.sampling = sampling
+ return [
+ OmniRequestOutput.from_diffusion(
+ request_id="req",
+ images=[np.zeros((1, 2, 2, 3), dtype=np.float32)],
+ )
+ ]
+
+ config = next(config for config in QUALITY_CONFIGS if config.id == "fp8_ltx2")
+ _generate_video(DummyOmni(), config)
+
+ assert (config.width, config.height, config.num_frames, config.num_inference_steps) == (512, 384, 73, 10)
+ assert config.max_lpips == 0.20
+ assert config.sigmas is None
+ assert config.guidance_scale is None
+ assert config.negative_prompt is None
+ assert "negative_prompt" not in captured.prompt
+ assert captured.sampling.sigmas is None
+ assert captured.sampling.guidance_scale is None
+
+
_marks = hardware_marks(res={"cuda": "H100"})
_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
diff --git a/tests/diffusion/quantization/test_svdquant_config.py b/tests/diffusion/quantization/test_svdquant_config.py
new file mode 100644
index 000000000..adbb88ed0
--- /dev/null
+++ b/tests/diffusion/quantization/test_svdquant_config.py
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+ LinearBase,
+ UnquantizedLinearMethod,
+)
+
+from vllm_omni.quantization import svdquant_config
+from vllm_omni.quantization.factory import (
+ SUPPORTED_QUANTIZATION_METHODS,
+ build_quant_config,
+)
+from vllm_omni.quantization.svdquant_config import (
+ DiffusionSVDQuantConfig,
+ DiffusionSVDQuantLinearMethod,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_config_from_checkpoint_dict() -> None:
+ config = DiffusionSVDQuantConfig.from_config(
+ {
+ "rank": 32,
+ "precision": "nvfp4",
+ "act_unsigned": False,
+ "modules_to_not_convert": ["blocks.0.adaln_proj.linear"],
+ }
+ )
+
+ assert config.rank == 32
+ assert config.precision == "nvfp4"
+ assert config.modules_to_not_convert == ["blocks.0.adaln_proj.linear"]
+ assert config.get_supported_act_dtypes() == [torch.bfloat16]
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "message"),
+ [
+ ({"rank": 0}, "rank must be positive"),
+ ({"precision": "int4"}, "NVFP4 checkpoints only"),
+ ({"act_unsigned": True}, "unsigned activations"),
+ ],
+)
+def test_config_rejects_unsupported_phase1_options(
+ kwargs: dict[str, object],
+ message: str,
+) -> None:
+ with pytest.raises(ValueError, match=message):
+ DiffusionSVDQuantConfig(**kwargs)
+
+
+def test_factory_builds_svdquant_with_extra_checkpoint_fields() -> None:
+ config = build_quant_config(
+ {
+ "quant_method": "svdquant",
+ "rank": 16,
+ "precision": "nvfp4",
+ "w4a16_modules": ["adaln_proj.linear"],
+ "w4a16_group_size": 64,
+ }
+ )
+
+ assert isinstance(config, DiffusionSVDQuantConfig)
+ assert config.rank == 16
+ assert "svdquant" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_skipped_and_non_linear_modules_remain_unquantized() -> None:
+ config = DiffusionSVDQuantConfig(modules_to_not_convert=["blocks.0.adaln_proj.linear"])
+ linear = Mock(spec=LinearBase)
+
+ assert isinstance(
+ config.get_quant_method(linear, "blocks.0.adaln_proj.linear"),
+ UnquantizedLinearMethod,
+ )
+ assert config.get_quant_method(torch.nn.ReLU(), "blocks.0.act") is None
+
+
+def test_linear_method_creates_canonical_partitioned_weights(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+ method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig(rank=32))
+ layer = torch.nn.Module()
+
+ method.create_weights(
+ layer,
+ input_size_per_partition=128,
+ output_partition_sizes=[64, 64],
+ input_size=128,
+ output_size=128,
+ params_dtype=torch.bfloat16,
+ )
+
+ assert layer.qweight.shape == (128, 64)
+ assert layer.qweight.dtype == torch.int8
+ assert layer.wscales.shape == (8, 128)
+ assert layer.wscales.dtype == torch.float8_e4m3fn
+ assert layer.proj_down.shape == (128, 32)
+ assert layer.proj_up.shape == (128, 32)
+ assert layer.smooth_factor.shape == (128,)
+ assert layer.wcscales.shape == (128,)
+ assert layer.wtscale.shape == (1,)
+
+ assert layer.qweight.input_dim == 1
+ assert layer.qweight.output_dim == 0
+ assert layer.wscales.input_dim == 0
+ assert layer.wscales.output_dim == 1
+ assert layer.proj_down.input_dim == 0
+ assert layer.proj_up.output_dim == 0
+ assert layer.wcscales.output_dim == 0
+
+
+def test_linear_method_rejects_misaligned_input_partition(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+ method = DiffusionSVDQuantLinearMethod(DiffusionSVDQuantConfig())
+
+ with pytest.raises(ValueError, match="divisible by the block size 16"):
+ method.create_weights(
+ torch.nn.Module(),
+ input_size_per_partition=24,
+ output_partition_sizes=[32],
+ input_size=24,
+ output_size=32,
+ params_dtype=torch.bfloat16,
+ )
+
+
+def test_active_linear_uses_svdquant_method(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(svdquant_config, "_assert_supported", lambda: None)
+ config = DiffusionSVDQuantConfig()
+ linear = Mock(spec=LinearBase)
+
+ method = config.get_quant_method(linear, "blocks.0.attn.out_proj")
+
+ assert isinstance(method, DiffusionSVDQuantLinearMethod)
diff --git a/tests/diffusion/quantization/test_svdquant_linear.py b/tests/diffusion/quantization/test_svdquant_linear.py
new file mode 100644
index 000000000..bdca88d9d
--- /dev/null
+++ b/tests/diffusion/quantization/test_svdquant_linear.py
@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _FakeNvfp4Kernel:
+ def __init__(self, base_weight: torch.Tensor | None = None) -> None:
+ self.base_weight = base_weight
+ self.processed = False
+ self.last_input: torch.Tensor | None = None
+
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+ del layer
+ self.processed = True
+
+ def apply_weights(
+ self,
+ layer: torch.nn.Module,
+ x: torch.Tensor,
+ bias: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ del layer
+ assert bias is None
+ assert self.base_weight is not None
+ self.last_input = x.clone()
+ return torch.mm(x, self.base_weight)
+
+
+def _register_parameter(
+ layer: torch.nn.Module,
+ name: str,
+ value: torch.Tensor,
+) -> None:
+ layer.register_parameter(
+ name,
+ torch.nn.Parameter(value, requires_grad=False),
+ )
+
+
+def test_supports_only_validated_datacenter_blackwell() -> None:
+ assert not svdquant._supports_capability(SimpleNamespace(major=10, minor=0))
+ assert svdquant._supports_capability(SimpleNamespace(major=10, minor=3))
+ assert not svdquant._supports_capability(SimpleNamespace(major=11, minor=0))
+ assert not svdquant._supports_capability(SimpleNamespace(major=12, minor=0))
+ assert not svdquant._supports_capability(None)
+
+
+def test_rejects_nvfp4_backend_with_incompatible_layer_contract(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ svdquant._nvfp4_kernel.cache_clear()
+ monkeypatch.setattr(
+ svdquant,
+ "init_nvfp4_linear_kernel",
+ lambda: _FakeNvfp4Kernel(),
+ )
+
+ with pytest.raises(RuntimeError, match="incompatible with the SVDQuant checkpoint"):
+ svdquant._nvfp4_kernel()
+
+ svdquant._nvfp4_kernel.cache_clear()
+
+
+def test_prepare_weights_uses_existing_nvfp4_layout_contract(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ kernel = _FakeNvfp4Kernel()
+ monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+ monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+ method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+ layer = torch.nn.Module()
+ qweight = torch.arange(16, dtype=torch.int8).reshape(4, 4)
+ wscales = torch.arange(8, dtype=torch.float32).to(torch.float8_e4m3fn).reshape(2, 4)
+ _register_parameter(layer, "qweight", qweight)
+ _register_parameter(layer, "wscales", wscales)
+ _register_parameter(layer, "wtscale", torch.tensor([2.0], dtype=torch.bfloat16))
+ _register_parameter(
+ layer,
+ "wcscales",
+ torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16),
+ )
+
+ method.process_weights_after_loading(layer)
+
+ assert kernel.processed
+ assert not hasattr(layer, "qweight")
+ assert not hasattr(layer, "wscales")
+ assert layer.weight.dtype == torch.uint8
+ assert torch.equal(layer.weight.view(torch.int8), qweight)
+ assert torch.equal(layer.weight_scale, wscales.transpose(0, 1))
+ torch.testing.assert_close(layer.alpha, torch.tensor([2.0]))
+ assert layer.input_global_scale_inv.dtype == torch.float32
+ assert layer.output_channel_scale.shape == (4,)
+
+
+def test_prepare_weights_drops_identity_output_scale(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ kernel = _FakeNvfp4Kernel()
+ monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+ monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+ method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+ layer = torch.nn.Module()
+ _register_parameter(layer, "qweight", torch.zeros(4, 4, dtype=torch.int8))
+ _register_parameter(
+ layer,
+ "wscales",
+ torch.ones(2, 4, dtype=torch.float8_e4m3fn),
+ )
+ _register_parameter(layer, "wtscale", torch.ones(1, dtype=torch.bfloat16))
+ _register_parameter(layer, "wcscales", torch.ones(4, dtype=torch.bfloat16))
+
+ method.process_weights_after_loading(layer)
+
+ assert layer.output_channel_scale is None
+
+
+def test_compatibility_path_applies_qkv_scale_and_rank_correction(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ base_weight = torch.tensor(
+ [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+ dtype=torch.bfloat16,
+ )
+ kernel = _FakeNvfp4Kernel(base_weight)
+ monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+ monkeypatch.setattr(svdquant, "_nvfp4_kernel", lambda: kernel)
+ method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+
+ layer = torch.nn.Module()
+ layer.smooth_factor = torch.tensor([2.0, 4.0], dtype=torch.bfloat16)
+ layer.proj_down = torch.tensor([[1.0], [2.0]], dtype=torch.bfloat16)
+ layer.proj_up = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.bfloat16)
+ layer.output_channel_scale = torch.tensor([1.0, 2.0, 4.0], dtype=torch.bfloat16)
+ layer.output_size_per_partition = 3
+ x = torch.tensor([[2.0, 4.0], [4.0, 8.0]], dtype=torch.bfloat16)
+ bias = torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16)
+
+ actual = method.apply(layer, x, bias)
+
+ smoothed = x / layer.smooth_factor
+ expected = torch.mm(smoothed, base_weight)
+ expected = expected * layer.output_channel_scale
+ expected = expected + torch.mm(torch.mm(x, layer.proj_down), layer.proj_up.T)
+ expected = expected + bias
+ torch.testing.assert_close(actual, expected)
+ torch.testing.assert_close(kernel.last_input, smoothed)
+
+
+def test_compatibility_path_rejects_fp16_activations(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+ method = svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig())
+ layer = torch.nn.Module()
+ with pytest.raises(ValueError, match="requires BF16 activations"):
+ method.apply(
+ layer,
+ torch.zeros(1, 16, dtype=torch.float16),
+ )
diff --git a/tests/diffusion/quantization/test_svdquant_tp_loading.py b/tests/diffusion/quantization/test_svdquant_tp_loading.py
new file mode 100644
index 000000000..14138dc1c
--- /dev/null
+++ b/tests/diffusion/quantization/test_svdquant_tp_loading.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ QKVParallelLinear,
+ RowParallelLinear,
+)
+
+from vllm_omni.quantization import svdquant_config as svdquant
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _linear_without_init(linear_type: type[torch.nn.Module], tp_rank: int):
+ layer = object.__new__(linear_type)
+ torch.nn.Module.__init__(layer)
+ layer.tp_rank = tp_rank
+ return layer
+
+
+def _values(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor:
+ values = torch.arange(torch.Size(shape).numel(), dtype=torch.int64).remainder(127)
+ return values.reshape(shape).to(dtype)
+
+
+def _method(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(svdquant, "_assert_supported", lambda: None)
+ return svdquant.DiffusionSVDQuantLinearMethod(svdquant.DiffusionSVDQuantConfig(rank=8))
+
+
+def test_column_parallel_loader_shards_svdquant_outputs(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ layer = _linear_without_init(ColumnParallelLinear, tp_rank=1)
+ method = _method(monkeypatch)
+ method.create_weights(
+ layer,
+ input_size_per_partition=128,
+ output_partition_sizes=[32],
+ input_size=128,
+ output_size=64,
+ params_dtype=torch.bfloat16,
+ weight_loader=layer.weight_loader,
+ )
+
+ loaded_tensors = {
+ "qweight": _values((64, 64), torch.int8),
+ "wscales": _values((8, 64), torch.float8_e4m3fn),
+ "proj_up": _values((64, 8), torch.bfloat16),
+ "wcscales": _values((64,), torch.bfloat16),
+ }
+ for name, loaded in loaded_tensors.items():
+ parameter = getattr(layer, name)
+ parameter.weight_loader(parameter, loaded)
+ expected = loaded.narrow(parameter.output_dim, 32, 32)
+ torch.testing.assert_close(parameter, expected)
+
+
+def test_row_parallel_loader_shards_svdquant_inputs(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ layer = _linear_without_init(RowParallelLinear, tp_rank=1)
+ method = _method(monkeypatch)
+ method.create_weights(
+ layer,
+ input_size_per_partition=64,
+ output_partition_sizes=[64],
+ input_size=128,
+ output_size=64,
+ params_dtype=torch.bfloat16,
+ weight_loader=layer.weight_loader,
+ )
+
+ loaded_tensors = {
+ "qweight": _values((64, 64), torch.int8),
+ "wscales": _values((8, 64), torch.float8_e4m3fn),
+ "proj_down": _values((128, 8), torch.bfloat16),
+ "smooth_factor": _values((128,), torch.bfloat16),
+ }
+ for name, loaded in loaded_tensors.items():
+ parameter = getattr(layer, name)
+ parameter.weight_loader(parameter, loaded)
+ shard_size = parameter.shape[parameter.input_dim]
+ expected = loaded.narrow(parameter.input_dim, shard_size, shard_size)
+ torch.testing.assert_close(parameter, expected)
+
+
+def test_qkv_parallel_loader_shards_fused_svdquant_outputs(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ layer = _linear_without_init(QKVParallelLinear, tp_rank=1)
+ layer.total_num_heads = 4
+ layer.total_num_kv_heads = 2
+ layer.num_heads = 2
+ layer.num_kv_heads = 1
+ layer.num_kv_head_replicas = 1
+ layer.head_size = 8
+ layer.v_head_size = 8
+ method = _method(monkeypatch)
+ method.create_weights(
+ layer,
+ input_size_per_partition=128,
+ output_partition_sizes=[16, 8, 8],
+ input_size=128,
+ output_size=64,
+ params_dtype=torch.bfloat16,
+ weight_loader=layer.weight_loader,
+ )
+
+ loaded_tensors = {
+ "qweight": _values((64, 64), torch.int8),
+ "wscales": _values((8, 64), torch.float8_e4m3fn),
+ "proj_up": _values((64, 8), torch.bfloat16),
+ "wcscales": _values((64,), torch.bfloat16),
+ }
+ expected_global_outputs = [slice(16, 32), slice(40, 48), slice(56, 64)]
+ for name, loaded in loaded_tensors.items():
+ parameter = getattr(layer, name)
+ parameter.weight_loader(parameter, loaded)
+ expected = torch.cat(
+ [
+ loaded.narrow(parameter.output_dim, part.start, part.stop - part.start)
+ for part in expected_global_outputs
+ ],
+ dim=parameter.output_dim,
+ )
+ torch.testing.assert_close(parameter, expected)
diff --git a/tests/diffusion/quantization/test_wan_autoround_mxfp4.py b/tests/diffusion/quantization/test_wan_autoround_mxfp4.py
new file mode 100644
index 000000000..89c317970
--- /dev/null
+++ b/tests/diffusion/quantization/test_wan_autoround_mxfp4.py
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests."""
+
+from unittest.mock import patch
+
+import pytest
+from torch import nn
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.model_loader.utils import configure_quant_config
+
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
+ WanTransformer3DModel,
+)
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.inc_config import OmniINCConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _autoround_mxfp4_config() -> dict:
+ return {
+ "bits": 4,
+ "act_bits": 4,
+ "data_type": "mx_fp",
+ "act_data_type": "mx_fp",
+ "group_size": 32,
+ "act_group_size": 32,
+ "sym": True,
+ "act_sym": True,
+ "act_dynamic": True,
+ "enable_quanted_input": False,
+ "block_name_to_quantize": "blocks",
+ "quant_method": "auto-round",
+ "packing_format": "auto_round:llm_compressor",
+ "autoround_version": "0.14.2",
+ }
+
+
+def test_build_wan_autoround_mxfp4_config():
+ config = build_quant_config(_autoround_mxfp4_config())
+
+ assert isinstance(config, OmniINCConfig)
+ assert config.get_name() == "inc"
+ assert config.weight_bits == 4
+ assert config.data_type == "mx_fp"
+ assert config.group_size == 32
+ assert config.packing_format == "auto_round:llm_compressor"
+ assert config.block_name_to_quantize == ["blocks"]
+
+
+def test_wan_autoround_config_uses_runtime_layer_names():
+ config = build_quant_config(_autoround_mxfp4_config())
+ configure_quant_config(config, WanTransformer3DModel)
+
+ assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping
+ assert config.block_name_to_quantize == ["blocks"]
+
+ layer = object.__new__(LinearBase)
+ nn.Module.__init__(layer)
+ for prefix in (
+ "blocks.0.ffn.net_0.proj",
+ "blocks.0.ffn.net_2",
+ "blocks.0.attn1.to_out",
+ ):
+ layer_config = config.config_parser.resolve(layer, prefix)
+ assert layer_config.quantized
+ assert layer_config.is_mxfp4
+
+
+def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4():
+ config = build_quant_config(_autoround_mxfp4_config())
+ configure_quant_config(config, WanTransformer3DModel)
+
+ layer = object.__new__(LinearBase)
+ nn.Module.__init__(layer)
+ layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj")
+ outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj")
+
+ assert layer_config.quantized
+ assert layer_config.is_mxfp4
+ assert not outside_config.quantized
+
+ fake_method = object()
+ with patch(
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method",
+ return_value=fake_method,
+ ):
+ method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj")
+ assert method is fake_method
diff --git a/tests/diffusion/sched/test_dmd2_sigma_schedule.py b/tests/diffusion/sched/test_dmd2_sigma_schedule.py
index 108cbe68c..8fbad959b 100644
--- a/tests/diffusion/sched/test_dmd2_sigma_schedule.py
+++ b/tests/diffusion/sched/test_dmd2_sigma_schedule.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
from vllm_omni.diffusion.sched import DMD2SigmaSchedule
@@ -18,7 +18,8 @@ def test_step_count_is_intervals_not_boundaries():
[
[],
[1.0],
- [0.9, 0.5, 0.0],
+ [1.5, 0.5, 0.0],
+ [0.0, 0.5, 0.0],
[1.0, 0.5, 0.1],
[1.0, 0.4, 0.6, 0.0],
[1.0, 0.5, 0.5, 0.0],
@@ -31,6 +32,15 @@ def test_malformed_positions_are_rejected(positions):
DMD2SigmaSchedule.from_positions(positions)
+def test_a_capped_first_position_is_a_schedule_not_a_malformed_one():
+ # A DMD2 student trained with max_timestep_ratio < 1 never sees pure noise,
+ # so its opening rung sits just below 1.0 (FastH3: 999/1000).
+ schedule = DMD2SigmaSchedule.from_positions([0.999, 0.749, 0.5, 0.25, 0.0])
+
+ assert schedule.base_schedule[0] == 0.999
+ assert schedule.num_inference_steps == 4
+
+
def test_one_base_schedule_drives_several_shift_scales():
schedule = DMD2SigmaSchedule.from_positions([1.0, 0.7, 0.4, 0.15, 0.0])
diff --git a/tests/diffusion/test_async_output_timeout.py b/tests/diffusion/test_async_output_timeout.py
new file mode 100644
index 000000000..2875ed02f
--- /dev/null
+++ b/tests/diffusion/test_async_output_timeout.py
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the async-output wait bound (_async_output_timeout)."""
+
+import logging
+
+import pytest
+
+from vllm_omni.diffusion.diffusion_engine import (
+ _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+ _ASYNC_OUTPUT_TIMEOUT_ENV,
+ _async_output_timeout,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestAsyncOutputTimeoutIsConfigurable:
+ def test_default_is_generous(self, monkeypatch):
+ """The bound covers a copy queued behind a full denoise step, so the
+ default must not be tight enough to abort a slow but healthy render.
+ """
+ monkeypatch.delenv(_ASYNC_OUTPUT_TIMEOUT_ENV, raising=False)
+ assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT == 600.0
+
+ def test_env_var_overrides_the_default(self, monkeypatch):
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "1800")
+ assert _async_output_timeout() == 1800.0
+
+ def test_env_var_accepts_a_float(self, monkeypatch):
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "45.5")
+ assert _async_output_timeout() == 45.5
+
+ def test_value_is_not_frozen_at_import(self, monkeypatch):
+ """Resolved per call rather than captured in a module constant, so the
+ value never goes stale relative to os.environ and no module reload is
+ needed to exercise it.
+ """
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "60")
+ assert _async_output_timeout() == 60.0
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, "120")
+ assert _async_output_timeout() == 120.0
+
+
+class TestAsyncOutputTimeoutRejectsBadValues:
+ """This runs on the request path, so a typo in the environment must degrade
+ to the default rather than fail the generation.
+ """
+
+ @pytest.mark.parametrize("value", ["", "abc", "30s", "None"])
+ def test_non_numeric_falls_back_to_the_default(self, monkeypatch, caplog, value):
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+ with caplog.at_level(logging.WARNING):
+ assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+
+ @pytest.mark.parametrize("value", ["0", "-1", "-0.5"])
+ def test_non_positive_falls_back_to_the_default(self, monkeypatch, caplog, value):
+ monkeypatch.setenv(_ASYNC_OUTPUT_TIMEOUT_ENV, value)
+ with caplog.at_level(logging.WARNING):
+ assert _async_output_timeout() == _ASYNC_OUTPUT_TIMEOUT_DEFAULT
diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py
index da684a4f5..a4bff4625 100644
--- a/tests/diffusion/test_async_output_worker.py
+++ b/tests/diffusion/test_async_output_worker.py
@@ -3,6 +3,8 @@
"""Unit tests for WorkerProc async output logic."""
import queue
+import threading
+import time
from unittest.mock import MagicMock
import pytest
@@ -29,9 +31,59 @@ def _make_worker_proc(step_execution=False):
proc._running = True
proc._async_output_queue = queue.Queue() if not step_execution else None
proc._async_output_thread = None
+ proc._async_output_done = threading.Condition()
+ proc._async_output_pending = 0
+ proc._result_mq_lock = threading.Lock()
return proc
+class _OverlapDetectingQueue:
+ """Fails if two threads are inside enqueue() at the same time.
+
+ Mirrors vLLM's MessageQueue, whose acquire_write() reads current_idx,
+ writes that block, then advances it. Overlapping writers can target the
+ same ring block and silently drop a message.
+ """
+
+ def __init__(self):
+ self._active = 0
+ self._guard = threading.Lock()
+ self.overlaps = 0
+ self.messages = []
+
+ def enqueue(self, msg, timeout=None):
+ with self._guard:
+ self._active += 1
+ if self._active > 1:
+ self.overlaps += 1
+ time.sleep(0.001) # widen the window a real ring-buffer write would have
+ with self._guard:
+ self.messages.append(msg)
+ self._active -= 1
+
+
+class TestResultQueueSingleWriter:
+ """Worker must serialize writes to the single-writer result queue."""
+
+ def test_concurrent_enqueues_do_not_overlap(self):
+ proc = _make_worker_proc(step_execution=False)
+ detector = _OverlapDetectingQueue()
+ proc.result_mq = detector
+
+ def writer(tag):
+ for i in range(20):
+ proc._enqueue_result(f"{tag}-{i}")
+
+ threads = [threading.Thread(target=writer, args=(tag,)) for tag in ("compute", "output")]
+ for t in threads:
+ t.start()
+ for t in threads:
+ t.join()
+
+ assert detector.overlaps == 0, f"{detector.overlaps} overlapping writes to result_mq"
+ assert len(detector.messages) == 40
+
+
class TestGenerateAsyncOutputId:
"""Test _generate_async_output_id uniqueness."""
@@ -192,6 +244,87 @@ class TestAsyncOutputLoopLogic:
assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
+class TestDrainAsyncOutputs:
+ """Async outputs must finish packing before device memory is released."""
+
+ def test_drain_returns_immediately_when_nothing_in_flight(self):
+ proc = _make_worker_proc(step_execution=False)
+
+ assert proc.drain_async_outputs(timeout=0.01) is True
+
+ def test_return_result_keeps_output_pending_until_loop_completes(self, mocker):
+ """_return_result marks work pending; the bg loop clears it."""
+ proc = _make_worker_proc(step_execution=False)
+
+ mock_platform = mocker.MagicMock()
+ mock_platform.record_device_event.return_value = MagicMock()
+ mocker.patch(
+ "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
+ mock_platform,
+ )
+ mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm")
+ mock_torch = mocker.MagicMock()
+ mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
+ mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
+
+ proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
+
+ # Packing has not run yet, so a sleep here would free tensors in flight.
+ assert proc._async_output_pending == 1
+ assert proc.drain_async_outputs(timeout=0.05) is False
+
+ thread = threading.Thread(target=proc._async_output_loop, daemon=True)
+ thread.start()
+ try:
+ assert proc.drain_async_outputs(timeout=5.0) is True
+ finally:
+ proc._async_output_queue.put(None)
+ thread.join(timeout=2.0)
+
+ def test_sleep_rpc_drains_before_releasing_memory(self, mocker):
+ proc = _make_worker_proc(step_execution=False)
+ proc.worker = MagicMock()
+
+ call_order: list[str] = []
+ mocker.patch.object(
+ proc,
+ "drain_async_outputs",
+ side_effect=lambda *a, **kw: call_order.append("drain") or True,
+ )
+ proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append("sleep")
+
+ proc._execute_rpc(
+ {
+ "method": "handle_sleep_task",
+ "args": (),
+ "kwargs": {},
+ "output_rank": 0,
+ "exec_all_ranks": False,
+ "collect_rank_status": False,
+ }
+ )
+
+ assert call_order == ["drain", "sleep"]
+
+ def test_non_sleep_rpc_does_not_drain(self, mocker):
+ proc = _make_worker_proc(step_execution=False)
+ proc.worker = MagicMock()
+ drain = mocker.patch.object(proc, "drain_async_outputs", return_value=True)
+
+ proc._execute_rpc(
+ {
+ "method": "execute_model",
+ "args": (),
+ "kwargs": {},
+ "output_rank": 0,
+ "exec_all_ranks": False,
+ "collect_rank_status": False,
+ }
+ )
+
+ drain.assert_not_called()
+
+
class TestWorkerProcShutdown:
def test_shutdown_stops_async_thread_and_releases_queues(self):
import threading
diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py
index 665aa6825..3cf8f1428 100644
--- a/tests/diffusion/test_compile.py
+++ b/tests/diffusion/test_compile.py
@@ -6,6 +6,7 @@ import torch.nn as nn
import vllm_omni.diffusion.compile as compile_module
from vllm_omni.diffusion.compile import regionally_compile
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -76,6 +77,37 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
+def test_regionally_compile_keeps_hook_dispatch_outside_compiled_graph(monkeypatch):
+ model = _ModelWithWrappedRepeatedBlocks()
+ block = model.transformer_blocks[0]
+ registry = HookRegistry.get_or_create(block)
+ hook = ModelHook()
+ registry.register_hook("test", hook)
+
+ wrapped_forward = block.forward
+ original_forward = block._omni_original_forward
+ compile_calls = []
+
+ def _compile(fn, *args, **kwargs):
+ compile_calls.append(fn)
+
+ def _compiled(*fn_args, **fn_kwargs):
+ return f"compiled:{fn(*fn_args, **fn_kwargs)}"
+
+ return _compiled
+
+ monkeypatch.setattr(compile_module.torch, "compile", _compile)
+
+ regionally_compile(model)
+
+ assert compile_calls[0] is original_forward
+ assert compile_calls[0] is not wrapped_forward
+ assert block.forward is wrapped_forward
+ assert block._omni_original_forward is not original_forward
+ assert hook.fn_ref.original_forward is block._omni_original_forward
+ assert block("ok") == "compiled:ok"
+
+
def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch):
"""cache-dit matches blocks via inspect.signature(block.forward).
diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py
index e25d6497c..60d0c63b8 100644
--- a/tests/diffusion/test_diffusion_config_propagation.py
+++ b/tests/diffusion/test_diffusion_config_propagation.py
@@ -92,6 +92,30 @@ class TestParallelConfigPropagation:
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
assert od.num_gpus == 2
+ def test_dp_is_inferred_from_num_gpus(self):
+ pc = DiffusionParallelConfig(tensor_parallel_size=2, ulysses_degree=2)
+ od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=8)
+ assert od.parallel_config.data_parallel_size == 2
+ assert od.parallel_config.world_size == 8
+
+ def test_explicit_dp_is_validated_against_num_gpus(self):
+ pc = DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2)
+ od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+ assert od.parallel_config.data_parallel_size == 2
+
+ with pytest.raises(ValueError, match="does not match WORLD-derived value"):
+ OmniDiffusionConfig.from_kwargs(
+ model="x",
+ parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, data_parallel_size=2),
+ num_gpus=8,
+ )
+
+ def test_hsdp_does_not_infer_ordinary_dp(self):
+ pc = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=4)
+ od = OmniDiffusionConfig.from_kwargs(model="x", parallel_config=pc, num_gpus=4)
+ assert od.parallel_config.data_parallel_size == 1
+ assert od.parallel_config.world_size == 4
+
class TestCreateDefaultDiffusion:
"""Verify engine_args structure from create_default_diffusion."""
@@ -125,10 +149,33 @@ class TestCreateDefaultDiffusion:
assert ea["enforce_eager"] is True
assert ea["lora_path"] == "/tmp/lora"
- def test_diffusion_kv_mode_roundtrip(self):
- od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
+ def test_diffusion_kv_mode_roundtrip(self, monkeypatch):
+ from vllm_omni.platforms import current_omni_platform
+
+ monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+ od = _roundtrip_diffusion_config(
+ model="x",
+ diffusion_kv_mode="paged_scheduler",
+ diffusion_kv_max_rows_per_request=2,
+ )
assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+ assert od.diffusion_kv_max_rows_per_request == 2
+
+ def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
+ monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
+ od = _roundtrip_diffusion_config(
+ model="x",
+ kv_cache_memory_bytes=4096,
+ gpu_memory_utilization=0.75,
+ max_num_batched_tokens=2048,
+ max_model_len=4096,
+ )
+
+ assert od.kv_cache_memory_bytes == 4096
+ assert od.gpu_memory_utilization == 0.75
+ assert od.max_num_batched_tokens == 2048
+ assert od.max_model_len == 4096
def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py
index a0c6acb02..04d528af4 100644
--- a/tests/diffusion/test_diffusion_engine.py
+++ b/tests/diffusion/test_diffusion_engine.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import queue
@@ -197,6 +197,59 @@ class TestRequestBatchCapability:
assert diffusion_engine_module.supports_request_batch(od_config) is True
+ def test_engine_rejects_multi_seq_diffusers_using_adapter_capability(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ mocker: MockerFixture,
+ ) -> None:
+ od_config = SimpleNamespace(
+ model_class_name="QwenImagePipeline",
+ custom_pipeline_args=None,
+ diffusion_load_format="diffusers",
+ streaming_output=False,
+ max_num_seqs=2,
+ )
+ registry_load = mocker.Mock(
+ side_effect=lambda model_class_name: (
+ _SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline
+ )
+ )
+ monkeypatch.setattr(
+ diffusion_engine_module.DiffusionModelRegistry,
+ "_try_load_model_cls",
+ registry_load,
+ )
+
+ engine = DiffusionEngine.__new__(DiffusionEngine)
+ with pytest.raises(ValueError, match="max_num_seqs=1"):
+ engine._resolve_execution_mode(od_config)
+
+ registry_load.assert_called_once_with("DiffusersAdapterPipeline")
+
+ def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ mocker: MockerFixture,
+ ) -> None:
+ od_config = SimpleNamespace(
+ model_class_name="QwenImagePipeline",
+ custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
+ diffusion_load_format="diffusers",
+ streaming_output=False,
+ max_num_seqs=2,
+ )
+ registry_load = mocker.Mock(return_value=_SingleRequestPipeline)
+ monkeypatch.setattr(
+ diffusion_engine_module.DiffusionModelRegistry,
+ "_try_load_model_cls",
+ registry_load,
+ )
+
+ engine = DiffusionEngine.__new__(DiffusionEngine)
+
+ assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH
+ registry_load.assert_not_called()
+
def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
od_config = SimpleNamespace(
model_class_name="SinglePipeline",
diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py
index fb8858a48..ba7d7d5d2 100644
--- a/tests/diffusion/test_diffusion_engine_cleanup.py
+++ b/tests/diffusion/test_diffusion_engine_cleanup.py
@@ -12,6 +12,7 @@ import pytest
from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
assert engine.scheduler.get_request_state(request_id) is None
+def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]:
+ engine = DiffusionEngine.__new__(DiffusionEngine)
+ engine.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+ cleanup = Mock()
+ engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup)
+ return engine, cleanup
+
+
+def test_paged_terminal_output_clears_worker_rows() -> None:
+ engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+ engine._finalize_finished_request = lambda request_id, *_args: request_id
+ engine._put_output = lambda *_args: None
+
+ engine._emit_finished_outputs({"req-0", "req-1"})
+
+ cleanup.assert_called_once()
+ assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"}
+
+
+def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None:
+ engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+ engine.scheduler = SimpleNamespace(
+ get_request_state=lambda _request_id: None,
+ finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")),
+ )
+
+ engine._abort_requests(["req-idle", "req-idle"])
+
+ cleanup.assert_called_once_with(["req-idle"])
+ engine.scheduler.finish_requests.assert_not_called()
+
+
+def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None:
+ engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY)
+ engine._finalize_finished_request = lambda request_id, *_args: request_id
+ engine._put_output = lambda *_args: None
+ engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None)
+
+ engine._emit_finished_outputs({"req-0"})
+ engine._abort_requests(["req-idle"])
+
+ cleanup.assert_not_called()
+
+
def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
od_config = SimpleNamespace(
custom_pipeline_args=None,
@@ -161,9 +206,83 @@ def test_scheduler_initialization_failure_closes_scheduler_and_executor(monkeypa
assert exc_info.value is initialization_error
custom_scheduler.initialize.assert_called_once_with(od_config)
custom_scheduler.close.assert_called_once_with()
+
+
+def test_init_shuts_down_executor_when_kv_control_plane_initialization_fails(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ fake_executor = SimpleNamespace(shutdown=Mock())
+ monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+ monkeypatch.setattr(
+ DiffusionEngine,
+ "_resolve_execution_mode",
+ lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+ )
+ monkeypatch.setattr(
+ DiffusionEngine, "_init_executor", lambda self, config: setattr(self, "executor", fake_executor)
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+ Mock(side_effect=RuntimeError("KV initialization failed")),
+ )
+
+ with pytest.raises(RuntimeError, match="KV initialization failed"):
+ DiffusionEngine(SimpleNamespace())
+
fake_executor.shutdown.assert_called_once_with()
+def test_paged_init_profiles_before_scheduler_initialization(monkeypatch: pytest.MonkeyPatch) -> None:
+ events = []
+ expected_profile_requests = [object()]
+ fake_executor = SimpleNamespace(shutdown=Mock())
+ kv_cache_config = object()
+ kv_vllm_config = object()
+ od_config = SimpleNamespace(diffusion_kv_mode="paged_scheduler")
+
+ monkeypatch.setattr(DiffusionEngine, "_init_process_hooks", lambda self, config: None)
+ monkeypatch.setattr(
+ DiffusionEngine,
+ "_resolve_execution_mode",
+ lambda self, config: DiffusionExecutionMode.REQUEST_BATCH,
+ )
+ monkeypatch.setattr(
+ DiffusionEngine,
+ "_init_executor",
+ lambda self, config: setattr(self, "executor", fake_executor),
+ )
+ monkeypatch.setattr(
+ DiffusionEngine,
+ "_prepare_diffusion_kv_profile_requests",
+ lambda self: events.append("prepare-profile") or expected_profile_requests,
+ )
+
+ def initialize(executor, config, *, profile_requests):
+ events.append("profile-workers")
+ assert executor is fake_executor
+ assert config is od_config
+ assert profile_requests is expected_profile_requests
+ return kv_cache_config, 16, 16, kv_vllm_config
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+ initialize,
+ )
+
+ def init_scheduler(self, config, scheduler, *args, **kwargs):
+ events.append("initialize-scheduler")
+ self.scheduler = SimpleNamespace(close=Mock())
+
+ monkeypatch.setattr(DiffusionEngine, "_init_scheduler", init_scheduler)
+ monkeypatch.setattr(DiffusionEngine, "_init_runtime_state", lambda self: None)
+ monkeypatch.setattr(DiffusionEngine, "_init_execute_fn", lambda self: None)
+ monkeypatch.setattr(DiffusionEngine, "_log_execution_mode", lambda self, config: None)
+
+ DiffusionEngine(od_config)
+
+ assert events == ["prepare-profile", "profile-workers", "initialize-scheduler"]
+
+
@pytest.mark.asyncio
async def test_step_compatibility_wrapper_returns_final_batch() -> None:
engine = _make_engine()
diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py
index 50dde3083..247fee029 100644
--- a/tests/diffusion/test_diffusion_engine_dummy_run.py
+++ b/tests/diffusion/test_diffusion_engine_dummy_run.py
@@ -1,9 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+from unittest.mock import Mock
+
import pytest
from vllm_omni.diffusion import io_support
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -55,3 +63,151 @@ def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.Monk
)
assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2
+
+
+def test_dummy_run_image_count_resolves_hunyuan_architecture_alias() -> None:
+ assert io_support.get_dummy_run_num_image_inputs("HunyuanImage3ForCausalMM") == 3
+ assert io_support.get_dummy_run_num_image_inputs("unknown") == 1
+
+
+def test_dense_mode_does_not_build_kv_profile_request() -> None:
+ engine = object.__new__(DiffusionEngine)
+ engine.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+ engine._make_dummy_request = Mock(side_effect=AssertionError("dense mode must not profile"))
+
+ assert engine._prepare_diffusion_kv_profile_requests() is None
+ engine._make_dummy_request.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
+ [
+ (DiffusionExecutionMode.STEP_BATCH, False, 2, 2),
+ (DiffusionExecutionMode.REQUEST_BATCH, True, 4, 1),
+ (DiffusionExecutionMode.STEP_BATCH, True, 3, 3),
+ ],
+)
+def test_paged_kv_profile_requests_match_per_rank_batch(
+ execution_mode: DiffusionExecutionMode,
+ uses_dlo_dp: bool,
+ max_num_seqs: int,
+ expected_profile_requests: int,
+) -> None:
+ engine = object.__new__(DiffusionEngine)
+ engine.execution_mode = execution_mode
+ engine.od_config = SimpleNamespace(
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ model_class_name="HunyuanImage3ForCausalMM",
+ max_num_seqs=max_num_seqs,
+ parallel_config=SimpleNamespace(data_parallel_size=max_num_seqs if uses_dlo_dp else 1),
+ enable_distributed_layerwise_offload=uses_dlo_dp,
+ dlo_use_allgather=True,
+ )
+ request = OmniDiffusionRequest(
+ prompt="profile",
+ request_id="profile-request",
+ sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+ )
+ engine._make_dummy_request = Mock(return_value=request)
+ prepared_layout = object()
+ scheduler_kv_state = (
+ DiffusionKVRequest(
+ "profile-request/diffusion-kv/0",
+ sequence_id=0,
+ prefix_len=3,
+ target_len=8,
+ seq_len=12,
+ ),
+ DiffusionKVRequest(
+ "profile-request/diffusion-kv/1",
+ sequence_id=1,
+ prefix_len=4,
+ target_len=8,
+ seq_len=13,
+ ),
+ )
+
+ def preprocess(req: OmniDiffusionRequest) -> OmniDiffusionRequest:
+ req.prepared_layout = prepared_layout
+ req.diffusion_kv_requests = scheduler_kv_state
+ return req
+
+ engine._prepare_request_for_admission = Mock(side_effect=preprocess)
+
+ result = engine._prepare_diffusion_kv_profile_requests()
+
+ assert result is not None
+ assert len(result) == expected_profile_requests
+ assert len({profile_request.request_id for profile_request in result}) == expected_profile_requests
+ assert all(profile_request.is_dummy_run() for profile_request in result)
+ assert all(profile_request.prepared_layout is prepared_layout for profile_request in result)
+ assert all(profile_request.diffusion_kv_requests is None for profile_request in result)
+ assert len({id(profile_request.sampling_params) for profile_request in result}) == expected_profile_requests
+ assert engine._diffusion_kv_profile_limits == (2, 13, 8)
+ assert len(result) * engine._diffusion_kv_profile_limits[0] == expected_profile_requests * 2
+ engine._make_dummy_request.assert_called_once_with(
+ height=1024,
+ width=1024,
+ guidance_scale=5.0,
+ num_image_inputs=3,
+ )
+ engine._prepare_request_for_admission.assert_called_once_with(request)
+
+
+@pytest.mark.parametrize(
+ ("num_sequences", "seq_len", "target_len"),
+ [
+ (3, 16, 8),
+ (1, 17, 8),
+ (1, 16, 9),
+ ],
+)
+def test_paged_kv_admission_rejects_shape_beyond_profile_envelope(
+ num_sequences: int,
+ seq_len: int,
+ target_len: int,
+) -> None:
+ engine = object.__new__(DiffusionEngine)
+ engine._diffusion_kv_profile_limits = (2, 16, 8)
+ engine.pre_process_func = lambda request: request
+ request = OmniDiffusionRequest(
+ prompt="too large",
+ request_id="too-large",
+ sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+ diffusion_kv_requests=tuple(
+ DiffusionKVRequest(
+ f"too-large/diffusion-kv/{sequence_id}",
+ sequence_id=sequence_id,
+ prefix_len=0,
+ target_len=target_len,
+ seq_len=seq_len,
+ )
+ for sequence_id in range(num_sequences)
+ ),
+ )
+
+ with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+ engine._prepare_request_for_admission(request)
+
+
+def test_paged_kv_admission_accepts_shape_at_profile_envelope() -> None:
+ engine = object.__new__(DiffusionEngine)
+ engine._diffusion_kv_profile_limits = (2, 16, 8)
+ engine.pre_process_func = lambda request: request
+ request = OmniDiffusionRequest(
+ prompt="fits",
+ request_id="fits",
+ sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+ diffusion_kv_requests=tuple(
+ DiffusionKVRequest(
+ f"fits/diffusion-kv/{sequence_id}",
+ sequence_id=sequence_id,
+ prefix_len=8,
+ target_len=8,
+ seq_len=16,
+ )
+ for sequence_id in range(2)
+ ),
+ )
+
+ assert engine._prepare_request_for_admission(request) is request
diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py
index 8f434c0e2..dc909a2b9 100644
--- a/tests/diffusion/test_diffusion_engine_metrics.py
+++ b/tests/diffusion/test_diffusion_engine_metrics.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Source-level regression tests for diffusion output/engine helpers.
These tests verify naming conventions and patterns by inspecting source code
@@ -37,10 +37,50 @@ _FORMATTER_PATH = os.path.normpath(
"output_formatter.py",
)
)
+_METRICS_PATH = os.path.normpath(
+ os.path.join(
+ os.path.dirname(__file__),
+ os.pardir,
+ os.pardir,
+ "vllm_omni",
+ "metrics",
+ "utils.py",
+ )
+)
+_INLINE_CLIENT_PATH = os.path.normpath(
+ os.path.join(
+ os.path.dirname(__file__),
+ os.pardir,
+ os.pardir,
+ "vllm_omni",
+ "diffusion",
+ "inline_stage_diffusion_client.py",
+ )
+)
+_STAGE_PROC_PATH = os.path.normpath(
+ os.path.join(
+ os.path.dirname(__file__),
+ os.pardir,
+ os.pardir,
+ "vllm_omni",
+ "diffusion",
+ "stage_diffusion_proc.py",
+ )
+)
+_ORCHESTRATOR_PATH = os.path.normpath(
+ os.path.join(
+ os.path.dirname(__file__),
+ os.pardir,
+ os.pardir,
+ "vllm_omni",
+ "engine",
+ "orchestrator.py",
+ )
+)
def _read_source(path: str) -> str:
- with open(path) as f:
+ with open(path, encoding="utf-8") as f:
return f.read()
@@ -100,9 +140,116 @@ class TestMetricKeys:
raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
- assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
+ # step_total_ms is no longer emitted as a metric (no family consumes
+ # it); it lives only in the debug log breakdown below.
+ assert '"diffusion_engine_total_time_ms"' not in step_streaming_source
assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
+ def test_step_streaming_piggybacks_scheduler_waiting_snapshot(self) -> None:
+ source = _read_source(_ENGINE_PATH)
+ step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+ assert "diffusion_scheduler_waiting_metrics(" in step_streaming_source
+ assert "_scheduler_num_waiting_reqs" in step_streaming_source
+
+ def test_step_streaming_preserves_snapshot_on_postprocess_error(self) -> None:
+ source = _read_source(_ENGINE_PATH)
+ step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+ assert "scheduler_metrics = diffusion_scheduler_waiting_metrics(" in step_streaming_source
+ assert 'setattr(exc, "diffusion_metrics", scheduler_metrics)' in step_streaming_source
+ assert step_streaming_source.index("scheduler_metrics =") < step_streaming_source.index(
+ "self.postprocess_output("
+ )
+
+ def test_abort_keeps_output_consumers_alive_for_terminal_snapshot(self) -> None:
+ inline_source = _read_source(_INLINE_CLIENT_PATH)
+ abort_source = _get_function_source(inline_source, "InlineStageDiffusionClient", "abort_requests_async")
+ assert ".cancel(" not in abort_source
+
+ proc_source = _read_source(_STAGE_PROC_PATH)
+ run_loop_source = _get_function_source(proc_source, "StageDiffusionProc", "run_loop")
+ abort_branch = run_loop_source[run_loop_source.index('elif msg_type == "abort":') :]
+ abort_branch = abort_branch[: abort_branch.index('elif msg_type == "collective_rpc":')]
+ assert ".cancel(" not in abort_branch
+
+ def test_orchestrator_consumes_metrics_only_output_without_routing(self) -> None:
+ """A metrics-only sentinel contributes its queue depth and is not routed.
+
+ Absorption lives in ``_absorb_diffusion_metrics`` rather than inline in a
+ loop, so both orchestration loops share one implementation. The ordering
+ is asserted where it now lives: the snapshot inside the helper, and the
+ absorb-before-route guard inside every loop that polls diffusion output.
+ """
+ source = _read_source(_ORCHESTRATOR_PATH)
+
+ # Snapshot before the verdict, so a sentinel still reports its waiting
+ # depth on the way out instead of being dropped unaccounted.
+ absorb_source = _get_function_source(source, "Orchestrator", "_absorb_diffusion_metrics")
+ snapshot_pos = absorb_source.index("_update_stage_replica_waiting(")
+ sentinel_pos = absorb_source.index("diffusion_output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID")
+ assert snapshot_pos < sentinel_pos
+
+ # Absorb before routing, or a sentinel reaches the downstream consumer
+ # as if it were a real request output.
+ for loop_name in ("_orchestration_loop", "_orchestration_loop_event_driven"):
+ loop_source = _get_function_source(source, "Orchestrator", loop_name)
+ absorb_pos = loop_source.index("self._absorb_diffusion_metrics(")
+ route_pos = loop_source.index("record_output_timestamps(")
+ assert absorb_pos < route_pos, loop_name
+
+
+class TestVaeDecodeEmit:
+ """VAE decode timing is sourced from ``DiffusionOutput.stage_durations``
+ (populated by the diffusion pipeline profiler) and emitted as
+ ``vae_decode_time_ms`` so the ``_MS_TO_S`` accumulator picks it up.
+ """
+
+ def test_extract_vae_decode_ms_helper_exists(self) -> None:
+ source = _read_source(_METRICS_PATH)
+ helper_src = _get_function_source(source, None, "extract_diffusion_vae_decode_ms")
+ assert ".vae.decode" in helper_src, "helper must key on the '.vae.decode' suffix"
+
+ def test_step_streaming_emits_vae_decode_time_ms(self) -> None:
+ source = _read_source(_ENGINE_PATH)
+ step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+ assert "extract_diffusion_vae_decode_ms(" in step_streaming_source, (
+ "step_streaming must call extract_diffusion_vae_decode_ms to source VAE decode timing"
+ )
+ assert '"vae_decode_time_ms"' in step_streaming_source, (
+ "step_streaming must emit vae_decode_time_ms key for the _MS_TO_S accumulator"
+ )
+
+
+class TestDiffuseForwardEmit:
+ """Forward-only timing → forward_time_ms, mirroring TestVaeDecodeEmit."""
+
+ def test_extract_diffuse_ms_helper_exists(self) -> None:
+ source = _read_source(_METRICS_PATH)
+ helper_src = _get_function_source(source, None, "extract_diffusion_denoise_ms")
+ assert ".diffuse" in helper_src, "helper must key on the '.diffuse' suffix"
+
+ def test_step_streaming_emits_forward_time_ms(self) -> None:
+ source = _read_source(_ENGINE_PATH)
+ step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+ assert "extract_diffusion_denoise_ms(" in step_streaming_source, (
+ "step_streaming must call extract_diffusion_denoise_ms to source forward-only timing"
+ )
+ assert '"forward_time_ms"' in step_streaming_source, (
+ "step_streaming must emit forward_time_ms key for the _MS_TO_S accumulator"
+ )
+
+
+class TestKvRecvEmit:
+ """KV-recv timing → kv_recv_time_ms, sourced from DiffusionOutput.kv_recv_ms
+ (set by the runner's _prepare_request_for_forward, not the profiler)."""
+
+ def test_step_streaming_emits_kv_recv_time_ms(self) -> None:
+ source = _read_source(_ENGINE_PATH)
+ step_streaming_source = _get_function_source(source, "DiffusionEngine", "step_streaming")
+ assert "kv_recv_ms" in step_streaming_source, "step_streaming must read output.kv_recv_ms"
+ assert '"kv_recv_time_ms"' in step_streaming_source, (
+ "step_streaming must emit kv_recv_time_ms key for the _MS_TO_S accumulator"
+ )
+
class TestDummyRunAllocation:
"""Verify _dummy_run generates exact-sized audio arrays."""
diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py
index 1d4c2d17e..119939ae2 100644
--- a/tests/diffusion/test_diffusion_engine_rpc_routing.py
+++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py
@@ -123,6 +123,7 @@ class _ConcurrencyTrackingExecutor:
def _make_request(tag: str):
sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
sampling_params["num_inference_steps"] = 1
+ sampling_params["extra_args"] = {}
return SimpleNamespace(
request_id=tag,
prompt=f"prompt_{tag}",
@@ -472,6 +473,38 @@ async def test_cancelled_future_does_not_kill_busy_loop():
_stop_engine(engine)
+@pytest.mark.asyncio
+async def test_native_kv_reservation_error_wakes_stream_without_killing_busy_loop():
+ loop = asyncio.get_running_loop()
+ engine = _make_engine_with_loop(loop)
+
+ def fail_reservation(*args, **kwargs):
+ raise RuntimeError("native allocation bug")
+
+ engine.scheduler._diffusion_kv_manager = SimpleNamespace(
+ has_request=lambda request_id: False,
+ reserve_request=fail_reservation,
+ free_request=lambda request_id: None,
+ )
+ request = _make_request("kv-error")
+ request.diffusion_kv_requests = (object(),)
+ try:
+ response_stream = engine.async_add_req_and_stream_response(request)
+ output_queue = engine._out_streams[request.request_id]
+ for _ in range(300):
+ if not output_queue.empty():
+ break
+ await asyncio.sleep(0.01)
+ assert not output_queue.empty(), "terminal KV allocation error was not delivered to the request stream"
+
+ output = await anext(response_stream)
+ assert output.error == "native allocation bug"
+ assert output.finished
+ assert engine.worker_thread.is_alive()
+ finally:
+ _stop_engine(engine)
+
+
# ─────────────────────────── _RpcTask basics ───────────────────────────────
diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py
index 3e4c587f1..aec062057 100644
--- a/tests/diffusion/test_diffusion_ipc.py
+++ b/tests/diffusion/test_diffusion_ipc.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import contextlib
import multiprocessing as mp
@@ -7,6 +7,7 @@ import queue
import threading
import time
from multiprocessing import shared_memory
+from multiprocessing.connection import Connection
import numpy as np
import pytest
@@ -29,7 +30,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def _large_numel(dtype: torch.dtype) -> int:
- return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1
+ return int(_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1
def _cleanup_shm_handle(value: object) -> None:
@@ -38,7 +39,7 @@ def _cleanup_shm_handle(value: object) -> None:
_unpack_if_shm_handle(value)
-def _result_queue_worker(connection, rank: int) -> None:
+def _result_queue_worker(connection: Connection, rank: int) -> None:
"""Send one nested NumPy result through a worker-owned MessageQueue."""
result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
try:
@@ -180,7 +181,7 @@ def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
- # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
+ # LTX2 returns (video, audio) tuples as DiffusionOutput.output.
video = torch.arange(300_000, dtype=torch.float32)
audio = torch.arange(300_000, dtype=torch.float32) * 3
output = DiffusionOutput(output=(video, audio))
@@ -305,6 +306,48 @@ def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
assert output["req-error"].result.error == "boom"
+def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None:
+ """Regression: per-request slices used to pickle the whole batch storage.
+
+ Each view is under the packing threshold, but pickle writes the entire
+ shared storage per view. Packing must key off storage nbytes so those
+ views become SHM handles instead of going on the wire.
+ """
+ num_requests = 16
+ # Each row stays under the threshold; the shared storage is well over it.
+ row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size())
+ batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel)
+ per_request = batch[0]
+ assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD
+ assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD
+
+ packed_view = _pack_value_if_large(per_request)
+ try:
+ assert isinstance(packed_view, dict)
+ assert packed_view["__tensor_shm__"] is True
+ finally:
+ _cleanup_shm_handle(packed_view)
+
+ output = BatchRunnerOutput.from_list(
+ [
+ RunnerOutput(
+ request_id=f"req-{i}",
+ finished=True,
+ result=DiffusionOutput(trajectory_latents=batch[i]),
+ )
+ for i in range(num_requests)
+ ]
+ )
+
+ pack_diffusion_output_shm(output)
+
+ assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True
+
+ unpack_diffusion_output_shm(output)
+ for i in range(num_requests):
+ torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i])
+
+
def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
tensor = torch.arange(
_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
@@ -337,31 +380,37 @@ def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_valu
large = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
small = torch.arange(8, dtype=torch.float32)
list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
+ original_list = [list_tensor]
payload = {
"media": {
"large": large,
"small": small,
},
- "list_value": [list_tensor],
+ "list_value": original_list,
"metadata": {"prompt": "keep inline"},
}
packed = _pack_value_if_large(payload)
try:
+ assert isinstance(packed, dict)
assert packed is not payload
- assert packed["media"] is not payload["media"]
- assert packed["media"]["large"]["__tensor_shm__"] is True
- assert packed["media"]["small"] is small
+ media = packed["media"]
+ list_value = packed["list_value"]
+ assert isinstance(media, dict)
+ assert isinstance(list_value, list)
+ assert media is not payload["media"]
+ assert media["large"]["__tensor_shm__"] is True
+ assert media["small"] is small
# Lists are recursed too: the large tensor inside is packed and a new
# list is returned, while the input payload is left untouched.
- assert packed["list_value"] is not payload["list_value"]
- assert packed["list_value"][0]["__tensor_shm__"] is True
- assert payload["list_value"][0] is list_tensor
+ assert list_value is not original_list
+ assert list_value[0]["__tensor_shm__"] is True
+ assert original_list[0] is list_tensor
assert packed["metadata"] == {"prompt": "keep inline"}
- torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large)
- torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor)
+ torch.testing.assert_close(_unpack_if_shm_handle(media["large"]), large)
+ torch.testing.assert_close(_unpack_if_shm_handle(list_value[0]), list_tensor)
finally:
if isinstance(packed, dict):
_cleanup_shm_handle(packed.get("media", {}).get("large"))
diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py
index 665196ad9..814450d92 100644
--- a/tests/diffusion/test_diffusion_model_runner.py
+++ b/tests/diffusion/test_diffusion_model_runner.py
@@ -1,8 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from contextlib import contextmanager
from types import SimpleNamespace
+from unittest.mock import Mock
import pytest
import torch
@@ -10,6 +11,7 @@ import torch
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
@@ -192,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
+ runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
+ vllm_config=runner.vllm_config,
+ od_config=runner.od_config,
+ device=runner.device,
+ )
runner.kv_transfer_manager = SimpleNamespace(
receive_kv_cache=lambda req, target_device=None: None,
receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
@@ -512,6 +519,102 @@ def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipelin
assert runner.pipeline.last_req.num_reqs == 1
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatch):
+ runner = _make_runner(cache_backend=None, cache_backend_name="none")
+ request = _make_request()
+ runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+ record_names = []
+ original_execute = runner._execute_request_list
+
+ def execute(*args, **kwargs):
+ record_names.append(kwargs["record_name"])
+ return original_execute(*args, **kwargs)
+
+ runner._execute_request_list = execute
+ forward_context_calls = []
+
+ @contextmanager
+ def record_forward_context(*args, **kwargs):
+ forward_context_calls.append((args, kwargs))
+ yield
+
+ monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
+ reset_peak_memory_stats = Mock()
+ monkeypatch.setattr(
+ model_runner_module.current_omni_platform,
+ "reset_peak_memory_stats",
+ reset_peak_memory_stats,
+ )
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+ synchronize = Mock()
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+ runner.profile_run([request])
+
+ assert runner.pipeline.forward_calls == 1
+ assert record_names == ["pipeline_memory_profile"]
+ assert len(forward_context_calls) == 1
+ assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
+ runner._validate_diffusion_kv_metadata.assert_not_called()
+ reset_peak_memory_stats.assert_not_called()
+ synchronize.assert_called_once_with()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypatch):
+ runner = _make_runner(cache_backend=None, cache_backend_name=None)
+ runner.pipeline = _FinalOnlyStepPipeline()
+ runner.od_config.step_execution = True
+ requests = [_make_request(), _make_request()]
+ requests[0].request_id = "profile-0"
+ requests[1].request_id = "profile-1"
+ for request in requests:
+ request.sampling_params.num_inference_steps = 1
+
+ observed_batch_rows = []
+ original_denoise_step = runner.pipeline.denoise_step
+
+ def denoise_step(input_batch, states):
+ observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
+ return original_denoise_step(input_batch, states)
+
+ runner.pipeline.denoise_step = denoise_step
+ runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
+ forward_context_calls = []
+
+ @contextmanager
+ def record_forward_context(*args, **kwargs):
+ forward_context_calls.append((args, kwargs))
+ yield
+
+ monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
+ reset_peak_memory_stats = Mock()
+ monkeypatch.setattr(
+ model_runner_module.current_omni_platform,
+ "reset_peak_memory_stats",
+ reset_peak_memory_stats,
+ )
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+ synchronize = Mock()
+ monkeypatch.setattr(model_runner_module.current_omni_platform, "synchronize", synchronize)
+
+ runner.profile_run(requests)
+
+ assert observed_batch_rows == [(2, 2)]
+ assert runner.state_cache == {}
+ assert runner.input_batch is None
+ assert len(forward_context_calls) == 1
+ assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
+ runner._validate_diffusion_kv_metadata.assert_not_called()
+ reset_peak_memory_stats.assert_not_called()
+ synchronize.assert_called_once_with()
+
+
class _BatchPipeline:
"""Pipeline returning a configurable list of outputs from forward()."""
@@ -779,7 +882,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
- monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
+ monkeypatch.setattr(
+ model_runner_module,
+ "enable_offload_backend",
+ lambda od_config, pipeline, device: (pipeline, None),
+ )
monkeypatch.setattr(
model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
)
@@ -791,8 +898,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
assert dummy_cache_backend.enabled is False
-@pytest.mark.core_model
-@pytest.mark.cpu
def test_set_forward_context_enters_vllm_config_contexts(monkeypatch):
"""Ensure `with set_forward_context(...):` enters vllm's context managers internally and calls desired vllm functions."""
import vllm.config.vllm as vllm_config_module
diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py
index dc21351ff..4578d8368 100644
--- a/tests/diffusion/test_diffusion_output_formatter.py
+++ b/tests/diffusion/test_diffusion_output_formatter.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -76,7 +76,7 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
result = results[0]
assert result.images == ["frame-0"]
assert result.prompt == "prompt-0"
- assert result.final_output_type == "image"
+ assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 48000},
@@ -315,7 +315,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
assert result.images == ["frame-0"]
assert result.prompt == "watch and listen"
- assert result.final_output_type == "image"
+ assert result.final_output_type == "video"
assert result.multimodal_output == {
"metadata": {
"audio": {"sample_rate": 16000},
@@ -328,6 +328,25 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
}
+def test_formatter_uses_native_model_identity_for_diffusers_video_frames(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+ frames = [["frame-0", "frame-1"]]
+ postprocess_output = normalize_diffusion_postprocess_output(frames)
+
+ [result] = format_diffusion_outputs(
+ request=_request("make a video"),
+ od_config=_config("WanImageToVideoPipeline"),
+ diffusion_output=DiffusionOutput(output=frames),
+ output_data=frames,
+ postprocess_output=postprocess_output,
+ )
+
+ assert result.images == frames
+ assert result.final_output_type == "video"
+
+
def test_formatter_preserves_audio_only_postprocess_dict(
monkeypatch: pytest.MonkeyPatch,
) -> None:
diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py
index ee3559f47..b1228cebd 100644
--- a/tests/diffusion/test_diffusion_output_metadata.py
+++ b/tests/diffusion/test_diffusion_output_metadata.py
@@ -25,7 +25,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None:
},
"audio": {
"sample_rate": 48000,
- "audiox_task": "tts",
},
"actions": {
"raw_action_dim": 7,
@@ -55,7 +54,6 @@ def test_validate_accepts_canonical_metadata_fields() -> None:
({"video": {"video_fps_multiplier": 0}}, ValueError),
({"video": {"shape": [1, 0, 64]}}, ValueError),
({"audio": {"sample_rate": 0}}, ValueError),
- ({"audio": {"audiox_task": ""}}, ValueError),
({"actions": {"raw_action_dim": 0}}, ValueError),
({"actions": {"action_mode": ""}}, ValueError),
({"actions": {"domain_id": "7"}}, TypeError),
diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py
index a437a572a..65ad9ccc7 100644
--- a/tests/diffusion/test_diffusion_plugin_hooks.py
+++ b/tests/diffusion/test_diffusion_plugin_hooks.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for diffusion engine plugin extensibility hooks.
@@ -20,6 +20,8 @@ from vllm_omni.diffusion.registry import (
_DIFFUSION_MODELS,
_DIFFUSION_POST_PROCESS_FUNCS,
_DIFFUSION_PRE_PROCESS_FUNCS,
+ get_diffusion_post_process_func,
+ get_diffusion_pre_process_func,
register_diffusion_model,
)
from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
@@ -40,6 +42,12 @@ class TestPlatformDiffusionHooks:
result = OmniPlatform.get_diffusion_model_runner_cls()
assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
+ def test_get_diffusion_kv_block_tables_cls_default(self):
+ """Test the default paged-KV BlockTables implementation."""
+ from vllm.v1.worker.gpu.block_table import BlockTables
+
+ assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables
+
def test_oot_enum_exists(self):
"""Test OOT is a valid platform enum value."""
assert OmniPlatformEnum.OOT.value == "oot"
@@ -106,6 +114,49 @@ class TestRegisterDiffusionModel:
assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"
+@pytest.mark.parametrize(
+ ("model_class_name", "get_process_func"),
+ [
+ ("QwenImagePipeline", get_diffusion_post_process_func),
+ ("WanImageToVideoPipeline", get_diffusion_post_process_func),
+ ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
+ ],
+)
+def test_diffusers_backend_skips_native_process_hooks(model_class_name, get_process_func):
+ od_config = SimpleNamespace(
+ model_class_name=model_class_name,
+ custom_pipeline_args=None,
+ diffusion_load_format="diffusers",
+ )
+
+ with patch("vllm_omni.diffusion.registry._load_process_func") as load_process_func:
+ assert get_process_func(od_config) is None
+
+ load_process_func.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("model_class_name", "get_process_func"),
+ [
+ ("QwenImagePipeline", get_diffusion_post_process_func),
+ ("WanImageToVideoPipeline", get_diffusion_post_process_func),
+ ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
+ ],
+)
+def test_custom_pipeline_preserves_native_process_hooks(model_class_name, get_process_func):
+ process_func = Mock()
+ od_config = SimpleNamespace(
+ model_class_name=model_class_name,
+ custom_pipeline_args={"pipeline_class": "custom.Pipeline"},
+ diffusion_load_format="diffusers",
+ )
+
+ with patch("vllm_omni.diffusion.registry._load_process_func", return_value=process_func) as load_process_func:
+ assert get_process_func(od_config) is process_func
+
+ load_process_func.assert_called_once()
+
+
class TestWorkerUsesHook:
"""Test that DiffusionWorker resolves model runner via platform hook."""
diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py
index 259492c1e..7476ec9c1 100644
--- a/tests/diffusion/test_diffusion_request.py
+++ b/tests/diffusion/test_diffusion_request.py
@@ -52,6 +52,7 @@ def test_dummy_run_request_is_identified_by_reserved_request_id():
assert req.is_dummy_run()
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
+ assert OmniDiffusionRequest.is_dummy_run_request_id(f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-1")
def test_non_dummy_request_is_not_identified_as_dummy_run():
diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py
index ce9f0ebe1..ae9b86f80 100644
--- a/tests/diffusion/test_diffusion_scheduler.py
+++ b/tests/diffusion/test_diffusion_scheduler.py
@@ -8,12 +8,14 @@ from types import SimpleNamespace
import pytest
import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
from pytest_mock import MockerFixture
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import (
BaseScheduler,
@@ -23,7 +25,7 @@ from vllm_omni.diffusion.sched import (
StepScheduler,
)
from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -87,6 +89,52 @@ def _cached_ids(sched_output) -> list[str]:
return list(sched_output.scheduled_cached_reqs.request_ids)
+def _initialize_paged_scheduler(
+ scheduler: BaseScheduler,
+ *,
+ num_blocks: int = 64,
+ max_num_seqs: int = 1,
+) -> None:
+ native_kv_managers.register_all_kvcache_specs(None)
+ spec = FullAttentionSpec(
+ block_size=4,
+ num_kv_heads=2,
+ head_size=8,
+ dtype=torch.bfloat16,
+ )
+ config = KVCacheConfig(
+ num_blocks=num_blocks,
+ kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])],
+ kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+ )
+ scheduler.initialize(
+ SimpleNamespace(
+ diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+ max_model_len=64,
+ max_num_seqs=max_num_seqs,
+ ),
+ kv_cache_config=config,
+ scheduler_block_size=4,
+ hash_block_size=4,
+ kv_vllm_config=SimpleNamespace(
+ model_config=SimpleNamespace(max_model_len=64),
+ max_in_flight_tokens=64,
+ ),
+ )
+
+
+def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+ request.diffusion_kv_requests = (
+ DiffusionKVRequest(
+ f"{request.request_id}/diffusion-kv/0",
+ sequence_id=0,
+ prefix_len=4,
+ target_len=4,
+ seq_len=seq_len,
+ ),
+ )
+
+
class _StubScheduler:
def __init__(self, request: OmniDiffusionRequest, output) -> None:
self._request = request
@@ -217,13 +265,21 @@ class TestGetRequestBatchSamplingParamsKey:
num_inference_steps: int = 2,
seed: int | None = 123,
generator: torch.Generator | None = None,
+ extra_args: dict | None = None,
+ condition_key: tuple | None = None,
) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(
num_inference_steps=num_inference_steps,
seed=seed,
generator=generator,
+ extra_args=extra_args or {},
+ )
+ return OmniDiffusionRequest(
+ prompt="prompt",
+ sampling_params=sp,
+ request_id=f"req-{num_inference_steps}",
+ batch_compatibility_key=condition_key,
)
- return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}")
def test_distinguishes_num_inference_steps(self) -> None:
scheduler = RequestScheduler()
@@ -240,6 +296,57 @@ class TestGetRequestBatchSamplingParamsKey:
self._make(seed=1, generator=gen_a)
) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b))
+ @pytest.mark.parametrize(
+ ("first_extra_args", "second_extra_args"),
+ [
+ ({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
+ ({"flow_shift": 3.0}, {"flow_shift": 5.0}),
+ ],
+ )
+ def test_distinguishes_wan_scheduler_structure(
+ self,
+ first_extra_args: dict,
+ second_extra_args: dict,
+ ) -> None:
+ scheduler = RequestScheduler()
+
+ assert scheduler._build_sampling_params_key(
+ self._make(extra_args=first_extra_args)
+ ) != scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
+
+ @pytest.mark.parametrize(
+ ("first_extra_args", "second_extra_args"),
+ [
+ ({"sample_solver": " Euler "}, {"sample_solver": "euler"}),
+ ({"flow_shift": "5.0"}, {"flow_shift": 5.0}),
+ ],
+ )
+ def test_normalizes_equivalent_wan_scheduler_structure(
+ self,
+ first_extra_args: dict,
+ second_extra_args: dict,
+ ) -> None:
+ scheduler = RequestScheduler()
+
+ assert scheduler._build_sampling_params_key(
+ self._make(extra_args=first_extra_args)
+ ) == scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args))
+
+ def test_uses_none_for_unspecified_wan_scheduler_structure(self) -> None:
+ scheduler = RequestScheduler()
+
+ key = scheduler._build_sampling_params_key(self._make())
+
+ assert key.sample_solver is None
+ assert key.flow_shift is None
+
+ def test_distinguishes_pipeline_condition_structure(self) -> None:
+ scheduler = RequestScheduler()
+
+ assert scheduler._build_sampling_params_key(
+ self._make(condition_key=("wan22_s2v_condition", True))
+ ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False)))
+
class TestRequestScheduler:
def setup_method(self) -> None:
@@ -341,18 +448,16 @@ class TestRequestScheduler:
assert self.scheduler.has_requests() is False
def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
- self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+ _initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
prepared_layout = object()
request.prepared_layout = prepared_layout
- text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
kv_request = DiffusionKVRequest(
"diffusion-kv/diffusion-kv/0",
sequence_id=0,
prefix_len=4,
target_len=8,
seq_len=16,
- kv_contexts=(text_context,),
)
request.diffusion_kv_requests = (kv_request,)
@@ -362,12 +467,179 @@ class TestRequestScheduler:
assert state is not None
assert state.diffusion_kv_requests == (kv_request,)
- assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
assert request.diffusion_kv_requests is None
assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+ metadata = scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata
+ assert metadata is not None
+ assert metadata.request_id == request.request_id
+ assert len(metadata.sequences[0].block_ids[0]) == 4
+
+ def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
+ _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
+ first_request = _make_request("first")
+ second_request = _make_request("second")
+ _attach_diffusion_kv(first_request)
+ _attach_diffusion_kv(second_request)
+ self.scheduler.add_request(first_request)
+ self.scheduler.add_request(second_request)
+
+ first_output = self.scheduler.schedule()
+
+ assert _new_ids(first_output) == ["first"]
+ assert first_output.num_waiting_reqs == 1
+
+ self.scheduler.update_from_output(first_output, _make_request_output("first"))
+ second_output = self.scheduler.schedule()
+
+ assert _new_ids(second_output) == ["second"]
+ assert second_output.num_waiting_reqs == 0
+ # The same output carries both the release and the newly admitted block
+ # table; the future Worker data plane must process finished ids first.
+ assert second_output.finished_req_ids == {"first"}
+ assert second_output.scheduled_new_reqs[0].diffusion_kv_metadata is not None
+
+ def test_impossible_diffusion_kv_capacity_finishes_only_that_request(self) -> None:
+ _initialize_paged_scheduler(self.scheduler, num_blocks=2, max_num_seqs=2)
+ impossible = _make_request("impossible")
+ schedulable = _make_request("schedulable")
+ _attach_diffusion_kv(impossible, seq_len=8)
+ schedulable.diffusion_kv_requests = (
+ DiffusionKVRequest(
+ "schedulable/diffusion-kv/0",
+ sequence_id=0,
+ prefix_len=0,
+ target_len=4,
+ seq_len=4,
+ ),
+ )
+ self.scheduler.add_request(impossible)
+ self.scheduler.add_request(schedulable)
+
+ sched_output = self.scheduler.schedule()
+
+ failed_state = self.scheduler.get_request_state("impossible")
+ assert failed_state is not None
+ assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+ assert failed_state.error is not None
+ assert "cannot fit even when the block pool is empty" in failed_state.error
+ assert sched_output.finished_req_ids == {"impossible"}
+ assert _new_ids(sched_output) == ["schedulable"]
+
+ finished = self.scheduler.update_from_output(
+ sched_output,
+ _make_request_output("schedulable"),
+ )
+
+ assert finished == {"impossible", "schedulable"}
+
+ def test_impossible_diffusion_kv_capacity_does_not_block_waiters_under_load(self) -> None:
+ _initialize_paged_scheduler(self.scheduler, num_blocks=4, max_num_seqs=2)
+ running = _make_request("running")
+ _attach_diffusion_kv(running)
+ self.scheduler.add_request(running)
+ assert _new_ids(self.scheduler.schedule()) == ["running"]
+
+ impossible = _make_request("impossible")
+ impossible.diffusion_kv_requests = tuple(
+ DiffusionKVRequest(
+ f"impossible/diffusion-kv/{sequence_id}",
+ sequence_id=sequence_id,
+ prefix_len=0,
+ target_len=4,
+ seq_len=4,
+ )
+ for sequence_id in range(4)
+ )
+ schedulable = _make_request("schedulable")
+ schedulable.diffusion_kv_requests = (
+ DiffusionKVRequest(
+ "schedulable/diffusion-kv/0",
+ sequence_id=0,
+ prefix_len=0,
+ target_len=4,
+ seq_len=4,
+ ),
+ )
+ self.scheduler.add_request(impossible)
+ self.scheduler.add_request(schedulable)
+
+ sched_output = self.scheduler.schedule()
+
+ failed_state = self.scheduler.get_request_state("impossible")
+ assert failed_state is not None
+ assert failed_state.status == DiffusionRequestStatus.FINISHED_ERROR
+ assert failed_state.error is not None
+ assert "required_blocks=4, available_blocks=3" in failed_state.error
+ assert sched_output.finished_req_ids == {"impossible"}
+ assert _new_ids(sched_output) == ["schedulable"]
+
+ def test_diffusion_kv_internal_allocation_error_finishes_request(self, monkeypatch) -> None:
+ _initialize_paged_scheduler(self.scheduler)
+ request = _make_request("native-error")
+ _attach_diffusion_kv(request)
+ self.scheduler.add_request(request)
+ manager = self.scheduler._diffusion_kv_manager
+ assert manager is not None
+
+ def raise_native_error(*args, **kwargs):
+ raise ValueError("native allocation bug")
+
+ monkeypatch.setattr(
+ manager,
+ "reserve_request",
+ raise_native_error,
+ )
+
+ sched_output = self.scheduler.schedule()
+
+ state = self.scheduler.get_request_state("native-error")
+ assert state is not None
+ assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+ assert state.error == "native allocation bug"
+ assert sched_output.finished_req_ids == {"native-error"}
+ assert sched_output.scheduled_request_ids == []
+
+ def test_diffusion_kv_preemption_retains_allocation(self) -> None:
+ _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+ request = _make_request("preempted")
+ _attach_diffusion_kv(request)
+ self.scheduler.add_request(request)
+ self.scheduler.schedule()
+ manager = self.scheduler._diffusion_kv_manager
+ assert manager is not None
+ metadata = manager.get_metadata("preempted")
+
+ assert self.scheduler.preempt_request("preempted") is True
+ resumed_output = self.scheduler.schedule()
+
+ assert resumed_output.scheduled_new_reqs == []
+ assert _cached_ids(resumed_output) == ["preempted"]
+ assert manager.get_metadata("preempted") == metadata
+
+ @pytest.mark.parametrize(
+ "status",
+ [
+ DiffusionRequestStatus.FINISHED_COMPLETED,
+ DiffusionRequestStatus.FINISHED_ABORTED,
+ DiffusionRequestStatus.FINISHED_ERROR,
+ ],
+ )
+ def test_diffusion_kv_terminal_status_releases_allocation(self, status: DiffusionRequestStatus) -> None:
+ _initialize_paged_scheduler(self.scheduler, num_blocks=3)
+ request = _make_request("terminal")
+ _attach_diffusion_kv(request)
+ self.scheduler.add_request(request)
+ self.scheduler.schedule()
+ manager = self.scheduler._diffusion_kv_manager
+ assert manager is not None
+ assert manager.has_request("terminal") is True
+
+ self.scheduler.finish_requests("terminal", status)
+
+ assert manager.has_request("terminal") is False
def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
- self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+ _initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
@@ -384,7 +656,7 @@ class TestRequestScheduler:
],
)
def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
- self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+ _initialize_paged_scheduler(self.scheduler)
request = _make_request("diffusion-kv")
owner = request if owner_name == "request" else request.sampling_params
setattr(owner, field_name, object())
@@ -501,6 +773,19 @@ class TestRequestScheduler:
assert third.num_running_reqs == 1
assert third.num_waiting_reqs == 0
+ def test_records_initial_scheduler_queue_wait(self, mocker: MockerFixture) -> None:
+ perf_counter = mocker.patch(
+ "vllm_omni.diffusion.sched.base_scheduler.time.perf_counter",
+ side_effect=[10.0, 10.125],
+ )
+ request = _make_step_request("queue-wait", num_inference_steps=2)
+
+ self.scheduler.add_request(request)
+ self.scheduler.schedule()
+
+ assert request.scheduler_queue_wait_ms == pytest.approx(125.0)
+ assert perf_counter.call_count == 2
+
def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -602,6 +887,38 @@ class TestRequestScheduler:
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
+ def test_batches_incompatible_pipeline_conditions_separately(self) -> None:
+ scheduler = RequestScheduler()
+ scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+ request_a = OmniDiffusionRequest(
+ prompt="a",
+ sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123),
+ request_id="a",
+ batch_compatibility_key=("wan22_s2v_condition", True),
+ )
+ request_b = OmniDiffusionRequest(
+ prompt="b",
+ sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456),
+ request_id="b",
+ batch_compatibility_key=("wan22_s2v_condition", False),
+ )
+ scheduler.add_request(request_a)
+ scheduler.add_request(request_b)
+
+ first = scheduler.schedule()
+
+ assert _new_ids(first) == ["a"]
+ assert first.num_running_reqs == 1
+ assert first.num_waiting_reqs == 1
+
+ scheduler.update_from_output(first, _make_request_output("a"))
+ second = scheduler.schedule()
+
+ assert _new_ids(second) == ["b"]
+ assert second.num_running_reqs == 1
+ assert second.num_waiting_reqs == 0
+
def test_batches_different_request_local_seed_together(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -625,6 +942,47 @@ class TestRequestScheduler:
assert first.num_running_reqs == 2
assert first.num_waiting_reqs == 0
+ @pytest.mark.parametrize(
+ ("first_extra_args", "second_extra_args"),
+ [
+ ({"sample_solver": "unipc"}, {"sample_solver": "euler"}),
+ ({"flow_shift": 3.0}, {"flow_shift": 5.0}),
+ ({"sample_solver": "unipc"}, {}),
+ ({"flow_shift": 3.0}, {}),
+ ],
+ )
+ def test_batches_incompatible_wan_scheduler_structure_separately(
+ self,
+ first_extra_args: dict,
+ second_extra_args: dict,
+ ) -> None:
+ scheduler = RequestScheduler()
+ scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+ first_id = scheduler.add_request(
+ _make_step_request(
+ "a",
+ sampling_params=OmniDiffusionSamplingParams(
+ num_inference_steps=2,
+ extra_args=first_extra_args,
+ ),
+ )
+ )
+ scheduler.add_request(
+ _make_step_request(
+ "b",
+ sampling_params=OmniDiffusionSamplingParams(
+ num_inference_steps=2,
+ extra_args=second_extra_args,
+ ),
+ )
+ )
+
+ first = scheduler.schedule()
+
+ assert _new_ids(first) == [first_id]
+ assert first.num_running_reqs == 1
+ assert first.num_waiting_reqs == 1
+
def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
@@ -789,6 +1147,45 @@ class TestDiffusionEngine:
assert engine.scheduler.max_num_running_reqs == 1
fake_executor_cls.assert_called_once_with(od_config)
+ def test_initializes_paged_scheduler_from_native_control_plane(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ mocker: MockerFixture,
+ ) -> None:
+ od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
+ fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
+ scheduler = mocker.Mock()
+ kv_cache_config = object()
+ kv_vllm_config = object()
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+ lambda *args, **kwargs: None,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+ lambda *args, **kwargs: None,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+ lambda *args, **kwargs: fake_executor_cls,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.diffusion_engine.initialize_diffusion_kv_control_plane",
+ lambda *args, **kwargs: (kv_cache_config, 16, 16, kv_vllm_config),
+ )
+
+ engine = DiffusionEngine(od_config, scheduler=scheduler)
+
+ assert engine.scheduler is scheduler
+ scheduler.initialize.assert_called_once_with(
+ od_config,
+ kv_cache_config=kv_cache_config,
+ scheduler_block_size=16,
+ hash_block_size=16,
+ kv_vllm_config=kv_vllm_config,
+ )
+
def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
scheduler = Scheduler()
scheduler.initialize(SimpleNamespace())
@@ -855,6 +1252,21 @@ class TestDiffusionEngine:
assert output.aborted is True
assert output.abort_message == "Request req-finalize aborted."
+ def test_finalize_finished_request_returns_scheduler_admission_error(self) -> None:
+ engine = DiffusionEngine.__new__(DiffusionEngine)
+ engine.scheduler = StepScheduler()
+ engine.scheduler.initialize(SimpleNamespace())
+
+ req_id = engine.scheduler.add_request(_make_request("req-error"))
+ engine.scheduler._finish_requests(
+ {req_id: DiffusionRequestStatus.FINISHED_ERROR},
+ {req_id: "KV request cannot fit"},
+ )
+
+ output = engine._finalize_finished_request(req_id)
+
+ assert output.error == "KV request cannot fit"
+
@pytest.mark.asyncio
async def test_streaming_runner_output_notifies_each_chunk(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
@@ -1160,6 +1572,39 @@ class TestStepScheduler:
assert sched_output.num_running_reqs == 2
assert sched_output.num_waiting_reqs == 0
+ def test_mixed_batch_error_finishes_only_failed_request(self) -> None:
+ scheduler = StepScheduler()
+ scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+ req_bad = scheduler.add_request(_make_step_request("bad", num_inference_steps=2))
+ req_good = scheduler.add_request(_make_step_request("good", num_inference_steps=2))
+ first = scheduler.schedule()
+
+ finished = scheduler.update_from_output(
+ first,
+ BatchRunnerOutput.from_list(
+ [
+ _make_step_output(
+ req_bad,
+ step_index=0,
+ finished=True,
+ error="invalid request input",
+ ),
+ _make_step_output(req_good, step_index=1),
+ ]
+ ),
+ )
+
+ assert finished == {req_bad}
+ assert scheduler.get_request_state(req_bad).status == DiffusionRequestStatus.FINISHED_ERROR
+ assert scheduler.get_request_state(req_good).status == DiffusionRequestStatus.RUNNING
+ assert scheduler.get_request_state(req_good).req.sampling_params.step_index == 1
+
+ second = scheduler.schedule()
+ assert second.finished_req_ids == {req_bad}
+ assert _new_ids(second) == []
+ assert _cached_ids(second) == [req_good]
+
def test_step_batch_rejects_different_sampling_key(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py
index 0964833d7..94c3eb48f 100644
--- a/tests/diffusion/test_diffusion_step_pipeline.py
+++ b/tests/diffusion/test_diffusion_step_pipeline.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for step-level diffusion execution across runner / worker / executor / engine."""
import contextlib
@@ -17,6 +17,8 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
from vllm_omni.diffusion.distributed.parallel_state import (
@@ -45,6 +47,7 @@ from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.errors import OmniClientError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -100,6 +103,22 @@ class _StepPipeline:
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
+class _PerRequestErrorStepPipeline(_StepPipeline):
+ def prepare_encode(self, state, **kwargs):
+ if state.prompt == "fail-prepare":
+ raise OmniClientError(
+ "invalid request input",
+ status_code=422,
+ error_type="UnprocessableEntityError",
+ )
+ return super().prepare_encode(state, **kwargs)
+
+ def denoise_step(self, input_batch, **kwargs):
+ del kwargs
+ self.denoise_calls += 1
+ return torch.ones_like(input_batch.latents)
+
+
class _ProfilingStepPipeline(_StepPipeline):
enable_diffusion_pipeline_profiler = True
@@ -306,11 +325,14 @@ def _make_vllm_config():
)
-def _make_runner():
+def _make_runner(
+ diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
+):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = _make_vllm_config()
runner.od_config = SimpleNamespace(
cache_backend=None,
+ diffusion_kv_mode=diffusion_kv_mode,
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
@@ -330,6 +352,7 @@ def _make_distributed_runner(mode: str, device: torch.device):
runner.vllm_config = _make_vllm_config()
runner.od_config = SimpleNamespace(
cache_backend=None,
+ diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
parallel_config=SimpleNamespace(use_hsdp=False),
streaming_output=False,
)
@@ -502,6 +525,81 @@ def test_step_profiler_reports_denoise_step_as_diffuse():
class TestRunner:
"""DiffusionModelRunner.execute_stepwise"""
+ @pytest.mark.parametrize(
+ ("diffusion_kv_mode", "should_cleanup"),
+ [
+ (DiffusionKVCacheMode.DENSE_LEGACY, False),
+ (DiffusionKVCacheMode.PAGED_SCHEDULER, True),
+ ],
+ )
+ def test_scheduler_finished_cleanup_only_runs_for_paged_kv(
+ self,
+ mocker: MockerFixture,
+ diffusion_kv_mode: DiffusionKVCacheMode,
+ should_cleanup: bool,
+ ):
+ runner = _make_runner(diffusion_kv_mode)
+ cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
+ expected_output = object()
+ mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output)
+ scheduler_output = DiffusionSchedulerOutput(
+ step_id=1,
+ scheduled_new_reqs=[],
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
+ finished_req_ids={"finished-request"},
+ num_running_reqs=0,
+ num_waiting_reqs=0,
+ )
+
+ output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+ assert output is expected_output
+ if should_cleanup:
+ cleanup.assert_called_once_with(["finished-request"])
+ else:
+ cleanup.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("diffusion_kv_mode", "should_cleanup"),
+ [
+ (DiffusionKVCacheMode.DENSE_LEGACY, False),
+ (DiffusionKVCacheMode.PAGED_SCHEDULER, True),
+ ],
+ )
+ def test_terminal_cleanup_only_runs_for_paged_kv(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ mocker: MockerFixture,
+ diffusion_kv_mode: DiffusionKVCacheMode,
+ should_cleanup: bool,
+ ):
+ runner = _make_runner(diffusion_kv_mode)
+ cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
+ install = mocker.patch.object(runner, "install_diffusion_kv_metadata")
+ req = _make_step_request()
+ scheduler_output = _make_scheduler_output(req, step_id=0)
+ if should_cleanup:
+ scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata(
+ request_id=req.request_id,
+ allocation_generation=1,
+ sequences=(),
+ )
+ monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+ first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+ assert first_output.get_request_output(req.request_id).finished is False
+ cleanup.assert_not_called()
+
+ output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+
+ assert output.get_request_output(req.request_id).finished is True
+ if should_cleanup:
+ install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata)
+ cleanup.assert_called_once_with([req.request_id])
+ else:
+ install.assert_not_called()
+ cleanup.assert_not_called()
+
def test_completes_request_and_clears_state(self, monkeypatch):
runner = _make_runner()
req = _make_step_request()
@@ -616,6 +714,72 @@ class TestRunner:
result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
assert len(result) == 2
+ def test_prepare_error_isolated_from_healthy_request(self, monkeypatch):
+ runner = _make_runner()
+ runner.pipeline = _PerRequestErrorStepPipeline()
+ bad_req = _make_step_request()
+ bad_req.request_id = "req-bad"
+ bad_req.prompt = "fail-prepare"
+ good_req = _make_step_request()
+ good_req.request_id = "req-good"
+ good_req.prompt = "valid"
+ monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+ result = DiffusionModelRunner.execute_stepwise(
+ runner,
+ _make_batch_scheduler_output([bad_req, good_req]),
+ )
+
+ bad_output = result.get_request_output("req-bad")
+ assert bad_output.finished is True
+ assert bad_output.result is not None
+ assert bad_output.result.error == "invalid request input"
+ assert bad_output.result.error_status_code == 422
+ assert bad_output.result.error_type == "UnprocessableEntityError"
+ good_output = result.get_request_output("req-good")
+ assert good_output.finished is False
+ assert good_output.step_index == 1
+ assert good_output.result is None
+ assert set(runner.state_cache) == {"req-good"}
+ assert runner.pipeline.denoise_calls == 1
+ assert runner.pipeline.scheduler_calls == 1
+
+ def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch):
+ """P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep)
+ can raise on rank 0 while non-zero ranks return cleanly and then hang
+ on a downstream ``dist.broadcast``. Beyond the pipeline-side guard,
+ the runner all-reduces a per-request failure flag across the DiT
+ group so a rank that ``prepare_encode`` returned cleanly on still
+ skips the request and does not enter the denoise step batch.
+ """
+ runner = _make_runner()
+ runner.pipeline = _StepPipeline() # locally always succeeds
+ monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+ # Simulate a peer DiT rank having reported failure. We patch the
+ # helper directly (rather than ``torch.distributed`` broadly) because
+ # ``execute_stepwise`` also queries ``torch.distributed.get_rank``.
+ monkeypatch.setattr(
+ model_runner_module,
+ "_dit_any_rank_failed",
+ lambda local_failed: True,
+ )
+
+ req = _make_step_request(num_inference_steps=1)
+ result = DiffusionModelRunner.execute_stepwise(
+ runner,
+ _make_scheduler_output(req),
+ )
+
+ output = result.get_request_output("req-1")
+ assert output.finished is True
+ assert output.result is not None
+ assert output.result.error is not None
+ assert "another DiT rank" in output.result.error
+ # A rank that skipped the request must not have driven a denoise step.
+ assert runner.pipeline.denoise_calls == 0
+ assert runner.state_cache == {}
+
def test_receives_kv_payload_before_prepare_encode(self, monkeypatch):
runner = _make_runner()
captured: dict[str, object] = {}
@@ -715,7 +879,11 @@ class TestRunner:
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
- monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+ monkeypatch.setattr(
+ model_runner_module,
+ "enable_offload_backend",
+ lambda _config, pipeline, **_kwargs: (pipeline, None),
+ )
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
with pytest.raises(ValueError, match="RequestOnlyPipeline"):
diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py
index d84f49b31..c833ae7a1 100644
--- a/tests/diffusion/test_diffusion_streaming_output.py
+++ b/tests/diffusion/test_diffusion_streaming_output.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Integration tests for diffusion pipeline streaming output mode.
Covers:
@@ -311,8 +311,8 @@ class TestPipelineStreamingOutputToEntrypoint:
lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(),
)
mocker.patch(
- "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
- return_value="diffusion",
+ "vllm_omni.entrypoints.openai.serving_video_output_stream.is_video_generation_pipeline",
+ return_value=True,
)
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
@@ -441,7 +441,6 @@ class TestPipelineStreamingOutputToEntrypoint:
model="test_model",
od_config=od_config,
metadata=metadata,
- batch_size=1,
)
@staticmethod
@@ -526,6 +525,7 @@ class TestSupportedPipelines:
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = _make_vllm_config()
runner.od_config = SimpleNamespace(
+ max_num_seqs=1,
enable_cpu_offload=False,
enable_layerwise_offload=False,
enforce_eager=True,
@@ -544,7 +544,11 @@ class TestSupportedPipelines:
monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader)
monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler)
- monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+ monkeypatch.setattr(
+ model_runner_module,
+ "enable_offload_backend",
+ lambda _config, pipeline, **_kwargs: (pipeline, None),
+ )
monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
with pytest.raises(ValueError, match="NoStepPipeline"):
diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py
index 6a49d0909..0d8d5e92a 100644
--- a/tests/diffusion/test_diffusion_worker.py
+++ b/tests/diffusion/test_diffusion_worker.py
@@ -13,11 +13,7 @@ import pytest
import torch
from pytest_mock import MockerFixture
-from vllm_omni.diffusion.worker.diffusion_worker import (
- DiffusionWorker,
- _create_diffusion_worker_vllm_config,
- _make_diffusion_vllm_model_config,
-)
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
@@ -58,34 +54,6 @@ def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
return worker
-def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
- from types import SimpleNamespace
-
- from vllm_omni.quantization import build_quant_config
-
- od_config = SimpleNamespace(
- model="dummy",
- dtype=torch.bfloat16,
- quantization_config=build_quant_config(
- {
- "quant_method": "modelopt",
- "quant_algo": "FP8",
- "ignore": [],
- }
- ),
- tf_model_config=SimpleNamespace(),
- enforce_eager=True,
- is_moe=False,
- )
-
- model_config = _make_diffusion_vllm_model_config(od_config)
-
- assert model_config.dtype is torch.bfloat16
- assert model_config.quantization == "modelopt"
- assert model_config.quantization_config is od_config.quantization_config
- assert model_config.is_quantized()
-
-
class TestDiffusionWorkerSleep:
"""Test DiffusionWorker.sleep method."""
@@ -329,49 +297,3 @@ class TestDiffusionWorkerWakeUp:
mock_buffer2.data.copy_.assert_not_called()
assert result is True
-
-
-class TestWorkerVllmConfigAdditionalConfig:
- """Test worker-side VllmConfig construction with additional_config."""
-
- def test_create_diffusion_worker_vllm_config_passes_additional_config(self, mocker: MockerFixture):
- mock_vllm_config = mocker.Mock()
- mock_vllm_cls = mocker.patch(
- "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
- return_value=mock_vllm_config,
- )
- od_config = mocker.Mock(additional_config={"torchair_graph_config": {"enabled": True}})
-
- result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
- assert result is mock_vllm_config
- assert mock_vllm_cls.call_args.kwargs["additional_config"] == od_config.additional_config
-
- def test_create_diffusion_worker_vllm_config_falls_back_when_constructor_rejects_additional_config(
- self, mocker: MockerFixture
- ):
- mock_vllm_config = mocker.Mock()
- mock_vllm_cls = mocker.patch(
- "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
- side_effect=[
- TypeError("VllmConfig.__init__() got an unexpected keyword argument 'additional_config'"),
- mock_vllm_config,
- ],
- )
- od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
- result = _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
-
- assert result is mock_vllm_config
- assert mock_vllm_cls.call_count == 2
- assert getattr(mock_vllm_config, "additional_config") == od_config.additional_config
-
- def test_create_diffusion_worker_vllm_config_reraises_other_type_errors(self, mocker: MockerFixture):
- mocker.patch(
- "vllm_omni.diffusion.worker.diffusion_worker.VllmConfig",
- side_effect=TypeError("additional_config values must be hashable"),
- )
- od_config = mocker.Mock(additional_config={"ascend_scheduler_config": {"foo": "bar"}})
-
- with pytest.raises(TypeError, match="hashable"):
- _create_diffusion_worker_vllm_config(torch.device("cpu"), od_config)
diff --git a/tests/diffusion/test_diffusion_worker_broadcast.py b/tests/diffusion/test_diffusion_worker_broadcast.py
new file mode 100644
index 000000000..750d4c30d
--- /dev/null
+++ b/tests/diffusion/test_diffusion_worker_broadcast.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+import time
+from unittest.mock import Mock
+
+import pytest
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
+
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
+from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _DelayedMultipartSocket:
+ def __init__(
+ self,
+ socket,
+ send_started: threading.Event,
+ release_send: threading.Event,
+ ):
+ self._socket = socket
+ self._send_started = send_started
+ self._release_send = release_send
+ self._delay_next_send = True
+
+ def __getattr__(self, name):
+ return getattr(self._socket, name)
+
+ def send_multipart(self, *args, **kwargs):
+ if self._delay_next_send:
+ self._delay_next_send = False
+ self._send_started.set()
+ if not self._release_send.wait(timeout=2):
+ raise TimeoutError("Timed out waiting to release delayed overflow payload")
+ return self._socket.send_multipart(*args, **kwargs)
+
+
+class _ObservedPollSocket:
+ def __init__(self, socket, poll_started: threading.Event):
+ self._socket = socket
+ self._poll_started = poll_started
+
+ def __getattr__(self, name):
+ return getattr(self._socket, name)
+
+ def poll(self, *args, **kwargs):
+ self._poll_started.set()
+ return self._socket.poll(*args, **kwargs)
+
+
+class _RecordingWorker:
+ def __init__(self):
+ self.calls: list[str] = []
+
+ def execute_method(self, method, *args, **kwargs):
+ assert method == "record"
+ self.calls.append(args[0])
+
+
+def test_recv_message_waits_indefinitely():
+ worker_proc = WorkerProc.__new__(WorkerProc)
+ worker_proc.mq = Mock()
+ worker_proc.mq.dequeue.return_value = {"type": "rpc"}
+
+ result = worker_proc.recv_message()
+
+ assert result == {"type": "rpc"}
+ worker_proc.mq.dequeue.assert_called_once_with(indefinite=True)
+
+
+def test_busy_loop_releases_failed_rpc_traceback_and_device_cache(monkeypatch):
+ original = RuntimeError("model execution failed")
+ worker_proc = WorkerProc.__new__(WorkerProc)
+ worker_proc.gpu_id = 0
+ worker_proc.mq = Mock()
+ worker_proc.mq.dequeue.side_effect = [
+ {
+ "type": "rpc",
+ "method": "execute_model",
+ "args": (),
+ "kwargs": {},
+ "output_rank": 0,
+ "exec_all_ranks": False,
+ "collect_rank_status": False,
+ },
+ {"type": "shutdown"},
+ ]
+ worker_proc.worker = Mock()
+ worker_proc.worker.execute_method.side_effect = original
+ worker_proc.result_mq = object()
+ worker_proc._return_result = Mock()
+ worker_proc.wake_event = None
+ worker_proc._running = True
+
+ gc_collect = Mock()
+ mock_platform = Mock()
+ monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect)
+ monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform)
+
+ worker_proc._worker_busy_loop()
+
+ worker_proc._return_result.assert_called_once_with(
+ {"status": "error", "error": "model execution failed", "wave_id": None}
+ )
+ assert original.__traceback__ is None
+ gc_collect.assert_called_once_with()
+ mock_platform.empty_cache.assert_called_once_with()
+
+
+def test_busy_loop_preserves_order_when_overflow_payload_is_delayed(monkeypatch):
+ """A consumed overflow marker must not be abandoned on a socket timeout."""
+ writer = MessageQueue(
+ n_reader=1,
+ n_local_reader=1,
+ local_reader_ranks=[0],
+ max_chunk_bytes=1024,
+ max_chunks=4,
+ )
+ reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
+ writer.wait_until_ready()
+ reader.wait_until_ready()
+
+ overflow_send_started = threading.Event()
+ release_overflow_send = threading.Event()
+ overflow_poll_started = threading.Event()
+ writer.local_socket = _DelayedMultipartSocket(
+ writer.local_socket,
+ overflow_send_started,
+ release_overflow_send,
+ )
+ reader.local_socket = _ObservedPollSocket(reader.local_socket, overflow_poll_started)
+
+ recording_worker = _RecordingWorker()
+ worker_proc = WorkerProc.__new__(WorkerProc)
+ worker_proc.gpu_id = 0
+ worker_proc.mq = reader
+ worker_proc.worker = recording_worker
+ worker_proc.result_mq = None
+ worker_proc.wake_event = None
+ worker_proc._running = True
+
+ # Make the pre-fix implementation fail quickly: it read this module-level
+ # value directly from _worker_busy_loop. The fixed implementation ignores it
+ # and delegates to recv_message(), which waits indefinitely.
+ monkeypatch.setattr(
+ diffusion_worker_module,
+ "_BROADCAST_DEQUEUE_TIMEOUT_S",
+ 0.02,
+ raising=False,
+ )
+
+ overflow_message = {
+ "type": "rpc",
+ "method": "record",
+ "args": ("overflow-1",),
+ "padding": b"x" * 2048,
+ }
+ inline_message = {
+ "type": "rpc",
+ "method": "record",
+ "args": ("inline-2",),
+ }
+ shutdown_message = {"type": "shutdown"}
+ thread_errors: list[BaseException] = []
+
+ def send_messages():
+ try:
+ writer.enqueue(overflow_message)
+ writer.enqueue(inline_message)
+ writer.enqueue(shutdown_message)
+ except BaseException as exc:
+ thread_errors.append(exc)
+
+ def run_busy_loop():
+ try:
+ worker_proc._worker_busy_loop()
+ except BaseException as exc:
+ thread_errors.append(exc)
+
+ writer_thread = threading.Thread(target=send_messages, daemon=True)
+ worker_thread = threading.Thread(target=run_busy_loop, daemon=True)
+
+ try:
+ writer_thread.start()
+ worker_thread.start()
+
+ assert overflow_send_started.wait(timeout=2)
+ assert overflow_poll_started.wait(timeout=2)
+
+ # Keep the multipart unavailable beyond the pre-fix 20 ms timeout.
+ time.sleep(0.1)
+ release_overflow_send.set()
+
+ writer_thread.join(timeout=2)
+ worker_thread.join(timeout=2)
+
+ assert not writer_thread.is_alive()
+ assert not worker_thread.is_alive()
+ assert thread_errors == []
+ assert recording_worker.calls == ["overflow-1", "inline-2"]
+ finally:
+ release_overflow_send.set()
+ worker_proc._running = False
+ reader.shutdown()
+ writer.shutdown()
+ writer_thread.join(timeout=1)
+ worker_thread.join(timeout=1)
diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py
new file mode 100644
index 000000000..83863bb86
--- /dev/null
+++ b/tests/diffusion/test_diffusion_worker_process_title.py
@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+import shutil
+import subprocess
+import sys
+import textwrap
+import time
+from dataclasses import dataclass
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@dataclass(frozen=True)
+class _FakeGroup:
+ world_size: int = 1
+ rank_in_group: int = 0
+
+
+_GROUP_GETTERS = (
+ "get_dp_group",
+ "get_pp_group",
+ "get_sp_group",
+ "get_cfg_group",
+ "get_tp_group",
+ "get_fs_group",
+ "get_hsdp_replicate_group",
+ "get_ep_group",
+)
+
+
+def _patch_groups(
+ mocker: MockerFixture,
+ groups: dict[str, _FakeGroup],
+) -> dict[str, Mock]:
+ patched_getters = {}
+ for getter_name in _GROUP_GETTERS:
+ patched_getters[getter_name] = mocker.patch.object(
+ worker_module,
+ getter_name,
+ return_value=groups.get(getter_name, _FakeGroup()),
+ )
+ return patched_getters
+
+
+def test_setup_uses_default_name_before_model_parallel_init(
+ mocker: MockerFixture,
+) -> None:
+ mocker.patch.object(
+ worker_module,
+ "model_parallel_is_initialized",
+ return_value=False,
+ )
+ group_getters = _patch_groups(mocker, {})
+ set_process_title = mocker.patch.object(worker_module, "set_process_title")
+ decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+ worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True)
+
+ set_process_title.assert_called_once_with(
+ name="DiffusionWorker",
+ prefix="vLLM-Omni",
+ )
+ decorate_logs.assert_called_once_with("DiffusionWorker")
+ for group_getter in group_getters.values():
+ group_getter.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
+ [
+ ({}, False, False, 1, "DiffusionWorker"),
+ ({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"),
+ ({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"),
+ ({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"),
+ (
+ {
+ "get_cfg_group": _FakeGroup(2, 1),
+ "get_tp_group": _FakeGroup(2, 1),
+ },
+ False,
+ False,
+ 1,
+ "DiffusionWorker_CFG1_TP1",
+ ),
+ (
+ {
+ "get_dp_group": _FakeGroup(2, 1),
+ "get_sp_group": _FakeGroup(4, 2),
+ "get_tp_group": _FakeGroup(2, 1),
+ },
+ False,
+ False,
+ 1,
+ "DiffusionWorker_DP1_SP2_TP1",
+ ),
+ (
+ {"get_fs_group": _FakeGroup(4, 3)},
+ False,
+ True,
+ 1,
+ "DiffusionWorker_FS3",
+ ),
+ (
+ {
+ "get_sp_group": _FakeGroup(2, 1),
+ "get_fs_group": _FakeGroup(4, 3),
+ },
+ False,
+ True,
+ 1,
+ "DiffusionWorker_SP1_FS3",
+ ),
+ (
+ {
+ "get_fs_group": _FakeGroup(1, 0),
+ "get_hsdp_replicate_group": _FakeGroup(2, 1),
+ },
+ False,
+ True,
+ 2,
+ "DiffusionWorker_RP1",
+ ),
+ ({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"),
+ ({}, True, False, 1, "DiffusionWorker"),
+ ],
+ ids=[
+ "all-singleton",
+ "tp-only",
+ "dp-only",
+ "pp-only",
+ "tp-cfg",
+ "dp-sp-tp",
+ "hsdp-only",
+ "hsdp-sp",
+ "hsdp-replicated",
+ "expert-parallel",
+ "singleton-ep",
+ ],
+)
+def test_setup_uses_initialized_parallel_groups(
+ mocker: MockerFixture,
+ groups: dict[str, _FakeGroup],
+ enable_ep: bool,
+ use_hsdp: bool,
+ hsdp_replicate_size: int,
+ expected_name: str,
+) -> None:
+ mocker.patch.object(
+ worker_module,
+ "model_parallel_is_initialized",
+ return_value=True,
+ )
+ group_getters = _patch_groups(mocker, groups)
+ set_process_title = mocker.patch.object(worker_module, "set_process_title")
+ decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+ worker_module._setup_diffusion_worker_proc_title_and_log_prefix(
+ enable_ep=enable_ep,
+ use_hsdp=use_hsdp,
+ hsdp_replicate_size=hsdp_replicate_size,
+ )
+
+ set_process_title.assert_called_once_with(
+ name=expected_name,
+ prefix="vLLM-Omni",
+ )
+ decorate_logs.assert_called_once_with(expected_name)
+ if use_hsdp:
+ group_getters["get_fs_group"].assert_called_once_with()
+ if hsdp_replicate_size > 1:
+ group_getters["get_hsdp_replicate_group"].assert_called_once_with()
+ else:
+ group_getters["get_hsdp_replicate_group"].assert_not_called()
+ else:
+ group_getters["get_fs_group"].assert_not_called()
+ group_getters["get_hsdp_replicate_group"].assert_not_called()
+ if enable_ep:
+ group_getters["get_ep_group"].assert_called_once_with()
+ else:
+ group_getters["get_ep_group"].assert_not_called()
+
+
+def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
+ mocker.patch.object(
+ worker_module,
+ "model_parallel_is_initialized",
+ return_value=False,
+ )
+ mocker.patch.dict(sys.modules, {"setproctitle": None})
+ decorate_logs = mocker.patch.object(worker_module, "decorate_logs")
+
+ worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
+
+ decorate_logs.assert_called_once_with("DiffusionWorker")
+
+
+@pytest.mark.skipif(
+ sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None,
+ reason="requires Linux ps and setproctitle",
+)
+def test_process_title_is_visible_through_ps() -> None:
+ expected_name = "vLLM-Omni::DiffusionWorker_TP1"
+ child_script = textwrap.dedent(
+ """
+ import time
+
+ from vllm.utils.system_utils import set_process_title
+
+
+ set_process_title(
+ name="DiffusionWorker_TP1",
+ prefix="vLLM-Omni",
+ )
+ time.sleep(30)
+ """
+ )
+ process = subprocess.Popen(
+ [sys.executable, "-c", child_script],
+ stderr=subprocess.PIPE,
+ text=True,
+ )
+
+ observed_title = ""
+ try:
+ for _ in range(100):
+ if process.poll() is not None:
+ stderr = process.stderr.read() if process.stderr is not None else ""
+ pytest.fail(f"child process exited before title check: {stderr}")
+ result = subprocess.run(
+ ["ps", "-o", "args=", "-p", str(process.pid)],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+ observed_title = result.stdout.strip()
+ if expected_name in observed_title:
+ break
+ time.sleep(0.1)
+ else:
+ pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}")
+ finally:
+ process.terminate()
+ try:
+ process.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ process.wait(timeout=5)
diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py
index 824c1483a..e39c18ff4 100644
--- a/tests/diffusion/test_inline_stage_diffusion_client.py
+++ b/tests/diffusion/test_inline_stage_diffusion_client.py
@@ -47,11 +47,16 @@ def client(mock_engine):
with patch.object(InlineStageDiffusionClient, "_enrich_config"):
od_config = MagicMock(spec=OmniDiffusionConfig)
od_config.streaming_output = False
- c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1)
+ c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata)
yield c
c.shutdown()
+def test_inline_client_implements_shared_prompt_hook_protocol(client):
+ assert client.prompt_transform_func is None
+ assert client.prompt_expand_func is None
+
+
@pytest.mark.asyncio
async def test_inline_dispatch_request_success(client, mock_engine):
mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
@@ -229,5 +234,4 @@ def test_inline_client_requires_replica_id(mock_engine):
model="test_model",
od_config=od_config,
metadata=metadata,
- batch_size=1,
)
diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py
index aa40d0770..855f38c4a 100644
--- a/tests/diffusion/test_multiproc_engine_concurrency.py
+++ b/tests/diffusion/test_multiproc_engine_concurrency.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import multiprocessing as mp
@@ -14,6 +14,7 @@ import torch
import zmq
from vllm.v1.engine.exceptions import EngineDeadError
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
@@ -413,6 +414,26 @@ class TestRequestModeDispatch:
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
+ def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
+ executor, _, _ = _make_executor(num_gpus=2)
+ executor.od_config = SimpleNamespace(
+ step_execution=False,
+ parallel_config=SimpleNamespace(data_parallel_size=2),
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=True,
+ )
+ executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+ scheduler_output = _make_sched_output("A", "B")
+ for new_req in scheduler_output.scheduled_new_reqs:
+ new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
+
+ result = executor.execute_request(scheduler_output)
+
+ assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+ forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
+ assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
+ assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
+
@pytest.mark.parametrize(
("field", "value"),
[
@@ -802,6 +823,31 @@ class TestWorkerProcRpcRankStatus:
assert status["bool_result"] is None
assert "local boom" in status["traceback"]
+ def test_execute_rpc_collect_exception_releases_traceback_and_device_cache(self, monkeypatch):
+ proc = self._make_worker_proc()
+ original = RuntimeError("local boom")
+ proc.worker.execute_method = Mock(side_effect=original)
+ gc_collect = Mock()
+ mock_platform = Mock()
+ monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+ monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect)
+ monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform)
+
+ proc._execute_rpc(
+ {
+ "method": "add_lora",
+ "args": (),
+ "kwargs": {},
+ "output_rank": 0,
+ "exec_all_ranks": True,
+ "collect_rank_status": True,
+ }
+ )
+
+ assert original.__traceback__ is None
+ gc_collect.assert_called_once_with()
+ mock_platform.empty_cache.assert_called_once_with()
+
def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
proc = self._make_worker_proc()
@@ -1033,6 +1079,42 @@ class TestStageDiffusionClientErrorPropagation:
assert client.get_diffusion_output_nowait() is None
+ def test_error_response_preserves_scheduler_metrics(self):
+ from vllm_omni.metrics import definitions as metric_defs
+
+ client = self._make_client()
+ client._response_socket.recv.side_effect = [b"message", zmq.Again()]
+ client._decoder.decode.return_value = {
+ "type": "error",
+ "request_id": "req-error",
+ "error": "gpu fault",
+ "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
+ }
+
+ output = client.get_diffusion_output_nowait()
+
+ assert output is not None
+ assert output.error == "gpu fault"
+ assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
+
+ def test_metrics_response_creates_metrics_only_output(self):
+ from vllm_omni.metrics import definitions as metric_defs
+ from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
+
+ client = self._make_client()
+ client._response_socket.recv.side_effect = [b"message", zmq.Again()]
+ client._decoder.decode.return_value = {
+ "type": "metrics",
+ "metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
+ }
+
+ output = client.get_diffusion_output_nowait()
+
+ assert output is not None
+ assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID
+ assert output.error is None
+ assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
+
def test_check_health_raises_when_proc_dead(self):
"""``check_health`` detects a dead subprocess via the manager's proc
and raises ``EngineDeadError``, setting ``_engine_dead`` as a
@@ -1089,7 +1171,6 @@ class TestStageDiffusionClientErrorPropagation:
metadata,
"tcp://req",
"tcp://resp",
- batch_size=1,
)
@pytest.mark.asyncio
diff --git a/tests/diffusion/test_omni_diffusion_config_master_port.py b/tests/diffusion/test_omni_diffusion_config_master_port.py
index 8addd1c8c..0fe2dd08d 100644
--- a/tests/diffusion/test_omni_diffusion_config_master_port.py
+++ b/tests/diffusion/test_omni_diffusion_config_master_port.py
@@ -1,22 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for OmniDiffusionConfig master_port resolution (issue #3794)."""
import socket
import pytest
+from tests.helpers.runtime import get_open_port
from vllm_omni.diffusion.data import OmniDiffusionConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-def _get_available_port() -> int:
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
- sock.bind(("", 0))
- return sock.getsockname()[1]
-
-
def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket:
for _ in range(100):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -36,27 +31,27 @@ def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket:
class TestOmniDiffusionConfigMasterPort:
def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
- port = _get_available_port()
+ port = get_open_port()
monkeypatch.setenv("MASTER_PORT", str(port))
config = OmniDiffusionConfig(model="test")
assert config.master_port == port
def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("MASTER_PORT", raising=False)
- port = _get_available_port()
+ port = get_open_port()
config = OmniDiffusionConfig(model="test", master_port=port)
assert config.master_port == port
def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
- env_port = _get_available_port()
- kwarg_port = _get_available_port()
+ env_port = get_open_port()
+ kwarg_port = get_open_port()
monkeypatch.setenv("MASTER_PORT", str(env_port))
config = OmniDiffusionConfig(model="test", master_port=kwarg_port)
assert config.master_port == env_port
def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("MASTER_PORT", raising=False)
- port = _get_available_port()
+ port = get_open_port()
ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)}
assert ports == {port}
diff --git a/tests/diffusion/test_prompt_update.py b/tests/diffusion/test_prompt_update.py
index cdee35979..f08825d1a 100644
--- a/tests/diffusion/test_prompt_update.py
+++ b/tests/diffusion/test_prompt_update.py
@@ -387,7 +387,6 @@ class TestPromptUpdateIntegration:
model="test_model",
od_config=od_config,
metadata=metadata,
- batch_size=1,
)
@staticmethod
diff --git a/tests/diffusion/test_result_pump.py b/tests/diffusion/test_result_pump.py
index cf00cb6cc..fa172fd0c 100644
--- a/tests/diffusion/test_result_pump.py
+++ b/tests/diffusion/test_result_pump.py
@@ -6,6 +6,7 @@ import concurrent.futures
import queue
import threading
import time
+from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
@@ -264,6 +265,119 @@ class TestResultPumpDispatch:
fut = executor.wait_output_ready(async_output_id)
assert fut.done()
assert fut.result(timeout=1.0) is output
+ assert async_output_id not in executor._output_futures
+
+ def test_output_ready_atomic_resolution_when_future_already_waiting(self):
+ """When OUTPUT_READY arrives and a future is already waiting, resolve it directly."""
+ executor = _make_executor()
+ async_output_id = "abc123"
+ output = DiffusionOutput(output="waiting")
+ fut = executor.wait_output_ready(async_output_id)
+ assert not fut.done()
+
+ msg = AsyncDiffusionOutput(
+ kind=AsyncOutputKind.OUTPUT_READY,
+ async_output_id=async_output_id,
+ output=output,
+ )
+ _feed_one_msg_to_pump(executor, msg)
+
+ assert fut.done()
+ assert fut.result(timeout=1.0) is output
+ assert async_output_id not in executor._output_futures
+ assert async_output_id not in executor._completed_outputs
+
+
+class _FakeBatchOutput:
+ """Batch-level output exposing per-request results."""
+
+ def __init__(self, results):
+ self._results = results
+
+ def get_request_output(self, req_id):
+ result = self._results.get(req_id)
+ if result is None:
+ return None
+ return SimpleNamespace(result=result)
+
+
+def _make_scheduler_output(req_ids):
+ return SimpleNamespace(
+ scheduled_new_reqs=[SimpleNamespace(request_id=rid, req=SimpleNamespace()) for rid in req_ids]
+ )
+
+
+class TestBatchSplitDelivery:
+ """execute_batch must resolve per-request futures in either arrival order."""
+
+ @staticmethod
+ def _run(executor, req_ids, batch_id, deliver_early):
+ # Keep execute_batch on the fused request-batch path (not DLO DP).
+ executor.od_config.parallel_config.data_parallel_size = 1
+ executor.od_config.enable_distributed_layerwise_offload = False
+ executor._ensure_open = lambda: None
+
+ outputs = {rid: DiffusionOutput(output=f"img-{rid}") for rid in req_ids}
+ ready = AsyncDiffusionOutput(
+ kind=AsyncOutputKind.OUTPUT_READY,
+ async_output_id=batch_id,
+ output=_FakeBatchOutput(outputs),
+ )
+
+ def fake_collective_rpc(*args, **kwargs):
+ if deliver_early:
+ # Worker's background D2H/SHM thread wins the race: OUTPUT_READY
+ # is pumped before execute_batch registers the split map.
+ _feed_one_msg_to_pump(executor, ready)
+ return AsyncDiffusionOutput(
+ kind=AsyncOutputKind.COMPUTE_DONE,
+ rpc_id="1",
+ async_output_id=batch_id,
+ )
+
+ executor.collective_rpc = fake_collective_rpc
+ batch = executor.execute_batch(_make_scheduler_output(req_ids))
+ if not deliver_early:
+ _feed_one_msg_to_pump(executor, ready)
+ return batch, outputs
+
+ def test_output_ready_after_split_map(self):
+ executor = _make_executor()
+ req_ids = ["r0", "r1", "r2"]
+ _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=False)
+
+ for rid in req_ids:
+ fut = executor.wait_output_ready(f"batch-1/{rid}")
+ assert fut.done()
+ assert fut.result(timeout=1.0) is outputs[rid]
+
+ def test_output_ready_before_split_map(self):
+ """Regression: the whole batch used to be lost, hanging every request."""
+ executor = _make_executor()
+ req_ids = ["r0", "r1", "r2"]
+ _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=True)
+
+ for rid in req_ids:
+ fut = executor.wait_output_ready(f"batch-1/{rid}")
+ assert fut.done(), f"request {rid} never resolved"
+ assert fut.result(timeout=1.0) is outputs[rid]
+
+ # No stale batch-level state left behind.
+ assert executor._batch_split_map == {}
+ assert executor._completed_outputs == {}
+
+ def test_engine_waiter_before_output_is_resolved(self):
+ """Consumers already blocked in wait_output_ready must be woken."""
+ executor = _make_executor()
+ req_ids = ["r0", "r1"]
+ batch_id = "batch-2"
+ futures = {rid: executor.wait_output_ready(f"{batch_id}/{rid}") for rid in req_ids}
+
+ _, outputs = self._run(executor, req_ids, batch_id, deliver_early=True)
+
+ for rid in req_ids:
+ assert futures[rid].done(), f"request {rid} never resolved"
+ assert futures[rid].result(timeout=1.0) is outputs[rid]
class TestShutdownCleansUpFutures:
@@ -304,3 +418,77 @@ class TestShutdownCleansUpFutures:
assert len(executor._rpc_futures) == 0
assert len(executor._output_futures) == 0
+
+
+class _RacyFuture(concurrent.futures.Future):
+ # done() always lies and reports False, to deterministically force the
+ # pump's check-then-act race window without needing real thread timing.
+ def done(self) -> bool:
+ return False
+
+
+def _racy_cancelled_future() -> concurrent.futures.Future:
+ fut = _RacyFuture()
+ fut.cancel()
+ assert fut.cancelled()
+ return fut
+
+
+class TestResultPumpCancelledFutureRace:
+ """Regression for #5793: a future cancelled concurrently with the pump's
+ resolve call (e.g. an asyncio.wait_for timeout, or a request abort) must
+ be dropped, not raise InvalidStateError and kill the pump thread.
+ """
+
+ def test_compute_done_racing_cancel_does_not_crash_pump(self):
+ executor = _make_executor()
+ fut = _racy_cancelled_future()
+ with executor._futures_lock:
+ executor._rpc_futures["1"] = fut
+
+ msg = AsyncDiffusionOutput(kind=AsyncOutputKind.COMPUTE_DONE, rpc_id="1", async_output_id="abc")
+ _feed_one_msg_to_pump(executor, msg)
+
+ assert fut.cancelled()
+
+ def test_output_ready_racing_cancel_does_not_crash_pump(self):
+ executor = _make_executor()
+ fut = _racy_cancelled_future()
+ with executor._futures_lock:
+ executor._output_futures["abc123"] = fut
+
+ msg = AsyncDiffusionOutput(
+ kind=AsyncOutputKind.OUTPUT_READY,
+ async_output_id="abc123",
+ output=DiffusionOutput(output="late"),
+ )
+ _feed_one_msg_to_pump(executor, msg)
+
+ assert fut.cancelled()
+
+ def test_batch_split_racing_cancel_does_not_crash_pump(self):
+ executor = _make_executor()
+ cancelled_fut = _racy_cancelled_future()
+ healthy_fut = concurrent.futures.Future()
+ with executor._futures_lock:
+ executor._output_futures["batch-1/r-aborted"] = cancelled_fut
+ executor._output_futures["batch-1/r-healthy"] = healthy_fut
+ executor._batch_split_map["batch-1"] = {
+ "batch-1/r-aborted": "r-aborted",
+ "batch-1/r-healthy": "r-healthy",
+ }
+
+ outputs = {
+ "r-aborted": DiffusionOutput(output="late"),
+ "r-healthy": DiffusionOutput(output="ok"),
+ }
+ msg = AsyncDiffusionOutput(
+ kind=AsyncOutputKind.OUTPUT_READY,
+ async_output_id="batch-1",
+ output=_FakeBatchOutput(outputs),
+ )
+ _feed_one_msg_to_pump(executor, msg)
+
+ assert cancelled_fut.cancelled()
+ assert healthy_fut.done()
+ assert healthy_fut.result(timeout=1.0) is outputs["r-healthy"]
diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py
index 697465097..90e753c9b 100644
--- a/tests/diffusion/test_stage_diffusion_proc.py
+++ b/tests/diffusion/test_stage_diffusion_proc.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import asyncio
import time
@@ -7,6 +7,8 @@ from dataclasses import asdict, dataclass
import pytest
+import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc
+import vllm_omni.plugins as omni_plugins
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -26,6 +28,33 @@ BASE_INFER_STEPS = 10
DELAY_BASE = 0.01
+def test_run_diffusion_proc_sets_lifecycle_before_loading_plugins(monkeypatch: pytest.MonkeyPatch) -> None:
+ events: list[str] = []
+
+ class StopProcessError(Exception):
+ pass
+
+ class TestStageDiffusionProc(StageDiffusionProc):
+ def __init__(self, model, od_config):
+ events.append("proc")
+ raise StopProcessError
+
+ monkeypatch.setattr(omni_plugins, "load_omni_general_plugins", lambda: events.append("plugins"))
+ monkeypatch.setattr(stage_diffusion_proc, "set_death_signal", lambda _: events.append("death_signal"))
+ monkeypatch.setattr(stage_diffusion_proc.signal, "signal", lambda *_: events.append("signal_handler"))
+
+ with pytest.raises(StopProcessError):
+ TestStageDiffusionProc.run_diffusion_proc(
+ model="test-model",
+ od_config=None,
+ handshake_address="test-address",
+ local_client=True,
+ headless=False,
+ )
+
+ assert events == ["death_signal", "signal_handler", "signal_handler", "plugins", "proc"]
+
+
class MockDiffusionEngine:
async def step_streaming(self, request):
def simulate_step_delay(height, width, num_inference_steps) -> float:
diff --git a/tests/diffusion/test_uniproc_executor.py b/tests/diffusion/test_uniproc_executor.py
new file mode 100644
index 000000000..dfc1d2eb1
--- /dev/null
+++ b/tests/diffusion/test_uniproc_executor.py
@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Single-GPU (uniproc) diffusion executor: selection and RPC behaviour."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeODConfig:
+ def __init__(self, num_gpus: int = 1, backend: str | None = None) -> None:
+ self.num_gpus = num_gpus
+ self.distributed_executor_backend = backend
+ self.worker_extension_cls = None
+ self.custom_pipeline_args = None
+
+
+def _od_config(num_gpus: int = 1, backend: str | None = None) -> _FakeODConfig:
+ return _FakeODConfig(num_gpus=num_gpus, backend=backend)
+
+
+def test_single_gpu_defaults_to_uniproc():
+ assert DiffusionExecutor.get_class(_od_config(num_gpus=1)) is UniProcDiffusionExecutor
+
+
+def test_explicit_mp_is_honored_on_single_gpu():
+ assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="mp")) is MultiprocDiffusionExecutor
+
+
+def test_multi_gpu_defaults_to_multiproc():
+ assert DiffusionExecutor.get_class(_od_config(num_gpus=2)) is MultiprocDiffusionExecutor
+
+
+def test_explicit_uni_backend_is_honored():
+ assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="uni")) is UniProcDiffusionExecutor
+
+
+@pytest.fixture
+def executor(monkeypatch: pytest.MonkeyPatch):
+ """A ``UniProcDiffusionExecutor`` with a mocked worker (no model load)."""
+ worker = MagicMock()
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+ MagicMock(return_value="vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker"),
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname",
+ MagicMock(return_value=object),
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase",
+ MagicMock(return_value=worker),
+ )
+ return UniProcDiffusionExecutor(_od_config(num_gpus=1)), worker
+
+
+def test_rejects_multi_gpu_config(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+ MagicMock(return_value="unused"),
+ )
+ with pytest.raises(ValueError, match="single GPU only"):
+ UniProcDiffusionExecutor(_od_config(num_gpus=4))
+
+
+def test_shutdown_is_safe_on_partially_constructed_executor(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+ MagicMock(return_value="unused"),
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname",
+ MagicMock(return_value=object),
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase",
+ MagicMock(side_effect=RuntimeError("model load blew up")),
+ )
+ constructed = UniProcDiffusionExecutor.__new__(UniProcDiffusionExecutor)
+ with pytest.raises(RuntimeError, match="model load blew up"):
+ constructed.__init__(_od_config(num_gpus=1))
+
+ constructed.shutdown()
+
+
+def test_collective_rpc_calls_worker_directly(executor):
+ exec_, worker = executor
+ worker.execute_method.return_value = "result"
+
+ out = exec_.collective_rpc("some_method", args=(1, 2), kwargs={"k": "v"}, unique_reply_rank=0)
+
+ assert out == "result"
+ worker.execute_method.assert_called_once_with("some_method", 1, 2, k="v")
+
+
+def test_collective_rpc_returns_list_when_no_reply_rank(executor):
+ exec_, worker = executor
+ worker.execute_method.return_value = "result"
+
+ assert exec_.collective_rpc("some_method") == ["result"]
+
+
+def test_collective_rpc_propagates_worker_exceptions(executor):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+
+ with pytest.raises(RuntimeError, match="boom"):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+
+def test_recoverable_worker_failure_keeps_the_executor_alive(executor, monkeypatch):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("CUDA out of memory")
+ monkeypatch.setattr(exec_, "_device_is_usable", lambda: True)
+ died = MagicMock()
+ exec_.register_failure_callback(died)
+
+ with pytest.raises(RuntimeError):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ assert exec_.is_dead is False
+ died.assert_not_called()
+ exec_.check_health()
+
+
+def test_poisoned_cuda_context_marks_the_executor_dead(executor, monkeypatch):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("an illegal memory access was encountered")
+ monkeypatch.setattr(exec_, "_device_is_usable", lambda: False)
+ died = MagicMock()
+ exec_.register_failure_callback(died)
+
+ with pytest.raises(RuntimeError):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ assert exec_.is_dead is True
+ died.assert_called_once_with()
+ with pytest.raises(EngineDeadError):
+ exec_.check_health()
+
+
+def test_failure_is_latched_and_callbacks_fire_once(executor, monkeypatch):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+ monkeypatch.setattr(exec_, "_device_is_usable", lambda: False)
+ died = MagicMock()
+ exec_.register_failure_callback(died)
+
+ for _ in range(3):
+ with pytest.raises(RuntimeError):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ died.assert_called_once_with()
+
+
+def test_device_probe_is_skipped_when_cuda_was_never_initialized(executor):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+
+ with pytest.raises(RuntimeError):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ if not torch.cuda.is_initialized():
+ assert exec_.is_dead is False
+
+
+def test_device_probe_skipped_when_accelerator_unavailable(executor, monkeypatch):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+ monkeypatch.setattr(torch.accelerator, "is_available", lambda: False)
+ sync = MagicMock()
+ monkeypatch.setattr(torch.accelerator, "synchronize", sync)
+
+ with pytest.raises(RuntimeError, match="boom"):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ sync.assert_not_called()
+ assert exec_.is_dead is False
+
+
+def test_device_probe_skipped_when_device_not_initialized(executor, monkeypatch):
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+ monkeypatch.setattr(torch.accelerator, "is_available", lambda: True)
+ monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu"))
+ monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: False), raising=False)
+ sync = MagicMock()
+ monkeypatch.setattr(torch.accelerator, "synchronize", sync)
+
+ with pytest.raises(RuntimeError, match="boom"):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ sync.assert_not_called()
+ assert exec_.is_dead is False
+
+
+def test_device_probe_latches_sticky_fault_on_npu(executor, monkeypatch):
+ """NPU must not short-circuit the probe the way a CUDA-only guard did."""
+ exec_, worker = executor
+ worker.execute_method.side_effect = RuntimeError("boom")
+ monkeypatch.setattr(torch.accelerator, "is_available", lambda: True)
+ monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu"))
+ monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: True), raising=False)
+ monkeypatch.setattr(
+ torch.accelerator,
+ "synchronize",
+ MagicMock(side_effect=RuntimeError("NPU context poisoned")),
+ )
+ died = MagicMock()
+ exec_.register_failure_callback(died)
+
+ with pytest.raises(RuntimeError, match="boom"):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+ assert exec_.is_dead is True
+ died.assert_called_once_with()
+
+
+def test_check_health_ok_then_dead(executor):
+ exec_, _ = executor
+ exec_.check_health()
+
+ exec_._is_failed = True
+ with pytest.raises(EngineDeadError):
+ exec_.check_health()
+
+
+def test_shutdown_is_idempotent_and_closes_executor(executor, monkeypatch):
+ exec_, worker = executor
+ collect = MagicMock()
+ empty_cache = MagicMock()
+ monkeypatch.setattr("vllm_omni.diffusion.executor.uniproc_executor.gc.collect", collect)
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.is_available",
+ lambda: True,
+ )
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.empty_cache",
+ empty_cache,
+ )
+ exec_.register_failure_callback(MagicMock())
+
+ exec_.shutdown()
+ exec_.shutdown()
+
+ worker.shutdown.assert_called_once()
+ assert exec_.driver_worker is None
+ assert exec_._failure_callbacks == []
+ collect.assert_called_once_with()
+ empty_cache.assert_called_once_with()
+ with pytest.raises(RuntimeError, match="closed"):
+ exec_.collective_rpc("some_method", unique_reply_rank=0)
diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py
index 170faec80..709fc7283 100644
--- a/tests/diffusion/test_worker_wrapper_base.py
+++ b/tests/diffusion/test_worker_wrapper_base.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Unit tests for WorkerWrapperBase class.
@@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation:
- Dynamic worker class extension
"""
+from types import SimpleNamespace
from typing import Any
import pytest
@@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation:
wrapper.worker.shutdown.assert_called_once()
assert result is None
+ def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
+ events: list[str] = []
+ offload_backend = mocker.Mock()
+ offload_backend.disable.side_effect = lambda: events.append("offload")
+ kv_manager = mocker.Mock()
+ kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
+ destroy = mocker.patch(
+ "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
+ side_effect=lambda: events.append("distributed"),
+ )
+ worker = DiffusionWorker.__new__(DiffusionWorker)
+ worker.model_runner = SimpleNamespace(
+ offload_backend=offload_backend,
+ kv_transfer_manager=kv_manager,
+ )
+
+ worker.shutdown()
+
+ assert events == ["offload", "kv", "distributed"]
+ destroy.assert_called_once_with()
+
# -------------------------------------------------------------------------
# Tests: execute_method
diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
index 57997905c..c31b80a6f 100644
--- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
+++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end test for Bagel text2img with the Mooncake inter-stage connector.
@@ -22,7 +22,7 @@ import pytest
from PIL import Image
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
+from tests.helpers.runtime import OmniRunner, get_open_port
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
@@ -55,15 +55,6 @@ PIXEL_TOLERANCE = 10
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
-def _find_free_port() -> int:
- """Find and return a free ephemeral port by binding to port 0."""
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- s.listen(1)
- port = s.getsockname()[1]
- return port
-
-
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
"""Configure sampling parameters for Bagel text2img generation.
@@ -277,9 +268,9 @@ def test_bagel_text2img_mooncake_connector(run_level):
"mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)"
)
MOONCAKE_HOST = "127.0.0.1"
- MOONCAKE_RPC_PORT = _find_free_port()
- MOONCAKE_HTTP_PORT = _find_free_port()
- MOONCAKE_METRICS_PORT = _find_free_port()
+ MOONCAKE_RPC_PORT = get_open_port()
+ MOONCAKE_HTTP_PORT = get_open_port()
+ MOONCAKE_METRICS_PORT = get_open_port()
mooncake_master_proc = None
temp_config_file = None
@@ -314,7 +305,7 @@ def test_bagel_text2img_mooncake_connector(run_level):
temp_config_file = _resolve_deploy_config(temp_config_file, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
- stage_configs_path=temp_config_file,
+ deploy_config=temp_config_file,
stage_init_timeout=300,
) as runner:
generated_image = _generate_bagel_image(runner.omni)
diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
index b97161335..6b23ba321 100644
--- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
+++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end tests for Bagel with shared memory connector: img2img and text2img.
@@ -13,7 +13,6 @@ End-to-end tests for Bagel with shared memory connector: img2img and text2img.
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-import socket
from typing import Any
import pytest
@@ -82,15 +81,6 @@ def _load_input_image() -> Image.Image:
return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
-def _find_free_port() -> int:
- """Find and return a free ephemeral port by binding to port 0."""
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind(("127.0.0.1", 0))
- s.listen(1)
- port = s.getsockname()[1]
- return port
-
-
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
"""Configure sampling parameters for Bagel img2img generation.
@@ -251,7 +241,7 @@ def test_bagel_img2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
- stage_configs_path=config_path,
+ deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_img2img(runner.omni, input_image)
if run_level == "advanced_model":
@@ -266,7 +256,7 @@ def test_bagel_text2img_shared_memory_connector(run_level):
config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level)
with OmniRunner(
"ByteDance-Seed/BAGEL-7B-MoT",
- stage_configs_path=config_path,
+ deploy_config=config_path,
) as runner:
generated_image = _generate_bagel_text2img(runner.omni)
if run_level == "advanced_model":
diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
index da4e81995..0e3194190 100644
--- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
@@ -1,18 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import threading
from collections import deque
-from types import SimpleNamespace
-from unittest.mock import patch
+from types import MethodType, SimpleNamespace
+from unittest.mock import Mock, patch
import pytest
import torch
from pytest_mock import MockerFixture
+from vllm.sampling_params import SamplingParams
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
-from vllm.v1.request import RequestStatus
+from vllm.v1.request import Request, RequestStatus
+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
from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
@@ -32,9 +34,14 @@ class DummyWaitingQueue(list):
def add_request(self, request):
self.append(request)
+ def remove_requests(self, requests):
+ remove = set(requests)
+ self[:] = [request for request in self if request not in remove]
+
def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None):
- return SimpleNamespace(
+ request = Mock(
+ client_index=0,
request_id=req_id,
external_req_id=external_req_id or req_id,
status=status,
@@ -44,7 +51,22 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
num_output_placeholders=0,
prefill_stats=None,
additional_information=None,
- is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
+ resumable=False,
+ )
+ request.is_finished = lambda: RequestStatus.is_finished(request.status)
+ return request
+
+
+def _streaming_request(mocker: MockerFixture, num_computed_tokens: int) -> SimpleNamespace:
+ prompt_token_ids = [0] * num_computed_tokens
+ return SimpleNamespace(
+ _all_token_ids=prompt_token_ids.copy(),
+ _output_token_ids=[],
+ prompt_token_ids=prompt_token_ids,
+ num_computed_tokens=num_computed_tokens,
+ num_prompt_tokens=num_computed_tokens,
+ num_output_placeholders=0,
+ update_block_hashes=mocker.Mock(),
)
@@ -75,30 +97,219 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture)
request.update_block_hashes.assert_called_once_with()
-def test_streaming_payload_can_append_exact_prompt_length(mocker: MockerFixture) -> None:
+def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=4064)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "replace_streaming_prompt": True,
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ replaced = construct_next_stage_streaming_input_prompt(
+ payload,
+ request,
+ max_model_len=4096,
+ previous_condition_len=10,
+ previous_condition_seq=7,
+ condition_seq=8,
+ recompute_previous_chunks=1,
+ )
+
+ assert replaced is True
+ assert request.prompt_token_ids == [0] * 10
+ assert payload["meta"]["streaming_prompt_recompute"] is False
+ assert "streaming_prompt_previous_codes" not in payload["ids"]
+
+
+@pytest.mark.parametrize("reserve", [True, 0, -1, 1.5, "26", None])
+def test_streaming_prompt_rejects_invalid_generation_reserve(
+ mocker: MockerFixture,
+ reserve: object,
+) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=10)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": reserve,
+ },
+ }
+
+ with pytest.raises(ValueError, match="next_stage_generation_tokens must be a positive integer"):
+ construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
+
+
+def test_capacity_managed_streaming_prompt_at_limit_still_appends(mocker: MockerFixture) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=4060)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ replaced = construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
+
+ assert replaced is False
+ assert request.num_computed_tokens == 4060
+ assert request.num_prompt_tokens == 4070
+
+
+def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None:
+ previous_condition_len = 10
+ previous_codes = list(range(25))
request = SimpleNamespace(
- _all_token_ids=[0, 0, 7, 8],
- _output_token_ids=[7, 8],
- prompt_token_ids=[0, 0],
- num_computed_tokens=4,
- num_prompt_tokens=2,
+ _all_token_ids=[0] * 4039 + previous_codes + [999],
+ _output_token_ids=previous_codes + [999],
+ prompt_token_ids=[0] * 4039,
+ num_computed_tokens=4065,
+ num_prompt_tokens=4039,
+ num_output_placeholders=1,
update_block_hashes=mocker.Mock(),
)
payload = {
- "ids": {"prompt": [1, 2, 3]},
- "meta": {"next_stage_prompt_len": 3},
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ },
}
- construct_next_stage_streaming_input_prompt(payload, request)
+ replaced = construct_next_stage_streaming_input_prompt(
+ payload,
+ request,
+ max_model_len=4096,
+ previous_condition_len=previous_condition_len,
+ previous_condition_seq=7,
+ condition_seq=8,
+ recompute_previous_chunks=1,
+ )
- assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0]
- assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0]
- assert request._output_token_ids == []
- assert request.num_computed_tokens == 4
- assert request.num_prompt_tokens == 7
+ assert replaced is True
+ assert request.prompt_token_ids == [0] * 45
+ assert request.num_prompt_tokens == 45
+ assert request.num_computed_tokens == 0
+ assert payload["ids"]["streaming_prompt_previous_codes"] == previous_codes
+ assert payload["meta"] == {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ "streaming_prompt_recompute": True,
+ "streaming_condition_seq": 8,
+ }
request.update_block_hashes.assert_called_once_with()
+def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: MockerFixture) -> None:
+ request = SimpleNamespace(
+ _all_token_ids=[0] * 10 + [101, 102],
+ _output_token_ids=[101, 102],
+ prompt_token_ids=[0] * 10,
+ num_computed_tokens=12,
+ num_prompt_tokens=10,
+ num_output_placeholders=0,
+ update_block_hashes=mocker.Mock(),
+ )
+ second = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 12,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ assert construct_next_stage_streaming_input_prompt(
+ second,
+ request,
+ max_model_len=4096,
+ previous_condition_len=10,
+ previous_condition_seq=0,
+ condition_seq=1,
+ recompute_previous_chunks=1,
+ )
+ assert request.num_prompt_tokens == 24
+ assert second["ids"]["streaming_prompt_previous_codes"] == [101, 102]
+
+ request._all_token_ids.extend([201, 202, 203])
+ request._output_token_ids.extend([201, 202, 203])
+ request.num_computed_tokens = 27
+ third = {
+ "ids": {"prompt": [2]},
+ "meta": {
+ "next_stage_prompt_len": 8,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ assert construct_next_stage_streaming_input_prompt(
+ third,
+ request,
+ max_model_len=4096,
+ previous_condition_len=12,
+ previous_condition_seq=1,
+ condition_seq=2,
+ recompute_previous_chunks=1,
+ )
+ assert request.num_prompt_tokens == 23
+ assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203]
+ assert third["meta"]["streaming_condition_seq"] == 2
+
+
+def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=4064)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ with pytest.raises(ValueError, match="requires window_size=1"):
+ construct_next_stage_streaming_input_prompt(
+ payload,
+ request,
+ max_model_len=4096,
+ previous_condition_len=10,
+ previous_condition_seq=0,
+ condition_seq=1,
+ )
+
+
+@pytest.mark.parametrize("next_stage_prompt_len", [0, -1])
+def test_capacity_managed_streaming_prompt_rejects_nonpositive_length(
+ mocker: MockerFixture, next_stage_prompt_len: int
+) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=0)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": next_stage_prompt_len,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ with pytest.raises(ValueError, match="positive next_stage_prompt_len"):
+ construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
+
+
+def test_capacity_managed_streaming_prompt_rejects_fresh_overflow(mocker: MockerFixture) -> None:
+ request = _streaming_request(mocker, num_computed_tokens=0)
+ payload = {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 4071,
+ "next_stage_generation_tokens": 26,
+ },
+ }
+
+ with pytest.raises(ValueError, match="exceeds max_model_len"):
+ construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096)
+
+
@pytest.fixture
def build_adapter(monkeypatch, mocker: MockerFixture):
def _build(
@@ -106,6 +317,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
stage_id: int = 1,
model_mode: str = "ar",
max_num_seqs: int = 2,
+ max_model_len: int = 0,
+ tts_max_model_len: int = 0,
+ tts_attention_type: str = "full_attention",
+ flat_tts_config: bool = False,
active_stream_window: int = 0,
connector_extra: dict | None = None,
):
@@ -122,6 +337,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
self._cancelled_load_reqs = set()
self._pending_save_reqs = deque()
self._finished_save_reqs = set()
+ self._send_failures = {}
+ self._send_failure_lock = threading.Lock()
+ self._receive_failures = {}
+ self._receive_failure_lock = threading.Lock()
self.stop_event = threading.Event()
self._recv_cond = threading.Condition()
self._save_cond = threading.Condition()
@@ -133,9 +352,23 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
classmethod(lambda cls, _model_config: connector),
)
+ hf_config = SimpleNamespace(
+ model_type="minicpmtts",
+ max_position_embeddings=tts_max_model_len,
+ attention_type=tts_attention_type,
+ )
+ if not flat_tts_config:
+ hf_config = SimpleNamespace(
+ tts_config=SimpleNamespace(
+ max_position_embeddings=tts_max_model_len,
+ attention_type=tts_attention_type,
+ )
+ )
model_config = SimpleNamespace(
worker_type=model_mode,
max_num_seqs=max_num_seqs,
+ max_model_len=max_model_len,
+ hf_config=hf_config,
active_stream_window=active_stream_window,
stage_connector_config={
"name": "SharedMemoryConnector",
@@ -151,6 +384,27 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
return _build
+def _dequeue_load_entry(adapter, request):
+ """Register and take the entry as the background recv loop would."""
+ adapter.load_async(request)
+ entry = adapter._registered_load_entries[request.request_id]
+ try:
+ adapter._pending_load_reqs.remove(entry)
+ except ValueError:
+ pass
+ return entry
+
+
+@pytest.mark.parametrize(
+ ("attention_type", "expected_previous_chunks"),
+ [("full_attention", 0), ("sliding_recompute", 1)],
+)
+def test_talker_attention_policy_controls_streaming_recompute(build_adapter, attention_type, expected_previous_chunks):
+ adapter, _ = build_adapter(tts_attention_type=attention_type)
+
+ assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks
+
+
@pytest.mark.parametrize(
("raw_cfg", "expected_name", "expected_extra"),
[
@@ -191,15 +445,258 @@ def test_load_poll(build_adapter):
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
}
connector.get.return_value = (payload, 16)
- adapter._poll_single_request(request)
+ adapter._poll_single_request(_dequeue_load_entry(adapter, request))
+ assert request.additional_information is None
+ assert "req-1" in adapter._finished_load_reqs
+ request.status = RequestStatus.WAITING_FOR_CHUNK
+ adapter.process_pending_chunks(
+ DummyWaitingQueue(),
+ [request],
+ scheduler_requests={request.request_id: request},
+ )
assert request.additional_information == payload
assert adapter.get_req_chunk["req-1"] == 1
- assert "req-1" in adapter._finished_load_reqs
+ assert "req-1" not in adapter._finished_load_reqs
assert "req-1" in adapter.upstream_exhausted_requests
assert "req-1" not in adapter._pending_load_reqs
+def test_load_async_does_not_requeue_registered_inflight_request(build_adapter):
+ adapter, connector = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("req-inflight", RequestStatus.RUNNING, external_req_id="ext-inflight")
+
+ entry = _dequeue_load_entry(adapter, request)
+
+ # The recv loop owns the popped request until polling completes. A
+ # scheduler-side duplicate registration must not append a second copy.
+ adapter.load_async(request)
+ assert adapter._pending_load_reqs == deque()
+
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {"finished": True},
+ },
+ 1,
+ )
+ assert adapter._poll_single_request(entry) is True
+
+ # Successful polling releases the registration for the next chunk.
+ adapter.load_async(request)
+ assert [queued.request for queued in adapter._pending_load_reqs] == [request]
+
+
+def test_cleanup_then_reregister_drops_stale_queued_entry(build_adapter):
+ adapter, connector = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("req-resume", RequestStatus.WAITING, external_req_id="ext-resume")
+
+ adapter.load_async(request)
+ stale_entry = adapter._pending_load_reqs[0]
+ adapter.cleanup_receiver(request.request_id)
+ adapter.load_async(request)
+ fresh_entry = adapter._pending_load_reqs[-1]
+ assert len(adapter._pending_load_reqs) == 2
+ assert adapter._pending_load_reqs.popleft() is stale_entry
+ assert adapter._pending_load_reqs.popleft() is fresh_entry
+
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {"finished": True},
+ },
+ 1,
+ )
+ assert adapter._poll_single_request(stale_entry) is True
+ connector.get.assert_not_called()
+ assert adapter._registered_load_entries[request.request_id] is fresh_entry
+
+ assert adapter._poll_single_request(fresh_entry) is True
+ connector.get.assert_called_once()
+ assert request.request_id not in adapter._registered_load_entries
+ assert adapter.get_req_chunk[request.request_id] == 1
+
+
+def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter):
+ adapter, connector = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace")
+ request.resumable = True
+ request.prompt_token_ids = [0] * 4
+ request._all_token_ids = [0] * 4
+ request._output_token_ids = []
+ request.num_prompt_tokens = 4
+ request.num_computed_tokens = 4
+ request.update_block_hashes = Mock()
+ adapter.get_req_chunk[request.request_id] = 1
+ adapter.requests_num_chunks_sent[request.external_req_id] = 4
+ adapter.request_ids_mapping[request.request_id] = request.external_req_id
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "replace_streaming_prompt": True,
+ "finished": False,
+ },
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+ request.status = RequestStatus.WAITING_FOR_CHUNK
+ running_queue = [request]
+ waiting_queue = DummyWaitingQueue()
+ adapter.process_pending_chunks(
+ waiting_queue,
+ running_queue,
+ scheduler_requests={request.request_id: request},
+ )
+ assert request.num_computed_tokens == 0
+ assert request.prompt_token_ids == [0] * 10
+ assert request.request_id in adapter.replaced_streaming_prompt_ids
+ assert running_queue == []
+ assert waiting_queue == [request]
+ assert request.status == RequestStatus.WAITING
+
+ adapter.requests_num_chunks_sent[request.external_req_id] = 9
+ adapter.postprocess_scheduler_output(
+ SimpleNamespace(
+ scheduled_new_reqs=[SimpleNamespace(req_id=request.request_id)],
+ scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
+ )
+ )
+ assert request.request_id not in adapter.replaced_streaming_prompt_ids
+ assert request.external_req_id not in adapter.requests_num_chunks_sent
+
+
+@pytest.mark.parametrize("flat_tts_config", [False, True])
+def test_load_poll_ar_recomputes_native_duplex_prompt_each_condition(build_adapter, flat_tts_config):
+ adapter, connector = build_adapter(
+ stage_id=1,
+ model_mode="ar",
+ max_model_len=8192,
+ tts_max_model_len=4096,
+ tts_attention_type="sliding_recompute",
+ flat_tts_config=flat_tts_config,
+ )
+ request = _req("req-rollover", RequestStatus.RUNNING, external_req_id="external-rollover")
+ request.resumable = True
+ previous_codes = list(range(25))
+ request.prompt_token_ids = [0] * 10
+ request._all_token_ids = [0] * 10 + previous_codes
+ request._output_token_ids = previous_codes.copy()
+ request.num_prompt_tokens = 10
+ request.num_computed_tokens = 35
+ request.num_output_placeholders = 0
+ request.update_block_hashes = Mock()
+ adapter.get_req_chunk[request.request_id] = 1
+ adapter.request_ids_mapping[request.request_id] = request.external_req_id
+ adapter._streaming_condition_lengths[request.request_id] = 10
+ adapter._streaming_condition_seqs[request.request_id] = 0
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ "finished": True,
+ },
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+ request.status = RequestStatus.WAITING_FOR_CHUNK
+ running_queue = [request]
+ waiting_queue = DummyWaitingQueue()
+ adapter.process_pending_chunks(
+ waiting_queue,
+ running_queue,
+ scheduler_requests={request.request_id: request},
+ )
+
+ assert request.num_computed_tokens == 0
+ assert request.prompt_token_ids == [0] * 45
+ assert request.additional_information["ids"]["streaming_prompt_previous_codes"] == previous_codes
+ assert request.additional_information["meta"]["streaming_prompt_recompute"] is True
+ assert request.resumable is False
+ assert request.request_id in adapter.replaced_streaming_prompt_ids
+ assert adapter._max_model_len == 4096
+ assert running_queue == []
+ assert waiting_queue == [request]
+
+
+def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None:
+ adapter, connector = build_adapter(
+ stage_id=1,
+ model_mode="ar",
+ tts_attention_type="sliding_recompute",
+ )
+ request = _req("req-condition-seq", RequestStatus.RUNNING)
+
+ def receive(payload):
+ connector.get.return_value = (payload, 1)
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+ adapter.requests_with_ready_chunks.add(request.request_id)
+ adapter._apply_pending_ar_prompt_updates({request.request_id: request})
+ adapter.requests_with_ready_chunks.discard(request.request_id)
+ return request.additional_information
+
+ first = receive(
+ {
+ "ids": {"prompt": [1]},
+ "meta": {"next_stage_prompt_len": 10},
+ }
+ )
+ assert first["meta"]["streaming_condition_seq"] == 0
+ assert adapter._streaming_condition_seqs[request.request_id] == 0
+
+ control = receive({"meta": {"is_segment_finished": True}})
+ assert "streaming_condition_seq" not in control["meta"]
+ assert adapter._streaming_condition_seqs[request.request_id] == 0
+
+ second = receive(
+ {
+ "ids": {"prompt": [2]},
+ "meta": {"next_stage_prompt_len": 12},
+ }
+ )
+ assert second["meta"]["streaming_condition_seq"] == 1
+ assert adapter._streaming_condition_lengths[request.request_id] == 12
+
+
+def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter):
+ adapter, connector = build_adapter(stage_id=1, model_mode="ar", max_model_len=4096)
+ request = _req("req-invalid-capacity", RequestStatus.RUNNING)
+ request.resumable = True
+ adapter.get_req_chunk[request.request_id] = 1
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 4096,
+ "next_stage_generation_tokens": 26,
+ "finished": False,
+ },
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+ request.status = RequestStatus.WAITING_FOR_CHUNK
+ running_queue = [request]
+ adapter.process_pending_chunks(
+ DummyWaitingQueue(),
+ running_queue,
+ scheduler_requests={request.request_id: request},
+ )
+ assert adapter.collect_failed_receive_request_ids() == {
+ request.request_id: "fresh streaming prompt plus generation reserve exceeds max_model_len: "
+ "prompt=4096, reserve=26, limit=4096"
+ }
+ assert request.request_id not in adapter._finished_load_reqs
+
+
def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
adapter, connector = build_adapter(stage_id=1, model_mode="generation")
request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
@@ -214,7 +711,7 @@ def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter)
}
connector.get.return_value = (payload, 16)
- adapter._poll_single_request(request)
+ adapter._poll_single_request(_dequeue_load_entry(adapter, request))
assert request.prompt_token_ids == [0]
assert request.num_computed_tokens == 0
@@ -244,12 +741,13 @@ def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapte
}
connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]
- assert adapter._poll_single_request(request) is False
+ entry = _dequeue_load_entry(adapter, request)
+ assert adapter._poll_single_request(entry) is False
assert request.request_id not in adapter._finished_load_reqs
assert request.request_id not in adapter.requests_with_ready_chunks
assert adapter.get_req_chunk[request.request_id] == 1
- assert adapter._poll_single_request(request) is True
+ assert adapter._poll_single_request(entry) is True
assert request.request_id in adapter._finished_load_reqs
assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
assert adapter.get_req_chunk[request.request_id] == 2
@@ -280,6 +778,126 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
assert len(adapter._pending_save_reqs) == 1
+def test_segment_boundary_starts_new_send_watermark_before_background_flush(build_adapter):
+ """A queued boundary owns the end of its deduplication generation.
+
+ The next segment can start before the save thread sends the old boundary.
+ Sending that old task later must not erase the new segment's watermark.
+ """
+ adapter, _ = build_adapter(stage_id=1)
+ request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+ request.resumable = True
+ request._omni_segment_generation = 1
+ request.num_computed_tokens = 0
+ adapter.requests_num_chunks_sent["ext-stream"] = 26
+
+ adapter.save_async(
+ multimodal_output=None,
+ request=request,
+ is_segment_finished=True,
+ confirmed_num_computed_tokens=26,
+ segment_generation=0,
+ )
+
+ assert len(adapter._pending_save_reqs) == 1
+ boundary_task = adapter._pending_save_reqs.popleft()
+ assert boundary_task["segment_generation"] == 0
+ assert "ext-stream" not in adapter.requests_num_chunks_sent
+
+ request.num_computed_tokens = 3
+ request._omni_segment_generation = 1
+ adapter.save_async(
+ multimodal_output=None,
+ request=request,
+ is_segment_finished=False,
+ )
+
+ assert len(adapter._pending_save_reqs) == 1
+ assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+
+ adapter._send_single_request(boundary_task)
+
+ assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+
+
+def test_background_send_uses_enqueued_request_snapshot(build_adapter):
+ """A queued segment must not observe later in-place request mutations."""
+ adapter, _ = build_adapter(stage_id=1)
+ request = Request(
+ request_id="req-stream",
+ prompt_token_ids=[1, 2],
+ sampling_params=SamplingParams(max_tokens=8),
+ pooling_params=None,
+ resumable=True,
+ )
+ request.external_req_id = "ext-stream"
+ ref_audio = torch.tensor([0.1, -0.1])
+ request.additional_information = {"codes": {"ref": ref_audio}, "meta": {"segment": "old"}}
+ request.append_output_token_ids([7])
+ seen_requests = []
+
+ def recording_processor(**kwargs):
+ queued = kwargs["request"]
+ seen_requests.append(
+ (
+ queued.additional_information["meta"]["segment"],
+ list(queued.prompt_token_ids),
+ list(queued.output_token_ids),
+ list(queued.all_token_ids),
+ queued.additional_information["codes"]["ref"] is ref_audio,
+ )
+ )
+ return OmniPayloadStruct()
+
+ adapter.custom_process_next_stage_input_func = recording_processor
+ adapter.save_async(
+ multimodal_output=None,
+ request=request,
+ is_segment_finished=True,
+ )
+ request.additional_information["meta"]["segment"] = "next"
+ request.prompt_token_ids.append(3)
+ request.append_output_token_ids([8])
+ request._omni_segment_generation = 1
+ adapter.save_async(
+ multimodal_output=None,
+ request=request,
+ is_segment_finished=True,
+ )
+ request.additional_information["meta"]["segment"] = "later"
+ request.prompt_token_ids.append(4)
+ request.append_output_token_ids([9])
+
+ assert len(adapter._pending_save_reqs) == 2
+ first_task = adapter._pending_save_reqs.popleft()
+ second_task = adapter._pending_save_reqs.popleft()
+ assert first_task["request"].additional_information is not request.additional_information
+ assert first_task["request"].additional_information["meta"] is not request.additional_information["meta"]
+
+ adapter._send_single_request(first_task)
+ adapter._send_single_request(second_task)
+
+ assert seen_requests == [
+ ("old", [1, 2], [7], [1, 2, 7], True),
+ ("next", [1, 2, 3], [7, 8], [1, 2, 7, 8], True),
+ ]
+
+
+def test_save_without_custom_processor_does_not_snapshot_request(build_adapter, mocker):
+ adapter, _ = build_adapter(stage_id=1)
+ request = _req("req-direct", RequestStatus.RUNNING, external_req_id="ext-direct")
+ snapshot = mocker.patch.object(
+ adapter,
+ "_snapshot_processor_request",
+ side_effect=AssertionError("unexpected snapshot"),
+ )
+
+ adapter.save_async(multimodal_output=None, request=request)
+
+ snapshot.assert_not_called()
+ assert adapter._pending_save_reqs.popleft()["request"] is request
+
+
def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
"""A terminal stop is not a segment boundary (#5383), but the producer-side
processor must still receive the flush signal on the terminal chunk.
@@ -489,7 +1107,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
assert second.codes is None
-def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
+def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
request.resumable = True
@@ -508,6 +1126,51 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
assert adapter.requests_num_chunks_sent["ext-stream"] == 0
+def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter):
+ adapter, _ = build_adapter(stage_id=1)
+ request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+ request.resumable = True
+ request._omni_segment_generation = 0
+ adapter.requests_num_chunks_sent["ext-stream"] = 111
+
+ adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+ late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream")
+ late.resumable = True
+ late.num_computed_tokens = 111
+ late._omni_segment_generation = 0
+ adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False)
+
+ next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream")
+ next_segment.resumable = True
+ next_segment.num_computed_tokens = 3
+ next_segment._omni_segment_generation = 1
+ adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False)
+
+ assert len(adapter._pending_save_reqs) == 2
+ assert adapter.requests_num_chunks_sent["ext-stream"] == 3
+ assert adapter._pending_save_reqs[-1]["request"] is next_segment
+
+
+def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter):
+ adapter, connector = build_adapter(stage_id=1)
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct())
+ request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old")
+ request.resumable = True
+ request._omni_segment_generation = 0
+
+ adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
+ adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+
+ old_task = adapter._pending_save_reqs.popleft()
+ adapter._send_single_request(old_task)
+ connector.put.assert_called_once()
+
+ boundary_task = adapter._pending_save_reqs.popleft()
+ adapter._send_single_request(boundary_task)
+ assert connector.put.call_count == 2
+
+
def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
@@ -528,6 +1191,77 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
assert args[1] == "ext-finished"
+def test_old_turn_terminal_send_preserves_new_turn_state_and_sender(build_adapter):
+ adapter, connector = build_adapter(
+ stage_id=1,
+ model_mode="ar",
+ max_model_len=4096,
+ tts_attention_type="sliding_recompute",
+ )
+ request = _req("req-live", RequestStatus.RUNNING, external_req_id="ext-live")
+ request.resumable = True
+ request._all_token_ids = []
+ request._output_token_ids = []
+ request.update_block_hashes = Mock()
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+ meta=MetaStruct(
+ finished=torch.tensor(True, dtype=torch.bool),
+ codec_streaming=True,
+ )
+ )
+
+ put_started = threading.Event()
+ release_put = threading.Event()
+ put_keys = []
+
+ def blocking_first_put(*_args, **kwargs):
+ put_keys.append(kwargs["put_key"])
+ if len(put_keys) == 1:
+ put_started.set()
+ assert release_put.wait(timeout=5)
+ return True, 1, {}
+
+ connector.put.side_effect = blocking_first_put
+ adapter.save_async(multimodal_output=None, request=request)
+ first_task = adapter._pending_save_reqs.popleft()
+ first_sender_token = first_task["sender_token"]
+ send_thread = threading.Thread(target=lambda: adapter._send_single_request(first_task))
+ send_thread.start()
+ assert put_started.wait(timeout=5)
+
+ entry = _dequeue_load_entry(adapter, request)
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "replace_streaming_prompt": True,
+ "next_stage_prompt_len": 10,
+ "next_stage_generation_tokens": 26,
+ },
+ },
+ 1,
+ )
+ assert adapter._poll_single_request(entry) is True
+ adapter.requests_with_ready_chunks.add(request.request_id)
+ adapter._apply_pending_ar_prompt_updates({request.request_id: request})
+ assert adapter._streaming_condition_seqs[request.request_id] == 0
+
+ request.num_computed_tokens = 1
+ adapter.save_async(multimodal_output=None, request=request)
+ second_task = adapter._pending_save_reqs.popleft()
+ assert second_task["sender_token"] is first_sender_token
+
+ release_put.set()
+ send_thread.join(timeout=5)
+ assert not send_thread.is_alive()
+ assert adapter._streaming_condition_seqs[request.request_id] == 0
+ assert adapter._sender_tokens[request.external_req_id] is first_sender_token
+ assert first_sender_token.cancelled is False
+
+ adapter._send_single_request(second_task)
+ assert put_keys == ["ext-live_1_0", "ext-live_1_1"]
+
+
def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="diffusion")
request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar")
@@ -547,7 +1281,7 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap
}
connector.get.return_value = (payload, 8)
- assert adapter._poll_single_request(request) is True
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert request.prompt_token_ids == [7, 8]
assert request.num_computed_tokens == 0
@@ -565,6 +1299,92 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap
assert "req-non-ar" in adapter.upstream_exhausted_requests
+def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapter):
+ adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+ request = _req("req-marker", RequestStatus.WAITING, external_req_id="external-marker")
+ request.additional_information = {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+ }
+ connector.get.return_value = (
+ {
+ "codes": {
+ "audio": torch.tensor([7, 8], dtype=torch.long),
+ "ref": torch.tensor([0.1, -0.1]),
+ },
+ "meta": {
+ "finished": torch.tensor(False, dtype=torch.bool),
+ "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+ "request_finished": torch.tensor(False, dtype=torch.bool),
+ "replace_runtime_additional_information": True,
+ },
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+
+ assert request.prompt_token_ids == [7, 8]
+ assert "audio" not in request.additional_information["codes"]
+ torch.testing.assert_close(
+ request.additional_information["codes"]["ref"],
+ torch.tensor([0.1, -0.1]),
+ )
+ assert not {"cache_epoch", "chunk_seq", "last_chunk"}.intersection(request.additional_information["meta"])
+ assert request.request_id in adapter.segment_finished_requests
+
+
+def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter):
+ adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+ request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker")
+ request.additional_information = {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+ }
+ connector.get.return_value = (
+ {
+ "meta": {
+ "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+ "replace_runtime_additional_information": True,
+ }
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
+
+ assert request.prompt_token_ids == [0]
+ assert "codes" not in request.additional_information
+ assert request.additional_information["meta"]["replace_runtime_additional_information"] is True
+ assert request.request_id in adapter.segment_finished_requests
+ assert request.request_id in adapter._finished_load_reqs
+
+
+def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(build_adapter):
+ adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+ request = _req("req-incremental", RequestStatus.WAITING, external_req_id="external-incremental")
+ request.additional_information = {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"cache_epoch": 3, "chunk_seq": 2},
+ }
+ connector.get.return_value = (
+ {
+ "meta": {
+ "finished": torch.tensor(False, dtype=torch.bool),
+ "phase": "decode",
+ }
+ },
+ 1,
+ )
+
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is False
+
+ assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2]))
+ assert request.additional_information["meta"]["cache_epoch"] == 3
+ assert request.additional_information["meta"]["chunk_seq"] == 2
+ assert request.additional_information["meta"]["phase"] == "decode"
+
+
def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
adapter, connector = build_adapter(stage_id=2, model_mode="ar")
request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
@@ -581,9 +1401,17 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt
}
connector.get.return_value = (payload, 8)
- adapter._poll_single_request(request)
+ adapter._poll_single_request(_dequeue_load_entry(adapter, request))
- # AR mode now forwards the latest payload directly.
+ # The recv thread records the latest payload; scheduler-side queue
+ # restoration publishes it onto the Request.
+ assert torch.equal(request.additional_information["hidden_states"]["output"], torch.tensor([[1.0]]))
+ request.status = RequestStatus.WAITING_FOR_CHUNK
+ adapter.process_pending_chunks(
+ DummyWaitingQueue(),
+ [request],
+ scheduler_requests={request.request_id: request},
+ )
assert request.additional_information == payload
assert request.additional_information["meta"]["finished"].item() is True
@@ -601,7 +1429,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter)
},
8,
)
- assert adapter._poll_single_request(request) is True
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
OmniGenerationScheduler._record_prefill_stats(request)
first_chunk_stats = request.prefill_stats
request.prefill_stats = None
@@ -613,7 +1441,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter)
},
8,
)
- assert adapter._poll_single_request(request) is True
+ assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True
assert isinstance(request.prefill_stats, PrefillStats)
OmniGenerationScheduler._record_prefill_stats(request)
second_chunk_stats = request.prefill_stats
@@ -712,7 +1540,7 @@ def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
assert waiting_queue == []
assert active.status == RequestStatus.WAITING_FOR_CHUNK
assert non_active.status == RequestStatus.WAITING
- assert list(adapter._pending_load_reqs) == [active]
+ assert [entry.request for entry in adapter._pending_load_reqs] == [active]
assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active]
@@ -863,8 +1691,19 @@ def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
adapter.get_req_chunk[req_id] = 3
adapter.requests_with_ready_chunks.add(req_id)
adapter.request_ids_mapping[req_id] = ext_id
- adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id))
+ entry = SimpleNamespace(request_id=req_id)
+ adapter._pending_load_reqs.append(entry)
+ adapter._registered_load_entries[req_id] = entry
adapter._finished_load_reqs.add(req_id)
+ adapter._pending_ar_prompt_updates[req_id] = (
+ SimpleNamespace(request_id=req_id),
+ {},
+ False,
+ False,
+ False,
+ )
+ adapter._streaming_condition_lengths[req_id] = 10
+ adapter._streaming_condition_seqs[req_id] = 2
adapter.put_req_chunk[ext_id] = 5
adapter.request_payload[ext_id] = {"hidden": [1, 2]}
@@ -886,12 +1725,71 @@ def test_cleanup_clears_all_state(build_adapter):
assert req_id not in adapter.request_ids_mapping
assert req_id in adapter._cancelled_load_reqs
assert req_id not in adapter._finished_load_reqs
+ assert req_id not in adapter._pending_ar_prompt_updates
+ assert req_id not in adapter._streaming_condition_lengths
+ assert req_id not in adapter._streaming_condition_seqs
+ assert req_id not in adapter._registered_load_entries
assert ext_id not in adapter.put_req_chunk
assert ext_id not in adapter.request_payload
assert ext_id not in adapter.code_prompt_token_ids
+def test_cleanup_and_reregister_discards_inflight_old_segment_chunk(build_adapter):
+ """A late old-segment chunk cannot overwrite a resumed registration."""
+ adapter, connector = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("req-late", RequestStatus.RUNNING, external_req_id="ext-late")
+ request.resumable = True
+ entry = _dequeue_load_entry(adapter, request)
+
+ get_started = threading.Event()
+ release_get = threading.Event()
+
+ def blocking_get(*_args):
+ get_started.set()
+ assert release_get.wait(timeout=5)
+ return (
+ {
+ "ids": {"prompt": [1]},
+ "meta": {
+ "next_stage_prompt_len": 10,
+ },
+ },
+ 1,
+ )
+
+ connector.get.side_effect = blocking_get
+ poll_results = []
+ poll_thread = threading.Thread(target=lambda: poll_results.append(adapter._poll_single_request(entry)))
+ poll_thread.start()
+ assert get_started.wait(timeout=5)
+
+ adapter.cleanup_receiver(request.request_id)
+ fresh_entry = _dequeue_load_entry(adapter, request)
+ release_get.set()
+ poll_thread.join(timeout=5)
+
+ assert not poll_thread.is_alive()
+ assert poll_results == [True]
+ assert request.additional_information is None
+ assert request.request_id not in adapter._pending_ar_prompt_updates
+ assert request.request_id not in adapter._streaming_condition_seqs
+ assert adapter.get_req_chunk[request.request_id] == 0
+ assert adapter._registered_load_entries[request.request_id] is fresh_entry
+
+ connector.get.side_effect = None
+ connector.get.return_value = (
+ {
+ "ids": {"prompt": [2]},
+ "meta": {"finished": True},
+ },
+ 1,
+ )
+ assert adapter._poll_single_request(fresh_entry) is True
+ assert adapter.get_req_chunk[request.request_id] == 1
+ assert request.request_id not in adapter._registered_load_entries
+
+
def test_cleanup_infers_external_id(build_adapter):
"""When external_req_id is None, cleanup should look it up from the mapping."""
adapter, _ = build_adapter(stage_id=1)
@@ -950,6 +1848,67 @@ def test_cleanup_preserves_pending_save(build_adapter):
assert len(adapter._pending_save_reqs) == 1
+def test_abort_invalidates_queued_sender_task_before_external_id_reuse(build_adapter):
+ adapter, connector = build_adapter(stage_id=1)
+ request = _req("req-old", RequestStatus.WAITING, external_req_id="ext-reused")
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
+ codes=CodesStruct(audio=torch.tensor([1], dtype=torch.long))
+ )
+
+ adapter.save_async(multimodal_output=None, request=request)
+ stale_task = adapter._pending_save_reqs.popleft()
+ adapter.cleanup(request.request_id, request.external_req_id)
+ adapter._send_single_request(stale_task)
+
+ connector.put.assert_not_called()
+ assert request.external_req_id not in adapter.put_req_chunk
+
+ replacement = _req("req-new", RequestStatus.WAITING, external_req_id=request.external_req_id)
+ adapter.save_async(multimodal_output=None, request=replacement)
+ current_task = adapter._pending_save_reqs.popleft()
+ adapter._send_single_request(current_task)
+
+ connector.put.assert_called_once()
+ assert connector.put.call_args.kwargs["put_key"] == "ext-reused_1_0"
+
+
+def test_finish_requests_does_not_wait_for_inflight_send(build_adapter):
+ adapter, connector = build_adapter(stage_id=1)
+ first = _req("req-first", RequestStatus.WAITING, external_req_id="ext-first")
+ adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+ put_started = threading.Event()
+ release_put = threading.Event()
+
+ def blocking_put(**kwargs):
+ put_started.set()
+ release_put.wait(timeout=2)
+ return True, 1, {}
+
+ connector.put.side_effect = blocking_put
+ adapter.save_async(multimodal_output=None, request=first)
+ task = adapter._pending_save_reqs.popleft()
+ sender = threading.Thread(target=adapter._send_single_request, args=(task,))
+ sender.start()
+ assert put_started.wait(timeout=1)
+ cleaner = threading.Thread(
+ target=adapter.finish_requests,
+ args=([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first}),
+ )
+ cleaner.start()
+ try:
+ cleaner.join(timeout=0.5)
+ assert not cleaner.is_alive()
+ assert task["sender_token"].cancelled
+ assert first.external_req_id in adapter._sender_tokens
+ finally:
+ release_put.set()
+ sender.join(timeout=1)
+ cleaner.join(timeout=1)
+ assert not sender.is_alive()
+ assert first.external_req_id not in adapter._sender_tokens
+ assert first.external_req_id not in adapter.put_req_chunk
+
+
def test_cleanup_only_affects_target_request(build_adapter):
"""Cleanup for one request must not affect another request's state."""
adapter, _ = build_adapter(stage_id=1)
@@ -979,7 +1938,7 @@ def test_cleanup_after_poll_flow(build_adapter):
"meta": {"finished": torch.tensor(True, dtype=torch.bool)},
}
connector.get.return_value = (payload, 8)
- adapter._poll_single_request(request)
+ adapter._poll_single_request(_dequeue_load_entry(adapter, request))
assert "req-flow" in adapter.upstream_exhausted_requests
assert adapter.get_req_chunk["req-flow"] == 1
@@ -1006,12 +1965,30 @@ def test_finish_requests_restores_status(build_adapter):
prior = RequestStatus.RUNNING
request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK)
adapter.requests_origin_status[req_id] = prior
+ adapter.waiting_for_chunk_running_requests.append(request)
requests_map = {req_id: request}
adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map)
assert request.status == prior
assert req_id not in adapter.requests_origin_status
+ assert not adapter.waiting_for_chunk_running_requests
+
+
+def test_finish_requests_does_not_restore_stale_status_without_connector_ownership(build_adapter):
+ adapter, _ = build_adapter(stage_id=1)
+ request = _req("req-restored", RequestStatus.WAITING_FOR_STREAMING_REQ)
+ request.resumable = True
+ adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
+
+ adapter.finish_requests(
+ [request.request_id],
+ RequestStatus.FINISHED_ABORTED,
+ {request.request_id: request},
+ )
+
+ assert request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+ assert request.request_id not in adapter.requests_origin_status
def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
@@ -1061,6 +2038,108 @@ def test_finish_requests_releases_active_stream_slot(build_adapter):
assert waiting.status == RequestStatus.WAITING_FOR_CHUNK
+@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
+@pytest.mark.parametrize(
+ ("placement", "origin_status", "initial_status", "streaming_counter_owned"),
+ [
+ ("hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, False),
+ ("running", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
+ ("waiting", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
+ ("skipped", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, True),
+ ("skipped_streaming", RequestStatus.RUNNING, RequestStatus.WAITING_FOR_STREAMING_REQ, True),
+ ("skipped_hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, True),
+ ("skipped_non_streaming", RequestStatus.WAITING, RequestStatus.WAITING, False),
+ ],
+)
+def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
+ build_adapter,
+ scheduler_cls,
+ placement,
+ origin_status,
+ initial_status,
+ streaming_counter_owned,
+):
+ adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
+ request = _req(
+ "req-segment-stop",
+ initial_status,
+ external_req_id="ext-segment-stop",
+ )
+ request.resumable = True
+ adapter.requests_origin_status[request.request_id] = origin_status
+ if placement in {"hidden", "skipped_hidden"}:
+ adapter.waiting_for_chunk_running_requests.append(request)
+ adapter._active_streams[request.request_id] = request
+
+ scheduler = scheduler_cls.__new__(scheduler_cls)
+ scheduler.chunk_transfer_adapter = adapter
+ scheduler.input_coordinator = None
+ scheduler.requests = {request.request_id: request}
+ scheduler.running = [request] if placement == "running" else []
+ scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
+ scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+ scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
+
+ freed = []
+
+ def fake_free_request(self, live, delay_free_blocks=False):
+ del delay_free_blocks
+ freed.append(live.request_id)
+ self.requests.pop(live.request_id)
+ return None, None
+
+ scheduler._free_request = MethodType(fake_free_request, scheduler)
+
+ first = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
+ second = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
+
+ assert len(first) == 1
+ assert second == []
+ assert freed == [request.request_id]
+ assert request.request_id not in scheduler.requests
+ assert request.request_id not in adapter._active_streams
+ assert not adapter.waiting_for_chunk_running_requests
+ assert scheduler.running == []
+ assert list(scheduler.waiting) == []
+ assert list(scheduler.skipped_waiting) == []
+ assert scheduler.num_waiting_for_streaming_input == 0
+
+ fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
+ fresh_b = _req("req-fresh-b", RequestStatus.WAITING)
+ assert adapter._ensure_active_stream(fresh_a)
+ assert adapter._ensure_active_stream(fresh_b)
+ assert set(adapter._active_streams) == {fresh_a.request_id, fresh_b.request_id}
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
+def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapter, scheduler_cls):
+ adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
+ request = _req("req-off-queue", RequestStatus.FINISHED_STOPPED)
+ request.resumable = True
+ adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
+
+ scheduler = scheduler_cls.__new__(scheduler_cls)
+ scheduler.chunk_transfer_adapter = adapter
+ scheduler.input_coordinator = None
+ scheduler.requests = {request.request_id: request}
+ scheduler.running = []
+ scheduler.waiting = DummyWaitingQueue()
+ scheduler.skipped_waiting = DummyWaitingQueue()
+ scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
+
+ assert (
+ scheduler_cls.finish_requests(
+ scheduler,
+ [request.request_id],
+ RequestStatus.FINISHED_ABORTED,
+ )
+ == []
+ )
+ assert scheduler.requests == {request.request_id: request}
+ assert request.status == RequestStatus.FINISHED_STOPPED
+ assert request.request_id not in adapter.requests_origin_status
+
+
def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
adapter, _ = build_adapter(stage_id=1)
zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
@@ -1128,6 +2207,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
scheduler.structured_output_manager.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
+ scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
request = _HashableRequest(
@@ -1208,6 +2288,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager = mocker.MagicMock()
scheduler.kv_cache_manager.take_events.return_value = None
+ scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler.waiting_for_transfer_free = set()
scheduler.transfer_triggered_requests = set()
@@ -1363,6 +2444,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
scheduler.structured_output_manager.should_advance.return_value = False
scheduler.finished_req_ids_dict = {}
scheduler.kv_cache_manager.take_events.return_value = None
+ scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = mocker.MagicMock()
scheduler._pending_finish_reqs = list(pending_finish_reqs)
@@ -1728,3 +2810,183 @@ def test_purge_is_noop_on_empty_deques(build_adapter):
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
assert running_queue == []
assert waiting_queue == []
+
+
+# --------------------------------------------------------------------------- #
+# Chunk-wait deadline (RFC #4855 R1.1, issue #3833)
+#
+# Before this, the async-chunk path had no deadline of any kind, so a dropped
+# terminal chunk or an upstream stage that died mid-stream parked the request
+# in WAITING_FOR_CHUNK forever. None of these scenarios were covered.
+# --------------------------------------------------------------------------- #
+
+
+def _park_in_chunk_wait(adapter, request, *, waiting=True):
+ """Drive a request through one process_pending_chunks round into the wait."""
+ queue = DummyWaitingQueue([request]) if waiting else [request]
+ if waiting:
+ adapter.process_pending_chunks(queue, [])
+ else:
+ adapter.process_pending_chunks(DummyWaitingQueue(), queue)
+ return queue
+
+
+def test_chunk_wait_clock_starts_when_a_request_parks(build_adapter):
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("r1", RequestStatus.WAITING)
+
+ _park_in_chunk_wait(adapter, request)
+
+ assert request.status == RequestStatus.WAITING_FOR_CHUNK
+ assert "r1" in adapter._waiting_since
+
+
+def test_chunk_wait_does_not_expire_before_the_deadline(build_adapter):
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
+
+ assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
+ assert "r1" in adapter._waiting_since
+
+
+def test_dropped_terminal_chunk_expires_the_request(build_adapter):
+ """The #3833 scenario: upstream stops sending and never marks the stream done."""
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
+
+ adapter._waiting_since["r1"] -= 601.0
+
+ assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"r1"}
+ # Cleared, so a second sweep does not re-report the same request.
+ assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
+
+
+def test_arriving_chunk_resets_the_clock(build_adapter):
+ """A slow but healthy stream must never expire: the deadline measures stall
+ time between chunks, not the lifetime of the stream."""
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("r1", RequestStatus.WAITING)
+ queue = _park_in_chunk_wait(adapter, request)
+ adapter.restore_queues(queue, [], scheduler_requests={"r1": request})
+
+ # Age the wait to just under the deadline, then deliver a chunk.
+ adapter._waiting_since["r1"] -= 599.0
+ adapter._finished_load_reqs.add("r1")
+ adapter.process_pending_chunks(queue, [])
+
+ assert "r1" not in adapter._waiting_since
+ assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == set()
+
+
+def test_a_disabled_deadline_never_expires(build_adapter):
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
+ adapter._waiting_since["r1"] -= 10_000.0
+
+ assert adapter.collect_timed_out_request_ids(timeout_s=0.0) == set()
+ assert adapter.collect_timed_out_request_ids(timeout_s=-1.0) == set()
+
+
+def test_finished_request_leaves_no_stale_timestamp(build_adapter):
+ """Otherwise a completed request would be reported as timed out later."""
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ request = _req("r1", RequestStatus.WAITING)
+ _park_in_chunk_wait(adapter, request)
+
+ adapter.finish_requests(["r1"], RequestStatus.FINISHED_STOPPED, {"r1": request})
+
+ assert "r1" not in adapter._waiting_since
+ assert adapter.collect_timed_out_request_ids(timeout_s=0.001) == set()
+
+
+def test_aborted_request_leaves_no_stale_timestamp(build_adapter):
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar")
+ _park_in_chunk_wait(adapter, _req("r1", RequestStatus.WAITING))
+
+ adapter.cleanup_receiver("r1")
+
+ assert "r1" not in adapter._waiting_since
+
+
+def test_expiry_is_per_request(build_adapter):
+ """One stalled stream must not take down its healthy neighbours."""
+ adapter, _ = build_adapter(stage_id=1, model_mode="ar", max_num_seqs=4)
+ stalled = _req("stalled", RequestStatus.WAITING)
+ healthy = _req("healthy", RequestStatus.WAITING)
+ adapter.process_pending_chunks(DummyWaitingQueue([stalled, healthy]), [])
+
+ adapter._waiting_since["stalled"] -= 601.0
+
+ assert adapter.collect_timed_out_request_ids(timeout_s=600.0) == {"stalled"}
+ assert "healthy" in adapter._waiting_since
+
+
+def test_abort_clears_native_codec_state_before_external_id_reuse(build_adapter):
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ adapter, _ = build_adapter(stage_id=1, connector_extra={"codec_chunk_frames": 1})
+ external_req_id = "ext-native-abort"
+ first = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+ first.resumable = True
+ first.model_intermediate_buffer = None
+ first.additional_information = {"meta": {"codec_streaming": True}, "nvc_logical_prompt_len": 1}
+ first_frame = torch.full((1, 31), 101, dtype=torch.long)
+
+ talker2code2wav_async_chunk(
+ adapter,
+ {"codes": {"audio": first_frame}, "meta": {"codec_streaming": True}},
+ first,
+ is_finished=False,
+ )
+
+ adapter.requests_num_chunks_sent[external_req_id] = 1
+ adapter.finish_requests([first.request_id], RequestStatus.FINISHED_ABORTED, {first.request_id: first})
+
+ assert external_req_id not in adapter.request_payload
+ assert external_req_id not in adapter.requests_num_chunks_sent
+
+ replacement = _req("req-native-abort", RequestStatus.WAITING, external_req_id=external_req_id)
+ replacement.resumable = True
+ replacement.model_intermediate_buffer = None
+ replacement.additional_information = first.additional_information
+ replacement_payload = talker2code2wav_async_chunk(
+ adapter,
+ {"codes": {"audio": torch.full((1, 31), 202, dtype=torch.long)}, "meta": {"codec_streaming": True}},
+ replacement,
+ is_finished=False,
+ )
+
+ assert torch.equal(replacement_payload.codes.audio, torch.full((1, 31), 202, dtype=torch.long))
+
+
+def test_save_async_boundary_holds_generation_without_request_counter(build_adapter):
+ """A boundary must not advance past a generation nothing can catch up to.
+
+ When the producing stage maintains ``_omni_segment_generation`` on the
+ request, a boundary may start the next segment's watermark at
+ ``generation + 1`` because the following frame will carry it. On paths
+ where the request has no such counter, that speculative advance makes every
+ subsequent frame look stale and starves the stream permanently
+ (vllm-project/vllm-omni#6816).
+ """
+ adapter, _ = build_adapter(stage_id=1)
+
+ request = _req("req-nogen", RequestStatus.WAITING, external_req_id="ext-nogen")
+ request.resumable = True
+ request._omni_segment_generation = None
+
+ adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True)
+ assert adapter._segment_generation["ext-nogen"] == 0
+
+ follow_up = _req("req-nogen-next", RequestStatus.WAITING, external_req_id="ext-nogen")
+ follow_up.resumable = True
+ follow_up.num_computed_tokens = 3
+ follow_up._omni_segment_generation = None
+ queued_before = len(adapter._pending_save_reqs)
+
+ adapter.save_async(multimodal_output=None, request=follow_up, is_segment_finished=False)
+
+ assert len(adapter._pending_save_reqs) == queued_before + 1
+ assert adapter._pending_save_reqs[-1]["request"] is follow_up
diff --git a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py
index 60653b3cd..66f0327f7 100755
--- a/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py
+++ b/tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Integration tests for MooncakeTransferEngineConnector.
@@ -18,6 +18,7 @@ import pytest
import torch
from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import get_open_port
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
ManagedBuffer,
MooncakeTransferEngineConnector,
@@ -164,14 +165,6 @@ RDMA_HOST = get_rdma_host()
RDMA_DEVICE = _detect_rdma_device()
-def _free_port() -> int:
- import socket
-
- with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
- s.bind((RDMA_HOST, 0))
- return s.getsockname()[1]
-
-
def _connector_config(
zmq_port: int,
pool_size: int = 16 * 1024 * 1024,
@@ -204,7 +197,7 @@ class TestBasicConnector:
"""Verify connector initialization, put, cleanup, and health check."""
def test_initialization(self):
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c:
assert c.rpc_port != 0
assert c.pool_size == 1024 * 1024
@@ -214,7 +207,7 @@ class TestBasicConnector:
def test_put_tensor_bytes_object(self):
"""Put tensor / bytes / dict and verify metadata."""
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100))
assert ok
@@ -229,7 +222,7 @@ class TestBasicConnector:
assert not meta["is_fast_path"]
def test_cleanup_releases_buffer(self):
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
c.put("s0", "s1", "r1", torch.randn(100))
key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1")
@@ -239,7 +232,7 @@ class TestBasicConnector:
def test_pool_exhaustion_and_recovery(self):
"""Fill pool, verify failure, free, verify recovery."""
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024)) as c:
ids = []
for i in range(10):
@@ -264,8 +257,8 @@ class TestEndToEnd:
"""E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
def _pair(self, pool_size=16 * 1024 * 1024):
- p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
- c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+ p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size))
+ c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size))
return p, c
def test_tensor_e2e(self):
@@ -401,7 +394,7 @@ class TestEndToEnd:
def test_concurrent_put(self):
"""10 concurrent puts should all succeed."""
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
conn = MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024 * 1024))
errors: list[str] = []
lock = threading.Lock()
@@ -457,20 +450,20 @@ class TestLifecycle:
"""Close, context manager, double-close safety."""
def test_close_releases_resources(self):
- c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024))
+ c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size=1024 * 1024))
c.put("s0", "s1", "x", torch.randn(100))
c.close()
assert c._stop_event.is_set()
assert len(c._local_buffers) == 0
def test_context_manager(self):
- with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c:
+ with MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST))) as c:
ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50))
assert ok
assert c._stop_event.is_set()
def test_double_close_safe(self):
- c = MooncakeTransferEngineConnector(_connector_config(_free_port()))
+ c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST)))
c.close()
c.close()
@@ -489,12 +482,12 @@ class TestGPUPool:
return _connector_config(port, pool_size, pool_device="cuda:0")
def test_gpu_pool_init(self):
- with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
+ with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c:
assert c.pool_device == "cuda:0"
assert c.pool.is_cuda
def test_gpu_pool_put_cpu_and_gpu_tensor(self):
- with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
+ with MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST))) as c:
ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256))
assert ok
assert meta["is_fast_path"]
@@ -504,8 +497,8 @@ class TestGPUPool:
assert meta["is_fast_path"]
def test_gpu_e2e_transfer(self):
- p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
- c = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
+ p = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST)))
+ c = MooncakeTransferEngineConnector(self._gpu_cfg(get_open_port(RDMA_HOST)))
try:
orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0")
ok, _, meta = p.put("s0", "s1", "ge", orig)
@@ -536,8 +529,8 @@ class TestStressCorrectness:
"""
def _pair(self, pool_size=64 * 1024 * 1024):
- p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
- c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+ p = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size))
+ c = MooncakeTransferEngineConnector(_connector_config(get_open_port(RDMA_HOST), pool_size))
return p, c
# -- Concurrent put + get with data integrity --
@@ -712,7 +705,7 @@ class TestStressCorrectness:
def test_empty_bytes_rejected(self):
"""Connector should gracefully reject empty bytes payload."""
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
ok, sz, meta = c.put("s0", "s1", "empty_b", b"")
assert not ok, "Empty bytes should be rejected by connector"
@@ -744,7 +737,7 @@ class TestStressCorrectness:
def test_rapid_alloc_free_cycle(self):
"""Put + cleanup in tight loop to stress allocator under real connector."""
- port = _free_port()
+ port = get_open_port(RDMA_HOST)
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
for i in range(50):
rid = f"cycle_{i}"
diff --git a/tests/docs/test_example_output_accessors.py b/tests/docs/test_example_output_accessors.py
new file mode 100644
index 000000000..e7a618890
--- /dev/null
+++ b/tests/docs/test_example_output_accessors.py
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guard the public accessor that examples and docs use to read omni outputs.
+
+``OmniRequestOutput`` used to hold a nested ``RequestOutput`` in a
+``request_output`` field. Since #5146 it *inherits* ``RequestOutput``, so the
+generation content lives on the object itself and the nested field is gone.
+
+Call sites that kept the old accessor only fail once a real model runs, which
+means a GPU nightly job is the first thing to notice. These checks pin both
+halves of the contract on CPU instead: the field is really gone, and no doc,
+example, test or tool still reaches for it.
+"""
+
+import re
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+ROOT_DIR = Path(__file__).parents[2]
+SCANNED_DIRS = ("docs", "examples", "tests", "tools")
+SCANNED_SUFFIXES = (".py", ".md")
+
+# ``.request_outputs`` (plural, an OutputProcessor batch) is a different and
+# still-valid attribute, so require a non-identifier character after the name.
+NESTED_ACCESSOR = re.compile(r"\.request_output(?![\w])")
+
+
+def _iter_scanned_files():
+ for directory in SCANNED_DIRS:
+ for path in sorted((ROOT_DIR / directory).rglob("*")):
+ if path.suffix in SCANNED_SUFFIXES and path.resolve() != Path(__file__).resolve():
+ yield path
+
+
+def test_omni_request_output_has_no_nested_request_output():
+ assert not hasattr(OmniRequestOutput(request_id="req-accessor"), "request_output")
+
+
+def test_no_source_reads_the_nested_request_output():
+ offenders = []
+ for path in _iter_scanned_files():
+ for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
+ if NESTED_ACCESSOR.search(line):
+ offenders.append(f"{path.relative_to(ROOT_DIR)}:{lineno}: {line.strip()}")
+
+ assert not offenders, (
+ "OmniRequestOutput no longer nests a RequestOutput; read the content off the "
+ "output itself (e.g. `outputs[0].images`) instead of `.request_output`:\n" + "\n".join(offenders)
+ )
diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py
index 99880fc73..832d34557 100644
--- a/tests/e2e/accuracy/helpers.py
+++ b/tests/e2e/accuracy/helpers.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import json
import os
import re
@@ -14,7 +17,7 @@ from hashlib import sha1
from io import BytesIO
from mimetypes import guess_type
from pathlib import Path
-from typing import Any
+from typing import IO, Any
from urllib.parse import urlparse
import numpy as np
@@ -35,11 +38,14 @@ class IdentityFtfy:
def apply_ftfy_mock(*, wan_i2v_module: Any | None = None) -> None:
"""Mock ftfy library for test cases in the main process"""
- if wan_i2v_module is None:
- from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
+ module: Any = wan_i2v_module
+ if module is None:
+ from diffusers.pipelines.wan import pipeline_wan_i2v as wan_pipeline
- if not hasattr(wan_i2v_module, "ftfy"):
- wan_i2v_module.ftfy = IdentityFtfy()
+ module = wan_pipeline
+
+ if not hasattr(module, "ftfy"):
+ module.ftfy = IdentityFtfy()
print("ftfy (text encoding sanitizer) is not installed. Using mock ftfy implementation (identity function)")
else:
print("ftfy (text encoding sanitizer) is installed. Using actual ftfy implementation.")
@@ -118,6 +124,7 @@ def probe_binary(binary: str, *, test_name: str = "video similarity e2e test") -
resolved = shutil.which(binary)
if resolved is None:
pytest.skip(f"{binary} is required for {test_name}.")
+ assert resolved is not None
return resolved
@@ -235,7 +242,7 @@ def video_artifact_dir(result_root: Path, source: str) -> Path:
def send_video_request_with_timeout(
- openai_client,
+ online_client,
request_config: dict[str, Any],
*,
timeout_seconds: int,
@@ -266,7 +273,7 @@ def send_video_request_with_timeout(
normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
start_time = time.perf_counter()
- create_url = openai_client._build_url("/v1/videos")
+ create_url = online_client._build_url("/v1/videos")
response = requests.post(
create_url,
data=normalized_form_data,
@@ -280,8 +287,8 @@ def send_video_request_with_timeout(
raise requests.HTTPError(f"{exc}; response body: {response.text}", response=response) from exc
job_data = response.json()
video_id = job_data["id"]
- openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
- video_content = openai_client._download_video_content(video_id)
+ online_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
+ video_content = online_client._download_video_content(video_id)
print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
return video_content
@@ -571,11 +578,11 @@ class SemanticSimilarityScorer:
raise ImportError("FlagEmbedding not installed. Install with: pip install FlagEmbedding")
return self._text_model
- def text_similarity(self, text1: str, text2: str) -> float:
+ def text_similarity(self, text1: str, text2: str) -> dict[str, float | int]:
"""
Compute semantic similarity between two texts using BGE-M3
"""
- results = {}
+ results: dict[str, float | int] = {}
# 1. Text prefix matching (maximum matching character length)
match_count = 0
for char_vllm, char_ref in zip(text1, text2):
@@ -653,7 +660,7 @@ def dreamzero_cfg_parallel_size() -> int:
def _dreamzero_parity_vllm_env() -> dict[str, str]:
- from tests.helpers.env import pick_least_used_device_indices
+ from tests.helpers.clean import pick_least_used_device_indices
from tests.helpers.runtime import get_open_port
parallel_size = dreamzero_cfg_parallel_size()
@@ -683,8 +690,9 @@ def dreamzero_parity_vllm_params(model: str | None = None):
]
if parallel_size > 1:
server_args.extend(["--cfg-parallel-size", str(parallel_size)])
+ resolved_model = model or os.environ.get("VLLM_DREAMZERO_MODEL") or _DREAMZERO_DEFAULT_MODEL
return OmniServerParams(
- model=model or os.environ.get("VLLM_DREAMZERO_MODEL", _DREAMZERO_DEFAULT_MODEL),
+ model=resolved_model,
stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"),
server_args=server_args,
env_dict=_dreamzero_parity_vllm_env(),
@@ -725,9 +733,9 @@ def dreamzero_parity_vllm_client(
def _dreamzero_create_openpi_vllm_client(host: str, port: int) -> Any:
from openpi_client import msgpack_numpy
- _, WebsocketClientPolicy = load_dreamzero_upstream_modules()
+ _, websocket_client_policy_cls = load_dreamzero_upstream_modules()
- class _OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
+ class _OpenPIWebsocketClientPolicy(websocket_client_policy_cls): # type: ignore[misc, valid-type]
def __init__(
self,
host: str = "127.0.0.1",
@@ -869,7 +877,7 @@ class _DreamZeroParityServer(ABC):
self.port = get_open_port() if port is None else port
self.ready_timeout_s = ready_timeout_s
self.proc: subprocess.Popen[str] | None = None
- self._log_file = None
+ self._log_file: IO[str] | None = None
@abstractmethod
def _build_argv(self) -> list[str]:
@@ -932,7 +940,7 @@ class DreamZeroUpstreamServer(_DreamZeroParityServer):
]
def _build_env(self) -> dict[str, str]:
- from tests.helpers.env import pick_least_used_device_indices
+ from tests.helpers.clean import pick_least_used_device_indices
dreamzero_repo, _ = require_dreamzero_parity_env()
env = os.environ.copy()
diff --git a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py
index b830cdc6c..eb8f2d769 100644
--- a/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py
+++ b/tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import os
@@ -149,7 +152,7 @@ def _generate_offline_video(*, image_source: str) -> Path:
def _generate_online_video(
*,
omni_server,
- openai_client,
+ online_client,
image_source: str,
timeout_seconds: int,
) -> Path:
@@ -171,7 +174,7 @@ def _generate_online_video(
"image_reference": build_online_image_reference(online_image_source),
}
online_video_bytes = send_video_request_with_timeout(
- openai_client,
+ online_client,
request_config,
timeout_seconds=timeout_seconds,
)
@@ -204,7 +207,7 @@ def test_hunyuanvideo15_i2v_diffusers_offline_generates_video(
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
def test_hunyuanvideo15_i2v_online_serving_generates_video(
omni_server,
- openai_client,
+ online_client,
hunyuanvideo15_i2v_image_source: str | None,
hunyuanvideo15_online_timeout_seconds: int,
) -> None:
@@ -216,7 +219,7 @@ def test_hunyuanvideo15_i2v_online_serving_generates_video(
validate_image_source(image_source)
online_path = _generate_online_video(
omni_server=omni_server,
- openai_client=openai_client,
+ online_client=online_client,
image_source=image_source,
timeout_seconds=hunyuanvideo15_online_timeout_seconds,
)
diff --git a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py
index f2c14e2b0..f5b90822d 100644
--- a/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py
+++ b/tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import os
@@ -115,7 +118,7 @@ def _generate_offline_video() -> Path:
return offline_path
-def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) -> Path:
+def _generate_online_video(*, omni_server, online_client, timeout_seconds: int) -> Path:
online_path, _ = _artifact_paths()
request_config = {
"model": omni_server.model,
@@ -132,7 +135,7 @@ def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int)
},
}
online_video_bytes = send_video_request_with_timeout(
- openai_client,
+ online_client,
request_config,
timeout_seconds=timeout_seconds,
)
@@ -161,7 +164,7 @@ def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None:
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
def test_hunyuanvideo15_t2v_online_serving_generates_video(
omni_server,
- openai_client,
+ online_client,
hunyuanvideo15_online_timeout_seconds: int,
) -> None:
if not torch.cuda.is_available():
@@ -170,7 +173,7 @@ def test_hunyuanvideo15_t2v_online_serving_generates_video(
probe_binary("ffprobe")
online_path = _generate_online_video(
omni_server=omni_server,
- openai_client=openai_client,
+ online_client=online_client,
timeout_seconds=hunyuanvideo15_online_timeout_seconds,
)
assert online_path.exists(), f"Expected online video artifact at {online_path}"
diff --git a/tests/e2e/accuracy/ltx/run_ltx25_reference.py b/tests/e2e/accuracy/ltx/run_ltx25_reference.py
new file mode 100644
index 000000000..cbc8917a8
--- /dev/null
+++ b/tests/e2e/accuracy/ltx/run_ltx25_reference.py
@@ -0,0 +1,591 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Generate raw LTX video and audio outputs with the official or Omni runtime."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--backend", choices=("official", "omni"), required=True)
+ parser.add_argument("--request", type=Path, required=True)
+ parser.add_argument("--output-dir", type=Path, required=True)
+ parser.add_argument("--model")
+ parser.add_argument("--model-class-name")
+ parser.add_argument("--official-root", type=Path)
+ parser.add_argument("--checkpoint", type=Path)
+ parser.add_argument("--gemma-root", type=Path)
+ parser.add_argument("--transformer-path", type=Path)
+ parser.add_argument("--text-encoder-path", type=Path)
+ parser.add_argument("--video-vae-path", type=Path)
+ parser.add_argument("--audio-vae-path", type=Path)
+ parser.add_argument("--spatial-upsampler-path", type=Path)
+ parser.add_argument("--distilled-lora-path", type=Path)
+ parser.add_argument("--connector-model", type=Path)
+ parser.add_argument(
+ "--official-pipeline",
+ choices=("distilled_two_stage", "full_one_stage", "full_two_stage"),
+ default="distilled_two_stage",
+ )
+ parser.add_argument("--enable-layerwise-offload", action="store_true")
+ return parser.parse_args()
+
+
+def _save_outputs(
+ output_dir: Path,
+ *,
+ video: torch.Tensor,
+ audio: torch.Tensor,
+ audio_sample_rate: int,
+ fps: float,
+ metadata: dict[str, Any],
+) -> None:
+ output_dir.mkdir(parents=True, exist_ok=True)
+ video = video.detach().float().cpu().clamp(0.0, 1.0)
+ audio = audio.detach().float().cpu()
+ np.save(output_dir / "video.npy", video.numpy())
+ np.save(output_dir / "audio.npy", audio.numpy())
+
+ from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+
+ frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy()
+ audio_array = audio.numpy()
+ while audio_array.ndim > 2 and audio_array.shape[0] == 1:
+ audio_array = audio_array[0]
+ (output_dir / "output.mp4").write_bytes(
+ mux_video_audio_bytes(
+ frames_u8,
+ audio_array,
+ fps=float(fps),
+ audio_sample_rate=int(audio_sample_rate),
+ video_codec_options={"preset": "ultrafast"},
+ )
+ )
+
+ frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1})
+ for index in frame_indices:
+ frame = video[index].mul(255.0).round().to(torch.uint8).numpy()
+ Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png")
+
+ metadata.update(
+ {
+ "video_shape": list(video.shape),
+ "audio_shape": list(audio.shape),
+ "audio_sample_rate": int(audio_sample_rate),
+ "frame_indices": frame_indices,
+ "mp4": "output.mp4",
+ }
+ )
+ (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
+
+
+def _reset_accelerator_peak_memory_stats() -> None:
+ if torch.accelerator.is_available():
+ torch.accelerator.reset_peak_memory_stats()
+
+
+def _accelerator_peak_memory_stats() -> dict[str, float]:
+ if not torch.accelerator.is_available():
+ return {}
+ torch.accelerator.synchronize()
+ mib = 1024**2
+ return {
+ "peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib,
+ "peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib,
+ }
+
+
+def _insert_official_paths(official_root: Path) -> None:
+ for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"):
+ path = str((official_root / relative_path).resolve())
+ if path not in sys.path:
+ sys.path.insert(0, path)
+
+
+def _configure_official_sdpa(pipeline: Any) -> None:
+ """Pin official connector and denoiser attention to cuDNN, with MATH fallback."""
+ from ltx_core.loader.attention_ops import set_attention_module_op
+ from ltx_core.model.transformer.attention import PytorchAttention
+ from torch.nn.attention import SDPBackend
+
+ attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH])
+ module_op = set_attention_module_op(
+ attention=attention,
+ masked_attention=attention,
+ )
+ transformer_owners = [
+ owner
+ for attribute in ("stage", "stage_1", "stage_2")
+ if (owner := getattr(pipeline, attribute, None)) is not None
+ ]
+ owners_and_attributes = [
+ *((owner, "_transformer_builder") for owner in transformer_owners),
+ (pipeline.prompt_encoder, "_embeddings_processor_builder"),
+ ]
+ for owner, attribute in owners_and_attributes:
+ builder = getattr(owner, attribute)
+ setattr(
+ owner,
+ attribute,
+ builder.with_module_ops((*builder.module_ops, module_op)),
+ )
+
+
+def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None:
+ """Run official connector code with the checkpoint under test."""
+ from safetensors import safe_open
+
+ connector_root = model / "connectors"
+ shards = sorted(connector_root.glob("*.safetensors"))
+ if not shards:
+ raise ValueError(f"No connector safetensors found under {connector_root}")
+
+ original_build_embeddings_processor = prompt_encoder._build_embeddings_processor
+
+ def build_embeddings_processor():
+ processor = original_build_embeddings_processor()
+ parameters = dict(processor.named_parameters())
+ expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))}
+ loaded: set[str] = set()
+ with torch.no_grad():
+ for shard in shards:
+ with safe_open(str(shard), framework="pt", device="cpu") as weights:
+ for source_name in weights.keys():
+ if not source_name.startswith(("video_connector.", "audio_connector.")):
+ continue
+ target_name = (
+ source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.")
+ .replace(".attn1.norm_q.", ".attn1.q_norm.")
+ .replace(".attn1.norm_k.", ".attn1.k_norm.")
+ )
+ parameter = parameters[target_name]
+ parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype))
+ loaded.add(target_name)
+ if loaded != expected:
+ missing = sorted(expected - loaded)
+ unexpected = sorted(loaded - expected)
+ raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}")
+ return processor
+
+ prompt_encoder._build_embeddings_processor = build_embeddings_processor
+
+
+@torch.inference_mode()
+def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
+ if args.official_root is None:
+ raise ValueError("Official backend requires --official-root")
+ _insert_official_paths(args.official_root)
+ # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is
+ # still pinned independently to the explicit cuDNN backend below.
+ torch.backends.cuda.enable_cudnn_sdp(False)
+ _reset_accelerator_peak_memory_stats()
+
+ from ltx_pipelines.utils.args import ImageConditioningInput
+ from ltx_pipelines.utils.model_paths import ModelPaths
+ from ltx_pipelines.utils.types import OffloadMode
+
+ offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
+ image_path = request.get("image")
+ images = (
+ []
+ if image_path is None
+ else [
+ ImageConditioningInput(
+ path=str(image_path),
+ frame_idx=0,
+ strength=1.0,
+ crf=request.get("image_crf", 18),
+ )
+ ]
+ )
+
+ if args.transformer_path is not None:
+ required = {
+ "--text-encoder-path": args.text_encoder_path,
+ "--video-vae-path": args.video_vae_path,
+ "--audio-vae-path": args.audio_vae_path,
+ }
+ if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}:
+ required["--spatial-upsampler-path"] = args.spatial_upsampler_path
+ if args.official_pipeline == "full_two_stage":
+ required["--distilled-lora-path"] = args.distilled_lora_path
+ missing = [name for name, value in required.items() if value is None]
+ if missing:
+ raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}")
+
+ model_paths = ModelPaths.from_split(
+ transformer_path=str(args.transformer_path),
+ text_encoder_path=str(args.text_encoder_path),
+ video_vae_path=str(args.video_vae_path),
+ audio_vae_path=str(args.audio_vae_path),
+ )
+ if args.official_pipeline == "distilled_two_stage":
+ from ltx_pipelines.distilled import DistilledPipeline
+
+ pipeline = DistilledPipeline(
+ model_paths=model_paths,
+ spatial_upsampler_path=str(args.spatial_upsampler_path),
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ if args.connector_model is not None:
+ _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+ _configure_official_sdpa(pipeline)
+ video, audio, _, _ = pipeline(
+ prompt=request["prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ images=images,
+ )
+ pipeline_name = "DistilledPipeline"
+ elif args.official_pipeline == "full_two_stage":
+ from ltx_core.components.guiders import MultiModalGuiderParams
+ from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+ from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+
+ pipeline = TI2VidTwoStagesPipeline(
+ model_paths=model_paths,
+ distilled_lora=[
+ LoraPathStrengthAndSDOps(
+ path=str(args.distilled_lora_path),
+ strength=1.0,
+ sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+ )
+ ],
+ spatial_upsampler_path=str(args.spatial_upsampler_path),
+ loras=[],
+ offload_mode=offload_mode,
+ )
+ if args.connector_model is not None:
+ _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+ _configure_official_sdpa(pipeline)
+ video, audio, _, _ = pipeline(
+ prompt=request["prompt"],
+ negative_prompt=request["negative_prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ num_inference_steps=request["num_inference_steps"],
+ video_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["video_cfg_scale"],
+ stg_scale=request["video_stg_scale"],
+ rescale_scale=request["video_rescale_scale"],
+ modality_scale=request["video_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["video_stg_blocks"],
+ ),
+ audio_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["audio_cfg_scale"],
+ stg_scale=request["audio_stg_scale"],
+ rescale_scale=request["audio_rescale_scale"],
+ modality_scale=request["audio_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["audio_stg_blocks"],
+ ),
+ images=images,
+ max_batch_size=4,
+ stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32),
+ stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32),
+ )
+ pipeline_name = "TI2VidTwoStagesPipeline"
+ else:
+ from ltx_core.components.guiders import MultiModalGuiderParams
+ from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+
+ pipeline = TI2VidOneStagePipeline(
+ model_paths=model_paths,
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ if args.connector_model is not None:
+ _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+ _configure_official_sdpa(pipeline)
+ request_sigmas = request.get("sigmas")
+ video, audio, _ = pipeline(
+ prompt=request["prompt"],
+ negative_prompt=request["negative_prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ num_inference_steps=request["num_inference_steps"],
+ video_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["video_cfg_scale"],
+ stg_scale=request["video_stg_scale"],
+ rescale_scale=request["video_rescale_scale"],
+ modality_scale=request["video_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["video_stg_blocks"],
+ ),
+ audio_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["audio_cfg_scale"],
+ stg_scale=request["audio_stg_scale"],
+ rescale_scale=request["audio_rescale_scale"],
+ modality_scale=request["audio_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["audio_stg_blocks"],
+ ),
+ images=images,
+ max_batch_size=4,
+ sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32),
+ )
+ pipeline_name = "TI2VidOneStagePipeline"
+ checkpoint_metadata = {
+ "transformer_path": str(args.transformer_path),
+ "video_vae_path": str(args.video_vae_path),
+ "connector_model": None if args.connector_model is None else str(args.connector_model),
+ "pipeline": pipeline_name,
+ }
+ else:
+ if args.checkpoint is None or args.gemma_root is None:
+ raise ValueError("Official monolith backend requires --checkpoint and --gemma-root")
+
+ from ltx_core.components.guiders import MultiModalGuiderParams
+ from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+
+ pipeline = TI2VidOneStagePipeline(
+ model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)),
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ _configure_official_sdpa(pipeline)
+ video, audio, _ = pipeline(
+ prompt=request["prompt"],
+ negative_prompt=request["negative_prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ num_inference_steps=request["num_inference_steps"],
+ video_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["video_cfg_scale"],
+ stg_scale=request["video_stg_scale"],
+ rescale_scale=request["video_rescale_scale"],
+ modality_scale=request["video_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["video_stg_blocks"],
+ ),
+ audio_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["audio_cfg_scale"],
+ stg_scale=request["audio_stg_scale"],
+ rescale_scale=request["audio_rescale_scale"],
+ modality_scale=request["audio_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["audio_stg_blocks"],
+ ),
+ images=images,
+ max_batch_size=4,
+ )
+ checkpoint_metadata = {
+ "checkpoint": str(args.checkpoint),
+ "pipeline": "TI2VidOneStagePipeline",
+ }
+
+ video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
+ _save_outputs(
+ args.output_dir,
+ video=_canonical_video(video_tensor),
+ audio=audio.waveform,
+ audio_sample_rate=audio.sampling_rate,
+ fps=float(request["fps"]),
+ metadata={
+ "backend": "official",
+ "attention_backend": "torch_sdpa_cudnn_bf16",
+ "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
+ "seed": request["seed"],
+ "sigmas": request.get("sigmas"),
+ "stage_1_sigmas": request.get("stage_1_sigmas"),
+ "stage_2_sigmas": request.get("stage_2_sigmas"),
+ **_accelerator_peak_memory_stats(),
+ **checkpoint_metadata,
+ },
+ )
+
+
+def _omni_worker_peak_memory_mb(output: Any) -> float:
+ result = output[0] if isinstance(output, list) and output else output
+ return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
+ from vllm_omni.outputs import OmniRequestOutput
+
+ audio = None
+ audio_sample_rate = None
+ frames = output[0] if isinstance(output, list) and output else output
+ if isinstance(frames, OmniRequestOutput):
+ multimodal_output = frames.multimodal_output or {}
+ audio = multimodal_output.get("audio")
+ audio_sample_rate = multimodal_output.get("audio_sample_rate")
+ nested_output = getattr(frames, "request_output", None)
+ if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput):
+ frames = nested_output
+ multimodal_output = frames.multimodal_output or {}
+ audio = multimodal_output.get("audio", audio)
+ audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
+ if isinstance(frames, OmniRequestOutput):
+ if not frames.images:
+ raise ValueError("No video frames found in OmniRequestOutput")
+ frames = frames.images
+
+ if isinstance(frames, list) and len(frames) == 1:
+ frames = frames[0]
+ if isinstance(frames, tuple) and len(frames) == 2:
+ frames, audio = frames
+ elif isinstance(frames, dict):
+ audio = frames.get("audio", audio)
+ audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
+ frames = frames.get("frames", frames.get("video"))
+
+ if frames is None or audio is None or audio_sample_rate is None:
+ raise ValueError("Omni output did not contain video, audio, and audio_sample_rate")
+ return frames, audio, int(audio_sample_rate)
+
+
+def _canonical_video(video: Any) -> torch.Tensor:
+ if isinstance(video, list):
+ if len(video) == 1:
+ return _canonical_video(video[0])
+ tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video]
+ if all(item.ndim == 3 for item in tensors):
+ video = torch.stack(tensors)
+ elif all(item.ndim == 4 for item in tensors):
+ video = torch.cat(tensors)
+ else:
+ raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}")
+ tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu()
+ if tensor.ndim == 5:
+ tensor = tensor[0]
+ if tensor.ndim != 4:
+ raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
+ if tensor.shape[-1] in (3, 4):
+ tensor = tensor[..., :3]
+ elif tensor.shape[1] in (3, 4):
+ tensor = tensor[:, :3].permute(0, 2, 3, 1)
+ elif tensor.shape[0] in (3, 4):
+ tensor = tensor[:3].permute(1, 2, 3, 0)
+ else:
+ raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}")
+ tensor = tensor.float()
+ if tensor.numel() and tensor.max() > 1:
+ tensor = tensor / 255.0
+ elif tensor.numel() and tensor.min() < 0:
+ tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5)
+ return tensor.clamp(0.0, 1.0)
+
+
+@torch.inference_mode()
+def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
+ if args.model is None or args.model_class_name is None:
+ raise ValueError("Omni backend requires --model and --model-class-name")
+
+ # vLLM imports can otherwise change the process-wide Gemma dispatch.
+ torch.backends.cuda.enable_cudnn_sdp(False)
+ _reset_accelerator_peak_memory_stats()
+
+ attention_config = {"default": {"backend": "CUDNN_ATTN"}}
+
+ from vllm_omni.diffusion.data import DiffusionParallelConfig
+ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+ from vllm_omni.model_extras import get_extra_body_params
+ from vllm_omni.platforms import current_omni_platform
+
+ generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"])
+ omni = Omni(
+ model=args.model,
+ model_class_name=args.model_class_name,
+ enforce_eager=True,
+ enable_layerwise_offload=args.enable_layerwise_offload,
+ diffusion_attention_config=attention_config,
+ parallel_config=DiffusionParallelConfig(),
+ )
+ try:
+ # The explicit public class argument selects the worker pipeline. The
+ # parent Omni proxy can only expose the checkpoint's base class, which
+ # is insufficient to distinguish LTX execution modes.
+ model_class_name = args.model_class_name
+ sampling_params = OmniDiffusionSamplingParams(
+ height=request["height"],
+ width=request["width"],
+ generator=generator,
+ guidance_scale=None,
+ num_inference_steps=request["num_inference_steps"],
+ num_frames=request["num_frames"],
+ fps=request["fps"],
+ frame_rate=float(request["fps"]),
+ output_type="np",
+ )
+ extra_args = {
+ key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
+ }
+ for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"):
+ if key in request:
+ extra_args[key] = request[key]
+ apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args)
+ prompt: dict[str, Any] = {"prompt": request["prompt"]}
+ if "negative_prompt" in request:
+ prompt["negative_prompt"] = request["negative_prompt"]
+ image_path = request.get("image")
+ if image_path is not None:
+ with Image.open(str(image_path)) as source_image:
+ image = source_image.convert("RGB")
+ image.load()
+ prompt["multi_modal_data"] = {"image": image}
+ output = omni.generate(prompt, sampling_params)
+ video, audio, audio_sample_rate = _unwrap_omni_output(output)
+ audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio)
+ _save_outputs(
+ args.output_dir,
+ video=_canonical_video(video),
+ audio=audio_tensor,
+ audio_sample_rate=audio_sample_rate,
+ fps=float(request["fps"]),
+ metadata={
+ "backend": "omni",
+ "attention_backend": "torch_sdpa_cudnn_bf16",
+ "seed": request["seed"],
+ "sigmas": request.get("sigmas"),
+ "model": args.model,
+ "stage_1_sigmas": request.get("stage_1_sigmas"),
+ "stage_2_sigmas": request.get("stage_2_sigmas"),
+ "model_class_name": model_class_name,
+ "worker_peak_memory_mb": _omni_worker_peak_memory_mb(output),
+ **_accelerator_peak_memory_stats(),
+ },
+ )
+ finally:
+ omni.shutdown()
+
+
+def main() -> None:
+ args = _parse_args()
+ request = json.loads(args.request.read_text())
+ if args.backend == "official":
+ _run_official(args, request)
+ else:
+ _run_omni(args, request)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py
index a4a180047..ccabd4596 100644
--- a/tests/e2e/accuracy/ltx/run_ltx_reference.py
+++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py
@@ -20,6 +20,11 @@ from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
+ parser.add_argument(
+ "--pipeline-kind",
+ choices=("one_stage", "distilled", "two_stage"),
+ default="one_stage",
+ )
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
@@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
+ parser.add_argument("--spatial-upsampler", type=Path)
+ parser.add_argument("--distilled-lora", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
@@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None:
attention=attention,
masked_attention=attention,
)
- owners_and_attributes = (
- (pipeline.stage, "_transformer_builder"),
+ owners_and_attributes = [
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
+ ]
+ owners_and_attributes.extend(
+ (stage, "_transformer_builder")
+ for name in ("stage", "stage_1", "stage_2")
+ if (stage := getattr(pipeline, name, None)) is not None
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
@@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None:
)
+def _require_path(path: Path | None, description: str) -> str:
+ if path is None or not path.is_file():
+ raise ValueError(f"Official {description} is required and must exist: {path}")
+ return str(path)
+
+
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
@@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
+ from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+ from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+ from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
- pipeline = TI2VidOneStagePipeline(
- checkpoint_path=str(args.checkpoint),
- gemma_root=str(args.gemma_root),
- loras=(),
- offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
- )
+ if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
+ raise ValueError(
+ f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
+ )
+ checkpoint = _require_path(args.checkpoint, "checkpoint")
+ gemma_root = str(args.gemma_root)
+ if not Path(gemma_root).is_dir():
+ raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}")
+ offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
+ pipeline: Any
+ if args.pipeline_kind == "one_stage":
+ pipeline = TI2VidOneStagePipeline(
+ checkpoint_path=checkpoint,
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ elif args.pipeline_kind == "distilled":
+ pipeline = DistilledPipeline(
+ distilled_checkpoint_path=checkpoint,
+ spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ else:
+ distilled_lora = [
+ LoraPathStrengthAndSDOps(
+ path=_require_path(args.distilled_lora, "distilled LoRA"),
+ strength=1.0,
+ sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+ )
+ ]
+ pipeline = TI2VidTwoStagesPipeline(
+ checkpoint_path=checkpoint,
+ distilled_lora=distilled_lora,
+ spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
@@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
)
]
)
- video, audio = pipeline(
- prompt=request["prompt"],
- negative_prompt=request["negative_prompt"],
- seed=request["seed"],
- height=request["height"],
- width=request["width"],
- num_frames=request["num_frames"],
- frame_rate=request["fps"],
- num_inference_steps=request["num_inference_steps"],
- video_guider_params=MultiModalGuiderParams(
- cfg_scale=request["video_cfg_scale"],
- stg_scale=request["video_stg_scale"],
- rescale_scale=request["video_rescale_scale"],
- modality_scale=request["video_modality_scale"],
- skip_step=0,
- stg_blocks=request["video_stg_blocks"],
- ),
- audio_guider_params=MultiModalGuiderParams(
- cfg_scale=request["audio_cfg_scale"],
- stg_scale=request["audio_stg_scale"],
- rescale_scale=request["audio_rescale_scale"],
- modality_scale=request["audio_modality_scale"],
- skip_step=0,
- stg_blocks=request["audio_stg_blocks"],
- ),
- images=images,
- max_batch_size=4,
- )
+ if args.pipeline_kind == "distilled":
+ video, audio = pipeline(
+ prompt=request["prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ images=images,
+ )
+ else:
+ video, audio = pipeline(
+ prompt=request["prompt"],
+ negative_prompt=request["negative_prompt"],
+ seed=request["seed"],
+ height=request["height"],
+ width=request["width"],
+ num_frames=request["num_frames"],
+ frame_rate=request["fps"],
+ num_inference_steps=request["num_inference_steps"],
+ video_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["video_cfg_scale"],
+ stg_scale=request["video_stg_scale"],
+ rescale_scale=request["video_rescale_scale"],
+ modality_scale=request["video_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["video_stg_blocks"],
+ ),
+ audio_guider_params=MultiModalGuiderParams(
+ cfg_scale=request["audio_cfg_scale"],
+ stg_scale=request["audio_stg_scale"],
+ rescale_scale=request["audio_rescale_scale"],
+ modality_scale=request["audio_modality_scale"],
+ skip_step=0,
+ stg_blocks=request["audio_stg_blocks"],
+ ),
+ images=images,
+ max_batch_size=4,
+ )
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
@@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
+ "pipeline_kind": args.pipeline_kind,
},
)
@@ -193,15 +260,9 @@ def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
- if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
- frames = frames.request_output
- multimodal_output = frames.multimodal_output or {}
- audio = multimodal_output.get("audio", audio)
- audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
- if isinstance(frames, OmniRequestOutput):
- if not frames.images:
- raise ValueError("No video frames found in OmniRequestOutput")
- frames = frames.images
+ if not frames.images:
+ raise ValueError("No video frames found in OmniRequestOutput")
+ frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
@@ -253,6 +314,10 @@ def _canonical_video(video: Any) -> torch.Tensor:
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
+ if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
+ raise ValueError(
+ f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
+ )
attention_config = {"default": {"backend": "TORCH_SDPA"}}
@@ -273,7 +338,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
parallel_config=DiffusionParallelConfig(),
)
try:
- model_class_name = get_model_class_name(omni)
+ detected_model_class_name = get_model_class_name(omni)
+ model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
@@ -289,10 +355,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
- prompt: dict[str, Any] = {
- "prompt": request["prompt"],
- "negative_prompt": request["negative_prompt"],
- }
+ prompt: dict[str, Any] = {"prompt": request["prompt"]}
+ if request.get("negative_prompt"):
+ prompt["negative_prompt"] = request["negative_prompt"]
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
@@ -312,6 +377,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
+ "model_config_class_name": detected_model_class_name,
+ "pipeline_kind": args.pipeline_kind,
},
)
finally:
diff --git a/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py
new file mode 100644
index 000000000..c8ce29c37
--- /dev/null
+++ b/tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py
@@ -0,0 +1,748 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""E2E accuracy guards against pinned official Lightricks LTX references.
+
+The LTX-2/2.3 comparison runs both runtimes through PyTorch SDPA and uses
+``max_batch_size=4`` in the official reference to match Omni's fused guidance
+batch. Video and audio guidance use the official non-HQ one-stage defaults;
+only the generation shape and step count are reduced for CI runtime.
+LTX-2.5 runs the official split-artifact pipelines with connector weights from
+the same Diffusers checkpoint under test. The distilled case covers the fixed
+two-stage schedule; the Full/SFT case covers raw dev weights and an explicit
+shared one-stage schedule.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import subprocess
+import sys
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+
+from tests.e2e.accuracy.helpers import reset_artifact_dir
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_test
+
+OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git"
+LTX25_OFFICIAL_REVISION = "7954dcb0d986bdc36ef272564a9789ade07fcc65"
+LTX25_OMNI_MODEL_ID = "Lightricks/LTX-2.5-Diffusers"
+LTX25_OMNI_MODEL_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
+LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
+LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
+LTX25_OFFICIAL_COMMON_FILES = {
+ "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
+ "video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
+ "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
+}
+LTX25_OFFICIAL_VARIANT_FILES = {
+ "distilled": {
+ "transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
+ "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
+ },
+ "full": {
+ "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
+ },
+ "full_two_stage": {
+ "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
+ "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
+ "distilled_lora": "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
+ },
+}
+# Version selected by the pinned official source's uv.lock. Keep it isolated
+# from Omni's runtime and development dependencies.
+OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0"
+LTX25_PROMPT = "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside."
+LTX25_I2V_IMAGE_REPO = "huggingface/documentation-images"
+LTX25_I2V_IMAGE_FILENAME = "diffusers/svd/rocket.png"
+LTX25_I2V_IMAGE_REVISION = "645d8364f0c7a101180b364811b5a11a362e4010"
+NEGATIVE_PROMPT = (
+ "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
+ "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
+ "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
+ "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, "
+ "background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting "
+ "direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, "
+ "incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, "
+ "silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect "
+ "dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural "
+ "transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, "
+ "stylized filters, or AI artifacts."
+)
+
+# Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
+# Degenerate one-token attention falls back to PyTorch SDPA MATH.
+ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
+LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
+
+
+@dataclass(frozen=True)
+class LTXAccuracyThresholds:
+ video_ssim_mean: float
+ video_ssim_min: float
+ video_psnr_mean_db: float
+ audio_relative_l2: float
+ audio_cosine_similarity: float
+
+
+STRICT_THRESHOLDS = LTXAccuracyThresholds(
+ video_ssim_mean=0.99,
+ video_ssim_min=0.99,
+ video_psnr_mean_db=40.0,
+ audio_relative_l2=0.10,
+ audio_cosine_similarity=0.99,
+)
+
+
+def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
+ from vllm_omni.outputs import OmniRequestOutput
+
+ from .run_ltx25_reference import _omni_worker_peak_memory_mb, _unwrap_omni_output
+
+ frame = object()
+ audio = object()
+ output = OmniRequestOutput(
+ stage_id=0,
+ images=[frame],
+ _multimodal_output={"audio": audio, "audio_sample_rate": 48_000},
+ )
+
+ frames, actual_audio, sample_rate = _unwrap_omni_output(output)
+
+ assert frames is frame
+ assert actual_audio is audio
+ assert sample_rate == 48_000
+ output.peak_memory_mb = 12_345.0
+ assert _omni_worker_peak_memory_mb([output]) == 12_345.0
+
+
+def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
+ start = time.perf_counter()
+ subprocess.run(command, env=env, timeout=timeout, check=True)
+ print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s")
+
+
+def _clone_pinned_source(root: Path, repository: str, revision: str) -> None:
+ root.mkdir(parents=True)
+ subprocess.run(["git", "init", "-q", str(root)], check=True)
+ subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True)
+ last_error: subprocess.CalledProcessError | None = None
+ for attempt in range(3):
+ try:
+ subprocess.run(
+ ["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision],
+ check=True,
+ )
+ last_error = None
+ break
+ except subprocess.CalledProcessError as error:
+ last_error = error
+ if attempt < 2:
+ time.sleep(5 * (attempt + 1))
+ if last_error is not None:
+ raise last_error
+ subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
+
+
+def _git_revision(root: Path) -> str | None:
+ try:
+ return subprocess.check_output(
+ ["git", "-C", str(root), "rev-parse", "HEAD"],
+ stderr=subprocess.DEVNULL,
+ text=True,
+ ).strip()
+ except subprocess.CalledProcessError:
+ return None
+
+
+def _ltx25_official_source(artifact_root: Path) -> tuple[Path, str]:
+ repository = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY)
+ revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REVISION", LTX25_OFFICIAL_REVISION)
+ configured_root = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_ROOT")
+ root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}"
+ actual_revision = _git_revision(root) if root.exists() else None
+ if actual_revision != revision and configured_root:
+ raise AssertionError(f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}")
+ if actual_revision != revision:
+ if root.exists():
+ shutil.rmtree(root)
+ _clone_pinned_source(root, repository, revision)
+ actual_revision = _git_revision(root)
+ assert actual_revision == revision, f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}"
+ return root, revision
+
+
+def _official_runner_prefix() -> list[str]:
+ """Run the reference with its missing binary dependency isolated from CI."""
+ uv = shutil.which("uv")
+ assert uv is not None, "uv is required to run the pinned official LTX reference"
+ return [
+ uv,
+ "run",
+ "--no-project",
+ "--with",
+ f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}",
+ "--python",
+ sys.executable,
+ "python",
+ ]
+
+
+def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> tuple[Path, str, str | None]:
+ if transformer_subfolder not in {"transformer", "transformer_full"}:
+ raise ValueError(f"Unsupported LTX-2.5 transformer subfolder: {transformer_subfolder!r}")
+ configured_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
+ configured_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
+ if configured_model and Path(configured_model).exists():
+ model = Path(configured_model).resolve()
+ snapshot_revision = model.name if model.parent.name == "snapshots" else None
+ assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
+ assert (model / "diffusion_decoder" / "config.json").is_file(), (
+ f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}"
+ )
+ return model, configured_model, configured_revision or snapshot_revision
+ model_id = configured_model or LTX25_OMNI_MODEL_ID
+ revision = configured_revision
+ if revision is None and model_id == LTX25_OMNI_MODEL_ID:
+ revision = LTX25_OMNI_MODEL_REVISION
+ model = Path(
+ snapshot_download(
+ repo_id=model_id,
+ revision=revision,
+ allow_patterns=[
+ "model_index.json",
+ "audio_vae/*",
+ "connectors/config.json",
+ "connectors/diffusion_pytorch_model.safetensors.index.json",
+ "connectors/diffusion_pytorch_model-*.safetensors",
+ "diffusion_decoder/config.json",
+ "scheduler/*",
+ "text_encoder/config.json",
+ "text_encoder/generation_config.json",
+ "text_encoder/model*",
+ "tokenizer/*",
+ f"{transformer_subfolder}/*",
+ "vae/*",
+ "vocoder/*",
+ ],
+ )
+ ).resolve()
+ snapshot_revision = model.name if model.parent.name == "snapshots" else None
+ return model, model_id, snapshot_revision or revision
+
+
+def _resolve_ltx25_official_artifacts(
+ *, checkpoint_variant: str = "distilled"
+) -> tuple[dict[str, Path], str, str | None]:
+ try:
+ variant_files = LTX25_OFFICIAL_VARIANT_FILES[checkpoint_variant]
+ except KeyError as exc:
+ raise ValueError(f"Unsupported official LTX-2.5 checkpoint variant: {checkpoint_variant!r}") from exc
+ files = {**LTX25_OFFICIAL_COMMON_FILES, **variant_files}
+ configured_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
+ configured_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
+ if configured_model and Path(configured_model).exists():
+ root = Path(configured_model).resolve()
+ model_source = configured_model
+ snapshot_revision = root.name if root.parent.name == "snapshots" else None
+ resolved_revision = configured_revision or snapshot_revision
+ else:
+ model_source = configured_model or LTX25_OFFICIAL_MODEL_ID
+ revision = configured_revision
+ if revision is None and model_source == LTX25_OFFICIAL_MODEL_ID:
+ revision = LTX25_OFFICIAL_MODEL_REVISION
+ root = Path(
+ snapshot_download(
+ repo_id=model_source,
+ revision=revision,
+ allow_patterns=list(files.values()),
+ )
+ ).resolve()
+ snapshot_revision = root.name if root.parent.name == "snapshots" else None
+ resolved_revision = snapshot_revision or revision
+
+ artifacts = {name: root / relative_path for name, relative_path in files.items()}
+ missing = [f"{name}: {path}" for name, path in artifacts.items() if not path.is_file()]
+ assert not missing, f"Official LTX-2.5 artifacts are missing: {', '.join(missing)}"
+ return artifacts, model_source, resolved_revision
+
+
+def _resolve_gemma_root(model: Path) -> Path:
+ configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
+ if configured_root:
+ root = Path(configured_root)
+ assert root.is_dir(), f"Gemma root not found: {root}"
+ return root
+ return model
+
+
+def _resolve_ltx25_image() -> Path:
+ configured_image = os.environ.get("VLLM_TEST_LTX25_I2V_IMAGE")
+ if configured_image:
+ image = Path(configured_image)
+ assert image.is_file(), f"LTX-2.5 I2V conditioning image not found: {image}"
+ return image
+ return Path(
+ hf_hub_download(
+ repo_id=LTX25_I2V_IMAGE_REPO,
+ repo_type="dataset",
+ filename=LTX25_I2V_IMAGE_FILENAME,
+ revision=LTX25_I2V_IMAGE_REVISION,
+ )
+ )
+
+
+def _require_ltx25_model_access() -> None:
+ configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
+ if not (configured_omni_model and Path(configured_omni_model).exists()):
+ omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID
+ omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
+ if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID:
+ omni_revision = LTX25_OMNI_MODEL_REVISION
+ skip_if_gated_repo_inaccessible(
+ omni_model,
+ revision=omni_revision,
+ filename="model_index.json",
+ )
+
+ configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
+ if not (configured_official_model and Path(configured_official_model).exists()):
+ official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID
+ official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
+ if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID:
+ official_revision = LTX25_OFFICIAL_MODEL_REVISION
+ skip_if_gated_repo_inaccessible(
+ official_model,
+ revision=official_revision,
+ filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"],
+ )
+
+
+def _ltx25_request(image: Path | None = None) -> dict[str, object]:
+ request: dict[str, object] = {
+ "prompt": LTX25_PROMPT,
+ # Keep the nightly parity test small enough to run the pinned official
+ # pipeline and Omni sequentially on one H100. The public recipe covers
+ # the release-resolution 1920x1088x121 configuration.
+ "width": 1024,
+ "height": 576,
+ "num_frames": 25,
+ "fps": 24,
+ "num_inference_steps": 8,
+ "seed": 42,
+ }
+ if image is not None:
+ request.update(image=str(image.resolve()), image_crf=18)
+ return request
+
+
+def _ltx25_full_sigmas(num_inference_steps: int) -> list[float]:
+ """Materialize the official LTX-2.5 one-stage schedule in FP32."""
+ sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)
+ base_anchor = 1024
+ max_anchor = 4096
+ base_shift = 0.95
+ max_shift = 2.05
+ slope = (max_shift - base_shift) / (max_anchor - base_anchor)
+ sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
+ exp_shift = math.exp(sigma_shift)
+ sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
+
+ non_zero = sigmas != 0
+ one_minus_sigmas = 1.0 - sigmas[non_zero]
+ scale = one_minus_sigmas[-1] / (1.0 - 0.1)
+ sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
+ return [float(sigma) for sigma in sigmas]
+
+
+def _ltx25_full_request(image: Path | None = None) -> dict[str, object]:
+ num_inference_steps = 30
+ request: dict[str, object] = {
+ "prompt": LTX25_PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ # Full guidance expands to a four-way batch. Keep the decoded output
+ # small enough for the pinned official and Omni runs to share one H100.
+ "width": 512,
+ "height": 384,
+ "num_frames": 25,
+ "fps": 24,
+ "num_inference_steps": num_inference_steps,
+ "seed": 42,
+ "video_cfg_scale": 3.0,
+ "audio_cfg_scale": 7.0,
+ "video_stg_scale": 1.0,
+ "audio_stg_scale": 1.0,
+ "video_modality_scale": 3.0,
+ "audio_modality_scale": 3.0,
+ "video_rescale_scale": 0.7,
+ "audio_rescale_scale": 0.7,
+ "video_stg_blocks": [28],
+ "audio_stg_blocks": [28],
+ "sigmas": _ltx25_full_sigmas(num_inference_steps),
+ }
+ if image is not None:
+ request.update(image=str(image.resolve()), image_crf=18)
+ return request
+
+
+def test_ltx25_full_request_pins_official_schedule() -> None:
+ request = _ltx25_full_request()
+ sigmas = request["sigmas"]
+
+ assert isinstance(sigmas, list)
+ assert isinstance(request["num_inference_steps"], int)
+ assert len(sigmas) == request["num_inference_steps"] + 1
+ assert math.isclose(sigmas[0], 1.0, rel_tol=0.0, abs_tol=1e-6)
+ assert sigmas[-1] == 0.0
+ assert all(sigmas[index] > sigmas[index + 1] for index in range(len(sigmas) - 1))
+
+
+def test_ltx25_i2v_requests_pin_official_crf() -> None:
+ image = Path("conditioning.png")
+ for request in (_ltx25_request(image), _ltx25_full_request(image)):
+ assert request["image"] == str(image.resolve())
+ assert request["image_crf"] == 18
+
+
+def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
+ assert reference.shape == prediction.shape
+ assert reference.ndim == 4 and reference.shape[-1] == 3
+ ssim_scores: list[float] = []
+ psnr_scores: list[float] = []
+ for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
+ reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
+ prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
+ ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
+ psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
+ difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
+ return {
+ "ssim_mean": float(np.mean(ssim_scores)),
+ "ssim_min": float(np.min(ssim_scores)),
+ "psnr_mean_db": float(np.mean(psnr_scores)),
+ "max_abs": float(difference.max()),
+ "mean_abs": float(difference.mean()),
+ }
+
+
+def _canonical_audio(audio: np.ndarray) -> np.ndarray:
+ while audio.ndim > 2 and audio.shape[0] == 1:
+ audio = audio[0]
+ return audio.astype(np.float64)
+
+
+def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]:
+ reference = _canonical_audio(reference)
+ prediction = _canonical_audio(prediction)
+ assert reference.shape == prediction.shape
+ difference = reference - prediction
+ reference_norm = max(float(np.linalg.norm(reference)), 1e-12)
+ prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12)
+ return {
+ "bitwise_equal": bool(np.array_equal(reference, prediction)),
+ "max_abs": float(np.abs(difference).max()),
+ "mean_abs": float(np.abs(difference).mean()),
+ "relative_l2": float(np.linalg.norm(difference) / reference_norm),
+ "cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)),
+ }
+
+
+def _assert_strict_similarity(
+ video_metrics: dict[str, float],
+ audio_metrics: dict[str, float | bool],
+) -> None:
+ assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+ assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+ assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+ assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+ assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
+
+
+@pytest.mark.slow
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
+ """Compare Omni with official canonical DiffVAE and matching connector weights."""
+ _require_ltx25_model_access()
+ artifact_parent = accuracy_artifact_root / "ltx_official"
+ output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
+ official_root, official_revision = _ltx25_official_source(artifact_parent)
+ official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts()
+ omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model()
+ image = _resolve_ltx25_image() if task == "i2v" else None
+ request_path = output_root / "request.json"
+ request_path.write_text(json.dumps(_ltx25_request(image), indent=2) + "\n")
+
+ runner = Path(__file__).with_name("run_ltx25_reference.py")
+ runner_args = [
+ str(runner),
+ "--request",
+ str(request_path),
+ "--enable-layerwise-offload",
+ ]
+ env = os.environ.copy()
+ env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
+ env["PYTHONUNBUFFERED"] = "1"
+ repository_root = Path(__file__).resolve().parents[4]
+ existing_pythonpath = env.get("PYTHONPATH")
+ env["PYTHONPATH"] = (
+ str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
+ )
+
+ official_output = output_root / "official"
+ _run(
+ _official_runner_prefix()
+ + runner_args
+ + [
+ "--backend",
+ "official",
+ "--output-dir",
+ str(official_output),
+ "--official-root",
+ str(official_root),
+ "--transformer-path",
+ str(official_artifacts["transformer"]),
+ "--text-encoder-path",
+ str(official_artifacts["text_encoder"]),
+ "--video-vae-path",
+ str(official_artifacts["video_vae"]),
+ "--audio-vae-path",
+ str(official_artifacts["audio_vae"]),
+ "--spatial-upsampler-path",
+ str(official_artifacts["spatial_upsampler"]),
+ "--connector-model",
+ str(omni_model),
+ ],
+ env=env,
+ )
+
+ omni_output = output_root / "omni"
+ _run(
+ [sys.executable]
+ + runner_args
+ + [
+ "--backend",
+ "omni",
+ "--output-dir",
+ str(omni_output),
+ "--model",
+ str(omni_model),
+ "--model-class-name",
+ "LTX2DistilledTwoStagePipeline",
+ ],
+ env=env,
+ )
+
+ official_metadata = json.loads((official_output / "metadata.json").read_text())
+ omni_metadata = json.loads((omni_output / "metadata.json").read_text())
+ assert official_metadata["official_revision"] == official_revision
+ assert official_metadata["pipeline"] == "DistilledPipeline"
+ assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
+ assert official_metadata["attention_backend"] == ATTENTION_BACKEND
+ assert official_metadata["connector_model"] == str(omni_model)
+ assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
+ assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
+ assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
+ video_metrics = _video_metrics(
+ np.load(official_output / "video.npy"),
+ np.load(omni_output / "video.npy"),
+ )
+ audio_metrics = _audio_metrics(
+ np.load(official_output / "audio.npy"),
+ np.load(omni_output / "audio.npy"),
+ )
+ result = {
+ "case": "ltx2_5_distilled",
+ "task": task,
+ "attention_backend": ATTENTION_BACKEND,
+ "official_revision": official_revision,
+ "official_model": official_model_source,
+ "official_model_revision": official_model_revision,
+ "official_connector_model": omni_model_source,
+ "official_connector_revision": omni_model_revision,
+ "omni_model": omni_model_source,
+ "omni_model_revision": omni_model_revision,
+ "resolved_omni_model_path": str(omni_model),
+ "thresholds": asdict(STRICT_THRESHOLDS),
+ "peak_memory_mb": {
+ "official_allocated": official_metadata["peak_memory_allocated_mb"],
+ "official_reserved": official_metadata["peak_memory_reserved_mb"],
+ "omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
+ "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
+ "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
+ },
+ "video": video_metrics,
+ "audio": audio_metrics,
+ }
+ (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+ print(json.dumps(result, indent=2))
+
+ _assert_strict_similarity(video_metrics, audio_metrics)
+
+
+@pytest.mark.slow
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
+def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
+ """Compare raw official Full/SFT and DiffVAE weights with Omni."""
+ _require_ltx25_model_access()
+ artifact_parent = accuracy_artifact_root / "ltx_official"
+ output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
+ model_class_name = {
+ "full_one_stage": "LTX2Pipeline",
+ "full_two_stage": "LTX2TwoStagePipeline",
+ }[pipeline_mode]
+ official_root, official_revision = _ltx25_official_source(artifact_parent)
+ official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts(
+ checkpoint_variant="full" if pipeline_mode == "full_one_stage" else "full_two_stage"
+ )
+ omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model(
+ transformer_subfolder="transformer_full"
+ )
+ image = _resolve_ltx25_image() if task == "i2v" else None
+ request = _ltx25_full_request(image)
+ if pipeline_mode == "full_two_stage":
+ request["stage_1_sigmas"] = request.pop("sigmas")
+ request["stage_2_sigmas"] = LTX25_STAGE_2_SIGMAS
+ request_path = output_root / "request.json"
+ request_path.write_text(json.dumps(request, indent=2) + "\n")
+
+ runner = Path(__file__).with_name("run_ltx25_reference.py")
+ runner_args = [
+ str(runner),
+ "--request",
+ str(request_path),
+ "--enable-layerwise-offload",
+ ]
+ env = os.environ.copy()
+ env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
+ env["PYTHONUNBUFFERED"] = "1"
+ repository_root = Path(__file__).resolve().parents[4]
+ existing_pythonpath = env.get("PYTHONPATH")
+ env["PYTHONPATH"] = (
+ str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
+ )
+
+ official_output = output_root / "official"
+ official_sidecar_args: list[str] = []
+ if pipeline_mode == "full_two_stage":
+ official_sidecar_args = [
+ "--spatial-upsampler-path",
+ str(official_artifacts["spatial_upsampler"]),
+ "--distilled-lora-path",
+ str(official_artifacts["distilled_lora"]),
+ ]
+
+ _run(
+ _official_runner_prefix()
+ + runner_args
+ + [
+ "--backend",
+ "official",
+ "--output-dir",
+ str(official_output),
+ "--official-root",
+ str(official_root),
+ "--official-pipeline",
+ pipeline_mode,
+ "--transformer-path",
+ str(official_artifacts["transformer"]),
+ "--text-encoder-path",
+ str(official_artifacts["text_encoder"]),
+ "--video-vae-path",
+ str(official_artifacts["video_vae"]),
+ "--audio-vae-path",
+ str(official_artifacts["audio_vae"]),
+ "--connector-model",
+ str(omni_model),
+ ]
+ + official_sidecar_args,
+ env=env,
+ )
+
+ omni_output = output_root / "omni"
+ _run(
+ [sys.executable]
+ + runner_args
+ + [
+ "--backend",
+ "omni",
+ "--output-dir",
+ str(omni_output),
+ "--model",
+ str(omni_model),
+ "--model-class-name",
+ model_class_name,
+ ],
+ env=env,
+ )
+
+ official_metadata = json.loads((official_output / "metadata.json").read_text())
+ omni_metadata = json.loads((omni_output / "metadata.json").read_text())
+ assert official_metadata["official_revision"] == official_revision
+ assert (
+ official_metadata["pipeline"]
+ == {
+ "full_one_stage": "TI2VidOneStagePipeline",
+ "full_two_stage": "TI2VidTwoStagesPipeline",
+ }[pipeline_mode]
+ )
+ assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
+ assert official_metadata["attention_backend"] == ATTENTION_BACKEND
+ assert official_metadata["connector_model"] == str(omni_model)
+ assert omni_metadata["model_class_name"] == model_class_name
+ assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
+ assert official_metadata["seed"] == omni_metadata["seed"]
+ for schedule_name in ("sigmas", "stage_1_sigmas", "stage_2_sigmas"):
+ assert official_metadata[schedule_name] == omni_metadata[schedule_name]
+ assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
+ video_metrics = _video_metrics(
+ np.load(official_output / "video.npy"),
+ np.load(omni_output / "video.npy"),
+ )
+ audio_metrics = _audio_metrics(
+ np.load(official_output / "audio.npy"),
+ np.load(omni_output / "audio.npy"),
+ )
+ result = {
+ "case": f"ltx2_5_{pipeline_mode}",
+ "task": task,
+ "attention_backend": ATTENTION_BACKEND,
+ "official_revision": official_revision,
+ "official_model": official_model_source,
+ "official_model_revision": official_model_revision,
+ "official_connector_model": omni_model_source,
+ "official_connector_revision": omni_model_revision,
+ "omni_model": omni_model_source,
+ "omni_model_revision": omni_model_revision,
+ "resolved_omni_model_path": str(omni_model),
+ "thresholds": asdict(STRICT_THRESHOLDS),
+ "peak_memory_mb": {
+ "official_allocated": official_metadata["peak_memory_allocated_mb"],
+ "official_reserved": official_metadata["peak_memory_reserved_mb"],
+ "omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
+ "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
+ "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
+ },
+ "video": video_metrics,
+ "audio": audio_metrics,
+ }
+ (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+ print(json.dumps(result, indent=2))
+
+ _assert_strict_similarity(video_metrics, audio_metrics)
diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
index 647ece7a7..63ce73923 100644
--- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
+++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
@@ -3,10 +3,8 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
-The comparison runs both runtimes through PyTorch SDPA and uses
-``max_batch_size=4`` in the official reference to match Omni's fused guidance
-batch. Video and audio guidance use the official non-HQ one-stage defaults;
-only the generation shape and step count are reduced for CI runtime.
+Four complementary default-shape cases cover both model versions, both
+two-stage weight families, and T2V/I2V without duplicating one-stage coverage.
"""
from __future__ import annotations
@@ -17,8 +15,9 @@ import shutil
import subprocess
import sys
import time
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass, replace
from pathlib import Path
+from typing import Literal
import numpy as np
import pytest
@@ -53,73 +52,198 @@ NEGATIVE_PROMPT = (
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
-VIDEO_SSIM_MEAN_THRESHOLD = 0.95
-VIDEO_SSIM_MIN_THRESHOLD = 0.90
-VIDEO_PSNR_MEAN_THRESHOLD = 30.0
-AUDIO_RELATIVE_L2_THRESHOLD = 0.2
-AUDIO_COSINE_THRESHOLD = 0.95
+
+
+@dataclass(frozen=True)
+class LTXArtifact:
+ repo_id: str
+ filename: str
+ revision: str
+ env: str
+ repo_type: Literal["model", "dataset"] = "model"
+
+
+@dataclass(frozen=True)
+class LTXAccuracyThresholds:
+ video_ssim_mean: float
+ video_ssim_min: float
+ video_psnr_mean_db: float
+ audio_relative_l2: float
+ audio_cosine_similarity: float
+
+
+STRICT_THRESHOLDS = LTXAccuracyThresholds(
+ video_ssim_mean=0.99,
+ video_ssim_min=0.99,
+ video_psnr_mean_db=40.0,
+ audio_relative_l2=0.10,
+ audio_cosine_similarity=0.99,
+)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
+ pipeline_kind: Literal["distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
model_class_name: str
- checkpoint_repo: str
- checkpoint_filename: str
- checkpoint_revision: str
- checkpoint_env: str
- stg_block: int
- prompt: str = PROMPT
- image_repo: str | None = None
- image_filename: str | None = None
- image_revision: str | None = None
+ checkpoint: LTXArtifact
+ width: int
+ height: int
+ num_frames: int
+ num_inference_steps: int
+ seed: int
+ stg_block: int | None
+ gemma_model_id: str | None = None
+ gemma_model_revision: str | None = None
+ gemma_model_env: str | None = None
+ image: LTXArtifact | None = None
+ spatial_upsampler: LTXArtifact | None = None
+ distilled_lora: LTXArtifact | None = None
+ # Official uses block streaming; Omni uses layerwise offload. Both preserve
+ # bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
+ enable_layerwise_offload: bool = False
-CASES = (
+LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
+LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25"
+LTX2_CHECKPOINT = LTXArtifact(
+ repo_id="Lightricks/LTX-2",
+ filename="ltx-2-19b-dev.safetensors",
+ revision=LTX2_REVISION,
+ env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
+)
+LTX2_DISTILLED_CHECKPOINT = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-19b-distilled.safetensors",
+ env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT",
+)
+LTX2_UPSAMPLER = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
+ env="VLLM_TEST_LTX2_UPSAMPLER",
+)
+LTX2_DISTILLED_LORA = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-19b-distilled-lora-384.safetensors",
+ env="VLLM_TEST_LTX2_DISTILLED_LORA",
+)
+LTX23_CHECKPOINT = LTXArtifact(
+ repo_id="Lightricks/LTX-2.3",
+ filename="ltx-2.3-22b-dev.safetensors",
+ revision=LTX23_REVISION,
+ env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
+)
+LTX23_DISTILLED_CHECKPOINT = replace(
+ LTX23_CHECKPOINT,
+ # The pinned Diffusers checkpoint contains the original merged distilled
+ # Transformer. Keep the reference Transformer on the same version; the
+ # spatial upsampler is versioned independently and remains on 1.1 below.
+ filename="ltx-2.3-22b-distilled.safetensors",
+ env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT",
+)
+LTX23_UPSAMPLER = replace(
+ LTX23_CHECKPOINT,
+ filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+ env="VLLM_TEST_LTX23_UPSAMPLER",
+)
+LTX23_DISTILLED_LORA = replace(
+ LTX23_CHECKPOINT,
+ filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
+ env="VLLM_TEST_LTX23_DISTILLED_LORA",
+)
+I2V_IMAGE = LTXArtifact(
+ repo_id="huggingface/documentation-images",
+ filename="diffusers/svd/rocket.png",
+ revision="645d8364f0c7a101180b364811b5a11a362e4010",
+ env="VLLM_TEST_LTX_I2V_IMAGE",
+ repo_type="dataset",
+)
+
+
+TWO_STAGE_CASES = (
LTXAccuracyCase(
- name="ltx2",
+ name="ltx2_distilled_t2v",
+ pipeline_kind="distilled",
+ model_id="rootonchair/LTX-2-19b-distilled",
+ model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5",
+ model_env="VLLM_TEST_LTX2_DISTILLED_MODEL",
+ model_class_name="LTX2DistilledPipeline",
+ checkpoint=LTX2_DISTILLED_CHECKPOINT,
+ spatial_upsampler=LTX2_UPSAMPLER,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=8,
+ seed=10,
+ stg_block=None,
+ gemma_model_id="Lightricks/LTX-2",
+ gemma_model_revision=LTX2_REVISION,
+ gemma_model_env="VLLM_TEST_LTX2_MODEL",
+ ),
+ LTXAccuracyCase(
+ name="ltx23_distilled_i2v",
+ pipeline_kind="distilled",
+ model_id="diffusers/LTX-2.3-Distilled-Diffusers",
+ model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee",
+ model_env="VLLM_TEST_LTX23_DISTILLED_MODEL",
+ model_class_name="LTX2DistilledPipeline",
+ checkpoint=LTX23_DISTILLED_CHECKPOINT,
+ spatial_upsampler=LTX23_UPSAMPLER,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=8,
+ seed=10,
+ stg_block=None,
+ gemma_model_id="diffusers/LTX-2.3-Diffusers",
+ gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
+ gemma_model_env="VLLM_TEST_LTX23_MODEL",
+ image=I2V_IMAGE,
+ ),
+ LTXAccuracyCase(
+ name="ltx2_two_stage_layer_fused_t2v",
+ pipeline_kind="two_stage",
model_id="Lightricks/LTX-2",
- model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
+ model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
- model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2",
- checkpoint_filename="ltx-2-19b-dev.safetensors",
- checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
- checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
+ model_class_name="LTX2TwoStagePipeline",
+ checkpoint=LTX2_CHECKPOINT,
+ spatial_upsampler=LTX2_UPSAMPLER,
+ distilled_lora=LTX2_DISTILLED_LORA,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=40,
+ seed=10,
stg_block=29,
+ enable_layerwise_offload=True,
),
LTXAccuracyCase(
- name="ltx2_3",
+ name="ltx23_two_stage_layer_fused_i2v",
+ pipeline_kind="two_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
- model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2.3",
- checkpoint_filename="ltx-2.3-22b-dev.safetensors",
- checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
- checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
+ model_class_name="LTX2TwoStagePipeline",
+ checkpoint=LTX23_CHECKPOINT,
+ spatial_upsampler=LTX23_UPSAMPLER,
+ distilled_lora=LTX23_DISTILLED_LORA,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=30,
+ seed=10,
stg_block=28,
- ),
- LTXAccuracyCase(
- name="ltx2_3_i2v",
- model_id="diffusers/LTX-2.3-Diffusers",
- model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
- model_env="VLLM_TEST_LTX23_MODEL",
- model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2.3",
- checkpoint_filename="ltx-2.3-22b-dev.safetensors",
- checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
- checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
- stg_block=28,
- image_repo="huggingface/documentation-images",
- image_filename="diffusers/svd/rocket.png",
- image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
+ enable_layerwise_offload=True,
+ image=I2V_IMAGE,
),
)
+WEEKLY_CASES = TWO_STAGE_CASES
+
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
@@ -209,6 +333,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
"model_index.json",
"audio_vae/*",
"connectors/*",
+ "latent_upsampler/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
@@ -223,74 +348,124 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
)
-def _resolve_gemma_root(model: Path) -> Path:
+def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
+ if case.gemma_model_id is not None:
+ configured_model = os.environ.get(case.gemma_model_env or "")
+ if configured_model and Path(configured_model).is_dir():
+ return Path(configured_model)
+ return Path(
+ snapshot_download(
+ repo_id=case.gemma_model_id,
+ revision=case.gemma_model_revision,
+ allow_patterns=[
+ "processor/*",
+ "text_encoder/*",
+ "tokenizer/*",
+ ],
+ )
+ )
return model
-def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
- configured_checkpoint = os.environ.get(case.checkpoint_env)
- if configured_checkpoint:
- checkpoint = Path(configured_checkpoint)
- assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
- return checkpoint
- model_checkpoint = model / case.checkpoint_filename
- if model_checkpoint.is_file():
- return model_checkpoint
+def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
+ configured_path = os.environ.get(artifact.env)
+ if configured_path:
+ path = Path(configured_path)
+ assert path.is_file(), f"Configured LTX artifact not found: {path}"
+ return path
+ if model is not None:
+ model_path = model / artifact.filename
+ if model_path.is_file():
+ return model_path
return Path(
hf_hub_download(
- repo_id=case.checkpoint_repo,
- filename=case.checkpoint_filename,
- revision=case.checkpoint_revision,
+ repo_id=artifact.repo_id,
+ repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
+ filename=artifact.filename,
+ revision=artifact.revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
- if case.image_filename is None:
+ if case.image is None:
return None
- if case.image_repo is None or case.image_revision is None:
- raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
- configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
- if configured_image:
- image = Path(configured_image)
- assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
- return image
- return Path(
- hf_hub_download(
- repo_id=case.image_repo,
- repo_type="dataset",
- filename=case.image_filename,
- revision=case.image_revision,
- )
- )
+ return _resolve_artifact(case.image)
+
+
+def _omni_model_with_pinned_artifacts(
+ model: Path,
+ output_root: Path,
+ artifacts: tuple[Path | None, ...],
+) -> Path:
+ """Expose the already-resolved sidecars to Omni without copying the model."""
+ resolved_artifacts = tuple(artifact for artifact in artifacts if artifact is not None)
+ if not resolved_artifacts:
+ return model
+
+ overlay = output_root / "omni-model"
+ overlay.mkdir()
+ for source in model.iterdir():
+ (overlay / source.name).symlink_to(source.resolve(), target_is_directory=source.is_dir())
+ for artifact in resolved_artifacts:
+ target = overlay / artifact.name
+ if target.exists():
+ assert target.samefile(artifact), f"Model sidecar does not match pinned artifact: {target} != {artifact}"
+ else:
+ target.symlink_to(artifact.resolve())
+ return overlay
+
+
+def test_omni_model_overlay_pins_artifacts(tmp_path: Path) -> None:
+ model = tmp_path / "model"
+ model.mkdir()
+ component = model / "model_index.json"
+ component.write_text("{}\n")
+ artifact = tmp_path / "sidecar.safetensors"
+ artifact.write_bytes(b"pinned")
+ output_root = tmp_path / "output"
+ output_root.mkdir()
+
+ overlay = _omni_model_with_pinned_artifacts(model, output_root, (artifact,))
+
+ assert (overlay / component.name).samefile(component)
+ assert (overlay / artifact.name).samefile(artifact)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
- "prompt": case.prompt,
- "negative_prompt": NEGATIVE_PROMPT,
- "width": 512,
- "height": 384,
- "num_frames": 25,
+ "pipeline_kind": case.pipeline_kind,
+ "prompt": PROMPT,
+ "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
+ "width": case.width,
+ "height": case.height,
+ "num_frames": case.num_frames,
"fps": 24,
- "num_inference_steps": 20,
- "seed": 42,
- "video_cfg_scale": 3.0,
- "audio_cfg_scale": 7.0,
- "video_stg_scale": 1.0,
- "audio_stg_scale": 1.0,
- "video_modality_scale": 3.0,
- "audio_modality_scale": 3.0,
- "video_rescale_scale": 0.7,
- "audio_rescale_scale": 0.7,
- "video_stg_blocks": [case.stg_block],
- "audio_stg_blocks": [case.stg_block],
+ "num_inference_steps": case.num_inference_steps,
+ "seed": case.seed,
}
+ if case.pipeline_kind != "distilled":
+ request.update(
+ video_cfg_scale=3.0,
+ audio_cfg_scale=7.0,
+ video_modality_scale=3.0,
+ audio_modality_scale=3.0,
+ )
+ if case.pipeline_kind == "two_stage":
+ assert case.stg_block is not None
+ request.update(
+ video_stg_scale=1.0,
+ audio_stg_scale=1.0,
+ video_rescale_scale=0.7,
+ audio_rescale_scale=0.7,
+ video_stg_blocks=[case.stg_block],
+ audio_stg_blocks=[case.stg_block],
+ )
if image is not None:
request["image"] = str(image.resolve())
return request
@@ -299,20 +474,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
+ device = torch.device("cpu")
+ ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
+ psnr = PeakSignalNoiseRatio(data_range=1.0).to(device)
ssim_scores: list[float] = []
psnr_scores: list[float] = []
- for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
- reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
- prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
- ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
- psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
- difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
+ max_abs = 0.0
+ absolute_error_sum = 0.0
+ with torch.inference_mode():
+ for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
+ reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+ prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+ ssim_scores.append(float(ssim(prediction_tensor, reference_tensor)))
+ psnr_scores.append(float(psnr(prediction_tensor, reference_tensor)))
+ ssim.reset()
+ psnr.reset()
+ difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64))
+ max_abs = max(max_abs, float(difference.max()))
+ absolute_error_sum += float(difference.sum())
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
- "max_abs": float(difference.max()),
- "mean_abs": float(difference.mean()),
+ "max_abs": max_abs,
+ "mean_abs": absolute_error_sum / reference.size,
}
@@ -338,29 +523,42 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
}
-@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
-@pytest.mark.full_model
+@pytest.mark.parametrize("case", WEEKLY_CASES, ids=lambda case: case.name)
+@pytest.mark.slow
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
- """Compare official and Omni raw AV outputs from the same E2E request."""
+def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
+ """Compare official and Omni raw AV outputs from the same single-H100 request."""
+ configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT")
+ if configured_artifact_root:
+ accuracy_artifact_root = Path(configured_artifact_root)
+ accuracy_artifact_root.mkdir(parents=True, exist_ok=True)
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
- gemma_root = _resolve_gemma_root(model)
- checkpoint = _resolve_checkpoint(case, model)
+ gemma_root = _resolve_gemma_root(case, model)
+ checkpoint = _resolve_artifact(case.checkpoint, model)
+ spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
+ distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
+ omni_model = _omni_model_with_pinned_artifacts(model, output_root, (spatial_upsampler, distilled_lora))
image = _resolve_image(case)
+ request = _request(case, image)
request_path = output_root / "request.json"
- request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
+ request_path.write_text(json.dumps(request, indent=2) + "\n")
runner = Path(__file__).with_name("run_ltx_reference.py")
runner_args = [
str(runner),
"--request",
str(request_path),
+ "--pipeline-kind",
+ case.pipeline_kind,
]
- if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
+ enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
+ "VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", ""
+ ).lower() in {"1", "true", "yes", "on"}
+ if enable_layerwise_offload:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
@@ -372,10 +570,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
)
official_output = output_root / "official"
- _run(
- _official_runner_prefix()
- + runner_args
- + [
+ official_args = _official_runner_prefix() + runner_args
+ official_args.extend(
+ [
"--backend",
"official",
"--output-dir",
@@ -386,12 +583,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(checkpoint),
"--gemma-root",
str(gemma_root),
- ],
- env=env,
+ ]
)
+ if spatial_upsampler is not None:
+ official_args.extend(["--spatial-upsampler", str(spatial_upsampler)])
+ if distilled_lora is not None:
+ official_args.extend(["--distilled-lora", str(distilled_lora)])
+ _run(official_args, env=env)
omni_output = output_root / "omni"
- _run(
+ omni_args = (
[sys.executable]
+ runner_args
+ [
@@ -400,12 +601,12 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
"--output-dir",
str(omni_output),
"--model",
- str(model),
+ str(omni_model),
"--model-class-name",
case.model_class_name,
- ],
- env=env,
+ ]
)
+ _run(omni_args, env=env)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
@@ -423,18 +624,25 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
+ "pipeline_kind": case.pipeline_kind,
+ "model_class_name": case.model_class_name,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
- "checkpoint_revision": case.checkpoint_revision,
+ "checkpoint_revision": case.checkpoint.revision,
+ "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
+ "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
+ "enable_layerwise_offload": enable_layerwise_offload,
+ "thresholds": asdict(STRICT_THRESHOLDS),
+ "request": request,
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
- assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
- assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
- assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
- assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
- assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
+ assert video_metrics["ssim_mean"] >= STRICT_THRESHOLDS.video_ssim_mean
+ assert video_metrics["ssim_min"] >= STRICT_THRESHOLDS.video_ssim_min
+ assert video_metrics["psnr_mean_db"] >= STRICT_THRESHOLDS.video_psnr_mean_db
+ assert audio_metrics["relative_l2"] <= STRICT_THRESHOLDS.audio_relative_l2
+ assert audio_metrics["cosine_similarity"] >= STRICT_THRESHOLDS.audio_cosine_similarity
diff --git a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
index 50474a644..e688b9360 100644
--- a/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
+++ b/tests/e2e/accuracy/minicpmo_4_5/test_minicpmo_4_5.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""MiniCPM-o 4.5 Daily-Omni + Seed-TTS accuracy regression coverage.
Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches
@@ -21,13 +22,25 @@ from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import (
build_acc_benchmark_cli_argv,
find_vllm_cli,
)
-from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import SERVER_PARAMS as DUPLEX_TEST_PARAMS
+from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+ DEPLOY_CONFIG as _DUPLEX_DEPLOY_CONFIG,
+)
+from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
+ SERVER_PARAMS as DUPLEX_TEST_PARAMS,
+)
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
_MODEL = os.environ.get("VLLM_TEST_MINICPMO_4_5_MODEL", "openbmb/MiniCPM-o-4_5")
_DEPLOY_CONFIG = get_deploy_config_path("minicpmo_4_5.yaml")
+_DUPLEX_BF16_DEPLOY_CONFIG = modify_stage_config(
+ _DUPLEX_DEPLOY_CONFIG,
+ updates={
+ "base_config": _DEPLOY_CONFIG,
+ "connectors.connector_of_shared_memory.extra.code2wav_bfloat16_attention_cache": True,
+ },
+)
_RESULT_DIR = Path(
os.environ.get(
"ACC_BENCH_RESULT_DIR",
@@ -91,6 +104,20 @@ seed_test_params = [
]
+duplex_accuracy_test_params = [
+ *DUPLEX_TEST_PARAMS,
+ pytest.param(
+ OmniServerParams(
+ model=_MODEL,
+ stage_config_path=_DUPLEX_BF16_DEPLOY_CONFIG,
+ use_stage_cli=False,
+ server_args=list(_SEED_TTS_SERVER_ARGS),
+ ),
+ id="three-stage-single-gpu-bf16-attention-cache",
+ ),
+]
+
+
def _require_vllm_cli() -> None:
try:
find_vllm_cli()
@@ -174,7 +201,7 @@ def test_minicpmo_4_5_seed_tts_wer_bench(omni_server) -> None:
@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", DUPLEX_TEST_PARAMS, indirect=True)
+@pytest.mark.parametrize("omni_server", duplex_accuracy_test_params, indirect=True)
def test_minicpmo_4_5_duplex_seed_tts_wer_bench(omni_server) -> None:
"""Gate Seed-TTS WER through the explicit Realtime TTS contract."""
_require_vllm_cli()
diff --git a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py
index b00e90d5e..60f829bee 100644
--- a/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py
+++ b/tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -119,6 +119,9 @@ def _model_name(model_env_var: str = MODEL_ENV_VAR, component: str = "FL2VA") ->
def _server_args() -> list[str]:
return [
"--trust-remote-code",
+ # Both official-reference requests use fixed shapes. Keep regional
+ # compilation static so cold Inductor codegen is stable across CI nodes.
+ "--no-diffusion-compile-dynamic",
"--num-gpus",
"4",
"--usp",
diff --git a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py
index dfc53a06e..d612b96a5 100644
--- a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py
+++ b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Shared helpers for Qwen3-Omni Daily-Omni / Seed-TTS ``vllm bench serve --omni`` accuracy runs.
Local dataset paths are **optional**. When ``VLLM_DAILY_OMNI_QA_JSON`` + ``VLLM_DAILY_OMNI_VIDEO_DIR``
@@ -28,7 +29,7 @@ DEFAULT_SEED_TTS_HF_REPO = "zhaochenyang20/seed-tts-eval"
class OmniBenchServerEndpoint(Protocol):
- """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.conftest.OmniServer`)."""
+ """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.helpers.runtime.OmniServer`)."""
host: str
port: int
diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py
index 121551664..7c48053e1 100644
--- a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py
+++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Qwen3-Omni accuracy benchmarks (Daily-Omni MCQ + Seed-TTS WER) via ``vllm bench serve --omni``.
Starts a **module-scoped** Omni OpenAI-compatible server (same pattern as ``tests/dfx/perf`` and
@@ -67,12 +68,8 @@ def get_chunk_config(config_path: str | None = None):
"""Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode."""
if config_path is None:
config_path = _CI_DEPLOY
- # TODO: remove this workaround once legacy `stage_args` path is deleted.
- # The pipeline (qwen3_omni/pipeline.py) already wires
- # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1,
- # so only async_chunk needs flipping. Writing nested `engine_args:` into
- # the new-schema overlay trips _parse_stage_deploy's legacy branch and
- # drops flat fields (load_format, max_num_seqs, ...).
+ # Pipeline already wires thinker2talker_async_chunk / talker2code2wav_async_chunk
+ # on stage 0/1; only the top-level async_chunk flag needs flipping here.
return modify_stage_config(config_path, updates={"async_chunk": True})
diff --git a/tests/e2e/accuracy/sana_video/README.md b/tests/e2e/accuracy/sana_video/README.md
new file mode 100644
index 000000000..cab38b7a2
--- /dev/null
+++ b/tests/e2e/accuracy/sana_video/README.md
@@ -0,0 +1,63 @@
+# SANA-Video golden outputs
+
+The v2 golden set covers both text-to-video and image-to-video for each
+checkpoint variant. Generate references with the Diffusers version pinned by
+this repository and the immutable Hub revisions embedded in
+`generate_goldens.py`:
+
+```bash
+python tests/e2e/accuracy/sana_video/generate_goldens.py \
+ --variant 480p --task t2v
+python tests/e2e/accuracy/sana_video/generate_goldens.py \
+ --variant 480p --task i2v --input-image /path/to/fixed-input.png
+python tests/e2e/accuracy/sana_video/generate_goldens.py \
+ --variant 720p --task t2v
+python tests/e2e/accuracy/sana_video/generate_goldens.py \
+ --variant 720p --task i2v --input-image /path/to/fixed-input.png
+```
+
+Use the same reviewed input image for both I2V variants. The generator
+normalizes it to an RGB PNG and includes that exact file and its SHA256 in the
+case manifest, so an I2V reference is self-contained.
+
+Each case is written under:
+
+```text
+///
+```
+
+It contains:
+
+- `transformer_case.safetensors`: deterministic transformer inputs and output.
+- `pipeline.mp4`: the encoded reference used by online similarity tests.
+- `pipeline_reference.safetensors`: final denoised latents and decoded,
+ pre-MP4 frames quantized to the exact uint8 video input.
+- `input.png`: the normalized conditioning image for I2V only.
+- `metadata.json`: model, generator, scheduler, input and sampling provenance.
+- `manifest.json`: the size and SHA256 of every artifact.
+
+## Publication requirements
+
+A canonical golden must be generated from the pinned Hugging Face model ID and
+immutable commit. The generator also requires a clean vLLM-Omni worktree and
+records the repository revision plus the generator script hash.
+
+`--model` is available for local development. It hashes every file in the local
+checkpoint directory and records optional `--local-source-model` and
+`--local-source-revision` claims, but marks the result `publishable: false`.
+The golden test rejects such a manifest. A local directory hash proves which
+bytes were used; it does not prove that those bytes match a claimed upstream
+commit.
+
+After reviewing the metadata, input image and generated videos, upload each
+publishable case directory without renaming files to:
+
+```text
+s3://vllm-public-assets/omni-assets/sana-video/v2///
+```
+
+Set `SANA_VIDEO_GOLDEN_BASE_URL` to the corresponding `v2` HTTP base URL when
+running the tests. Uploading is intentionally separate from generation so test
+runs cannot overwrite a frozen reference. The storage location must prevent
+in-place mutation; publish a new version prefix instead of replacing any v2
+file when a model, dependency, input or sampling parameter changes.
diff --git a/tests/e2e/accuracy/sana_video/generate_goldens.py b/tests/e2e/accuracy/sana_video/generate_goldens.py
new file mode 100644
index 000000000..8ab25dd18
--- /dev/null
+++ b/tests/e2e/accuracy/sana_video/generate_goldens.py
@@ -0,0 +1,380 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Generate frozen SANA-Video transformer and pipeline reference artifacts."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import subprocess
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from diffusers import SanaImageToVideoPipeline, SanaVideoPipeline
+from diffusers.utils import export_to_video
+from PIL import Image
+from safetensors.torch import save_file
+
+PROMPT = "A cat walking on the grass, facing the camera. motion score: 30."
+NEGATIVE_PROMPT = (
+ "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, "
+ "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, "
+ "jitter, and ghosting effects, creating a disorienting visual experience."
+)
+SCHEMA_VERSION = 2
+SEED = 42
+NUM_FRAMES = 81
+FPS = 16
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 6.0
+VARIANTS = {
+ "480p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+ "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f",
+ "height": 480,
+ "width": 832,
+ "vae_dtype": torch.float32,
+ "transformer_case_shape": (1, 16, 3, 8, 8),
+ },
+ "720p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+ "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e",
+ "height": 704,
+ "width": 1280,
+ "vae_dtype": torch.bfloat16,
+ "transformer_case_shape": (1, 128, 2, 4, 4),
+ },
+}
+
+
+def _sha256(path: Path) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _json_sha256(payload: Any) -> str:
+ encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode()
+ return hashlib.sha256(encoded).hexdigest()
+
+
+def _directory_fingerprint(root: Path) -> dict[str, Any]:
+ """Fingerprint every local checkpoint file without recording host paths."""
+ if not root.is_dir():
+ raise ValueError(f"Local checkpoint must be a directory: {root}")
+ files = []
+ digest = hashlib.sha256()
+ for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file()):
+ relative = path.relative_to(root).as_posix()
+ file_sha256 = _sha256(path)
+ size = path.stat().st_size
+ entry = f"{relative}\0{size}\0{file_sha256}\n".encode()
+ digest.update(entry)
+ files.append({"path": relative, "size": size, "sha256": file_sha256})
+ if not files:
+ raise ValueError(f"Local checkpoint contains no files: {root}")
+ return {
+ "algorithm": "sha256-tree-v1",
+ "sha256": digest.hexdigest(),
+ "file_count": len(files),
+ "total_size": sum(entry["size"] for entry in files),
+ }
+
+
+def _repository_provenance() -> dict[str, Any]:
+ repo_root = Path(__file__).resolve().parents[4]
+ revision = subprocess.run(
+ ["git", "rev-parse", "HEAD"],
+ cwd=repo_root,
+ check=True,
+ capture_output=True,
+ text=True,
+ ).stdout.strip()
+ dirty = bool(
+ subprocess.run(
+ ["git", "status", "--porcelain", "--untracked-files=no"],
+ cwd=repo_root,
+ check=True,
+ capture_output=True,
+ text=True,
+ ).stdout.strip()
+ )
+ return {
+ "repository": "vllm-project/vllm-omni",
+ "revision": revision,
+ "dirty": dirty,
+ "generator_sha256": _sha256(Path(__file__)),
+ }
+
+
+def _cpu(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.detach().cpu().contiguous()
+
+
+def _normalize_input_image(source: Path, destination: Path) -> Image.Image:
+ if not source.is_file():
+ raise ValueError(f"I2V input image does not exist: {source}")
+ with Image.open(source) as image:
+ normalized = image.convert("RGB")
+ normalized.load()
+ normalized.save(destination, format="PNG", optimize=False)
+ return normalized
+
+
+def _frames_to_uint8(frames: np.ndarray) -> torch.Tensor:
+ tensor = torch.from_numpy(np.asarray(frames))
+ if tensor.is_floating_point():
+ # Match diffusers.utils.export_to_video, which truncates rather than
+ # rounds when converting float frames to the encoder's uint8 input.
+ tensor = tensor.clamp(0, 1) * 255
+ return tensor.to(torch.uint8).contiguous()
+
+
+def _model_provenance(
+ config: dict[str, Any],
+ model: str | None,
+ local_source_model: str | None,
+ local_source_revision: str | None,
+) -> tuple[dict[str, Any], bool]:
+ if model is None:
+ return (
+ {
+ "kind": "huggingface_hub",
+ "model_id": config["model"],
+ "revision": config["revision"],
+ },
+ True,
+ )
+
+ local_path = Path(model).expanduser().resolve()
+ if bool(local_source_model) != bool(local_source_revision):
+ raise ValueError("--local-source-model and --local-source-revision must be specified together.")
+ provenance = {
+ "kind": "local_checkpoint",
+ "checkpoint_fingerprint": _directory_fingerprint(local_path),
+ "claimed_source_model": local_source_model,
+ "claimed_source_revision": local_source_revision,
+ "verification": "unverified-local-copy",
+ }
+ # A directory hash proves which local bytes were used, but not that those
+ # bytes match a claimed Hub commit. Such output is useful for development
+ # comparisons but must never be uploaded as the canonical frozen reference.
+ return provenance, False
+
+
+def generate(
+ variant: str,
+ task: str,
+ model: str | None,
+ input_image: Path | None,
+ output_root: Path,
+ local_source_model: str | None = None,
+ local_source_revision: str | None = None,
+) -> Path:
+ config = dict(VARIANTS[variant])
+ model_provenance, publishable = _model_provenance(
+ config,
+ model,
+ local_source_model,
+ local_source_revision,
+ )
+ repository_provenance = _repository_provenance()
+ if publishable and repository_provenance["dirty"]:
+ raise RuntimeError("Refusing to generate publishable golden assets from a dirty vllm-omni worktree.")
+ if model is not None:
+ config["model"] = str(Path(model).expanduser().resolve())
+ config["revision"] = None
+
+ output_dir = output_root / variant / task
+ output_dir.mkdir(parents=True, exist_ok=True)
+ normalized_image: Image.Image | None = None
+ input_path: Path | None = None
+ if task == "i2v":
+ if input_image is None:
+ raise ValueError("--input-image is required for task=i2v.")
+ input_path = output_dir / "input.png"
+ normalized_image = _normalize_input_image(input_image, input_path)
+ elif input_image is not None:
+ raise ValueError("--input-image is only valid for task=i2v.")
+
+ pipeline_class = SanaImageToVideoPipeline if task == "i2v" else SanaVideoPipeline
+ pipe = pipeline_class.from_pretrained(
+ config["model"],
+ revision=config["revision"],
+ torch_dtype=torch.bfloat16,
+ )
+ pipe.text_encoder.to(torch.bfloat16)
+ pipe.vae.to(config["vae_dtype"])
+ pipe.to("cuda")
+
+ torch.manual_seed(1234)
+ hidden_states = torch.randn(config["transformer_case_shape"], dtype=torch.bfloat16, device="cuda")
+ encoder_hidden_states = torch.randn(1, 8, 2304, dtype=torch.bfloat16, device="cuda")
+ encoder_attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 0]], device="cuda")
+ timestep = torch.tensor([500.0], device="cuda")
+ with torch.inference_mode():
+ transformer_output = pipe.transformer(
+ hidden_states,
+ encoder_hidden_states,
+ timestep,
+ encoder_attention_mask=encoder_attention_mask,
+ ).sample
+
+ transformer_path = output_dir / "transformer_case.safetensors"
+ save_file(
+ {
+ "hidden_states": _cpu(hidden_states),
+ "encoder_hidden_states": _cpu(encoder_hidden_states),
+ "encoder_attention_mask": _cpu(encoder_attention_mask),
+ "timestep": _cpu(timestep),
+ "output": _cpu(transformer_output),
+ },
+ transformer_path,
+ )
+
+ final_latents: dict[str, torch.Tensor] = {}
+
+ def capture_final_latents(_pipe, step, _timestep, callback_kwargs):
+ if step == NUM_INFERENCE_STEPS - 1:
+ final_latents["latents"] = _cpu(callback_kwargs["latents"])
+ return callback_kwargs
+
+ torch.accelerator.reset_peak_memory_stats()
+ generator = torch.Generator(device="cuda").manual_seed(SEED)
+ generation_kwargs = {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ "height": config["height"],
+ "width": config["width"],
+ "frames": NUM_FRAMES,
+ "guidance_scale": GUIDANCE_SCALE,
+ "num_inference_steps": NUM_INFERENCE_STEPS,
+ "generator": generator,
+ "output_type": "np",
+ "callback_on_step_end": capture_final_latents,
+ "callback_on_step_end_tensor_inputs": ["latents"],
+ }
+ if normalized_image is not None:
+ generation_kwargs["image"] = normalized_image
+
+ torch.accelerator.synchronize()
+ started = time.perf_counter()
+ with torch.inference_mode():
+ frames = pipe(**generation_kwargs).frames[0]
+ torch.accelerator.synchronize()
+ generation_seconds = time.perf_counter() - started
+ if "latents" not in final_latents:
+ raise RuntimeError("Diffusers did not expose the final denoised latents to the generation callback.")
+
+ video_path = output_dir / "pipeline.mp4"
+ export_to_video(frames, video_path, fps=FPS)
+ reference_path = output_dir / "pipeline_reference.safetensors"
+ save_file(
+ {
+ "final_latents": final_latents["latents"],
+ "decoded_frames_uint8": _frames_to_uint8(frames),
+ },
+ reference_path,
+ )
+
+ scheduler_config = dict(pipe.scheduler.config)
+ input_metadata = None
+ if input_path is not None:
+ input_metadata = {
+ "filename": input_path.name,
+ "sha256": _sha256(input_path),
+ "size": input_path.stat().st_size,
+ "format": "png-rgb",
+ }
+ metadata = {
+ "schema_version": SCHEMA_VERSION,
+ "variant": variant,
+ "task": task,
+ "publishable": publishable,
+ "reference_implementation": {
+ "library": "diffusers",
+ "version": __import__("diffusers").__version__,
+ "pipeline_class": pipeline_class.__name__,
+ "torch_version": torch.__version__,
+ },
+ "model_provenance": model_provenance,
+ "generator_provenance": repository_provenance,
+ "scheduler": {
+ "class": type(pipe.scheduler).__name__,
+ "config_sha256": _json_sha256(scheduler_config),
+ "config": scheduler_config,
+ },
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ "seed": SEED,
+ "generator_device": "cuda",
+ "dtype": "bfloat16",
+ "vae_dtype": str(config["vae_dtype"]).removeprefix("torch."),
+ "height": config["height"],
+ "width": config["width"],
+ "num_frames": NUM_FRAMES,
+ "fps": FPS,
+ "num_inference_steps": NUM_INFERENCE_STEPS,
+ "guidance_scale": GUIDANCE_SCALE,
+ "input_image": input_metadata,
+ "generation_seconds": generation_seconds,
+ "peak_reserved_gib": torch.accelerator.max_memory_reserved() / 1024**3,
+ }
+ metadata_path = output_dir / "metadata.json"
+ metadata_path.write_text(json.dumps(metadata, indent=2, default=str) + "\n")
+
+ artifact_paths = [transformer_path, video_path, reference_path, metadata_path]
+ if input_path is not None:
+ artifact_paths.append(input_path)
+ manifest = {
+ "schema_version": SCHEMA_VERSION,
+ "variant": variant,
+ "task": task,
+ "publishable": publishable,
+ "model_provenance": model_provenance,
+ "metadata_sha256": _sha256(metadata_path),
+ "files": {path.name: {"sha256": _sha256(path), "size": path.stat().st_size} for path in artifact_paths},
+ }
+ manifest_path = output_dir / "manifest.json"
+ manifest_path.write_text(json.dumps(manifest, indent=2) + "\n")
+ return manifest_path
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--variant", choices=VARIANTS, required=True)
+ parser.add_argument("--task", choices=("t2v", "i2v"), required=True)
+ parser.add_argument(
+ "--model",
+ help=(
+ "Optional local checkpoint. Local output is fingerprinted and marked development-only; "
+ "it cannot be published as the canonical golden."
+ ),
+ )
+ parser.add_argument("--input-image", type=Path, help="Fixed local RGB input for task=i2v.")
+ parser.add_argument("--local-source-model", help="Claimed upstream model ID for local checkpoint audit records.")
+ parser.add_argument("--local-source-revision", help="Claimed upstream commit for local checkpoint audit records.")
+ parser.add_argument("--output-root", type=Path, default=Path("/tmp/sana-video-goldens-v2"))
+ args = parser.parse_args()
+ print(
+ generate(
+ args.variant,
+ args.task,
+ args.model,
+ args.input_image,
+ args.output_root,
+ args.local_source_model,
+ args.local_source_revision,
+ )
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_golden.py b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py
new file mode 100644
index 000000000..c0c468315
--- /dev/null
+++ b/tests/e2e/accuracy/sana_video/test_sana_video_golden.py
@@ -0,0 +1,328 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+import torch
+from safetensors.torch import load_file
+from torch import nn
+
+from tests.e2e.accuracy.helpers import assert_video_metadata, assert_video_similarity_metrics, probe_video
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+GOLDEN_BASE_URL = os.getenv("SANA_VIDEO_GOLDEN_BASE_URL")
+SCHEMA_VERSION = 2
+PROMPT = "A cat walking on the grass, facing the camera. motion score: 30."
+NEGATIVE_PROMPT = (
+ "A chaotic sequence with misshapen, deformed limbs in heavy motion blur, sudden disappearance, jump cuts, "
+ "jerky movements, rapid shot changes, frames out of sync, inconsistent character shapes, temporal artifacts, "
+ "jitter, and ghosting effects, creating a disorienting visual experience."
+)
+VARIANTS = {
+ "480p": ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", 480, 832),
+ "720p": ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", 704, 1280),
+}
+REVISIONS = {
+ "480p": "fed3bce411c58a0f688a31afe8f52e61acc2b15f",
+ "720p": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e",
+}
+SSIM_THRESHOLDS = {
+ ("480p", "t2v"): 0.91,
+ ("480p", "i2v"): 0.93,
+ ("720p", "t2v"): 0.93,
+ ("720p", "i2v"): 0.93,
+}
+
+pytestmark = [
+ pytest.mark.full_model,
+ pytest.mark.diffusion,
+ pytest.mark.skipif(
+ not GOLDEN_BASE_URL,
+ reason="Set SANA_VIDEO_GOLDEN_BASE_URL after publishing the frozen v2 S3 assets.",
+ ),
+]
+
+
+class _TransformerCheckpoint(nn.Module):
+ def __init__(self, transformer: nn.Module, model: str, revision: str) -> None:
+ super().__init__()
+ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+
+ self.transformer = transformer
+ self.weights_sources = [
+ DiffusersPipelineLoader.ComponentSource(
+ model_or_path=model,
+ subfolder="transformer",
+ revision=revision,
+ prefix="transformer.",
+ fall_back_to_pt=True,
+ )
+ ]
+
+
+def _sha256(path: Path) -> str:
+ return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def _json_sha256(payload: Any) -> str:
+ encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str).encode()
+ return hashlib.sha256(encoded).hexdigest()
+
+
+def _assert_hex_digest(value: Any, *, length: int = 64) -> None:
+ assert isinstance(value, str) and len(value) == length
+ assert all(character in "0123456789abcdef" for character in value)
+
+
+def _validate_manifest(payload: Any, variant: str, task: str) -> dict[str, Any]:
+ assert isinstance(payload, dict)
+ assert payload["schema_version"] == SCHEMA_VERSION
+ assert payload["variant"] == variant
+ assert payload["task"] == task
+ # Local checkpoint generation is deliberately marked non-publishable. Do
+ # not permit such output to become the canonical reference by accident.
+ assert payload["publishable"] is True
+ assert payload["model_provenance"] == {
+ "kind": "huggingface_hub",
+ "model_id": VARIANTS[variant][0],
+ "revision": REVISIONS[variant],
+ }
+ _assert_hex_digest(payload["metadata_sha256"])
+
+ expected_files = {
+ "metadata.json",
+ "pipeline.mp4",
+ "pipeline_reference.safetensors",
+ "transformer_case.safetensors",
+ }
+ if task == "i2v":
+ expected_files.add("input.png")
+ files = payload["files"]
+ assert isinstance(files, dict)
+ assert set(files) == expected_files
+ for filename, expected in files.items():
+ # Prevent a compromised manifest from writing outside output_dir.
+ assert filename == Path(filename).name
+ assert isinstance(expected, dict) and set(expected) == {"sha256", "size"}
+ _assert_hex_digest(expected["sha256"])
+ assert isinstance(expected["size"], int) and expected["size"] > 0
+ return payload
+
+
+def _validate_metadata(payload: Any, manifest: dict[str, Any], variant: str, task: str) -> None:
+ assert isinstance(payload, dict)
+ _, height, width = VARIANTS[variant]
+ assert payload["schema_version"] == SCHEMA_VERSION
+ assert payload["variant"] == variant
+ assert payload["task"] == task
+ assert payload["publishable"] is True
+ assert payload["model_provenance"] == manifest["model_provenance"]
+ assert payload["prompt"] == PROMPT
+ assert payload["negative_prompt"] == NEGATIVE_PROMPT
+ assert payload["seed"] == 42
+ assert payload["generator_device"] == "cuda"
+ assert payload["dtype"] == "bfloat16"
+ assert payload["height"] == height
+ assert payload["width"] == width
+ assert payload["num_frames"] == 81
+ assert payload["fps"] == 16
+ assert payload["num_inference_steps"] == 50
+ assert payload["guidance_scale"] == 6.0
+
+ implementation = payload["reference_implementation"]
+ assert implementation["library"] == "diffusers"
+ assert implementation["pipeline_class"] == ("SanaImageToVideoPipeline" if task == "i2v" else "SanaVideoPipeline")
+ assert implementation["version"]
+ assert implementation["torch_version"]
+
+ generator = payload["generator_provenance"]
+ assert generator["repository"] == "vllm-project/vllm-omni"
+ _assert_hex_digest(generator["revision"], length=40)
+ _assert_hex_digest(generator["generator_sha256"])
+ assert generator["dirty"] is False
+
+ scheduler = payload["scheduler"]
+ assert scheduler["class"] == "DPMSolverMultistepScheduler"
+ _assert_hex_digest(scheduler["config_sha256"])
+ assert scheduler["config_sha256"] == _json_sha256(scheduler["config"])
+
+ input_image = payload["input_image"]
+ if task == "i2v":
+ assert input_image["filename"] == "input.png"
+ _assert_hex_digest(input_image["sha256"])
+ assert input_image["size"] == manifest["files"]["input.png"]["size"]
+ assert input_image["sha256"] == manifest["files"]["input.png"]["sha256"]
+ assert input_image["format"] == "png-rgb"
+ else:
+ assert input_image is None
+
+
+def _download_case(variant: str, task: str, output_dir: Path, filenames: set[str]) -> dict[str, Any]:
+ output_dir.mkdir(parents=True, exist_ok=True)
+ base_url = f"{GOLDEN_BASE_URL.rstrip('/')}/{variant}/{task}"
+ response = requests.get(f"{base_url}/manifest.json", timeout=60)
+ response.raise_for_status()
+ payload = _validate_manifest(response.json(), variant, task)
+ filenames = filenames | {"metadata.json"}
+ if task == "i2v":
+ filenames.add("input.png")
+ assert filenames <= payload["files"].keys()
+ for filename in filenames:
+ expected = payload["files"][filename]
+ path = output_dir / filename
+ response = requests.get(f"{base_url}/{filename}", timeout=300)
+ response.raise_for_status()
+ path.write_bytes(response.content)
+ assert path.stat().st_size == expected["size"]
+ assert _sha256(path) == expected["sha256"]
+ assert _sha256(output_dir / "metadata.json") == payload["metadata_sha256"]
+ metadata = json.loads((output_dir / "metadata.json").read_text())
+ _validate_metadata(metadata, payload, variant, task)
+ return payload
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("variant", VARIANTS)
+@pytest.mark.parametrize(
+ "_hardware",
+ [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))],
+)
+def test_sana_video_transformer_matches_frozen_golden(
+ variant: str,
+ _hardware,
+ accuracy_artifact_root: Path,
+) -> None:
+ del _hardware
+ from vllm.config import LoadConfig
+ from vllm.transformers_utils.config import get_hf_file_to_dict
+
+ from vllm_omni.diffusion.data import OmniDiffusionConfig
+ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+ from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+ model, _, _ = VARIANTS[variant]
+ revision = REVISIONS[variant]
+ output_dir = accuracy_artifact_root / "sana_video" / variant / "t2v"
+ _download_case(variant, "t2v", output_dir, {"transformer_case.safetensors"})
+ case = load_file(output_dir / "transformer_case.safetensors")
+ transformer_config = get_hf_file_to_dict("transformer/config.json", model, revision=revision)
+ assert transformer_config is not None
+ transformer = SanaVideoTransformer3DModel.from_config(transformer_config).to(
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ checkpoint = _TransformerCheckpoint(transformer, model, revision)
+ od_config = OmniDiffusionConfig(
+ model=model,
+ model_class_name="SanaVideoPipeline",
+ dtype=torch.bfloat16,
+ revision=revision,
+ )
+ loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
+ transformer.load_weights(
+ (name.removeprefix("transformer."), tensor)
+ for name, tensor in loader.get_all_weights(checkpoint)
+ if name.startswith("transformer.")
+ )
+ transformer.eval()
+ with torch.inference_mode():
+ actual = transformer(
+ case["hidden_states"].to("cuda"),
+ case["encoder_hidden_states"].to("cuda"),
+ case["timestep"].to("cuda"),
+ encoder_attention_mask=case["encoder_attention_mask"].to("cuda"),
+ ).sample.float()
+ expected = case["output"].to("cuda").float()
+ error = actual - expected
+ max_abs = error.abs().max()
+ relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected)
+ cosine = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0)
+ # CUDNN_ATTN BF16 reduction order is not bitwise deterministic on
+ # Blackwell. Bound rare outliers separately while keeping the aggregate
+ # relative-L2 and cosine checks strict.
+ max_abs_threshold = 0.2 if variant == "480p" else 0.1
+ relative_l2_threshold = 0.017 if variant == "480p" else 0.015
+ assert max_abs.item() <= max_abs_threshold
+ assert relative_l2.item() <= relative_l2_threshold
+ assert cosine.item() >= 0.999
+
+
+SERVER_CASES = [
+ pytest.param(
+ OmniServerParams(
+ model=model,
+ server_args=[
+ "--model-class-name",
+ pipeline_class,
+ ],
+ ),
+ id=f"{variant}-{task}",
+ marks=hardware_marks(res={"cuda": "H100"}),
+ )
+ for variant, (model, _, _) in VARIANTS.items()
+ for task, pipeline_class in (
+ ("t2v", "SanaVideoPipeline"),
+ ("i2v", "SanaImageToVideoPipeline"),
+ )
+]
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_sana_video_pipeline_matches_frozen_golden(
+ omni_server: OmniServer,
+ openai_client: OpenAIClientHandler,
+ accuracy_artifact_root: Path,
+) -> None:
+ variant = "720p" if "720p" in omni_server.model else "480p"
+ task = "i2v" if "SanaImageToVideoPipeline" in omni_server.serve_args else "t2v"
+ _, height, width = VARIANTS[variant]
+ output_dir = accuracy_artifact_root / "sana_video" / variant / task
+ _download_case(variant, task, output_dir, {"pipeline.mp4"})
+
+ request_config = {
+ "model": omni_server.model,
+ "form_data": {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ "height": height,
+ "width": width,
+ "num_frames": 81,
+ "fps": 16,
+ "num_inference_steps": 50,
+ "guidance_scale": 6.0,
+ "seed": 42,
+ },
+ }
+ if task == "i2v":
+ encoded = base64.b64encode((output_dir / "input.png").read_bytes()).decode()
+ request_config["image_reference"] = f"data:image/png;base64,{encoded}"
+ result = openai_client.send_video_diffusion_request(request_config)[0]
+ actual_path = output_dir / "actual.mp4"
+ actual_path.write_bytes(result.videos[0])
+ golden_path = output_dir / "pipeline.mp4"
+
+ metadata = probe_video(actual_path)
+ assert_video_metadata(metadata, width=width, height=height, fps=16, frame_count=81)
+ assert_video_similarity_metrics(
+ label=f"sana_video_{variant}_{task}",
+ online_path=actual_path,
+ offline_path=golden_path,
+ # The native prompt cross-attention uses vLLM-Omni Attention rather
+ # than Diffusers SDPA. Frozen stage checks prove exact inputs and
+ # scheduler state, while the BF16 backend reduction difference grows
+ # across CFG denoising. Calibrate only the affected 480p T2V case;
+ # its frozen result is 0.919861 SSIM / 28.7442 dB PSNR.
+ ssim_threshold=SSIM_THRESHOLDS[(variant, task)],
+ psnr_threshold=28.0,
+ )
diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py
new file mode 100644
index 000000000..8dd35eca2
--- /dev/null
+++ b/tests/e2e/accuracy/sana_video/test_sana_video_i2v_alignment.py
@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+ pytest.mark.full_model,
+ pytest.mark.diffusion,
+ pytest.mark.benchmark,
+]
+
+VARIANTS = {
+ "480p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+ "vae_class": "AutoencoderKLWan",
+ "native_vae_class": "DistributedAutoencoderKLWan",
+ "dtype": torch.float32,
+ "latent_channels": 16,
+ "spatial_scale": 8,
+ "temporal_scale": 4,
+ "frames": 5,
+ "max_abs": 1e-6,
+ "relative_l2": 1e-6,
+ },
+ "720p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+ "vae_class": "AutoencoderKLLTX2Video",
+ "native_vae_class": "DistributedAutoencoderKLLTX2Video",
+ "dtype": torch.bfloat16,
+ "latent_channels": 128,
+ "spatial_scale": 32,
+ "temporal_scale": 8,
+ "frames": 9,
+ "max_abs": 0.04,
+ "relative_l2": 0.004,
+ },
+}
+
+
+@pytest.mark.parametrize("variant", VARIANTS)
+@pytest.mark.parametrize(
+ "_hardware",
+ [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))],
+)
+def test_sana_video_i2v_first_frame_latents_match_diffusers(variant: str, _hardware) -> None:
+ del _hardware
+ from diffusers import (
+ AutoencoderKLLTX2Video,
+ AutoencoderKLWan,
+ )
+ from diffusers import (
+ SanaImageToVideoPipeline as DiffusersSanaImageToVideoPipeline,
+ )
+
+ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+ DistributedAutoencoderKLLTX2Video,
+ )
+ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+ DistributedAutoencoderKLWan,
+ )
+ from vllm_omni.diffusion.models.sana_video import SanaImageToVideoPipeline
+
+ config = VARIANTS[variant]
+ vae_classes = {
+ "AutoencoderKLLTX2Video": AutoencoderKLLTX2Video,
+ "AutoencoderKLWan": AutoencoderKLWan,
+ }
+ native_vae_classes = {
+ "DistributedAutoencoderKLLTX2Video": DistributedAutoencoderKLLTX2Video,
+ "DistributedAutoencoderKLWan": DistributedAutoencoderKLWan,
+ }
+ reference_vae = (
+ vae_classes[config["vae_class"]]
+ .from_pretrained(
+ config["model"],
+ subfolder="vae",
+ torch_dtype=config["dtype"],
+ )
+ .to("cuda")
+ )
+ native_vae = (
+ native_vae_classes[config["native_vae_class"]]
+ .from_config(dict(reference_vae.config))
+ .to(device="cuda", dtype=config["dtype"])
+ )
+ native_vae.load_state_dict(reference_vae.state_dict())
+ reference_vae.eval()
+ native_vae.eval()
+
+ height, width = 192, 320
+ num_latent_frames = (config["frames"] - 1) // config["temporal_scale"] + 1
+ image = torch.linspace(
+ -1,
+ 1,
+ 3 * height * width,
+ dtype=torch.float32,
+ device="cuda",
+ ).reshape(1, 3, height, width)
+ noise = torch.randn(
+ (
+ 1,
+ config["latent_channels"],
+ num_latent_frames,
+ height // config["spatial_scale"],
+ width // config["spatial_scale"],
+ ),
+ generator=torch.Generator(device="cuda").manual_seed(7),
+ device="cuda",
+ dtype=torch.float32,
+ )
+ reference_pipeline = SimpleNamespace(
+ vae=reference_vae,
+ vae_scale_factor_temporal=config["temporal_scale"],
+ vae_scale_factor_spatial=config["spatial_scale"],
+ )
+ native_pipeline = SimpleNamespace(
+ vae=native_vae,
+ vae_scale_factor_temporal=config["temporal_scale"],
+ vae_scale_factor_spatial=config["spatial_scale"],
+ device=torch.device("cuda"),
+ )
+
+ with torch.inference_mode():
+ expected = DiffusersSanaImageToVideoPipeline.prepare_latents(
+ reference_pipeline,
+ image,
+ 1,
+ config["latent_channels"],
+ height,
+ width,
+ config["frames"],
+ torch.float32,
+ torch.device("cuda"),
+ None,
+ noise.clone(),
+ )
+ actual = SanaImageToVideoPipeline._prepare_i2v_latents(
+ native_pipeline,
+ image,
+ 1,
+ config["latent_channels"],
+ height,
+ width,
+ config["frames"],
+ torch.float32,
+ None,
+ noise.clone(),
+ )
+
+ # Conditioning replaces only the first latent frame.
+ torch.testing.assert_close(actual[:, :, 1:], noise[:, :, 1:], rtol=0, atol=0)
+ error = (actual[:, :, :1] - expected[:, :, :1]).float()
+ expected_first_frame = expected[:, :, :1].float()
+ max_abs = error.abs().max()
+ relative_l2 = torch.linalg.vector_norm(error) / torch.linalg.vector_norm(expected_first_frame)
+ cosine = torch.nn.functional.cosine_similarity(
+ actual[:, :, :1].flatten(),
+ expected[:, :, :1].flatten(),
+ dim=0,
+ )
+ assert max_abs.item() <= config["max_abs"]
+ assert relative_l2.item() <= config["relative_l2"]
+ assert cosine.item() >= 0.9999
diff --git a/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py
new file mode 100644
index 000000000..a88aec7e2
--- /dev/null
+++ b/tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py
@@ -0,0 +1,432 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import gc
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+ pytest.mark.full_model,
+ pytest.mark.diffusion,
+ pytest.mark.benchmark,
+]
+
+PROMPT = "A small red panda walks through a bamboo forest."
+NEGATIVE_PROMPT = "blur, distortion, abrupt cuts"
+SEED = 42
+GUIDANCE_SCALE = 6.0
+NUM_INFERENCE_STEPS = 2
+HEIGHT = 192
+WIDTH = 320
+NUM_FRAMES = 9
+
+VARIANTS = {
+ "480p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+ "revision": "fed3bce411c58a0f688a31afe8f52e61acc2b15f",
+ "model_env": "SANA_VIDEO_480P_MODEL",
+ "max_abs": 0.2,
+ "relative_l2": 0.02,
+ "cosine": 0.999,
+ "cfg_max_abs": 0.5,
+ "cfg_relative_l2": 0.05,
+ "cfg_cosine": 0.999,
+ "step_max_abs": 0.075,
+ "step_relative_l2": 0.012,
+ "step_cosine": 0.9999,
+ },
+ "720p": {
+ "model": "Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+ "revision": "8bda5e623d0f48cd6da3b387b10ca35d15cf1c4e",
+ "model_env": "SANA_VIDEO_720P_MODEL",
+ "max_abs": 0.1,
+ "relative_l2": 0.015,
+ "cosine": 0.999,
+ "cfg_max_abs": 0.6,
+ "cfg_relative_l2": 0.065,
+ "cfg_cosine": 0.998,
+ "step_max_abs": 0.08,
+ "step_relative_l2": 0.014,
+ "step_cosine": 0.9999,
+ },
+}
+
+
+@dataclass(frozen=True)
+class ErrorMetrics:
+ shape: tuple[int, ...]
+ actual_dtype: torch.dtype
+ expected_dtype: torch.dtype
+ max_abs: float
+ mean_abs: float
+ relative_l2: float
+ cosine: float
+
+ def __str__(self) -> str:
+ return (
+ f"shape={self.shape}, actual_dtype={self.actual_dtype}, "
+ f"expected_dtype={self.expected_dtype}, max_abs={self.max_abs:.8g}, "
+ f"mean_abs={self.mean_abs:.8g}, relative_l2={self.relative_l2:.8g}, "
+ f"cosine={self.cosine:.8g}"
+ )
+
+
+def _metrics(actual: torch.Tensor, expected: torch.Tensor) -> ErrorMetrics:
+ assert actual.shape == expected.shape
+ actual_float = actual.detach().float()
+ expected_float = expected.detach().float()
+ error = actual_float - expected_float
+ expected_norm = torch.linalg.vector_norm(expected_float)
+ relative_l2 = torch.linalg.vector_norm(error) / expected_norm.clamp_min(torch.finfo(torch.float32).tiny)
+ if actual_float.numel() == 1:
+ cosine = torch.tensor(
+ 1.0 if bool(torch.equal(actual_float, expected_float)) else 0.0,
+ device=actual.device,
+ )
+ else:
+ cosine = F.cosine_similarity(actual_float.flatten(), expected_float.flatten(), dim=0)
+ return ErrorMetrics(
+ shape=tuple(actual.shape),
+ actual_dtype=actual.dtype,
+ expected_dtype=expected.dtype,
+ max_abs=error.abs().max().item(),
+ mean_abs=error.abs().mean().item(),
+ relative_l2=relative_l2.item(),
+ cosine=cosine.item(),
+ )
+
+
+def _assert_stage(
+ stage: str,
+ actual: torch.Tensor,
+ expected: torch.Tensor,
+ *,
+ max_abs: float,
+ relative_l2: float,
+ cosine: float,
+) -> ErrorMetrics:
+ result = _metrics(actual, expected)
+ print(f"\nSANA-Video alignment [{stage}]: {result}")
+ assert result.actual_dtype == result.expected_dtype, f"{stage}: {result}"
+ assert result.max_abs <= max_abs, f"{stage}: {result}"
+ assert result.relative_l2 <= relative_l2, f"{stage}: {result}"
+ assert result.cosine >= cosine, f"{stage}: {result}"
+ return result
+
+
+def _model_source(variant: str) -> tuple[str, str | None]:
+ config = VARIANTS[variant]
+ model = os.getenv(config["model_env"], config["model"])
+ revision = None if Path(model).exists() else config["revision"]
+ return model, revision
+
+
+def _native_method_proxy(reference_pipeline: Any):
+ """Build the minimum native-pipeline surface needed by stage helpers.
+
+ Sharing the tokenizer and text encoder makes this a comparison of pipeline
+ orchestration rather than a second copy of the same multi-billion-parameter
+ encoder. Transformer execution still uses separate Diffusers and native
+ modules below.
+ """
+ from vllm_omni.diffusion.models.sana_video import SanaVideoPipeline
+
+ proxy = object.__new__(SanaVideoPipeline)
+ torch.nn.Module.__init__(proxy)
+ proxy.tokenizer = reference_pipeline.tokenizer
+ proxy.text_encoder = reference_pipeline.text_encoder
+ proxy.device = torch.device("cuda")
+ proxy.vae_scale_factor_temporal = reference_pipeline.vae_scale_factor_temporal
+ proxy.vae_scale_factor_spatial = reference_pipeline.vae_scale_factor_spatial
+ return proxy
+
+
+def _cfg_prediction(
+ prediction: torch.Tensor,
+ *,
+ guidance_scale: float,
+ latent_channels: int,
+ out_channels: int,
+) -> torch.Tensor:
+ unconditional, conditional = prediction.float().chunk(2)
+ prediction = unconditional + guidance_scale * (conditional - unconditional)
+ if out_channels // 2 == latent_channels:
+ prediction = prediction.chunk(2, dim=1)[0]
+ return prediction
+
+
+@pytest.mark.parametrize("variant", VARIANTS)
+@pytest.mark.parametrize(
+ "_hardware",
+ [pytest.param(None, marks=hardware_marks(res={"cuda": "H100"}))],
+)
+def test_sana_video_quick_pipeline_stages_match_diffusers(
+ variant: str,
+ _hardware,
+) -> None:
+ """Compare the first denoising step without replacing native attention.
+
+ This deliberately stops after the first scheduler step. It keeps the test
+ useful as a stage-localizing PR/nightly check while the frozen 81-frame,
+ 50-step golden test covers the final latent and decoded-video boundary.
+ """
+ del _hardware
+ from diffusers import SanaVideoPipeline as DiffusersSanaVideoPipeline
+ from diffusers.pipelines.sana_video.pipeline_sana_video import (
+ retrieve_timesteps as diffusers_retrieve_timesteps,
+ )
+ from diffusers.schedulers import DPMSolverMultistepScheduler
+
+ from vllm_omni.diffusion.models.sana_video import (
+ SanaVideoPipeline,
+ SanaVideoTransformer3DModel,
+ )
+ from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+ retrieve_timesteps as native_retrieve_timesteps,
+ )
+
+ model, revision = _model_source(variant)
+ config = VARIANTS[variant]
+ load_kwargs: dict[str, Any] = {
+ "torch_dtype": torch.bfloat16,
+ # No stage in this test decodes latents. Avoid loading a multi-GB VAE
+ # that would add memory and runtime without increasing coverage.
+ "vae": None,
+ }
+ if revision is not None:
+ load_kwargs["revision"] = revision
+
+ reference = DiffusersSanaVideoPipeline.from_pretrained(model, **load_kwargs).to("cuda")
+ assert type(reference) is DiffusersSanaVideoPipeline
+ assert type(reference.transformer).__module__.startswith("diffusers.")
+ native_proxy = _native_method_proxy(reference)
+ assert type(native_proxy) is SanaVideoPipeline
+ assert type(native_proxy).__module__.startswith("vllm_omni.")
+ try:
+ with torch.inference_mode():
+ reference_prompt = reference.encode_prompt(
+ PROMPT,
+ True,
+ negative_prompt=NEGATIVE_PROMPT,
+ device=torch.device("cuda"),
+ clean_caption=False,
+ max_sequence_length=300,
+ complex_human_instruction=None,
+ )
+ native_prompt = SanaVideoPipeline.encode_prompt(
+ native_proxy,
+ PROMPT,
+ True,
+ negative_prompt=NEGATIVE_PROMPT,
+ device=torch.device("cuda"),
+ clean_caption=False,
+ max_sequence_length=300,
+ complex_human_instruction=None,
+ )
+ prompt_stage_names = (
+ "prompt_embeddings",
+ "prompt_attention_mask",
+ "negative_prompt_embeddings",
+ "negative_prompt_attention_mask",
+ )
+ for stage, actual, expected in zip(prompt_stage_names, native_prompt, reference_prompt, strict=True):
+ _assert_stage(
+ stage,
+ actual,
+ expected,
+ max_abs=0.0,
+ relative_l2=0.0,
+ cosine=0.999999,
+ )
+
+ reference_prompt_embeds, reference_prompt_mask, reference_negative_embeds, reference_negative_mask = (
+ reference_prompt
+ )
+ encoder_hidden_states = torch.cat([reference_negative_embeds, reference_prompt_embeds])
+ encoder_attention_mask = torch.cat([reference_negative_mask, reference_prompt_mask])
+
+ latent_channels = reference.transformer.config.in_channels
+ reference_latents = reference.prepare_latents(
+ 1,
+ latent_channels,
+ HEIGHT,
+ WIDTH,
+ NUM_FRAMES,
+ torch.float32,
+ torch.device("cuda"),
+ torch.Generator(device="cuda").manual_seed(SEED),
+ )
+ native_latents = SanaVideoPipeline.prepare_latents(
+ native_proxy,
+ 1,
+ latent_channels,
+ HEIGHT,
+ WIDTH,
+ NUM_FRAMES,
+ torch.float32,
+ torch.device("cuda"),
+ torch.Generator(device="cuda").manual_seed(SEED),
+ )
+ _assert_stage(
+ "initial_noise_latents",
+ native_latents,
+ reference_latents,
+ max_abs=0.0,
+ relative_l2=0.0,
+ cosine=0.999999,
+ )
+
+ reference_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config)
+ native_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config)
+ isolated_scheduler = DPMSolverMultistepScheduler.from_config(reference.scheduler.config)
+ reference_timesteps, reference_num_steps = diffusers_retrieve_timesteps(
+ reference_scheduler,
+ NUM_INFERENCE_STEPS,
+ torch.device("cuda"),
+ )
+ native_timesteps, native_num_steps = native_retrieve_timesteps(
+ native_scheduler,
+ NUM_INFERENCE_STEPS,
+ torch.device("cuda"),
+ )
+ isolated_retrieve_timesteps, isolated_num_steps = native_retrieve_timesteps(
+ isolated_scheduler,
+ NUM_INFERENCE_STEPS,
+ torch.device("cuda"),
+ )
+ assert reference_num_steps == native_num_steps == isolated_num_steps
+ _assert_stage(
+ "timesteps",
+ native_timesteps,
+ reference_timesteps,
+ max_abs=0.0,
+ relative_l2=0.0,
+ cosine=0.999999,
+ )
+ _assert_stage(
+ "sigmas",
+ native_scheduler.sigmas,
+ reference_scheduler.sigmas,
+ max_abs=0.0,
+ relative_l2=0.0,
+ cosine=0.999999,
+ )
+ torch.testing.assert_close(isolated_retrieve_timesteps, reference_timesteps, rtol=0, atol=0)
+
+ native_transformer = SanaVideoTransformer3DModel.from_config(dict(reference.transformer.config)).to(
+ device="cuda",
+ dtype=torch.bfloat16,
+ )
+ incompatible = native_transformer.load_state_dict(reference.transformer.state_dict(), strict=True)
+ assert not incompatible.missing_keys
+ assert not incompatible.unexpected_keys
+ assert type(native_transformer) is SanaVideoTransformer3DModel
+ assert type(native_transformer).__module__.startswith("vllm_omni.")
+ assert type(native_transformer) is not type(reference.transformer)
+ reference_parameter = next(reference.transformer.parameters())
+ native_parameter = next(native_transformer.parameters())
+ assert native_parameter.data_ptr() != reference_parameter.data_ptr()
+ torch.testing.assert_close(native_parameter, reference_parameter, rtol=0, atol=0)
+ native_transformer.eval()
+ reference.transformer.eval()
+
+ latent_model_input = torch.cat([reference_latents, reference_latents])
+ timestep = reference_timesteps[0].expand(latent_model_input.shape[0])
+ with torch.inference_mode():
+ reference_raw_prediction = reference.transformer(
+ latent_model_input.to(dtype=torch.bfloat16),
+ encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16),
+ encoder_attention_mask=encoder_attention_mask,
+ timestep=timestep,
+ return_dict=False,
+ )[0]
+ native_raw_prediction = native_transformer(
+ latent_model_input.to(dtype=torch.bfloat16),
+ encoder_hidden_states=encoder_hidden_states.to(dtype=torch.bfloat16),
+ encoder_attention_mask=encoder_attention_mask,
+ timestep=timestep,
+ return_dict=False,
+ )[0]
+
+ _assert_stage(
+ "first_transformer_raw_prediction",
+ native_raw_prediction,
+ reference_raw_prediction,
+ max_abs=config["max_abs"],
+ relative_l2=config["relative_l2"],
+ cosine=config["cosine"],
+ )
+ reference_cfg_prediction = _cfg_prediction(
+ reference_raw_prediction,
+ guidance_scale=GUIDANCE_SCALE,
+ latent_channels=latent_channels,
+ out_channels=reference.transformer.config.out_channels,
+ )
+ native_cfg_prediction = _cfg_prediction(
+ native_raw_prediction,
+ guidance_scale=GUIDANCE_SCALE,
+ latent_channels=latent_channels,
+ out_channels=native_transformer.config.out_channels,
+ )
+ _assert_stage(
+ "first_cfg_prediction",
+ native_cfg_prediction,
+ reference_cfg_prediction,
+ max_abs=config["cfg_max_abs"],
+ relative_l2=config["cfg_relative_l2"],
+ cosine=config["cfg_cosine"],
+ )
+
+ reference_first_step = reference_scheduler.step(
+ reference_cfg_prediction,
+ reference_timesteps[0],
+ reference_latents,
+ return_dict=False,
+ )[0]
+ # Isolate scheduler orchestration from native-attention error by
+ # feeding the Diffusers prediction through an independently configured
+ # scheduler initialized by the native helper.
+ isolated_first_step = isolated_scheduler.step(
+ reference_cfg_prediction,
+ isolated_retrieve_timesteps[0],
+ reference_latents,
+ return_dict=False,
+ )[0]
+ _assert_stage(
+ "first_scheduler_step_isolated",
+ isolated_first_step,
+ reference_first_step,
+ max_abs=0.0,
+ relative_l2=0.0,
+ cosine=0.999999,
+ )
+
+ native_first_step = native_scheduler.step(
+ native_cfg_prediction,
+ native_timesteps[0],
+ native_latents,
+ return_dict=False,
+ )[0]
+ _assert_stage(
+ "first_scheduler_step_end_to_end",
+ native_first_step,
+ reference_first_step,
+ max_abs=config["step_max_abs"],
+ relative_l2=config["step_relative_l2"],
+ cosine=config["step_cosine"],
+ )
+ finally:
+ del native_proxy
+ del reference
+ gc.collect()
+ torch.accelerator.empty_cache()
diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py
index 68f890c64..be3d7f84c 100644
--- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py
+++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Run vLLM-Omni with diffusers backend, then run diffusers directly. Compare the output similarity.
@@ -39,9 +39,9 @@ from tests.e2e.accuracy.helpers import (
from tests.e2e.accuracy.helpers import (
run_ffmpeg_similarity as _run_ffmpeg_similarity,
)
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OnlineOmniClient
from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import (
CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS,
)
@@ -128,7 +128,7 @@ def _run_vllm_omni_wan22_i2v(
"seed": SEED,
}
with OmniServer(model, server_args, env_dict=env_to_apply_ftfy_mock_in_subproc(), use_omni=True) as omni_server:
- client = OpenAIClientHandler(
+ client = OnlineOmniClient(
host=omni_server.host,
port=omni_server.port,
run_level="full_model",
@@ -139,16 +139,20 @@ def _run_vllm_omni_wan22_i2v(
"image_reference": f"data:image/png;base64,{pil_to_base64(conditioning_image, 'png')}",
}
result = client.send_video_diffusion_request(request_config)[0]
- video_bytes = result.videos[0] # pyright: ignore[reportOptionalSubscript] # Guaranteed not None
+ videos = result.videos
+ assert videos is not None
+ video_bytes = videos[0]
output_path.write_bytes(video_bytes)
- return result.e2e_latency # pyright: ignore[reportReturnType] # Guaranteed not None
+ latency = result.e2e_latency
+ assert latency is not None
+ return latency
def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_image: Image.Image) -> float:
from diffusers import WanImageToVideoPipeline # pyright: ignore[reportPrivateImportUsage]
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
- run_pre_test_cleanup()
+ cleanup_test_environment()
apply_ftfy_mock()
pipe: WanImageToVideoPipeline | None = None
try:
@@ -190,7 +194,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]:
@@ -248,7 +252,7 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.Image, float]:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: DiffusionPipeline | None = None
try:
pipe = DiffusionPipeline.from_pretrained(
@@ -288,7 +292,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
@pytest.mark.benchmark
diff --git a/tests/e2e/accuracy/test_hidream_o1_image.py b/tests/e2e/accuracy/test_hidream_o1_image.py
new file mode 100644
index 000000000..923e53a47
--- /dev/null
+++ b/tests/e2e/accuracy/test_hidream_o1_image.py
@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL_ID = "HiDream-ai/HiDream-O1-Image"
+PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'"
+BASELINE_PATH = Path(__file__).resolve().parents[2] / "assets" / "hidream_o1" / "hidream_o1_baseline.png"
+SSIM_THRESHOLD = 0.8
+PSNR_THRESHOLD = 19.0
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_hidream_o1_image_tp2_matches_baseline(
+ accuracy_artifact_root: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+ model = os.environ.get("HIDREAM_O1_IMAGE_MODEL", MODEL_ID)
+ with OmniRunner(
+ model,
+ parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+ enforce_eager=True,
+ ) as runner:
+ outputs = runner.omni.generate(
+ {"prompt": PROMPT},
+ OmniDiffusionSamplingParams(
+ height=2048,
+ width=2048,
+ num_inference_steps=50,
+ guidance_scale=5.0,
+ seed=42,
+ ),
+ )
+ images = extract_images_from_outputs(outputs)
+ assert images
+ image = images[0].convert("RGB")
+
+ output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+ image.save(output_dir / "vllm_omni_tp2.png")
+ assert BASELINE_PATH.exists()
+ baseline = Image.open(BASELINE_PATH).convert("RGB")
+ resized = image.resize(baseline.size, Image.Resampling.LANCZOS)
+ assert_similarity(
+ model_name=MODEL_ID,
+ vllm_image=resized,
+ diffusers_image=baseline,
+ ssim_threshold=SSIM_THRESHOLD,
+ psnr_threshold=PSNR_THRESHOLD,
+ )
diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py
index 6f3ed7e71..a09014c8d 100644
--- a/tests/e2e/accuracy/test_qwen_image.py
+++ b/tests/e2e/accuracy/test_qwen_image.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import base64
@@ -13,7 +16,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from PIL import Image
from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
@@ -29,14 +32,14 @@ HEIGHT = 512
NUM_INFERENCE_STEPS = 20
TRUE_CFG_SCALE = 4.0
SEED = 42
-SSIM_THRESHOLD = 0.97
-# The FA3-hub-matched comparison clears 0.97 with margin. The SDPA-matched
-# fallback (hub kernel unavailable for the image's torch) landed at 0.9687 in
-# build 2954 — the attention math matches but the surrounding stacks (vLLM
-# compiled path vs eager diffusers) still differ slightly. Keep the strict
-# gate for the FA3 path and a slightly looser one for the fallback only.
+SSIM_THRESHOLD = 0.94
+# With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8
+# vs Diffusers (see #5734 / #5963), so the primary gates are 0.94 / 27.0.
+# The SDPA-matched fallback (hub kernel unavailable for the image's torch)
+# landed at SSIM 0.9687 in build 2954 — keep a slightly tighter SSIM gate for
+# that path only.
SSIM_THRESHOLD_SDPA_FALLBACK = 0.96
-PSNR_THRESHOLD = 30.0
+PSNR_THRESHOLD = 27.0
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL"
@@ -149,7 +152,7 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: DiffusionPipeline | None = None
try:
pipe = DiffusionPipeline.from_pretrained(
@@ -179,7 +182,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
@@ -211,7 +214,7 @@ def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im
def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: DiffusionPipeline | None = None
try:
pipe = DiffusionPipeline.from_pretrained(
@@ -241,7 +244,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
@pytest.mark.benchmark
diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py
index 9ee9c2c87..9c3b9dd9c 100644
--- a/tests/e2e/accuracy/test_qwen_image_edit.py
+++ b/tests/e2e/accuracy/test_qwen_image_edit.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import gc
@@ -12,12 +15,12 @@ from PIL import Image
from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServer
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion]
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
@@ -109,7 +112,7 @@ def _run_diffusers_image_edit(
input_images: list[Image.Image],
output_path: Path,
) -> Image.Image:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
try:
images = input_images[0] if len(input_images) == 1 else input_images
@@ -140,7 +143,7 @@ def _run_diffusers_image_edit(
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
def _load_edit_2511_input_images() -> list[Image.Image]:
@@ -189,7 +192,7 @@ def _run_diffusers_image_edit_2511(
input_images: list[Image.Image],
output_path: Path,
) -> Image.Image:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: QwenImageEditPlusPipeline | None = None
try:
pipe = QwenImageEditPlusPipeline.from_pretrained(
@@ -222,7 +225,7 @@ def _run_diffusers_image_edit_2511(
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
def _vllm_omni_output_single_image(
@@ -253,6 +256,7 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
)
+@pytest.mark.slow
@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_qwen_image_edit_single_matches_diffusers(
@@ -278,6 +282,7 @@ def test_qwen_image_edit_single_matches_diffusers(
)
+@pytest.mark.slow
@pytest.mark.benchmark
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py
index f2bc1b6aa..53354a739 100644
--- a/tests/e2e/accuracy/test_qwen_image_layered.py
+++ b/tests/e2e/accuracy/test_qwen_image_layered.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import base64
@@ -13,11 +16,11 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from PIL import Image
from tests.e2e.accuracy.helpers import assert_image_sequence_similarity, model_output_dir
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.clean import cleanup_test_environment
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
MODEL_ID = "Qwen/Qwen-Image-Layered"
@@ -93,7 +96,7 @@ def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, o
def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]:
- run_pre_test_cleanup()
+ cleanup_test_environment()
pipe: DiffusionPipeline | None = None
try:
pipe = DiffusionPipeline.from_pretrained(
@@ -126,7 +129,7 @@ def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, o
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
- run_post_test_cleanup()
+ cleanup_test_environment()
@pytest.mark.benchmark
diff --git a/tests/e2e/accuracy/test_video_streaming_output_similarity.py b/tests/e2e/accuracy/test_video_streaming_output_similarity.py
index 799e6e125..cdb2c4a28 100644
--- a/tests/e2e/accuracy/test_video_streaming_output_similarity.py
+++ b/tests/e2e/accuracy/test_video_streaming_output_similarity.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Online serving smoke test for streaming video output similarity.
"""
@@ -9,9 +12,9 @@ import pytest
from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler
+from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, get_model_prefix
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "BestWishYsh/Helios-Distilled"
PROMPT = "A serene lakeside sunrise with mist over the water."
@@ -67,8 +70,8 @@ def _request_config(model: str, *, serialize_extra_params: bool) -> dict:
}
-def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler:
- return OpenAIClientHandler(
+def _client_for(server: OmniServer, *, run_level: str) -> OnlineOmniClient:
+ return OnlineOmniClient(
host=server.host,
port=server.port,
api_key="EMPTY",
@@ -95,12 +98,11 @@ def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Pat
)
def test_helios_streaming_output_matches_non_streaming(
model: str,
- model_prefix: str,
run_level: str,
tmp_path: Path,
record_property,
):
- model = model_prefix + model
+ model = get_model_prefix() + model
with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server:
client = _client_for(server, run_level=run_level)
diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
index a09d2330a..a995a9cf9 100644
--- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
+++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import os
@@ -338,7 +341,7 @@ def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
def _generate_online_video(
*,
omni_server,
- openai_client,
+ online_client,
image_source: str,
online_timeout_seconds: int,
) -> Path:
@@ -360,7 +363,7 @@ def _generate_online_video(
"image_reference": _build_online_image_reference(image_source),
}
online_video_bytes = _send_video_request_with_timeout(
- openai_client,
+ online_client,
request_config,
timeout_seconds=_online_timeout_seconds(online_timeout_seconds),
)
@@ -412,7 +415,7 @@ def test_wan22_i2v_diffusers_offline_generates_video(
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
def test_wan22_i2v_online_serving_generates_video(
omni_server,
- openai_client,
+ online_client,
wan22_i2v_image_source: str | None,
wan22_i2v_online_timeout_seconds: int,
) -> None:
@@ -424,7 +427,7 @@ def test_wan22_i2v_online_serving_generates_video(
_validate_image_source(image_source)
online_path = _generate_online_video(
omni_server=omni_server,
- openai_client=openai_client,
+ online_client=online_client,
image_source=image_source,
online_timeout_seconds=wan22_i2v_online_timeout_seconds,
)
diff --git a/tests/e2e/features/comfyui/test_comfyui_integration.py b/tests/e2e/features/comfyui/test_comfyui_integration.py
index 5dcb9c449..61cf76dcd 100644
--- a/tests/e2e/features/comfyui/test_comfyui_integration.py
+++ b/tests/e2e/features/comfyui/test_comfyui_integration.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Integration tests for ComfyUI nodes that use the Omni API client, with a mocked AsyncOmni and a real API server running in a background process.
These tests cover the integration between ComfyUI node and the API server, without actual model inference logic.
@@ -39,6 +42,7 @@ from pytest_mock import MockerFixture
from vllm import SamplingParams
from vllm.outputs import CompletionOutput, RequestOutput
+from tests.helpers.runtime import get_open_port
from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
from vllm_omni.entrypoints.cli.serve import OmniServeCommand
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
@@ -515,7 +519,7 @@ def mock_async_omni(
@pytest.fixture
-def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni, tmp_path):
+def api_server(server_case: ServerCase, mock_async_omni, tmp_path):
"""Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
# Override the STORAGE_MANAGER path to a writable temp directory before
# forking the server subprocess. The default /tmp/storage may be owned
@@ -533,7 +537,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
cmd = OmniServeCommand()
cmd.subparser_init(subparsers)
- port = unused_tcp_port_factory()
+ port = get_open_port(host="0.0.0.0")
args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
def run_server():
@@ -550,11 +554,12 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
wait_poll_interval = 1
for _ in range(wait_time // wait_poll_interval):
try:
- response = requests.get(f"http://127.0.0.1:{port}/health")
+ response = requests.get(f"http://127.0.0.1:{port}/health", timeout=wait_poll_interval)
if response.status_code == 200:
break
- except requests.ConnectionError:
- time.sleep(wait_poll_interval)
+ except requests.RequestException:
+ pass
+ time.sleep(wait_poll_interval)
else:
if server_process.is_alive():
server_process.terminate()
@@ -809,7 +814,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
stage_list=["diffusion"],
- stage_configs=[{"stage_type": "diffusion"}],
+ stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
@@ -820,7 +825,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
ServerCase(
served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
stage_list=["diffusion"],
- stage_configs=[{"stage_type": "diffusion"}],
+ stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
"Wan-AI/Wan2.2-I2V-A14B-Diffusers",
@@ -881,7 +886,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
- stage_configs=[{"stage_type": "diffusion"}],
+ stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
@@ -892,7 +897,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
ServerCase(
served_model="MiniMaxAI/MiniMax-H3",
stage_list=["diffusion"],
- stage_configs=[{"stage_type": "diffusion"}],
+ stage_configs=[{"stage_type": "diffusion", "final_output": True, "final_output_type": "video"}],
outputs=[_build_diffusion_video_output()],
),
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
index 04e83caeb..a11c9b7e4 100644
--- a/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
+++ b/tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py
@@ -166,8 +166,8 @@ async def test_sleep_memory_reclaimed_custom_pipeline():
The fix moves custom_pipeline init outside the CUDA context so all weights
go through the caching allocator and are therefore fully reclaimed by
- sleep(level=1). A non-zero ``CuMemAllocator.get_current_usage()`` after
- sleep is the direct signal that the bypass is still occurring.
+ sleep(level=1). The tracked allocations must transition to the allocator's
+ asleep state and report physical reclamation.
"""
with ExitStack() as after:
engine = AsyncOmni(
@@ -187,62 +187,72 @@ async def test_sleep_memory_reclaimed_custom_pipeline():
free_before, total = torch.cuda.mem_get_info()
used_before_gib = (total - free_before) / 1024**3
- # Measure CuMemAllocator-tracked usage before sleep. In inline mode
- # the worker runs in a thread pool inside this process, so the allocator
- # singleton is shared and can be read directly.
+ # Measure CuMemAllocator-tracked usage before sleep. In inline / uni
+ # mode the worker runs in this process, so the allocator singleton is
+ # shared and can be read directly.
allocator = None
tracked_before = 0
+ tracked_ptrs: set[int] = set()
try:
from vllm.device_allocator.cumem import CuMemAllocator
allocator = CuMemAllocator.get_instance()
- tracked_before = allocator.get_current_usage()
+ tracked_ptrs = {ptr for ptr, data in allocator.pointer_to_data.items() if not data.is_asleep}
+ tracked_before = sum(allocator.pointer_to_data[ptr].handle[1] for ptr in tracked_ptrs)
except Exception:
pass
- # Put the engine to sleep; all weights should be offloaded via the pool.
- acks = await engine.sleep(level=1)
- await asyncio.sleep(0.5) # allow the CUDA driver to settle
- torch.accelerator.synchronize()
+ slept = False
+ try:
+ # Put the engine to sleep; this engine's weights should be
+ # offloaded via the pool.
+ acks = await engine.sleep(level=1)
+ slept = True
+ await asyncio.sleep(0.5) # allow the CUDA driver to settle
+ torch.accelerator.synchronize()
- # Measure after sleep.
- free_after, _ = torch.cuda.mem_get_info()
- used_after_gib = (total - free_after) / 1024**3
- drop_gib = used_before_gib - used_after_gib
+ # Measure after sleep.
+ free_after, _ = torch.cuda.mem_get_info()
+ used_after_gib = (total - free_after) / 1024**3
+ drop_gib = used_before_gib - used_after_gib
- # --- Primary assertion: allocator reports zero tracked memory. ---
- # If this fails it means weights were allocated outside the CuMem pool
- # (safetensors direct-to-GPU bypass) — the exact regression this test
- # is designed to catch.
- if allocator is not None:
- tracked_after = allocator.get_current_usage()
- assert tracked_after == 0, (
- f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB "
- f"after sleep(level=1) on custom_pipeline path "
- f"(was {tracked_before / 1024**3:.3f} GiB before sleep). "
- "Weights were allocated outside the CuMem pool via the "
- "safetensors direct-to-GPU fast path — loader-context fix "
- "may not be applied."
+ # get_current_usage() includes unmapped sleeping handles, so it
+ # does not decrease on sleep. Verify that allocations mapped
+ # before the call instead transitioned to the asleep state.
+ if allocator is not None:
+ assert tracked_before > 0, (
+ "Expected custom_pipeline weights to use the CuMem pool, "
+ "but no mapped allocations were tracked before sleep."
+ )
+ still_awake = [
+ ptr
+ for ptr in tracked_ptrs
+ if ptr in allocator.pointer_to_data and not allocator.pointer_to_data[ptr].is_asleep
+ ]
+ assert not still_awake, (
+ f"{len(still_awake)} CuMem allocation(s) remained mapped "
+ "after sleep(level=1) on the custom_pipeline path."
+ )
+
+ # Physical VRAM or ACK freed_bytes confirms reclamation at the
+ # driver level.
+ total_freed_bytes = sum(
+ (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
+ for ack in acks
+ if ack is not None
+ )
+ freed_gib = total_freed_bytes / 1024**3
+ assert freed_gib > 0 or drop_gib > 0, (
+ f"Expected GPU memory to be reclaimed after sleep(level=1) on "
+ f"custom_pipeline + enable_sleep_mode=True. "
+ f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
+ f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
)
- # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms
- # reclamation at the driver level.
- total_freed_bytes = sum(
- (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
- for ack in acks
- if ack is not None
- )
- freed_gib = total_freed_bytes / 1024**3
- assert freed_gib > 0 or drop_gib > 0, (
- f"Expected GPU memory to be reclaimed after sleep(level=1) on "
- f"custom_pipeline + enable_sleep_mode=True. "
- f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
- f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
- )
-
- # Engine must report it is sleeping.
- assert await engine.is_sleeping()
-
- # Wake up and confirm the engine is functional again.
- await engine.wake_up()
- assert not await engine.is_sleeping()
+ assert await engine.is_sleeping()
+ await engine.wake_up()
+ slept = False
+ assert not await engine.is_sleeping()
+ finally:
+ if slept:
+ await engine.wake_up()
diff --git a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
index b24f364e4..12ad4f1b3 100644
--- a/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
+++ b/tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py
@@ -17,9 +17,12 @@ from functools import lru_cache
import numpy as np
import pytest
+import torch.distributed as dist
from transformers import AutoTokenizer
from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -131,6 +134,19 @@ async def _generate_once(
return last_output
+def _assert_live_uniproc_worker(engine: AsyncOmni) -> None:
+ """Default single-GPU path must use the in-process executor, not mp IPC."""
+ clients = engine.engine.stage_clients
+ assert len(clients) == 1
+ client = clients[0]
+ assert isinstance(client, InlineStageDiffusionClient)
+ executor = client._engine.executor
+ assert isinstance(executor, UniProcDiffusionExecutor)
+ assert executor.driver_worker is not None
+ assert dist.is_initialized()
+ assert dist.get_world_size() == 1
+
+
def _assert_valid_image_output(output: OmniRequestOutput) -> None:
assert output.final_output_type == "image"
assert output.images, "Expected at least one generated image"
@@ -158,6 +174,7 @@ async def test_async_omni_generate():
)
after.callback(engine.shutdown)
+ _assert_live_uniproc_worker(engine)
output = await _generate_once(
engine,
"a beautiful sunset over the ocean with vibrant orange and purple clouds "
diff --git a/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py
index abb387573..65d4f1ca9 100644
--- a/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py
+++ b/tests/e2e/features/fullduplex/engine/test_duplex_deploy_config.py
@@ -8,6 +8,7 @@ from dataclasses import FrozenInstanceError
import pytest
+from tests.helpers.stage_config import get_deploy_duplex_max_sessions
from vllm_omni.config.stage_config import (
DuplexSessionRuntimeConfig,
load_deploy_config,
@@ -66,3 +67,18 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str,
with pytest.raises(ValueError, match=rf"duplex_session\.{name} must be positive"):
load_deploy_config(deploy_path)
+
+
+@pytest.mark.parametrize(
+ ("deploy_yaml", "expected"),
+ [
+ ("minicpmo_4_5.yaml", 4),
+ ("minicpmo_4_5_8x4090.yaml", 1),
+ ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
+ ],
+)
+def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None:
+ # Guards the admission probe. A capacity edit, a config that declares none,
+ # or an overlay inheriting one from its base must surface here rather than
+ # as another nightly duplex admission timeout.
+ assert get_deploy_duplex_max_sessions(deploy_yaml) == expected
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py
new file mode 100644
index 000000000..5790b11f1
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_data_plane.py
@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.engine.duplex_runtime import duplex_resource_request_id
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.data_plane import (
+ NemotronVoiceChatDataPlaneContext,
+ NemotronVoiceChatDataPlaneSession,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_RUNTIME = {
+ "nvc_text_bos_id": 0,
+ "nvc_text_eos_id": 1,
+ "nvc_text_pad_id": 12,
+ "nvc_function_sotc_id": 20,
+ "nvc_function_eotc_id": 21,
+ "nvc_tokenizer_ref": "test-tokenizer",
+}
+
+
+def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
+ projector = NemotronVoiceChatDataPlaneSession(encode_audio)
+ projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
+ return projector
+
+
+def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
+ metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
+ if function_token is not None:
+ metadata["nvc_function_token"] = [function_token]
+ completion = SimpleNamespace(multimodal_output=metadata)
+ return SimpleNamespace(
+ stage_id=0,
+ request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
+ )
+
+
+def test_function_channel_projects_completed_call_without_ending_speech() -> None:
+ projector = _projector()
+ projector._decode = lambda token_ids: (
+ '[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
+ )
+ context = NemotronVoiceChatDataPlaneContext(epoch=0)
+
+ events = []
+ for function_token in (20, 99, 21):
+ events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
+
+ listen = [event for event in events if event.get("is_listen") is True]
+ function = [event for event in events if event.get("function_call") is True]
+ assert len(listen) == 3
+ assert len(function) == 1
+ assert function[0]["name"] == "weather"
+ assert function[0]["arguments"] == '{"city":"Shanghai"}'
+
+ projector._decode = lambda _token_ids: "not-json"
+ bad_events = []
+ for function_token in (20, 98, 21):
+ bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
+ assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
+
+
+def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
+ projector = _projector(lambda *_: "audio")
+ old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
+ new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
+
+ projector.begin_request(old_request)
+ list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
+ projector.begin_request(new_request)
+
+ # EOTC in the new epoch must not close the old epoch's partial function.
+ new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
+ assert not [event for event in new_events if event.get("function_call") is True]
+ assert not [event for event in new_events if event.get("end_of_turn") is True]
+
+ audio = SimpleNamespace(
+ stage_id=2,
+ request_output=SimpleNamespace(
+ request_id=new_request,
+ outputs=[
+ SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
+ ],
+ ),
+ )
+ audio_events = list(projector.project_output(audio))
+ assert sum(event.get("end_of_turn") is True for event in audio_events) == 1
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py
new file mode 100644
index 000000000..84c8b123a
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_input.py
@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import base64
+
+import numpy as np
+import pytest
+
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.input import (
+ NemotronVoiceChatPcmAppendBuffer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]:
+ return {
+ "type": "audio",
+ "audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"),
+ "format": fmt,
+ "sample_rate_hz": sample_rate_hz,
+ }
+
+
+def _append(buffer, samples, operation_id):
+ return buffer.prepare_append(
+ _payload(np.asarray(samples, dtype=np.float32)),
+ operation_id=operation_id,
+ chunk_period_ms=80,
+ allow_emit=True,
+ )
+
+
+def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
+ buffer = NemotronVoiceChatPcmAppendBuffer()
+ assert _append(buffer, np.arange(1000), "packet-1") is None
+
+ reservation = _append(buffer, np.arange(1000, 1280), "packet-2")
+
+ assert reservation is not None
+ assert reservation.operation_id == "packet-2"
+ decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32)
+ np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
+ assert buffer.pending_byte_count == 0
+
+
+def test_append_rejects_invalid_frame_contract() -> None:
+ cases = (
+ (1281, 16000, "pcm_f32le", 0.0, "at most 1280"),
+ (1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"),
+ (1, 16000, "audio/pcm", 0.0, "format"),
+ (1, 16000, "pcm_f32le", np.nan, "finite"),
+ )
+ for samples, sample_rate_hz, fmt, value, message in cases:
+ with pytest.raises(ValueError, match=message):
+ NemotronVoiceChatPcmAppendBuffer().prepare_append(
+ _payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt),
+ operation_id="invalid",
+ chunk_period_ms=80,
+ allow_emit=True,
+ )
+
+
+def test_terminal_commit_rejects_more_than_one_frame_tail() -> None:
+ buffer = NemotronVoiceChatPcmAppendBuffer()
+ for index, samples in enumerate((1000, 281)):
+ buffer.prepare_append(
+ _payload(np.zeros(samples, dtype=np.float32)),
+ operation_id=f"tail-{index}",
+ chunk_period_ms=80,
+ allow_emit=False,
+ )
+ with pytest.raises(ValueError, match="more than one 80 ms frame"):
+ buffer.prepare_commit(operation_id="commit", chunk_period_ms=80)
diff --git a/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py
new file mode 100644
index 000000000..0b32322c0
--- /dev/null
+++ b/tests/e2e/features/fullduplex/nemotron_voicechat/test_runtime_contract.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import base64
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexInputMode
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.runtime import (
+ NemotronVoiceChatDuplexRuntimeExtension,
+)
+from vllm_omni.experimental.fullduplex.nemotron_voicechat.serving_adapter import (
+ NemotronVoiceChatServingRuntimeAdapter,
+ _render_tool_response,
+)
+from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
+ NemotronVoiceChatThinkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _frame() -> dict[str, object]:
+ raw = np.zeros(1280, dtype=np.float32).tobytes()
+ return {
+ "type": "audio",
+ "audio": base64.b64encode(raw).decode("ascii"),
+ "format": "pcm_f32le",
+ "sample_rate_hz": 16000,
+ }
+
+
+def _plan(extension, *, input_seq: int):
+ return extension.plan_append(
+ request_id="req",
+ fence=DuplexFence("sid", incarnation=2, epoch=3),
+ session_config={},
+ runtime_config={
+ "nvc_prompt_token_ids": [0, 42, 1],
+ "nvc_text_pad_id": 12,
+ "nvc_max_model_len": 8192,
+ },
+ seq=input_seq,
+ turn_seq=input_seq,
+ mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
+ payload=_frame(),
+ final=False,
+ sampling_params=SamplingParams(),
+ )
+
+
+def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
+ extension = NemotronVoiceChatDuplexRuntimeExtension()
+
+ params = extension.configure_sampling_params(
+ runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
+ )[0]
+ assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
+
+ first = _plan(extension, input_seq=1)
+ later = _plan(extension, input_seq=2)
+
+ assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
+ assert later.prompt["prompt_token_ids"] == [12]
+ assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
+ assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
+
+
+def test_append_rejects_stage0_context_overflow() -> None:
+ with pytest.raises(ValueError, match="max_model_len"):
+ _plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
+
+
+def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
+ encoded: list[str] = []
+ adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
+ adapter._tokenizer = SimpleNamespace(
+ encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
+ )
+
+ first = adapter.runtime_config_for_function_output(
+ {"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
+ {},
+ )
+ second = adapter.runtime_config_for_function_output(
+ {"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
+ first,
+ )
+
+ assert _render_tool_response('{"result":20}') == '[{"result":20}]'
+ assert encoded == [
+ '[{"result":20}]',
+ '["plain text"]',
+ ]
+ assert first["nvc_function_response_generation"] == 1
+ assert second["nvc_function_response_generation"] == 2
+ assert second["nvc_function_response_token_ids"] == [31, 32, 33]
+ assert second["nvc_function_response_call_id"] == "call-2"
+ assert second["nvc_function_response_batches"] == [
+ {"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
+ {"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
+ ]
+
+ session = {"function_response_generation": 0, "forced_function_tokens": []}
+ NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
+
+ assert session["function_response_generation"] == 2
+ assert session["forced_function_token"] == 31
+ assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]
diff --git a/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py b/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py
index f4a3c994e..f2970a240 100644
--- a/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py
+++ b/tests/e2e/features/fullduplex/personaplex/test_unified_runtime.py
@@ -18,6 +18,10 @@ from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
DuplexInputMode,
)
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+from vllm_omni.experimental.fullduplex.openai.runtime_adapter import (
+ ServingRuntimeConfigError,
+)
+from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
PersonaPlexDataPlaneContext,
PersonaPlexDataPlaneSession,
@@ -242,6 +246,53 @@ def test_personaplex_capabilities_are_honest() -> None:
assert multi.supports_barge_in is False
+def test_personaplex_runtime_config_update_rejects_changed_persona() -> None:
+ adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
+ short_persona = "You are a helpful assistant."
+ longer_persona = (
+ "You are now a swashbuckling pirate captain who speaks entirely in "
+ "nautical slang, references buried treasure constantly, and never "
+ "breaks character no matter what the user asks."
+ )
+ current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"}
+
+ with pytest.raises(ServingRuntimeConfigError, match="persona"):
+ adapter.runtime_config_for_update(
+ SimpleNamespace(instructions=longer_persona, extra_body={}),
+ current,
+ )
+
+ unchanged = adapter.runtime_config_for_update(
+ SimpleNamespace(instructions=short_persona, extra_body={}),
+ current,
+ )
+ assert unchanged["personaplex_persona"] == short_persona
+
+
+def test_personaplex_runtime_config_update_rejects_changed_voice() -> None:
+ adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
+ current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"}
+
+ with pytest.raises(ServingRuntimeConfigError, match="voice"):
+ adapter.runtime_config_for_update(
+ SimpleNamespace(instructions=None, voice="OtherVoice.pt", extra_body={}),
+ current,
+ )
+
+ unchanged = adapter.runtime_config_for_update(
+ SimpleNamespace(instructions=None, voice="NATF2.pt", extra_body={}),
+ current,
+ )
+ assert unchanged["personaplex_voice_prompt"] == "NATF2.pt"
+
+
+def test_personaplex_rejects_minicpmo_only_opt_out_flag() -> None:
+ adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
+
+ with pytest.raises(ServingRuntimeConfigError, match="minicpmo45_native_duplex"):
+ adapter.validate_client_extra_body({"minicpmo45_native_duplex": False})
+
+
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
buffer = PersonaPlexPcmAppendBuffer()
reservation = buffer.prepare_append(
diff --git a/tests/e2e/features/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py
index d90c9e5aa..53965b89e 100644
--- a/tests/e2e/features/fullduplex/test_client.py
+++ b/tests/e2e/features/fullduplex/test_client.py
@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
import base64
import wave
from urllib.parse import parse_qs, urlsplit
@@ -14,13 +18,17 @@ from vllm_omni.experimental.fullduplex.client import (
from vllm_omni.experimental.fullduplex.minicpmo45.policy import (
MiniCPMO45DuplexPolicy,
)
+from vllm_omni.experimental.fullduplex.video_stacking import (
+ concat_frames,
+ unit_subframe_offsets,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_realtime_client_builds_explicit_native_duplex_url():
url = build_realtime_url(
- "ws://localhost:8099/v1/realtime?custom=1",
+ "ws://localhost:8099/v1/realtime?custom=1&duplex=0&model=stale&minicpmo45_native_duplex=0&session_id=stale",
"openbmb/MiniCPM-o-4_5",
session_id="session-a",
)
@@ -50,7 +58,7 @@ def test_seed_tts_initial_text_is_part_of_native_duplex_context():
def test_realtime_client_builds_resume_only_url_when_autostart_disabled():
url = build_realtime_url(
- "ws://localhost:8099/v1/realtime?duplex=1",
+ "ws://localhost:8099/v1/realtime?duplex=1&autostart=1",
"openbmb/MiniCPM-o-4_5",
autostart=False,
)
@@ -73,10 +81,15 @@ async def test_realtime_client_configure_omits_ref_audio_by_default():
client = Client()
- await client.configure("openbmb/MiniCPM-o-4_5", timeout_s=1)
+ await client.configure(
+ "openbmb/MiniCPM-o-4_5",
+ idle_timeout_s=900,
+ timeout_s=1,
+ )
session = client.sent[0]["session"]
assert "ref_audio" not in session
+ assert session["idle_timeout_s"] == 900
@pytest.mark.asyncio
@@ -100,6 +113,74 @@ async def test_realtime_client_configure_sends_explicit_ref_audio():
session = client.sent[0]["session"]
assert session["ref_audio"] == "data:audio/wav;base64,AAAA"
+ assert "idle_timeout_s" not in session
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_reports_reader_exit_to_event_waiters():
+ client = RealtimeDuplexClient("ws://unused")
+
+ async def stopped():
+ return None
+
+ client._reader_task = asyncio.create_task(stopped())
+ await client._reader_task
+
+ with pytest.raises(ConnectionError, match="closed before"):
+ client.raise_if_reader_stopped()
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_configure_reports_reader_failure_immediately():
+ class Client(RealtimeDuplexClient):
+ async def send(self, event):
+ async def fail():
+ raise RuntimeError("reader failed")
+
+ self._reader_task = asyncio.create_task(fail())
+ await asyncio.sleep(0)
+
+ client = Client("ws://unused")
+ with pytest.raises(ConnectionError, match="reader failed"):
+ await client.configure("model", timeout_s=10.0)
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_close_reports_reader_cancellation_immediately():
+ class Client(RealtimeDuplexClient):
+ async def send(self, event):
+ async def wait_forever():
+ await asyncio.Future()
+
+ self._reader_task = asyncio.create_task(wait_forever())
+ self._reader_task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await self._reader_task
+
+ client = Client("ws://unused")
+ with pytest.raises(ConnectionError, match="reader was cancelled"):
+ await client.close_session(timeout_s=10.0)
+
+
+@pytest.mark.asyncio
+async def test_realtime_client_close_waits_for_a_new_session_closed_event():
+ class Client(RealtimeDuplexClient):
+ async def send(self, event):
+ assert event["type"] == "session.close"
+
+ async def acknowledge_close():
+ await asyncio.sleep(0.02)
+ self.events.add({"type": "session.closed"})
+
+ asyncio.create_task(acknowledge_close())
+
+ client = Client("ws://unused")
+ client.events.add({"type": "session.closed", "reason": "old"})
+ started_at = asyncio.get_running_loop().time()
+
+ await client.close_session(timeout_s=0.2)
+
+ assert asyncio.get_running_loop().time() - started_at >= 0.01
@pytest.mark.asyncio
@@ -235,6 +316,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals():
"output_token_count": 4,
"ttft_ms": 120.0,
"tpot_ms": 15.0,
+ "itls_ms": [10.0, 14.0, 18.0],
"inter_token_interval_ms": {
"count": 3,
"mean": 14.0,
@@ -340,3 +422,62 @@ def test_realtime_client_pcm16_wav_round_trip(tmp_path):
assert wav_file.getnchannels() == 1
assert wav_file.getframerate() == 16_000
assert read_pcm16_wav(path) == pcm16
+
+
+def _stack_test_image(color, size=(40, 30)):
+ from PIL import Image
+
+ return Image.new("RGB", size, color)
+
+
+def test_video_stacking_tiles_a_units_subframes_into_one_image():
+ frames = [_stack_test_image((255, 0, 0)), _stack_test_image((0, 255, 0))]
+
+ composite = concat_frames(frames)
+
+ # Landscape cells stack into a column because that canvas is squarer, and
+ # the interior seam carries the separator band that lets the model tell the
+ # sub-frames apart.
+ assert composite.size == (40, 30 * 2 + 6)
+ assert composite.getpixel((20, 30 + 3)) == (0, 0, 0)
+
+
+def test_video_stacking_picks_the_squarest_grid():
+ portrait = [_stack_test_image((255, 0, 0), size=(20, 60)) for _ in range(2)]
+ four = [_stack_test_image((0, 0, 255)) for _ in range(4)]
+
+ # A portrait clip tiles side by side, four frames always tile 2x2.
+ assert concat_frames(portrait).size == (20 * 2 + 6, 60)
+ assert concat_frames(four).size == (40 * 2 + 6, 30 * 2 + 6)
+
+
+def test_video_stacking_skips_the_subframe_that_duplicates_the_base_frame():
+ # Official samples stack_frames=5 as 0.2/0.4/0.6/0.8 s: offset 0 is the base
+ # frame, already sent as frame_list[0].
+ assert unit_subframe_offsets(5) == pytest.approx([0.2, 0.4, 0.6, 0.8])
+ assert unit_subframe_offsets(1) == []
+
+
+def test_realtime_client_sends_each_units_composite_beside_its_base_frame():
+ sent: list[dict] = []
+
+ class _Client(RealtimeDuplexClient):
+ def __init__(self):
+ pass
+
+ async def send(self, event):
+ sent.append(event)
+
+ asyncio.run(
+ _Client().stream_pcm16(
+ b"\x01\x00" * (16_000 * 3),
+ chunk_ms=200,
+ realtime=False,
+ video_frames=["f0", "f1"],
+ stacked_video_frames=["s0", None],
+ )
+ )
+
+ # A composite belongs to the unit it was captured in, so it rides the same
+ # append as that unit's base frame; a unit without one sends the base alone.
+ assert [event["video_frames"] for event in sent if "video_frames" in event] == [["f0", "s0"], ["f1"]]
diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py
new file mode 100644
index 000000000..192859de5
--- /dev/null
+++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py
@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
+from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
+from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+
+
+def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
+ manifest = tmp_path / "manifest.json"
+ manifest.write_text(
+ json.dumps(
+ [
+ {
+ "id": "sample-1",
+ "split": "PR_correction",
+ "question_text": "What changed?",
+ "answer1": "The object moved.",
+ }
+ ]
+ ),
+ encoding="utf-8",
+ )
+ response_root = tmp_path / "responses"
+ score_root = tmp_path / "scores"
+
+ async def fake_generate(sample, *, output_root, **kwargs):
+ output = Path(output_root) / sample.split / f"{sample.id}.json"
+ output.parent.mkdir(parents=True, exist_ok=True)
+ output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
+ output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
+ return output
+
+ class FakeJudge:
+ def __init__(self, *args, **kwargs):
+ pass
+
+ def chat(self, *args, **kwargs):
+ return '{"success_score": 1, "is_relevant": 1}'
+
+ monkeypatch.setattr(cli, "generate_sample", fake_generate)
+ monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
+
+ parser = argparse.ArgumentParser()
+ OmniDuplexEvalSubcommand.add_cli_args(parser)
+
+ common = ["--dataset", str(manifest), "--family", "pr"]
+ generate = parser.parse_args(
+ [
+ "generate",
+ *common,
+ "--model",
+ "mock",
+ "--ref-audio",
+ str(manifest),
+ "--response-root",
+ str(response_root),
+ "--concurrency",
+ "2",
+ ]
+ )
+ OmniDuplexEvalSubcommand.cmd(generate)
+ evaluate = parser.parse_args(
+ [
+ "evaluate",
+ *common,
+ "--response-root",
+ str(response_root),
+ "--score-root",
+ str(score_root),
+ "--judge-model",
+ "mock-judge",
+ "--eval-workers",
+ "2",
+ ]
+ )
+ OmniDuplexEvalSubcommand.cmd(evaluate)
+ OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
+
+ summary = json.loads(capsys.readouterr().out)
+ assert summary["samples"] == 1
+ assert summary["pr"]["mean_all_success"] == 1.0
+
+
+@pytest.mark.asyncio
+async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+ import websockets
+
+ received = []
+
+ async def handler(websocket):
+ response_started = False
+ async for raw in websocket:
+ event = json.loads(raw)
+ received.append(event)
+ if event["type"] == "session.update":
+ await websocket.send(json.dumps({"type": "session.created"}))
+ elif event["type"] == "input_audio_buffer.append" and not response_started:
+ response_started = True
+ await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
+ await websocket.send(
+ json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
+ )
+ await websocket.send(
+ json.dumps(
+ {
+ "type": "response.audio.delta",
+ "response_id": "r1",
+ "delta": base64.b64encode(b"\0\0" * 240).decode(),
+ "sample_rate_hz": 24_000,
+ }
+ )
+ )
+ elif event["type"] == "input_audio_buffer.commit":
+ await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
+ elif event["type"] == "session.close":
+ await websocket.send(json.dumps({"type": "session.closed"}))
+ return
+
+ monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
+ monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
+ monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
+ audio = tmp_path / "question.wav"
+ video = tmp_path / "video.mp4"
+ ref = tmp_path / "ref.wav"
+ for path in (audio, video, ref):
+ path.write_bytes(b"data")
+ sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
+
+ async with websockets.serve(handler, "127.0.0.1", 0) as server:
+ port = server.sockets[0].getsockname()[1]
+ output = await runner.generate_sample(
+ sample,
+ url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
+ model="mock",
+ ref_audio=ref,
+ output_root=tmp_path / "responses",
+ )
+
+ assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
+ meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
+ assert meta["response_done"] is True
+ assert meta["drain_timeout"] is None
+ assert any(event["type"] == "playback.ack" for event in received)
+
+
+def test_judge_exercises_openai_http_schema():
+ requests = []
+
+ class Handler(BaseHTTPRequestHandler):
+ def do_POST(self):
+ length = int(self.headers["Content-Length"])
+ requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
+ payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
+ self.send_response(200)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(payload)))
+ self.end_headers()
+ self.wfile.write(payload)
+
+ def log_message(self, *args):
+ pass
+
+ server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
+ assert judge.chat("prompt") == '{"success_score": 1}'
+ finally:
+ server.shutdown()
+ thread.join()
+ server.server_close()
+
+ path, authorization, payload = requests[0]
+ assert path == "/v1/chat/completions"
+ assert authorization == "Bearer token"
+ assert payload["model"] == "judge"
+ assert payload["messages"] == [{"role": "user", "content": "prompt"}]
diff --git a/tests/e2e/features/fullduplex/test_personaplex_elastic.py b/tests/e2e/features/fullduplex/test_personaplex_elastic.py
index 453d90153..1b2f47984 100644
--- a/tests/e2e/features/fullduplex/test_personaplex_elastic.py
+++ b/tests/e2e/features/fullduplex/test_personaplex_elastic.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""GPU e2e for PersonaPlex elastic batching (bit-exact isolation + recycle).
Greedy decoding makes every cross-slot leak observable as a bit difference:
@@ -24,12 +24,17 @@ from __future__ import annotations
import numpy as np
import pytest
+from tests.helpers.mark import hardware_marks
+
torch = pytest.importorskip("torch")
-# slow + H100 + omni routes this module into the weekly Buildkite E2E group
-# ("Omni · H100" in .buildkite/cuda/test-weekly.yml: -m "slow and H100 and
-# omni" --run-level full_model); the ready/merge CPU sweeps exclude it by mark.
-pytestmark = [pytest.mark.full_model, pytest.mark.omni, pytest.mark.slow, pytest.mark.H100]
+# slow + H100 + omni + cards_1 routes this module into weekly
+# "Omni · H100 · Single-GPU" (test-weekly.yml).
+pytestmark = [
+ pytest.mark.omni,
+ pytest.mark.slow,
+ *hardware_marks(res={"cuda": "H100"}),
+]
if not torch.cuda.is_available(): # noqa: E402
pytest.skip("needs CUDA", allow_module_level=True)
diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py
deleted file mode 100644
index af944b144..000000000
--- a/tests/e2e/offline_inference/test_audiox_model_expansion.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import os
-
-import numpy as np
-import pytest
-import torch
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-# Prefer a tiny/random checkpoint for CI.
-# Override in CI if needed: AUDIOX_TEST_MODEL=
-models = [os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")]
-
-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
-_OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models]
-
-pytestmark = [
- pytest.mark.slow,
- pytest.mark.diffusion,
- pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True, ids=models),
-]
-
-
-@hardware_test(res={"cuda": "L4", "xpu": "B60"})
-def test_audiox_model(omni_runner: OmniRunner) -> None:
- # Keep runtime short for CI.
- seconds_total = 2.0
- # AudioXPipeline always emits 44.1 kHz stereo (advertised via class-level
- # ``audio_sample_rate``); the trimmed output should match this rate.
- sample_rate = 44100
-
- outputs = omni_runner.omni.generate(
- prompts={"prompt": "A dog barking in a quiet park."},
- sampling_params_list=OmniDiffusionSamplingParams(
- num_inference_steps=4,
- guidance_scale=6.0,
- generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
- num_outputs_per_prompt=1,
- extra_args={
- "audiox_task": "t2a",
- "seconds_start": 0.0,
- "seconds_total": seconds_total,
- },
- ),
- )
-
- assert outputs is not None
- first_output = outputs[0]
- assert first_output.final_output_type == "audio"
- assert isinstance(first_output, OmniRequestOutput) and first_output
-
- req_out = first_output
- assert isinstance(req_out, OmniRequestOutput)
- assert req_out.final_output_type == "audio"
- assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
-
- audio = req_out.multimodal_output.get("audio")
- assert isinstance(audio, np.ndarray)
- # audio shape: (batch, channels, samples)
- assert audio.ndim == 3
- assert audio.shape[0] == 1
- assert audio.shape[1] == 2
- assert audio.shape[2] > 0
- expected_samples = int(seconds_total * sample_rate)
- assert abs(audio.shape[2] - expected_samples) <= 2 * 1024
diff --git a/tests/e2e/offline_inference/test_bagel.py b/tests/e2e/offline_inference/test_bagel.py
index 9ba7bdac4..580a46742 100644
--- a/tests/e2e/offline_inference/test_bagel.py
+++ b/tests/e2e/offline_inference/test_bagel.py
@@ -46,7 +46,7 @@ REFERENCE_TEXT_IMG2TEXT = (
_OMNI_RUNNER_PARAM = (MODEL_NAME, STAGE_CONFIG)
pytestmark = [
- pytest.mark.advanced_model,
+ pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
diff --git a/tests/e2e/offline_inference/test_bagel_expansion.py b/tests/e2e/offline_inference/test_bagel_expansion.py
index 980fc633f..4c31be722 100644
--- a/tests/e2e/offline_inference/test_bagel_expansion.py
+++ b/tests/e2e/offline_inference/test_bagel_expansion.py
@@ -43,7 +43,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml")
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
-# (model, stage_configs_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, deploy_config_path) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_OMNI_RUNNER_PARAM = (MODEL, BAGEL_STAGE_CONFIG)
@@ -138,7 +138,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
pytestmark = [
- pytest.mark.full_model,
+ pytest.mark.slow,
pytest.mark.diffusion,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py
index fa1bdfead..7e57433ef 100644
--- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py
+++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py
@@ -126,7 +126,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s
]
-# (model, stage_configs_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
+# (model, deploy_config_path, extra_omni_kwargs) for ``@pytest.mark.parametrize("omni_runner", ..., indirect=True)``
_cosy_deployment = get_deploy_config_path("cosyvoice3.yaml")
_cosy_model_path = str(_resolve_model_dir())
_OMNI_RUNNER_PARAMS = [
diff --git a/tests/e2e/offline_inference/test_covo_audio_expansion.py b/tests/e2e/offline_inference/test_covo_audio_expansion.py
index 11672686e..e743814d3 100644
--- a/tests/e2e/offline_inference/test_covo_audio_expansion.py
+++ b/tests/e2e/offline_inference/test_covo_audio_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
E2E tests for Covo-Audio-Chat model with audio input and audio/text output.
"""
@@ -25,7 +28,7 @@ def get_question():
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_audio_to_audio(omni_runner, omni_runner_handler) -> None:
+def test_audio_to_audio(omni_runner, offline_client) -> None:
"""
Test audio input → text + audio output.
Deploy Setting: default yaml
@@ -44,4 +47,4 @@ def test_audio_to_audio(omni_runner, omni_runner_handler) -> None:
"audios": (audio, 16000),
"modalities": ["audio"],
}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
diff --git a/tests/e2e/offline_inference/test_dots_tts_expansion.py b/tests/e2e/offline_inference/test_dots_tts_expansion.py
new file mode 100644
index 000000000..1c9362a23
--- /dev/null
+++ b/tests/e2e/offline_inference/test_dots_tts_expansion.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for dots.tts native AR offline inference."""
+
+from collections.abc import Mapping
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+DOTS_TTS_MODEL = "rednote-hilab/dots.tts-soar"
+DEPLOY_CONFIG = get_deploy_config_path("dots_tts.yaml")
+SAMPLE_RATE = 48000
+
+# dots.tts ships a custom tokenizer carrying the added ``<|audio_gen_start|>``
+# token, so remote code must be explicitly enabled.
+_OMNI_RUNNER_PARAM = (DOTS_TTS_MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})
+
+pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
+
+
+def _build_prompt(omni_runner: OmniRunner, text: str) -> dict:
+ """Wrap ``text`` in the prefix scaffold dots.tts prefill expects.
+
+ Unlike most TTS models here, a bare ``{"prompt": ...}`` string does not
+ work: prefill needs ``[文本][文本对应语音]<|audio_gen_start|>`` as
+ real Qwen2 token ids.
+ """
+ from transformers import AutoTokenizer
+
+ from vllm_omni.model_executor.models.dots_tts.dots_tts_prompt import (
+ build_dots_tts_prompt,
+ )
+
+ tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True)
+ return build_dots_tts_prompt(tokenizer=tokenizer, text=text)
+
+
+def _extract_audio(multimodal_output: dict) -> torch.Tensor:
+ """Extract the final complete audio tensor from multimodal output."""
+ assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}"
+
+ # The talker emits delta chunks; the output processor accumulates them
+ # under "model_outputs" ("audio" kept for backwards compat).
+ audio = multimodal_output.get("model_outputs")
+ if audio is None:
+ audio = multimodal_output.get("audio")
+ assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}"
+
+ if isinstance(audio, list):
+ valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None]
+ assert valid, "No valid audio tensors in output list"
+ audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+
+ assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}"
+ return audio.float().cpu().reshape(-1)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_dots_tts_zero_shot_001(omni_runner: OmniRunner) -> None:
+ """Zero-shot synthesis produces audio of a plausible duration."""
+ prompt = _build_prompt(omni_runner, "Hello, this is a test of dots TTS running on vLLM Omni.")
+
+ outputs = omni_runner.omni.generate([prompt])
+ assert len(outputs) == 1
+
+ audio = _extract_audio(outputs[0].outputs[0].multimodal_output)
+ duration_s = audio.shape[0] / SAMPLE_RATE
+ assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
+ assert torch.isfinite(audio).all(), "Audio contains NaN/Inf"
+ assert audio.abs().max() > 1e-3, "Audio is silent"
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_dots_tts_prefill_decode_mixed_batch_002(omni_runner: OmniRunner) -> None:
+ """Mixed prefill+decode batches must keep per-request rows aligned.
+
+ ``compute_logits`` fills ``logits[i]`` positionally from
+ ``_results_queue``, so every request in the batch has to push exactly one
+ entry per step, prefills included. When the prefill branch pushed nothing,
+ a shorter request's stop signal could be read into another request's row.
+ Regression guard for that pairing (PR #4765 review).
+ """
+ long_text = (
+ "This is a deliberately long sentence that stays in the decode phase "
+ "for many steps, so the shorter requests behind it keep entering "
+ "prefill alongside it and reproduce the mixed batch pattern."
+ )
+ short_texts = ["Hello one.", "Hello two.", "Hello three.", "Hello four."]
+ prompts = [_build_prompt(omni_runner, long_text)] + [_build_prompt(omni_runner, t) for t in short_texts]
+
+ outputs = omni_runner.omni.generate(prompts)
+ assert len(outputs) == len(prompts)
+
+ for i, out in enumerate(outputs):
+ audio = _extract_audio(out.outputs[0].multimodal_output)
+ duration_s = audio.shape[0] / SAMPLE_RATE
+ assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_dots_tts_deterministic_flow_matching_noise_003(omni_runner: OmniRunner) -> None:
+ """The same text twice must yield the same audio.
+
+ The DiT flow-matching noise is drawn from a per-request generator seeded
+ from ``blake2b(seed:request_key:noise_step)`` rather than the global CUDA
+ RNG. Drawing from the global RNG made output non-reproducible and let
+ concurrent requests perturb each other's noise streams (PR #4765 review).
+ """
+ text = "Reading aloud every evening helps children build a lasting love of language."
+
+ first = _extract_audio(
+ omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output
+ )
+ second = _extract_audio(
+ omni_runner.omni.generate([_build_prompt(omni_runner, text)])[0].outputs[0].multimodal_output
+ )
+
+ assert first.shape == second.shape, f"Length differs across runs: {first.shape} vs {second.shape}"
+ assert torch.allclose(first, second, atol=1e-4), "Audio differs across runs; flow-matching noise is not seeded"
diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py
deleted file mode 100644
index e22aa1174..000000000
--- a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py
+++ /dev/null
@@ -1,290 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""E2E test for DreamID-Omni offline inference with Cache-Dit.
-
-Uses the official DreamID-Omni ``test_case/twoip`` example assets instead of
-synthetic inputs so the coverage matches the upstream demo structure.
-"""
-
-from __future__ import annotations
-
-import importlib.util
-import json
-import os
-import re
-import time
-from functools import lru_cache
-from pathlib import Path
-from typing import NamedTuple
-
-import numpy as np
-import pytest
-import soundfile as sf
-from PIL import Image
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow]
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni"
-MODEL_PATH_ENV = "DREAMID_OMNI_MODEL_PATH"
-EXAMPLE_ROOT_ENV = "DREAMID_OMNI_EXAMPLE_ROOT"
-NUM_INFERENCE_STEPS = 45
-
-CACHE_CONFIG = {
- "Fn_compute_blocks": 1,
- "Bn_compute_blocks": 0,
- "max_warmup_steps": 4,
- "max_cached_steps": 20,
- "residual_diff_threshold": 0.24,
- "max_continuous_cached_steps": 3,
- "enable_taylorseer": False,
- "taylorseer_order": 1,
- "scm_steps_mask_policy": None,
- "scm_steps_policy": "dynamic",
-}
-
-
-class DreamIDGenerationResult(NamedTuple):
- video: np.ndarray
- audio: np.ndarray
- fps: int
- duration_s: float
- elapsed_s: float
-
-
-class DreamIDExampleCase(NamedTuple):
- prompt: str
- image_paths: tuple[Path, ...]
- audio_paths: tuple[Path, ...]
-
-
-def _require_model_path() -> str:
- model_path = os.environ.get(MODEL_PATH_ENV, str(DEFAULT_MODEL_PATH))
- if not Path(model_path).exists():
- pytest.skip(f"DreamID-Omni model bundle not found at {model_path}. Set {MODEL_PATH_ENV}.")
-
- if importlib.util.find_spec("dreamid_omni") is None:
- pytest.skip(
- "dreamid_omni dependency package is not importable. "
- "Run examples/offline_inference/x_to_video_audio/download_dreamid_omni.py first."
- )
-
- return model_path
-
-
-def _dreamid_repo_root_candidates() -> list[Path]:
- candidates: list[Path] = []
- example_root = os.environ.get(EXAMPLE_ROOT_ENV)
- if example_root:
- candidates.append(Path(example_root).expanduser().resolve())
-
- spec = importlib.util.find_spec("dreamid_omni")
- if spec is not None and spec.origin:
- candidates.append(Path(spec.origin).resolve().parents[1] / "test_case" / "twoip")
-
- return candidates
-
-
-@lru_cache(maxsize=1)
-def _load_example_case() -> DreamIDExampleCase:
- for candidate_root in _dreamid_repo_root_candidates():
- caption_path = candidate_root / "captions" / "20.json"
- image_paths = (
- candidate_root / "imgs" / "20" / "0.png",
- candidate_root / "imgs" / "20" / "1.png",
- )
- audio_paths = (
- candidate_root / "audios" / "20" / "0.wav",
- candidate_root / "audios" / "20" / "1.wav",
- )
- if caption_path.exists() and all(path.exists() for path in image_paths + audio_paths):
- prompt = _load_prompt_from_json(caption_path)
- return DreamIDExampleCase(
- prompt=prompt,
- image_paths=image_paths,
- audio_paths=audio_paths,
- )
-
- searched = ", ".join(str(path) for path in _dreamid_repo_root_candidates()) or ""
- pytest.skip(
- "DreamID-Omni official example assets not found. "
- f"Set {EXAMPLE_ROOT_ENV} to the local 'test_case/twoip' directory. "
- f"Searched: {searched}"
- )
-
-
-def _load_prompt_from_json(path: Path) -> str:
- raw_text = path.read_text(encoding="utf-8").strip()
- parsed = json.loads(raw_text)
- if isinstance(parsed, str):
- prompt = parsed
- elif isinstance(parsed, dict) and isinstance(parsed.get("content"), str):
- prompt = parsed["content"]
- else:
- raise AssertionError(f"Unexpected DreamID caption format in {path}")
-
- prompt = re.sub(
- r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]",
- "",
- prompt,
- flags=re.DOTALL,
- ).strip()
- prompt = re.sub(
- r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]",
- "",
- prompt,
- flags=re.DOTALL,
- ).strip()
- prompt = re.sub(r"\[[A-Z_]+\]", "", prompt)
- prompt = re.sub(r"\n\s*\n", "\n", prompt).strip()
- return prompt
-
-
-def _load_example_images(example_case: DreamIDExampleCase) -> list[Image.Image]:
- images: list[Image.Image] = []
- for path in example_case.image_paths:
- with Image.open(path) as image:
- images.append(image.convert("RGB"))
- return images
-
-
-def _load_example_audios(example_case: DreamIDExampleCase) -> list[np.ndarray]:
- audios: list[np.ndarray] = []
- for path in example_case.audio_paths:
- audio_array, sample_rate = sf.read(str(path), dtype="float32")
- audio_np = np.asarray(audio_array, dtype=np.float32)
- if audio_np.ndim == 2:
- audio_np = np.mean(audio_np, axis=1)
- audio_np = audio_np[int(sample_rate * 1) : int(sample_rate * 3)]
- audios.append(audio_np.reshape(-1))
- return audios
-
-
-def _sampling_params(num_inference_steps: int) -> OmniDiffusionSamplingParams:
- return OmniDiffusionSamplingParams(
- height=704,
- width=1280,
- num_inference_steps=num_inference_steps,
- seed=42,
- extra_args={
- "solver_name": "unipc",
- "shift": 5.0,
- },
- )
-
-
-def _request_payload(example_case: DreamIDExampleCase) -> dict[str, object]:
- return {
- "prompt": example_case.prompt,
- "video_negative_prompt": "jitter, bad hands, blur, distortion",
- "audio_negative_prompt": "robotic, muffled, echo, distorted",
- "multi_modal_data": {
- "image": _load_example_images(example_case),
- "audio": _load_example_audios(example_case),
- },
- }
-
-
-def _extract_generation(outputs: list[object], elapsed_s: float) -> DreamIDGenerationResult:
- """Extract video and audio from raw Omni outputs"""
- assert outputs, "DreamID-Omni returned no outputs"
-
- result = outputs[0]
- assert hasattr(result, "images") and result.images, "not return video frames"
-
- video = result.images[0]
- assert isinstance(video, np.ndarray), f"Expected video ndarray, got {type(video)}"
- assert video.ndim == 4, f"Expected 4D DreamID video tensor, got {video.shape}"
-
- multimodal_output = getattr(result, "multimodal_output", None) or {}
- assert "audio" in multimodal_output, "not return audio"
- audio = np.asarray(multimodal_output["audio"], dtype=np.float32).reshape(-1)
- assert audio.size > 0, "returned empty audio"
-
- audio_sample_rate = int(multimodal_output.get("audio_sample_rate", 16000))
- fps = multimodal_output.get("fps", 24)
- duration_s = float(audio.size) / float(audio_sample_rate)
-
- generation = DreamIDGenerationResult(
- video=video.astype(np.float32, copy=False),
- audio=audio,
- fps=fps,
- duration_s=duration_s,
- elapsed_s=elapsed_s,
- )
- return generation
-
-
-def _run_generation(
- *,
- model_path: str,
- example_case: DreamIDExampleCase,
- use_cache_dit: bool,
-) -> DreamIDGenerationResult:
- runner_kwargs = {
- "model_type": "dreamid-omni",
- "parallel_config": DiffusionParallelConfig(cfg_parallel_size=1),
- }
- if use_cache_dit:
- runner_kwargs["cache_backend"] = "cache_dit"
- runner_kwargs["cache_config"] = CACHE_CONFIG
-
- with OmniRunner(model_path, **runner_kwargs) as runner:
- warmup_payload = _request_payload(example_case)
- _ = runner.omni.generate(warmup_payload, _sampling_params(2))
-
- elapsed_samples: list[float] = []
- extracted_result: DreamIDGenerationResult | None = None
-
- request_payload = _request_payload(example_case)
- start = time.perf_counter()
- outputs = runner.omni.generate(request_payload, _sampling_params(NUM_INFERENCE_STEPS))
- elapsed_s = time.perf_counter() - start
- elapsed_samples.append(elapsed_s)
- if extracted_result is None:
- extracted_result = _extract_generation(outputs, elapsed_s)
-
- assert extracted_result is not None
- median_elapsed_s = float(np.median(np.asarray(elapsed_samples, dtype=np.float64)))
- return extracted_result._replace(elapsed_s=median_elapsed_s)
-
-
-def _video_mae(a: np.ndarray, b: np.ndarray) -> float:
- assert a.shape == b.shape, f"Video shapes differ: {a.shape} vs {b.shape}"
- return float(np.mean(np.abs(a - b)))
-
-
-def _audio_rms(audio: np.ndarray) -> float:
- audio = np.asarray(audio, dtype=np.float32).reshape(-1)
- return float(np.sqrt(np.mean(audio**2)))
-
-
-def _audio_rms_ratio(a: DreamIDGenerationResult, b: DreamIDGenerationResult) -> float:
- rms_a = _audio_rms(a.audio)
- rms_b = _audio_rms(b.audio)
- return rms_b / max(rms_a, 1e-6)
-
-
-@hardware_test(res={"cuda": "H100"})
-def test_dreamid_omni_cache_dit_e2e() -> None:
- model_path = _require_model_path()
- example_case = _load_example_case()
-
- baseline = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=False)
- cached = _run_generation(model_path=model_path, example_case=example_case, use_cache_dit=True)
-
- video_mae = _video_mae(baseline.video, cached.video)
- audio_rms_ratio = _audio_rms_ratio(baseline, cached)
- speedup = baseline.elapsed_s / cached.elapsed_s if cached.elapsed_s > 0 else float("inf")
-
- assert video_mae <= 0.15, f"video drift too much: {video_mae:.6f}"
- assert 0.4 <= audio_rms_ratio <= 2.5, f"audio quality regressed too much: rms_ratio={audio_rms_ratio:.6f}"
- assert speedup >= 1.05, f"cache-dit improve not enough: {speedup:.3f}x"
diff --git a/tests/e2e/offline_inference/test_ernie_image_expansion.py b/tests/e2e/offline_inference/test_ernie_image_expansion.py
index 19793ff00..c986a6366 100644
--- a/tests/e2e/offline_inference/test_ernie_image_expansion.py
+++ b/tests/e2e/offline_inference/test_ernie_image_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end test for ErnieImage text-to-image generation.
@@ -11,13 +11,13 @@ Equivalent to running:
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "baidu/ERNIE-Image"
-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL,
None,
@@ -37,9 +37,9 @@ pytestmark = [
@hardware_test(res={"cuda": "L4"}, num_cards=2)
-def test_ernie_image_text_to_image(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_ernie_image_text_to_image(offline_client: OfflineOmniClient) -> None:
request_config = {
- "model": omni_runner_handler.runner.model_name,
+ "model": offline_client.runner.model_name,
"prompt": "A photo of a cat sitting on a laptop",
"sampling_params": OmniDiffusionSamplingParams(
height=512,
@@ -49,4 +49,4 @@ def test_ernie_image_text_to_image(omni_runner_handler: OmniRunnerHandler) -> No
seed=42,
),
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py
index ec48a49f7..92a710f6e 100644
--- a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py
+++ b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for FLUX AutoRound W4A16 quantized inference.
These tests require:
@@ -19,8 +19,8 @@ import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
diff --git a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py
index abf7bbc2c..61a47a3b0 100644
--- a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py
+++ b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Tests for FluxKontext model pipeline.
@@ -15,7 +15,7 @@ import pytest
from vllm.assets.image import ImageAsset
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -47,18 +47,18 @@ def _sampling_512() -> OmniDiffusionSamplingParams:
@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_flux_kontext_text_to_image(omni_runner_handler: OmniRunnerHandler):
+def test_flux_kontext_text_to_image(offline_client: OfflineOmniClient):
"""Test FluxKontext text-to-image generation with real model."""
request_config = {
"model": MODEL,
"prompt": "A photo of a cat sitting on a laptop",
"sampling_params": _sampling_512(),
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_flux_kontext_image_edit(omni_runner_handler: OmniRunnerHandler):
+def test_flux_kontext_image_edit(offline_client: OfflineOmniClient):
"""Test FluxKontext image-to-image editing with real model."""
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
request_config = {
@@ -67,4 +67,4 @@ def test_flux_kontext_image_edit(omni_runner_handler: OmniRunnerHandler):
"modalities": ["img2img"],
"sampling_params": _sampling_512(),
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py
index de6a758f6..9b61f876f 100644
--- a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py
+++ b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for Flux1 Schnell."""
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "black-forest-labs/FLUX.1-schnell"
@@ -21,7 +21,7 @@ pytestmark = [
@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler):
+def test_flux_schnell_text_to_image(offline_client: OfflineOmniClient):
"""Test FLUX.1-schnell text-to-image generation."""
request_config = {
"model": MODEL,
@@ -33,4 +33,4 @@ def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler):
seed=42,
),
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py
index 9656772c9..ed0cbaf91 100644
--- a/tests/e2e/offline_inference/test_flux2_klein.py
+++ b/tests/e2e/offline_inference/test_flux2_klein.py
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end test for Flux2 Klein inpainting.
-Inpainting uses ``omni_runner_handler.send_diffusion_request`` with
+Inpainting uses ``offline_client.send_diffusion_request`` with
``multi_modal_data`` containing ``image`` and ``mask_image``; see
-:meth:`OmniRunnerHandler.send_diffusion_request` in ``tests.helpers.runtime``.
+:meth:`OfflineOmniClient.send_diffusion_request` in ``tests.helpers.runtime``.
"""
import pytest
@@ -14,7 +14,7 @@ import torch
from PIL import Image, ImageDraw
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
+from tests.helpers.runtime import DiffusionResponse, OfflineOmniClient
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -47,15 +47,17 @@ def _create_test_inputs(color: tuple = (100, 150, 200)):
def _images_from_response(response: DiffusionResponse) -> list[Image.Image]:
- if isinstance(response.images[0], list):
- return [f for fr in response.images for f in fr]
- return list(response.images)
+ images = response.images
+ assert images is not None
+ if images and isinstance(images[0], list):
+ return [frame for frames in images for frame in frames]
+ return list(images)
def _send_inpaint_with_generator(
- omni_runner_handler: OmniRunnerHandler, prompt: str, input_image, mask_image, generator: torch.Generator
+ offline_client: OfflineOmniClient, prompt: str, input_image, mask_image, generator: torch.Generator
) -> DiffusionResponse:
- return omni_runner_handler.send_diffusion_request(
+ return offline_client.send_diffusion_request(
{
"model": MODEL,
"prompt": prompt,
@@ -72,18 +74,19 @@ def _send_inpaint_with_generator(
)
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
+@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandler):
+def test_flux2_klein_inpaint_deterministic(offline_client: OfflineOmniClient):
input_image, mask_image = _create_test_inputs()
seed = 12345
gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed)
- r1 = _send_inpaint_with_generator(omni_runner_handler, "A red flower in a field", input_image, mask_image, gen1)
- r2 = _send_inpaint_with_generator(omni_runner_handler, "A red flower in a field", input_image, mask_image, gen2)
+ r1 = _send_inpaint_with_generator(offline_client, "A red flower in a field", input_image, mask_image, gen1)
+ r2 = _send_inpaint_with_generator(offline_client, "A red flower in a field", input_image, mask_image, gen2)
images1 = _images_from_response(r1)
images2 = _images_from_response(r2)
@@ -93,17 +96,18 @@ def test_flux2_klein_inpaint_deterministic(omni_runner_handler: OmniRunnerHandle
)
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
+@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_flux2_klein_inpaint_different_seeds_different_output(omni_runner_handler: OmniRunnerHandler):
+def test_flux2_klein_inpaint_different_seeds_different_output(offline_client: OfflineOmniClient):
input_image, mask_image = _create_test_inputs()
gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42)
gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999)
- r1 = _send_inpaint_with_generator(omni_runner_handler, "A beautiful landscape", input_image, mask_image, gen1)
- r2 = _send_inpaint_with_generator(omni_runner_handler, "A beautiful landscape", input_image, mask_image, gen2)
+ r1 = _send_inpaint_with_generator(offline_client, "A beautiful landscape", input_image, mask_image, gen1)
+ r2 = _send_inpaint_with_generator(offline_client, "A beautiful landscape", input_image, mask_image, gen2)
images1 = _images_from_response(r1)
images2 = _images_from_response(r2)
diff --git a/tests/e2e/offline_inference/test_gepard_tts.py b/tests/e2e/offline_inference/test_gepard_tts.py
new file mode 100644
index 000000000..b201842b4
--- /dev/null
+++ b/tests/e2e/offline_inference/test_gepard_tts.py
@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E offline inference tests for the Gepard-1.0 single-stage native-AR TTS.
+
+Zero-shot only (the model's default learned voice). These tests need a GPU and
+the NeMo NanoCodec, so the whole file runs in the weekly tier -- see the comment
+above ``pytestmark``. The window arithmetic behind the streaming decode is
+covered on CPU in ``tests/model_executor/models/test_gepard_window.py``.
+"""
+
+from __future__ import annotations
+
+import functools
+import tempfile
+from pathlib import Path
+from typing import NamedTuple
+
+import pytest
+import soundfile as sf
+import torch
+from transformers import AutoTokenizer
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import convert_audio_file_to_text
+from tests.helpers.runtime import OmniRunner # noqa: F401 (indirect fixture)
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+
+# The codec decoder is an optional dependency; skip rather than fail the engine
+# inside the worker process when it is absent. The CI image does not carry it,
+# so the weekly step installs it first -- see .buildkite/cuda/test-weekly.yml.
+# Install recipe: examples/offline_inference/text_to_speech/README.md.
+pytest.importorskip("nemo.collections.tts.models")
+
+MODEL_NAME = "nineninesix/gepard-1.0"
+STAGE_CONFIG = get_deploy_config_path("gepard.yaml")
+SAMPLE_RATE = 22050
+# NanoCodec upsamples one FSQ frame to exactly this many samples.
+SAMPLES_PER_FRAME = 1024
+
+
+@functools.cache
+def _config() -> GepardConfig:
+ """The checkpoint's own layout, not the class defaults.
+
+ The worker builds its config from the ``gepard_config.json`` sidecar, so a
+ prompt assembled from the defaults would silently use a different layout
+ the day the sidecar moves — and a mismatched layout degrades the speech
+ rather than failing, which is the hardest way for this to be found.
+ """
+ return GepardConfig.from_checkpoint(MODEL_NAME)
+
+
+# Prompt -> a word that appears in no other prompt, so a clip can be traced
+# back to the request that asked for it.
+#
+# These are not interchangeable. Batching perturbs every request's sampling
+# trajectory, and on some trajectories this checkpoint's stop head never fires;
+# this set was screened over 10 batched runs with no runaway. Re-screen the
+# whole set when changing one, not just the row changed.
+_DISTINGUISHABLE_PROMPTS = {
+ "He drinks coffee every morning.": "coffee",
+ "Machine learning is interesting.": "learning",
+ "Please close the window before leaving.": "window",
+ "My favorite color is purple.": "purple",
+}
+
+# Above text_repetition.apply_below, so the layout is voiced once instead of
+# repeated. Every other prompt here takes the repeated branch.
+_LONG_TEXT = (
+ "The morning train was late again, so I walked the last two miles along the "
+ "river and reached the office just before nine."
+)
+
+_OMNI_RUNNER_PARAM = (
+ MODEL_NAME,
+ STAGE_CONFIG,
+ {"trust_remote_code": True},
+)
+
+# Only a common word is checked, which needs less ASR than the shared helper's
+# ``small`` default; a miss escalates to it.
+_ASR_MODEL = "base"
+_ASR_ESCALATION_MODEL = "small"
+
+# Transcription dominates this file's runtime -- Whisper runs on CPU on a
+# single-GPU runner -- so only two of the four concurrent clips are read as
+# speech. The pairwise waveform comparison covers all four.
+_TRANSCRIBED_CLIPS = 2
+
+# No level marker at module scope: each test carries the one matching how often
+# its claim needs re-checking, though the weekly step collects the file whole.
+# There is deliberately no core_model row -- under dummy weights the stop head is
+# random, so every request would run to gepard.yaml's max_tokens.
+pytestmark = [
+ pytest.mark.tts,
+ pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+
+def _build_request(text: str) -> dict:
+ """Assemble the [speaker slots, SOT, text, EOT, SOS] prompt."""
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+ prompt_token_ids = build_gepard_prompt_ids(
+ tokenizer(text, add_special_tokens=False)["input_ids"],
+ config=_config(),
+ )
+ return {
+ "prompt_token_ids": prompt_token_ids,
+ "additional_information": {"text": [text]},
+ }
+
+
+def _extract_audio(mm) -> torch.Tensor | None:
+ # The consolidation path renames "model_outputs" to "audio".
+ audio = mm.get("audio")
+ if audio is None:
+ audio = mm.get("model_outputs")
+ if isinstance(audio, list):
+ audio = audio[0] if len(audio) else None
+ return audio
+
+
+class _Clip(NamedTuple):
+ wav: torch.Tensor
+ # Counts the frames the codec was handed using only what the engine
+ # returned, never the model's internals.
+ committed_frames: int
+
+
+def _synthesize_all(omni: Omni, texts: list[str]) -> list[_Clip]:
+ """Submit every text in one call; return one clip per request, in
+ submission order.
+
+ ``Omni.generate`` yields in completion order and never sorts, so the pairing
+ is restored from the request id prefix it assigns (``f"{i}_{uuid4()}"``).
+ Without that, the only testable claim is that the batch as a whole looks
+ plausible, which a stable swap between two requests satisfies.
+
+ Nothing here authors a SamplingParams: a fresh one would replace the stage
+ defaults rather than merge over them, dropping the pipeline's stop token.
+ """
+ outputs = omni.generate([_build_request(t) for t in texts])
+ assert len(outputs) == len(texts), f"expected {len(texts)} outputs, got {len(outputs)}"
+
+ by_index: dict[int, _Clip] = {}
+ for stage_outputs in outputs:
+ index = int(str(stage_outputs.request_id).split("_", 1)[0])
+ waveform = None
+ frames = 0
+ for out in stage_outputs.outputs:
+ # Explicit None check: `x or []` would ask an array for a truth value.
+ token_ids = out.token_ids if out.token_ids is not None else []
+ # STOP sits just past head0's range, so a token equal to it is
+ # never a committed frame.
+ frames += sum(1 for t in token_ids if int(t) != _config().stop_token)
+ mm = out.multimodal_output
+ if mm is None:
+ continue
+ audio = _extract_audio(mm)
+ if audio is None:
+ continue
+ sr_t = mm.get("sr")
+ if hasattr(sr_t, "item"):
+ assert int(sr_t.item()) == SAMPLE_RATE
+ waveform = audio.cpu().float()
+ assert waveform is not None, "no audio output produced"
+ by_index[index] = _Clip(wav=waveform, committed_frames=frames)
+
+ assert sorted(by_index) == list(range(len(texts))), (
+ f"request ids did not cover every submitted prompt exactly once: got {sorted(by_index)}"
+ )
+ return [by_index[i] for i in range(len(texts))]
+
+
+def _synthesize(omni: Omni, text: str) -> _Clip:
+ return _synthesize_all(omni, [text])[0]
+
+
+def _assert_clip_is_sane(clip: _Clip) -> None:
+ """The checks every clip has to pass, whatever produced it."""
+ wav = clip.wav
+ assert wav.dim() == 1, f"expected mono 1-D, got shape {tuple(wav.shape)}"
+ assert wav.numel() > 0, "empty waveform"
+ assert torch.isfinite(wav).all(), "non-finite samples"
+
+ # Frame conservation against the engine's own token count, not just
+ # divisibility. This is what catches a payload carrying non-audio data or a
+ # dropped tail; both pass every other check here.
+ assert wav.numel() == clip.committed_frames * SAMPLES_PER_FRAME, (
+ f"{wav.numel()} samples for {clip.committed_frames} committed frames — "
+ f"expected exactly {clip.committed_frames * SAMPLES_PER_FRAME}"
+ )
+
+ # Anything far outside the waveform range is not codec output.
+ assert float(wav.abs().max()) <= 1.5, f"peak {float(wav.abs().max()):.2f} is out of range"
+
+
+def _assert_stopped_on_its_own(clip: _Clip) -> None:
+ """The request must stop on its own rather than run to ``max_tokens``.
+
+ Separate from the structural checks because it reads the stop head's
+ judgement, not the payload's shape: under ``load_format: dummy`` that head
+ is random and never fires, so every clip runs to the token cap.
+ """
+ seconds = clip.wav.numel() / SAMPLE_RATE
+ assert 0.5 < seconds < 30.0, f"implausible duration {seconds:.2f}s — did STOP fire?"
+
+
+def _transcribe(wav: torch.Tensor, tmp_dir: str, name: str, model_size: str = _ASR_MODEL) -> str:
+ path = str(Path(tmp_dir) / f"{name}.wav")
+ sf.write(path, wav.numpy(), SAMPLE_RATE)
+ return convert_audio_file_to_text(path, model_size=model_size, language="en").lower()
+
+
+def _transcribe_for(wav: torch.Tensor, tmp_dir: str, name: str, accepts) -> str:
+ """Transcribe with the fast model, retrying once with the stronger one.
+
+ ``accepts`` decides whether the fast transcript is usable. Escalating only
+ on a miss keeps the average clip at ``_ASR_MODEL`` while a genuine model
+ artifact still fails the caller's assertion, since the stronger ASR
+ mistranscribes it too.
+ """
+ transcript = _transcribe(wav, tmp_dir, name)
+ if accepts(transcript):
+ return transcript
+ return _transcribe(wav, tmp_dir, name, model_size=_ASR_ESCALATION_MODEL)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize(
+ ("text", "keyword"),
+ [
+ # The canonical path: one request, one transcript, the cheapest check
+ # that reads the audio as speech.
+ pytest.param(
+ "Hello, this is Gepard speaking.",
+ "hello",
+ marks=pytest.mark.advanced_model,
+ id="canonical",
+ ),
+ # No keyword: this row carries the structural and stopping checks over a
+ # second text and leaves the transcript to the row above.
+ pytest.param(
+ "The quick brown fox jumps over the lazy dog.",
+ None,
+ marks=pytest.mark.full_model,
+ id="second_text",
+ ),
+ ],
+)
+def test_gepard_offline_zero_shot(omni_runner, run_level: str, text: str, keyword: str | None) -> None:
+ """Zero-shot synthesis produces a finite mono waveform at 22.05 kHz.
+
+ The transcript check is the one assertion here that reads the audio as
+ speech rather than as a tensor. It stands in for a parity comparison, which
+ needs a reference implementation this PR cannot run — noise, a repeated
+ loop or the wrong clip all pass every numeric check above it.
+ """
+ clip = _synthesize(omni_runner, text)
+ _assert_clip_is_sane(clip)
+
+ if run_level not in {"advanced_model", "full_model"}:
+ return # dummy weights: the audio is structural only, so do not read it
+
+ _assert_stopped_on_its_own(clip)
+ if keyword is None:
+ return
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ transcript = _transcribe_for(clip.wav, tmp_dir, "zero_shot", lambda t: keyword in t)
+ assert keyword in transcript, f"expected {keyword!r} in transcript, got {transcript!r}"
+
+
+@pytest.mark.full_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_gepard_offline_long_text_skips_repetition(omni_runner, run_level: str) -> None:
+ """The other branch of the prompt layout, end to end.
+
+ A text at or above ``apply_below`` tokens is voiced once; shorter ones are
+ repeated to carry enough mass against the speaker prefix. The layout
+ assertion keeps this row honest — if a tokenizer change pushed this text
+ back under the threshold it would silently retest the repeated path.
+ """
+ prompt_ids = _build_request(_LONG_TEXT)["prompt_token_ids"]
+ assert prompt_ids.count(_config().start_of_text) == 1, "_LONG_TEXT is no longer above apply_below"
+
+ clip = _synthesize(omni_runner, _LONG_TEXT)
+ _assert_clip_is_sane(clip)
+
+ if run_level not in {"advanced_model", "full_model"}:
+ return
+
+ _assert_stopped_on_its_own(clip)
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ transcript = _transcribe_for(clip.wav, tmp_dir, "long_text", lambda t: "river" in t)
+ assert "river" in transcript, f"expected 'river' in transcript, got {transcript!r}"
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_gepard_offline_concurrent_requests_stay_isolated(omni_runner, run_level: str) -> None:
+ """Four requests in one batch, each holding its own generation state.
+
+ The prompts are deliberately unlike each other: cross-talk between the
+ per-request frame history, the codec window or the output routing shows up
+ as a transcript matching the wrong prompt, which similar prompts would
+ hide. ``max_num_seqs`` is 4, so these really do run together.
+ """
+ texts = list(_DISTINGUISHABLE_PROMPTS)
+ clips = _synthesize_all(omni_runner, texts)
+
+ for clip in clips:
+ _assert_clip_is_sane(clip)
+
+ if run_level not in {"advanced_model", "full_model"}:
+ return # dummy weights: the audio is structural only, so do not read it
+
+ for clip in clips:
+ _assert_stopped_on_its_own(clip)
+
+ # All four clips, not just the transcribed ones: delivering one request's
+ # audio for another leaves two byte-identical waveforms.
+ for i in range(len(clips)):
+ for j in range(i + 1, len(clips)):
+ assert not torch.equal(clips[i].wav, clips[j].wav), (
+ f"clips {i} and {j} are the same waveform — one request's audio was delivered for another"
+ )
+
+ # Each clip must say what its own prompt asked for: accepting any of the
+ # four would pass a stable permutation, which is what a mix-up produces.
+ keywords = set(_DISTINGUISHABLE_PROMPTS.values())
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ for i, clip in enumerate(clips[:_TRANSCRIBED_CLIPS]):
+ expected = _DISTINGUISHABLE_PROMPTS[texts[i]]
+ # Escalate on the expected word being absent, not on "no keyword
+ # at all" -- a neighbour's word is the failure under test.
+ transcript = _transcribe_for(clip.wav, tmp_dir, f"concurrent_{i}", lambda t: expected in t)
+ present = sorted(k for k in keywords if k in transcript)
+ assert present == [expected], (
+ f"clip for {texts[i]!r} transcribed as {transcript!r}: expected {expected!r} and "
+ f"no other prompt's keyword, found {present} — requests bled into each other"
+ )
diff --git a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
index 28eaab9b3..b38e26993 100644
--- a/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
+++ b/tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for GLM-Image AutoRound W4A16 quantized inference.
These tests cover text-to-image and image-to-image generation with
@@ -19,10 +19,10 @@ import pytest
from PIL import Image
from vllm import SamplingParams
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.monitor import DeviceMemoryMonitor
+from tests.helpers.runtime import OfflineOmniClient
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -109,7 +109,7 @@ def _diffusion_sampling_params(
pytestmark = [
- pytest.mark.full_model,
+ pytest.mark.slow,
pytest.mark.diffusion,
]
@@ -121,7 +121,7 @@ pytestmark = [
@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler):
+def test_glm_image_autoround_w4a16_generates_image(offline_client: OfflineOmniClient):
"""Load the W4A16 quantized GLM-Image model and verify it produces a valid image."""
gc.collect()
current_omni_platform.empty_cache()
@@ -152,7 +152,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn
seed=42,
)
- outputs = omni_runner_handler.runner.generate(
+ outputs = offline_client.runner.generate(
[prompt_dict],
[ar_params, diffusion_params],
)
@@ -186,7 +186,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn
@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler):
+def test_glm_image_autoround_w4a16_image_to_image(offline_client: OfflineOmniClient):
"""Load the W4A16 quantized GLM-Image and verify image-to-image generation works."""
ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"]
@@ -217,7 +217,7 @@ def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunne
seed=42,
)
- outputs = omni_runner_handler.runner.generate(
+ outputs = offline_client.runner.generate(
[prompt_dict],
[ar_params, diffusion_params],
)
diff --git a/tests/e2e/offline_inference/test_glm_tts_expansion.py b/tests/e2e/offline_inference/test_glm_tts_expansion.py
index e54dd5480..2bde38e6b 100644
--- a/tests/e2e/offline_inference/test_glm_tts_expansion.py
+++ b/tests/e2e/offline_inference/test_glm_tts_expansion.py
@@ -96,7 +96,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None:
with OmniRunner(
MODEL,
- stage_configs_path=_get_deploy_config(async_chunk=async_chunk),
+ deploy_config=_get_deploy_config(async_chunk=async_chunk),
stage_init_timeout=600,
trust_remote_code=True,
) as omni_runner:
diff --git a/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py
index 00482a132..53134446d 100644
--- a/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py
+++ b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Offline expansion tests for higgs-audio v2.
Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py:
build Stage-0 prompt_token_ids with the model's own processor + plain-text
chat template, drive the engine via ``Omni.generate``, and inspect the
-Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request
+Stage-1 PCM payload directly. The OfflineOmniClient.send_audio_speech_request
helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this
test bypasses it (same approach as test_voxtral_tts.py).
diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py
index 4bdb42c34..7e7c63900 100644
--- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py
+++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py
@@ -1,21 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Smoke test for HunyuanImage-3.0 Image-to-Text (I2T) pipeline."""
from collections.abc import Generator
-from pathlib import Path
import pytest
-import torch
from PIL import Image
+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 import Omni
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-REPO_ROOT = Path(__file__).resolve().parents[3]
-STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_i2t.yaml"
+DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3_ar.yaml")
# First 20 generated token IDs from the HF greedy reference on this input.
# vllm-omni AR output matches this prefix bitwise; the two implementations
@@ -52,13 +51,13 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
def omni() -> Generator[Omni, None, None]:
with OmniRunner(
MODEL_NAME,
- stage_configs_path=str(STAGE_CONFIG_PATH),
+ deploy_config=DEPLOY_CONFIG_PATH,
trust_remote_code=True,
) as runner:
yield runner.omni
-@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Need at least 4 CUDA GPUs.")
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
def test_i2t_generates_text(omni: Omni) -> None:
"""Verify I2T output's first 20 token IDs match the HF greedy baseline."""
# Solid-color image keeps the input self-contained and reproducible.
diff --git a/tests/e2e/offline_inference/test_indextts2_expansion.py b/tests/e2e/offline_inference/test_indextts2_expansion.py
index ba5b984f9..c95de98ea 100644
--- a/tests/e2e/offline_inference/test_indextts2_expansion.py
+++ b/tests/e2e/offline_inference/test_indextts2_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E offline inference tests for IndexTTS2 two-stage pipeline.
Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV.
@@ -8,12 +8,11 @@ Covers: basic TTS.
from __future__ import annotations
-from pathlib import Path
-
import pytest
import torch
from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni import Omni
@@ -29,8 +28,7 @@ _OMNI_RUNNER_PARAM = (
STAGE_CONFIG,
)
-ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2"
-REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav")
+REF_AUDIO = str(get_asset_path("indextts2/ref_audio.wav"))
pytestmark = [
pytest.mark.slow,
@@ -55,9 +53,12 @@ def _sample_rate_from_mm(mm: dict) -> int:
sr = mm.get("sr")
if isinstance(sr, list):
sr = sr[-1] if sr else None
- if hasattr(sr, "item"):
- return int(sr.item())
- return int(sr) if sr is not None else SAMPLE_RATE
+ if sr is None:
+ return SAMPLE_RATE
+ item = getattr(sr, "item", None)
+ if callable(item):
+ return int(item())
+ return int(sr)
def _build_request(text: str, ref_audio: str, **extra) -> dict:
diff --git a/tests/e2e/offline_inference/test_krea2_expansion.py b/tests/e2e/offline_inference/test_krea2_expansion.py
index a8f63032b..372a36afa 100644
--- a/tests/e2e/offline_inference/test_krea2_expansion.py
+++ b/tests/e2e/offline_inference/test_krea2_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``).
@@ -15,7 +15,7 @@ import numpy as np
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
@@ -49,9 +49,9 @@ def _sampling() -> OmniDiffusionSamplingParams:
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_text_to_image_001(offline_client: OfflineOmniClient) -> None:
"""Basic functional smoke: a single prompt produces a decoded image."""
- omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+ offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
@hardware_test(res={"cuda": "H100"})
@@ -60,22 +60,22 @@ def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None
[(MODEL, None, {"enable_layerwise_offload": True})],
indirect=True,
)
-def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_layerwise_offload(offline_client: OfflineOmniClient) -> None:
"""Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs)."""
- omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+ offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
@hardware_test(res={"cuda": "H100"})
@pytest.mark.cache
@pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True)
-def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_cache_dit(offline_client: OfflineOmniClient) -> None:
"""Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``).
Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware
(False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint
exercises the no-CFG branch.
"""
- omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+ offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@@ -88,13 +88,13 @@ def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None:
[(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})],
indirect=True,
)
-def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_hsdp(offline_client: OfflineOmniClient) -> None:
"""Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs.
- Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda)
- shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices.
+ Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU
+ shard by ``num_cards=2`` (``not cards_1``); skipped automatically on boxes with fewer than 2 CUDA devices.
"""
- omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+ offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@@ -107,27 +107,29 @@ def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None:
[(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})],
indirect=True,
)
-def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_vae_patch_parallel(offline_client: OfflineOmniClient) -> None:
"""Exercise VAE patch-parallel (decode) across 2 GPUs.
Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the
- multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with
+ multi-GPU shard by ``num_cards=2`` (``not cards_1``); skipped automatically on boxes with
fewer than 2 CUDA devices.
"""
- omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+ offline_client.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
-def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray:
+def _generate(handler: OfflineOmniClient, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray:
sp = _sampling()
sp.lora_request = lora_request
sp.lora_scale = lora_scale
resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp})
- return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16)
+ images = resp.images
+ assert images is not None
+ return np.asarray(images[0].convert("RGB"), dtype=np.int16)
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None:
+def test_krea2_lora(offline_client: OfflineOmniClient) -> None:
"""Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation.
Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``,
@@ -137,10 +139,10 @@ def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None:
"""
lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA)
- baseline = _generate(omni_runner_handler, None)
- img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0)
- img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0)
- restored = _generate(omni_runner_handler, None)
+ baseline = _generate(offline_client, None)
+ img_1x = _generate(offline_client, lora_request, lora_scale=1.0)
+ img_2x = _generate(offline_client, lora_request, lora_scale=2.0)
+ restored = _generate(offline_client, None)
diff_1x = np.abs(baseline - img_1x).mean()
diff_2x = np.abs(baseline - img_2x).mean()
diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py
new file mode 100644
index 000000000..39d9b9d95
--- /dev/null
+++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py
@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end smoke test for LingBot-World v2 generation."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunnerHandler
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = os.environ.get(
+ "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH",
+ "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
+)
+_IMAGE_PATH = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH")
+_ACTION_DIR_VALUE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR")
+_ACTION_DIR = Path(_ACTION_DIR_VALUE).expanduser().resolve() if _ACTION_DIR_VALUE else None
+_HAS_INPUT_ASSETS = _IMAGE_PATH is not None and _ACTION_DIR is not None
+
+_OMNI_RUNNER_PARAM = (
+ MODEL,
+ None,
+ {
+ "flow_shift": 5.0,
+ "parallel_config": DiffusionParallelConfig(tensor_parallel_size=1),
+ "model_config": {
+ "lingbot_action_root": str(_ACTION_DIR.parent) if _ACTION_DIR is not None else "",
+ },
+ },
+)
+
+pytestmark = [
+ pytest.mark.slow,
+ pytest.mark.diffusion,
+ pytest.mark.skipif(
+ not _HAS_INPUT_ASSETS,
+ reason=("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH and VLLM_OMNI_LINGBOT_WORLD_V2_ACTION_DIR are required"),
+ ),
+ pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) -> None:
+ """Load the real checkpoint and generate one causal block on one GPU."""
+
+ assert _IMAGE_PATH is not None
+ assert _ACTION_DIR is not None
+ response = omni_runner_handler.send_diffusion_request(
+ {
+ "model": omni_runner_handler.runner.model_name,
+ "prompt": "The camera moves slowly forward through the scene.",
+ "multi_modal_data": {"image": str(Path(_IMAGE_PATH).expanduser().resolve())},
+ "sampling_params": OmniDiffusionSamplingParams(
+ height=480,
+ width=832,
+ num_frames=9,
+ num_inference_steps=4,
+ max_sequence_length=512,
+ seed=42,
+ fps=16,
+ extra_args={
+ "action_path": _ACTION_DIR.name,
+ "flow_shift": 5.0,
+ },
+ ),
+ }
+ )
+
+ assert response.success
+ assert response.images is not None
+ assert len(response.images) == 1
+ video = np.asarray(response.images[0])
+ assert video.shape == (1, 9, 480, 832, 3)
+ assert np.isfinite(video).all()
+ assert float(np.ptp(video)) > 1e-6
+ assert float(np.mean(np.abs(np.diff(video, axis=1)))) > 1e-6
diff --git a/tests/e2e/offline_inference/test_longcat_video_avatar.py b/tests/e2e/offline_inference/test_longcat_video_avatar.py
index f0fa81eb3..404ad9485 100644
--- a/tests/e2e/offline_inference/test_longcat_video_avatar.py
+++ b/tests/e2e/offline_inference/test_longcat_video_avatar.py
@@ -84,7 +84,6 @@ def longcat_omni_runner() -> Generator[OmniRunner, None, None]:
with OmniRunner(
model,
seed=42,
- stage_configs_path=None,
model_class_name="LongCatVideoAvatarPipeline",
additional_config={
"model_type": "avatar-v1.5",
diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py
deleted file mode 100644
index e6b3c6a42..000000000
--- a/tests/e2e/offline_inference/test_magi_human_expansion.py
+++ /dev/null
@@ -1,162 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
-
-import io
-
-import av
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-MODEL = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
-
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
-_OMNI_RUNNER_PARAM = (
- MODEL,
- None,
- {
- "init_timeout": 1200,
- "tensor_parallel_size": 2,
- "vae_use_tiling": True,
- },
-)
-
-_OMNI_RUNNER_PARAM_OFFLOAD = (
- MODEL,
- None,
- {
- "init_timeout": 1200,
- "tensor_parallel_size": 2,
- "enable_cpu_offload": True,
- "vae_use_tiling": True,
- },
-)
-
-pytestmark = [
- pytest.mark.slow,
- pytest.mark.diffusion,
- pytest.mark.parametrize(
- "omni_runner",
- [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
- indirect=True,
- ),
-]
-
-
-def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
- """Validate that the MP4 contains meaningful video and audio tracks."""
- container = av.open(io.BytesIO(video_bytes))
-
- v_streams = [s for s in container.streams if s.type == "video"]
- assert len(v_streams) >= 1, "No video stream found in MP4"
-
- a_streams = [s for s in container.streams if s.type == "audio"]
- assert len(a_streams) >= 1, "No audio stream found in MP4"
-
- v_stream = v_streams[0]
- assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
- assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
-
- frame_count = 0
- for frame in container.decode(video=0):
- frame_count += 1
- if frame_count >= min_frames:
- break
- assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
-
- container.close()
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-def test_magi_human_e2e(omni_runner: OmniRunner):
- """End-to-end test for MagiHuman generating video and audio."""
- prompt = (
- "A young woman with long, wavy golden blonde hair and bright blue eyes, "
- "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
- "stationary in front of a softly lit, blurred warm-toned interior. Her "
- "overall disposition is warm, composed, and gently confident. The camera "
- "holds a static medium close-up, framing her from the shoulders up, "
- "with shallow depth of field keeping her face in sharp focus. Soft "
- "directional key light falls from the upper left, casting a gentle "
- "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
- "the levator labii superiors relaxing as her lips part. She speaks in "
- "clear, warm, unhurried American English: "
- "\"The most beautiful things in life aren't things at all — "
- "they're moments, feelings, and the people who make you feel truly alive.\" "
- "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
- "shapes each vowel with precision. A faint, genuine smile engages the "
- "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
- "in a soft, neutral arch throughout. She maintains steady, forward-facing "
- "eye contact. Head position remains level; no torso displacement occurs.\n\n"
- "Dialogue:\n"
- ": "
- "\"The most beautiful things in life aren't things at all — "
- "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
- "Background Sound:\n"
- ""
- )
-
- sampling_params = OmniDiffusionSamplingParams(
- height=256,
- width=448,
- num_inference_steps=8,
- seed=52,
- extra_args={
- "seconds": 5,
- "sr_height": 1080,
- "sr_width": 1920,
- "sr_num_inference_steps": 5,
- },
- )
-
- omni = omni_runner.omni
- outputs = list(
- omni.generate(
- prompts=[prompt],
- sampling_params_list=[sampling_params],
- )
- )
-
- assert len(outputs) > 0, "No outputs returned"
- first = outputs[0]
-
- assert hasattr(first, "images") and first.images, "No video frames in output"
- video_frames = first.images[0]
- assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
- assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
-
- mm = first.multimodal_output
- assert mm, "multimodal_output is empty or missing"
-
- audio_waveform = mm.get("audio")
- assert audio_waveform is not None, "No audio waveform in multimodal_output"
-
- audio_sample_rate = mm.get("audio_sample_rate")
- assert audio_sample_rate is not None, (
- "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
- )
- assert isinstance(audio_sample_rate, (int, float)), (
- f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
- )
- assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
-
- fps = mm.get("fps")
- assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
- assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
- assert int(fps) > 0, f"fps must be positive, got {fps}"
-
- video_bytes = mux_video_audio_bytes(
- video_frames,
- audio_waveform,
- fps=float(fps),
- audio_sample_rate=int(audio_sample_rate),
- )
- assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
- assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
-
- _validate_mp4(video_bytes)
diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
index ec1545f6a..f570a6ea7 100644
--- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
+++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
@@ -38,7 +38,7 @@ _AR_PATCH_SIZE = 16
MODEL_PATH = "bytedance-research/MammothModa2-Preview"
T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
-_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
+_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
# Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate.
_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
diff --git a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py
index 6d4bd197a..f8d43eaba 100644
--- a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py
+++ b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import os
@@ -53,7 +53,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]]
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
-def test_text_to_text(omni_runner, omni_runner_handler) -> None:
+def test_text_to_text(omni_runner, offline_client) -> None:
"""
Test text-only input processing and text output generation.
Input Modal: text
@@ -62,14 +62,14 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt("请详细介绍鹦鹉的生活习性。")
request_config = {"prompts": prompt, "modalities": ["text"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
-def test_image_to_text(omni_runner, omni_runner_handler) -> None:
+def test_image_to_text(omni_runner, offline_client) -> None:
"""
Test image understanding with text output.
Input Modal: image + text
@@ -79,14 +79,14 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.")
request_config = {"prompts": prompt, "images": image, "modalities": ["text"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
-def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
+def test_audio_to_text(omni_runner, offline_client) -> None:
"""
Test audio understanding with text output.
Input Modal: audio + text
@@ -98,14 +98,14 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt(f"{AUDIO_TOKEN}Please recognize the language of this speech and transcribe it. Format: oral.")
request_config = {"prompts": prompt, "audios": audio, "modalities": ["text"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
-def test_video_to_text(omni_runner, omni_runner_handler) -> None:
+def test_video_to_text(omni_runner, offline_client) -> None:
"""
Test video understanding with text output.
Input Modal: video + text
@@ -115,14 +115,14 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt(f"{VIDEO_TOKEN}Describe what is happening in this video.")
request_config = {"prompts": prompt, "videos": video, "modalities": ["text"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
-def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
+def test_mixed_to_text(omni_runner, offline_client) -> None:
"""
Test mixed modality input (image + audio) with text output.
Input Modal: image + audio + text
@@ -135,14 +135,14 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt(f"{IMAGE_TOKEN}{AUDIO_TOKEN}Describe the image and transcribe the audio.")
request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
+def test_text_to_audio(omni_runner, offline_client) -> None:
"""
Test text input with audio output via the thinker+talker pipeline.
Input Modal: text
@@ -151,14 +151,14 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt("请简单介绍一下北京。")
request_config = {"prompts": prompt, "modalities": ["audio"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
+def test_image_to_audio(omni_runner, offline_client) -> None:
"""
Test image + text input with audio output via the thinker+talker pipeline.
Input Modal: image + text
@@ -168,4 +168,4 @@ def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.")
request_config = {"prompts": prompt, "images": image, "modalities": ["audio"]}
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
diff --git a/tests/e2e/offline_inference/test_ming_tts_expansion.py b/tests/e2e/offline_inference/test_ming_tts_expansion.py
index 8dcefd182..e2e848ee0 100644
--- a/tests/e2e/offline_inference/test_ming_tts_expansion.py
+++ b/tests/e2e/offline_inference/test_ming_tts_expansion.py
@@ -35,7 +35,6 @@ MIN_AUDIO_SAMPLES = 1000
pytestmark = [
pytest.mark.slow,
pytest.mark.tts,
- pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"),
]
diff --git a/tests/e2e/offline_inference/test_minicpmo_4_5.py b/tests/e2e/offline_inference/test_minicpmo_4_5.py
index 6725ce2aa..5a518a7f7 100644
--- a/tests/e2e/offline_inference/test_minicpmo_4_5.py
+++ b/tests/e2e/offline_inference/test_minicpmo_4_5.py
@@ -17,7 +17,7 @@ models = ["openbmb/MiniCPM-o-4_5"]
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
-test_params = [(model, None, {"deploy_config": _CI_DEPLOY, "trust_remote_code": True}) for model in models]
+test_params = [(model, _CI_DEPLOY, {"trust_remote_code": True}) for model in models]
def get_question(prompt_type: str = "text") -> str:
diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py
index b894fe4ea..41e2229d3 100644
--- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py
+++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py
@@ -25,7 +25,7 @@ from vllm_omni import Omni
MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
STAGE_CONFIG = get_deploy_config_path("moss_tts_nano.yaml")
-# (model, stage_configs_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
+# (model, deploy_config_path, extra_omni_kwargs) for ``omni_runner`` indirect parametrize
_OMNI_RUNNER_PARAM = (
MODEL_NAME,
STAGE_CONFIG,
diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py
index 892a395ba..34d36118d 100644
--- a/tests/e2e/offline_inference/test_moss_tts_realtime.py
+++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py
@@ -96,7 +96,7 @@ def _get_test_config() -> str:
pytestmark = [
pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
- pytest.mark.full_model,
+ pytest.mark.slow,
pytest.mark.tts,
pytest.mark.parametrize(
"omni_runner",
@@ -271,7 +271,6 @@ def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor
# ---------------------------------------------------------------------------
-@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
@@ -283,7 +282,6 @@ def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str)
assert not torch.all(audio == 0), "Audio is silence"
-@pytest.mark.advanced_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None:
"""MossTTSRealtime: Chinese input produces non-empty audio."""
diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py
index 25c671e05..487a6bfa6 100644
--- a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py
+++ b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py
@@ -28,10 +28,14 @@ _OMNI_RUNNER_PARAM = (
{"stage_init_timeout": 600, "trust_remote_code": True},
)
+_SKIP_ISSUE_6417 = pytest.mark.skip(
+ reason="https://github.com/vllm-project/vllm-omni/issues/6417",
+)
+
pytestmark = [
- pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4643"),
pytest.mark.slow,
pytest.mark.tts,
+ _SKIP_ISSUE_6417,
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
]
diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py
index 5b4c7a935..dd65cbdcc 100644
--- a/tests/e2e/offline_inference/test_omnivoice_expansion.py
+++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py
@@ -20,13 +20,11 @@ from tests.helpers.stage_config import get_deploy_config_path
MODEL = "k2-fsa/OmniVoice"
DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
-# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
-# The migrated test passes deploy_config through extra Omni kwargs.
+# OmniRunner tuple: model, deploy config path, extra Omni kwargs.
_OMNI_RUNNER_PARAM = (
MODEL,
- None,
+ DEPLOY_CONFIG,
{
- "deploy_config": DEPLOY_CONFIG,
"trust_remote_code": True,
"log_stats": True,
},
diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py
index 6be54969c..41a84d677 100644
--- a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py
+++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for Qwen2.5-Omni AutoRound W4A16 quantized inference.
These tests cover text, audio, image, video, and mixed-modality inputs
@@ -36,8 +36,8 @@ BASELINE_MODEL = os.environ.get("QWEN2_5_OMNI_BASELINE_MODEL", BASELINE_MODEL)
_CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml")
-_SKIP_ISSUE_4663 = pytest.mark.skip(
- reason="https://github.com/vllm-project/vllm-omni/issues/4663",
+_SKIP_ISSUE_5652 = pytest.mark.skip(
+ reason="https://github.com/vllm-project/vllm-omni/issues/5652",
)
@@ -65,16 +65,16 @@ quant_params = [(QUANTIZED_MODEL, stage_config)]
# ------------------------------------------------------------------
-@_SKIP_ISSUE_4663
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_text_to_text(omni_runner, omni_runner_handler):
+def test_text_to_text(omni_runner, offline_client):
"""Text input → text output with W4A16 quantized Qwen2.5-Omni."""
request_config = {
"prompts": "What is the capital of China?",
"modalities": ["text"],
}
- response = omni_runner_handler.send_omni_request(request_config)
+ response = offline_client.send_omni_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -84,10 +84,10 @@ def test_text_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_4663
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_audio_to_text(omni_runner, omni_runner_handler):
+def test_audio_to_text(omni_runner, offline_client):
"""Audio input → text output with W4A16 quantized Qwen2.5-Omni."""
audio = generate_synthetic_audio(1, 1, 16000)["np_array"]
if len(audio.shape) == 2:
@@ -98,7 +98,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
"audios": (audio, 16000),
"modalities": ["text"],
}
- response = omni_runner_handler.send_omni_request(request_config)
+ response = offline_client.send_omni_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -108,10 +108,10 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_4663
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_image_to_text(omni_runner, omni_runner_handler):
+def test_image_to_text(omni_runner, offline_client):
"""Image input → text output with W4A16 quantized Qwen2.5-Omni."""
image = generate_synthetic_image(16, 16)["np_array"]
@@ -120,7 +120,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
"images": image,
"modalities": ["text"],
}
- response = omni_runner_handler.send_omni_request(request_config)
+ response = offline_client.send_omni_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -130,10 +130,10 @@ def test_image_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_4663
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_video_to_text(omni_runner, omni_runner_handler):
+def test_video_to_text(omni_runner, offline_client):
"""Video input → text output with W4A16 quantized Qwen2.5-Omni."""
video = generate_synthetic_video(16, 16, 30)["np_array"]
@@ -142,7 +142,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
"videos": video,
"modalities": ["text"],
}
- response = omni_runner_handler.send_omni_request(request_config)
+ response = offline_client.send_omni_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -152,10 +152,10 @@ def test_video_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_4663
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_mix_to_audio(omni_runner, omni_runner_handler):
+def test_mix_to_audio(omni_runner, offline_client):
"""Mixed-modality input → audio output with W4A16 quantized Qwen2.5-Omni."""
video = generate_synthetic_video(16, 16, 30)["np_array"]
image = generate_synthetic_image(16, 16)["np_array"]
@@ -170,5 +170,5 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
"audios": (audio, 16000),
"modalities": ["audio"],
}
- response = omni_runner_handler.send_omni_request(request_config)
+ response = offline_client.send_omni_request(request_config)
assert response.success, "Request failed"
diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
index 9d368a688..e2465234a 100644
--- a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
+++ b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
E2E tests for Qwen2.5-Omni model with mixed modality inputs, audio and text output.
"""
@@ -55,7 +58,7 @@ def get_question(prompt_type="mix"):
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
+def test_mix_to_audio(omni_runner, offline_client) -> None:
"""
Test multi-modal input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -79,14 +82,14 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
}
# Test single completion
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_text_to_text(omni_runner, omni_runner_handler) -> None:
+def test_text_to_text(omni_runner, offline_client) -> None:
"""
Test text input processing and text output generation via OpenAI API.
Deploy Setting: default yaml
@@ -99,4 +102,4 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
request_config = {"prompts": get_question("text_only"), "modalities": ["text"]}
# Test single completion
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py
index fd3a27d2e..32c345194 100644
--- a/tests/e2e/offline_inference/test_qwen3_omni.py
+++ b/tests/e2e/offline_inference/test_qwen3_omni.py
@@ -1,5 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
E2E offline tests for Omni model with video input and audio output.
+
+Abort / sleep-admission lives in ``test_qwen3_omni_colocate_async.py`` so it
+does not overlap this module's OmniRunners.
"""
import os
@@ -57,23 +63,23 @@ def get_question(prompt_type="video"):
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_video_to_audio(omni_runner, omni_runner_handler) -> None:
+def test_video_to_audio(omni_runner, offline_client) -> None:
"""Test processing video, generating audio output."""
video = generate_synthetic_video(224, 224, 300)["np_array"]
request_config = {"prompts": get_question(), "videos": video, "modalities": ["audio"]}
# Test single completion
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", thinker_test_params, indirect=True)
-def test_thinker_only_model_request(omni_runner, omni_runner_handler) -> None:
+def test_thinker_only_model_request(omni_runner, offline_client) -> None:
"""Test that we can load and run a request through a model that only has the thinker stage."""
request_config = {"prompts": "what color is the sky?", "modalities": ["text"]}
# Test single completion
- omni_runner_handler.send_omni_request(request_config)
+ offline_client.send_omni_request(request_config)
diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py
index 9111740ea..b1f17356b 100644
--- a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py
+++ b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for Qwen3-Omni AutoRound W4A16 quantized inference.
These tests cover text, audio, image, video, and mixed-modality inputs
@@ -27,8 +27,8 @@ pytestmark = [
pytest.mark.omni,
]
-_SKIP_ISSUE_3195 = pytest.mark.skip(
- reason="https://github.com/vllm-project/vllm-omni/issues/3195",
+_SKIP_ISSUE_5652 = pytest.mark.skip(
+ reason="https://github.com/vllm-project/vllm-omni/issues/5652",
)
QUANTIZED_MODEL = "Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound"
@@ -77,16 +77,16 @@ quant_params = [(QUANTIZED_MODEL, stage_config)]
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_text_to_text(omni_runner, omni_runner_handler):
+def test_text_to_text(omni_runner, offline_client):
"""Text input → text output with W4A16 quantized Qwen3-Omni."""
request_config = {
"prompts": "What is the capital of France?",
"modalities": ["text"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -96,10 +96,10 @@ def test_text_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_audio_to_text(omni_runner, omni_runner_handler):
+def test_audio_to_text(omni_runner, offline_client):
"""Audio input → text output with W4A16 quantized Qwen3-Omni."""
audio = generate_synthetic_audio(1, 1, 16000)["np_array"]
if len(audio.shape) == 2:
@@ -110,7 +110,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
"audios": (audio, 16000),
"modalities": ["text"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -120,10 +120,10 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_image_to_text(omni_runner, omni_runner_handler):
+def test_image_to_text(omni_runner, offline_client):
"""Image input → text output with W4A16 quantized Qwen3-Omni."""
image = generate_synthetic_image(16, 16)["np_array"]
@@ -132,7 +132,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
"images": image,
"modalities": ["text"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -142,10 +142,10 @@ def test_image_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_video_to_text(omni_runner, omni_runner_handler):
+def test_video_to_text(omni_runner, offline_client):
"""Video input → text output with W4A16 quantized Qwen3-Omni."""
video = generate_synthetic_video(224, 224, 300)["np_array"]
@@ -154,7 +154,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
"videos": video,
"modalities": ["text"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
assert response.text_content and len(response.text_content.strip()) > 0
@@ -164,10 +164,10 @@ def test_video_to_text(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_video_to_audio(omni_runner, omni_runner_handler):
+def test_video_to_audio(omni_runner, offline_client):
"""Video input → audio output with W4A16 quantized Qwen3-Omni."""
video = generate_synthetic_video(224, 224, 300)["np_array"]
@@ -176,7 +176,7 @@ def test_video_to_audio(omni_runner, omni_runner_handler):
"videos": video,
"modalities": ["audio"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
@@ -185,10 +185,10 @@ def test_video_to_audio(omni_runner, omni_runner_handler):
# ------------------------------------------------------------------
-@_SKIP_ISSUE_3195
+@_SKIP_ISSUE_5652
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_mix_to_audio(omni_runner, omni_runner_handler):
+def test_mix_to_audio(omni_runner, offline_client):
"""Mixed-modality input → audio output with W4A16 quantized Qwen3-Omni."""
video = generate_synthetic_video(224, 224, 300)["np_array"]
image = generate_synthetic_image(16, 16)["np_array"]
@@ -203,5 +203,5 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
"audios": (audio, 16000),
"modalities": ["audio"],
}
- response = omni_runner_handler.send_request(request_config)
+ response = offline_client.send_request(request_config)
assert response.success, "Request failed"
diff --git a/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py
new file mode 100644
index 000000000..54e9b8f62
--- /dev/null
+++ b/tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py
@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Isolated AsyncOmni abort / sleep-admission regression for Qwen3-Omni.
+
+This lives in its own module so the L3 merge job can run it *after*
+``test_qwen3_omni.py`` (offline) and the online serving file. Those modules keep
+module-scoped OmniRunner / OmniServer instances until teardown; overlapping a
+CuMem sleep-mode engine with them kills StageEngineCoreProc at handshake
+(``Failed core proc(s): {}``).
+"""
+
+import asyncio
+import os
+from contextlib import ExitStack
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+from vllm import SamplingParams
+from vllm.inputs import TokensPrompt
+
+from tests.helpers.clean import wait_for_gpu_memory_to_clear
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.platforms import current_omni_platform
+
+MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+
+def _output_token_ids(output) -> list[int]:
+ if not getattr(output, "outputs", None):
+ return []
+ completion = output.outputs[0]
+ token_ids = list(completion.token_ids or [])
+ if token_ids:
+ return token_ids
+ return list(getattr(completion, "cumulative_token_ids", None) or [])
+
+
+def _finish_reason(output) -> str | None:
+ if not getattr(output, "outputs", None):
+ return None
+ return getattr(output.outputs[0], "finish_reason", None)
+
+
+def _colocate_async_deploy() -> str:
+ """Thinker-only Instruct deploy with sleep mode on a single GPU.
+
+ Abort / sleep admission is a thinker control-plane contract. The 3-stage
+ CI yaml plus CuMem on every stage fails engine-core startup in the L3
+ merge job: Stage 1 resolves ``devices: "1"`` while only GPU 0 is visible.
+ """
+ return modify_stage_config(
+ get_deploy_config_path("qwen3_omni_moe_thinking.yaml"),
+ updates={
+ "stages": {
+ 0: {
+ "enable_sleep_mode": True,
+ "enforce_eager": True,
+ "trust_remote_code": True,
+ "tensor_parallel_size": 1,
+ "devices": "0",
+ "gpu_memory_utilization": 0.9,
+ "max_num_seqs": 1,
+ "max_model_len": 8192,
+ "max_num_batched_tokens": 8192,
+ "enable_prefix_caching": False,
+ "skip_mm_profiling": True,
+ },
+ },
+ },
+ )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
+@pytest.mark.asyncio
+async def test_colocate_async_abort_tokens_and_sleep_admission() -> None:
+ """**required** regression. Control-plane APIs (``abort`` / ``sleep`` /
+ ``wake_up`` / ``resume_generation``) are not in ``send_omni_request``.
+
+ On ``origin/main`` this test fails:
+ - abort: ``generate()`` never yields ``finish_reason="abort"`` with the
+ tokens produced so far, so resume has an empty prefix (or the stream
+ hangs after frontend state is dropped).
+ - sleep: ``generate()`` is admitted into EngineCore while weights/KV are
+ offloaded, so the task errors (asleep / corrupted ADD frame) instead of
+ waiting for ``resume_generation()``.
+
+ On this PR both legs pass: abort returns a resumeable prefix, and sleep
+ holds ``generate()`` until wake + resume.
+ """
+ prompt = "What color is the sky? Write a long, detailed explanation."
+ device_count = current_omni_platform.device_count()
+ if device_count > 0:
+ # Fail closed if a prior module-scoped runner is still on the GPU.
+ # ``clean_gpu_memory_between_tests`` only logs a note on timeout.
+ wait_for_gpu_memory_to_clear(
+ devices=list(range(device_count)),
+ threshold_ratio=0.15,
+ timeout_s=120,
+ )
+
+ with ExitStack() as after:
+ engine = AsyncOmni(
+ model=MODEL,
+ deploy_config=_colocate_async_deploy(),
+ enable_sleep_mode=True,
+ )
+ after.callback(engine.shutdown)
+
+ request_id = "qwen3-abort-partial"
+ outputs: list = []
+
+ async def _generate(max_tokens: int, req_id: str, gen_prompt) -> None:
+ async for output in engine.generate(
+ prompt=gen_prompt,
+ request_id=req_id,
+ sampling_params=SamplingParams(temperature=0.0, max_tokens=max_tokens),
+ output_modalities=["text"],
+ ):
+ outputs.append(output)
+
+ abort_task = asyncio.create_task(_generate(256, request_id, prompt))
+ prefix: list[int] = []
+ prompt_token_ids: list[int] = []
+ for _ in range(600):
+ if abort_task.done():
+ break
+ if outputs:
+ latest = outputs[-1]
+ prefix = _output_token_ids(latest)
+ prompt_token_ids = list(getattr(latest, "prompt_token_ids", None) or [])
+ if prefix:
+ break
+ await asyncio.sleep(0.1)
+
+ assert prefix, "generate produced no tokens before abort"
+ assert not abort_task.done(), "generate finished before abort; raise max_tokens"
+ await engine.abort(request_id)
+ await asyncio.wait_for(abort_task, timeout=90)
+
+ final = outputs[-1]
+ assert final.finished
+ assert _finish_reason(final) == "abort"
+ abort_prefix = _output_token_ids(final)
+ assert abort_prefix, "abort dropped the generated prefix (main-branch behavior)"
+
+ resume_prompt: str | TokensPrompt = (
+ TokensPrompt(prompt_token_ids=prompt_token_ids + abort_prefix) if prompt_token_ids else prompt
+ )
+ outputs.clear()
+ async for output in engine.generate(
+ prompt=resume_prompt,
+ request_id="qwen3-abort-resume",
+ sampling_params=SamplingParams(temperature=0.0, max_tokens=8),
+ output_modalities=["text"],
+ ):
+ outputs.append(output)
+ assert outputs
+ assert any(_output_token_ids(out) for out in outputs)
+
+ # Trainer order is pause → abort → sleep. Pause here after the
+ # resume generate so EngineCore is idle before CuMem offload.
+ await engine.pause_generation(mode="abort", clear_cache=True)
+ await engine.sleep(level=1)
+ outputs.clear()
+ sleep_task = asyncio.create_task(_generate(8, "qwen3-sleep-admission", prompt))
+ await asyncio.sleep(1.0)
+ assert not sleep_task.done(), "generate() ran while EngineCore was sleeping (main-branch admission race)"
+ await engine.wake_up()
+ await asyncio.sleep(0.5)
+ assert not sleep_task.done(), "generate() resumed before resume_generation()"
+ await engine.resume_generation()
+ await asyncio.wait_for(sleep_task, timeout=180)
+ assert outputs
+ assert _finish_reason(outputs[-1]) != "abort"
+ await asyncio.sleep(5)
diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py
index 5d1b3dda7..ad6935e71 100644
--- a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py
+++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for Qwen3-Omni ModelOpt NVFP4 W4A4 — Marlin fallback path.
Loads the published full-thinker NVFP4 W4A4 checkpoint on H100 (sm_90),
@@ -74,7 +74,7 @@ quant_params = [(QUANTIZED_MODEL, _stage_config())]
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
-def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler):
+def test_marlin_fallback_no_nan_collapse(omni_runner, offline_client):
"""H100 Marlin FP4 fallback: text in -> coherent text out.
On H100 the upstream NVFP4 PWAL routes through
@@ -89,7 +89,7 @@ def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler):
degenerate-output signature that indicates NaN propagation through
the FP4 GEMM.
"""
- response = omni_runner_handler.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]})
+ response = offline_client.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]})
assert response.success, "request failed"
text = (response.text_content or "").strip()
assert text, "empty response — likely server returned only EOS"
diff --git a/tests/e2e/offline_inference/test_qwen3_tts_base.py b/tests/e2e/offline_inference/test_qwen3_tts_base.py
index a076526a4..f800a1163 100644
--- a/tests/e2e/offline_inference/test_qwen3_tts_base.py
+++ b/tests/e2e/offline_inference/test_qwen3_tts_base.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E offline tests for Qwen3-TTS Base model with text input and audio output.
@@ -15,12 +15,12 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
# See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale.
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -70,7 +70,7 @@ def get_prompt():
@pytest.mark.tts
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
+def test_text_to_audio_001(omni_runner, offline_client) -> None:
"""
Test text input processing and audio output via offline Omni runner.
Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true
@@ -87,4 +87,4 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- omni_runner_handler.send_audio_speech_request(request_config)
+ offline_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
index 4fa2d4b6f..b303b1cda 100644
--- a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+++ b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E offline tests for Qwen3-TTS CustomVoice model with text input and audio output.
@@ -64,7 +64,7 @@ def get_prompt():
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
+def test_text_to_audio_001(omni_runner, offline_client) -> None:
"""
Test text input processing and audio output via offline Omni runner.
Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true
@@ -74,4 +74,4 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
Datasets: few requests
"""
request_config = {"input": get_prompt(), "voice": "vivian"}
- omni_runner_handler.send_audio_speech_request(request_config)
+ offline_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/offline_inference/test_sana_wm.py b/tests/e2e/offline_inference/test_sana_wm.py
new file mode 100644
index 000000000..8b1c460c2
--- /dev/null
+++ b/tests/e2e/offline_inference/test_sana_wm.py
@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import numpy as np
+import pytest
+
+pytestmark = [
+ pytest.mark.advanced_model,
+ pytest.mark.diffusion,
+ pytest.mark.gpu,
+]
+
+
+def _make_sana_wm_e2e_image() -> Any:
+ from PIL import Image
+
+ from vllm_omni.diffusion.models.sana_wm import SANA_WM_OUTPUT_HEIGHT, SANA_WM_OUTPUT_WIDTH
+
+ return Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+
+
+def _sana_wm_e2e_tensor_parallel_size() -> int:
+ return int(
+ os.environ.get(
+ "SANA_WM_E2E_TENSOR_PARALLEL_SIZE",
+ os.environ.get("SANA_WM_E2E_TP", "1"),
+ )
+ )
+
+
+def _coerce_video_array(video: Any) -> np.ndarray:
+ import torch
+
+ if isinstance(video, list):
+ assert video, "SANA-WM e2e produced an empty video list."
+ video = video[0]
+ if isinstance(video, torch.Tensor):
+ video = video.detach().cpu().float().numpy()
+ video_array = np.asarray(video)
+ if video_array.ndim == 5:
+ assert video_array.shape[0] == 1
+ video_array = video_array[0]
+ assert video_array.ndim == 4
+ return video_array
+
+
+def _assert_sana_wm_e2e_shape(
+ video: np.ndarray,
+ *,
+ output_type: str,
+ num_frames: int,
+) -> None:
+ if output_type == "latent":
+ # Latent output is channel-first: (C, T, H, W) after removing the
+ # batch dimension. Stage-1 emits LTX-2 latents with 128 channels.
+ assert video.shape[0] == 128
+ assert 0 < video.shape[1] <= num_frames
+ return
+ assert 0 < video.shape[0] <= num_frames
+
+
+def _run_sana_wm_e2e(*, output_type: str) -> np.ndarray:
+ import torch
+
+ from vllm_omni.diffusion.models.sana_wm import (
+ SANA_WM_MODEL_ID,
+ SANA_WM_OUTPUT_HEIGHT,
+ SANA_WM_OUTPUT_WIDTH,
+ )
+ from vllm_omni.entrypoints.omni import Omni
+ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+ from vllm_omni.outputs import OmniRequestOutput
+
+ if not torch.cuda.is_available():
+ pytest.skip("Sana-WM e2e requires CUDA.")
+
+ model = os.environ.get("SANA_WM_E2E_MODEL", SANA_WM_MODEL_ID)
+ model_class_name = os.environ.get("SANA_WM_E2E_MODEL_CLASS", "SanaWmPipeline")
+ num_frames = int(os.environ.get("SANA_WM_E2E_NUM_FRAMES", "9"))
+ stage1_steps = int(
+ os.environ.get(
+ "SANA_WM_E2E_STAGE1_STEPS",
+ os.environ.get("SANA_WM_E2E_NUM_INFERENCE_STEPS", "1"),
+ )
+ )
+ action = os.environ.get("SANA_WM_E2E_ACTION", f"w-{max(num_frames - 1, 1)}")
+ tensor_parallel_size = _sana_wm_e2e_tensor_parallel_size()
+ image = _make_sana_wm_e2e_image()
+ omni = Omni(
+ model=model,
+ model_class_name=model_class_name,
+ enforce_eager=True,
+ tensor_parallel_size=tensor_parallel_size,
+ )
+ extra_args: dict[str, Any] = {}
+ # The native Stage-1 path is latent-token capped; raise it for full-size
+ # e2e runs (704x1280 exceeds the default 4096-token cap).
+ native_max_tokens = os.environ.get("SANA_WM_E2E_NATIVE_MAX_TOKENS", "").strip()
+ if native_max_tokens:
+ extra_args["sana_wm_native_max_tokens"] = int(native_max_tokens)
+ # Return the requested output type so the shape assertion matches the
+ # pipeline output instead of the raw Stage-1 latent default.
+
+ output = omni.generate(
+ {
+ "prompt": "A slow forward camera move through a quiet city street.",
+ "multi_modal_data": {"image": image},
+ "sana_wm": {
+ "action": action,
+ "num_frames": num_frames,
+ "translation_speed": 0.055,
+ "rotation_speed_deg": 1.2,
+ # Avoid the optional Pi3X dependency by passing deterministic
+ # camera intrinsics directly.
+ "intrinsics": {
+ "fx": SANA_WM_OUTPUT_WIDTH / 2,
+ "fy": SANA_WM_OUTPUT_WIDTH / 2,
+ "cx": SANA_WM_OUTPUT_WIDTH / 2,
+ "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+ },
+ },
+ },
+ OmniDiffusionSamplingParams(
+ height=SANA_WM_OUTPUT_HEIGHT,
+ width=SANA_WM_OUTPUT_WIDTH,
+ num_frames=num_frames,
+ seed=0,
+ fps=16,
+ num_inference_steps=stage1_steps,
+ guidance_scale=1.0,
+ guidance_scale_provided=True,
+ output_type=output_type,
+ extra_args=extra_args,
+ ),
+ )
+
+ request_output = output[0] if isinstance(output, list) else output
+ assert isinstance(request_output, OmniRequestOutput)
+ assert request_output.error is None
+ assert request_output.images
+
+ frames = request_output.images[0]
+ video = _coerce_video_array(frames)
+ _assert_sana_wm_e2e_shape(video, output_type=output_type, num_frames=num_frames)
+ return video
+
+
+def test_sana_wm_native_generates_video() -> None:
+ output_type = os.environ.get("SANA_WM_E2E_OUTPUT_TYPE", "np")
+ video = _run_sana_wm_e2e(output_type=output_type)
+ if output_type == "latent":
+ assert video.ndim == 4
+ else:
+ assert video.shape[-1] in (3, 4)
diff --git a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py
index 4f5c97d11..42510d7c6 100644
--- a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py
+++ b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py
@@ -131,7 +131,7 @@ def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
- with OmniRunner(MODEL, stage_configs_path=None) as runner:
+ with OmniRunner(MODEL) as runner:
omni = runner.omni
baseline = _generate(omni)
img_1x = _generate(omni, lora_request, lora_scale=1.0)
diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py
index 1ac76da28..c2739844c 100644
--- a/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py
+++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py
@@ -164,10 +164,7 @@ def _generate_sensenova_u1_img2img(
def test_sensenova_u1_img2img(run_level):
"""Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML)."""
input_image = _load_input_image()
- with OmniRunner(
- "SenseNova/SenseNova-U1-8B-MoT",
- stage_configs_path=None,
- ) as runner:
+ with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -182,7 +179,6 @@ def test_sensenova_u1_img2img_cache_dit():
input_image = _load_input_image()
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
- stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_img2img(runner.omni, input_image)
diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py
index 092da16b9..2e07e2076 100644
--- a/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py
+++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py
@@ -151,10 +151,7 @@ def _generate_sensenova_u1_image(
@hardware_test(res={"cuda": "H100"})
def test_sensenova_u1_text2img(run_level):
"""Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML)."""
- with OmniRunner(
- "SenseNova/SenseNova-U1-8B-MoT",
- stage_configs_path=None,
- ) as runner:
+ with OmniRunner("SenseNova/SenseNova-U1-8B-MoT") as runner:
generated_image = _generate_sensenova_u1_image(runner.omni)
if run_level == "full_model" or run_level == "advanced_model":
_validate_pixels(generated_image)
@@ -168,7 +165,6 @@ def test_sensenova_u1_text2img_cache_dit():
"""Test SenseNova-U1 text2img with Cache-DiT enabled."""
with OmniRunner(
"SenseNova/SenseNova-U1-8B-MoT",
- stage_configs_path=None,
cache_backend="cache_dit",
) as runner:
_generate_sensenova_u1_image(runner.omni)
diff --git a/tests/e2e/offline_inference/test_soulxsinger_expansion.py b/tests/e2e/offline_inference/test_soulxsinger_expansion.py
deleted file mode 100644
index c3a6091d6..000000000
--- a/tests/e2e/offline_inference/test_soulxsinger_expansion.py
+++ /dev/null
@@ -1,241 +0,0 @@
-"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline)."""
-
-import functools
-import importlib
-import json
-import os
-import shutil
-from pathlib import Path
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-import numpy as np
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import get_asset_path
-from tests.helpers.runtime import OmniRunner
-from tests.helpers.stage_config import get_deploy_config_path
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3")
-TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3")
-PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json")
-TARGET_METADATA = get_asset_path("soulxsinger/music.json")
-SAMPLE_RATE = 24_000
-
-# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it
-# from the model dir. Stage it from this pinned upstream commit so the SVS DiT
-# can run in CI; the SVS tests skip gracefully if it cannot be fetched.
-_PHONE_SET_URL = (
- "https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/"
- "81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json"
-)
-
-if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file():
- pytest.skip(
- f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}",
- allow_module_level=True,
- )
-
-pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.tts]
-
-_CASES = (
- pytest.param(
- "SoulXSingerPipeline",
- "soulxsinger_svs.yaml",
- {
- "language": "Mandarin",
- "vocal_sep": False,
- "control": "score",
- "auto_shift": False,
- "pitch_shift": 0,
- },
- ("g2pM", "g2p_en"),
- id="svs",
- ),
- pytest.param(
- "SoulXSingerSVCPipeline",
- "soulxsinger_svc.yaml",
- {"vocal_sep": False, "auto_shift": False, "pitch_shift": 0},
- (),
- id="svc",
- ),
-)
-
-
-@functools.lru_cache(maxsize=1)
-def _resolve_weights() -> tuple[Path, Path, Path]:
- from huggingface_hub import snapshot_download
-
- base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"]))
-
- # phone_set.json is not on HF; best-effort stage it from pinned upstream so
- # SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent.
- phone_set = base / "phoneme" / "phone_set.json"
- if not phone_set.is_file() and not (base / "phone_set.json").is_file():
- try:
- from urllib.request import urlretrieve
-
- phone_set.parent.mkdir(exist_ok=True)
- urlretrieve(_PHONE_SET_URL, phone_set)
- except Exception:
- pass
-
- # make a temporary svc directory
- svc_dir = base.parent / "SoulX-Singer-SVC"
- svc_dir.mkdir(parents=True, exist_ok=True)
- shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml")
- shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt")
-
- if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"):
- pre = Path(raw).expanduser().resolve()
- if (pre / "rmvpe" / "rmvpe.pt").is_file():
- return base, svc_dir, pre
-
- from huggingface_hub import snapshot_download
-
- pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"]))
- return base, svc_dir, pre
-
-
-@pytest.fixture(scope="session")
-def soulx_weights() -> tuple[Path, Path, Path]:
- try:
- return _resolve_weights()
- except Exception as exc:
- pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})")
-
-
-def _flatten_audio(audio_val) -> np.ndarray:
- import torch
-
- if isinstance(audio_val, list):
- chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None]
- return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
- if isinstance(audio_val, torch.Tensor):
- return audio_val.detach().cpu().float().numpy().reshape(-1)
- return np.asarray(audio_val, dtype=np.float32).reshape(-1)
-
-
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES)
-def test_soulxsinger_multistage_from_audio(
- soulx_weights: tuple[Path, Path, Path],
- architecture: str,
- deploy_yaml: str,
- extra_args: dict,
- py_deps: tuple[str, ...],
-) -> None:
- for mod in py_deps:
- try:
- importlib.import_module(mod)
- except ImportError as exc:
- pytest.fail(f"SoulX SVS requires {mod}: {exc}")
-
- base_dir, svc_dir, preprocess_dir = soulx_weights
-
- # SVS mode requires phone_set.json in the model directory
- if architecture == "SoulXSingerPipeline":
- json.dump(
- {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
- open(base_dir / "config.json", "w"),
- indent=4,
- )
- model = str(base_dir)
- if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
- pytest.skip(
- "SoulX-Singer SVS test requires phoneme/phone_set.json. "
- "Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. "
- "See `examples/offline_inference/text_to_speech/README.md` for details."
- )
- if architecture == "SoulXSingerSVCPipeline":
- json.dump(
- {"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1},
- open(svc_dir / "config.json", "w"),
- indent=4,
- )
- model = str(svc_dir)
-
- with OmniRunner(
- model,
- stage_configs_path=get_deploy_config_path(deploy_yaml),
- async_chunk=False,
- ) as runner:
- sampling = OmniDiffusionSamplingParams(
- num_inference_steps=4,
- guidance_scale=3.0,
- seed=42,
- extra_args={
- "prompt_audio": str(PROMPT_AUDIO),
- "target_audio": str(TARGET_AUDIO),
- "preprocess_weights_dir": str(preprocess_dir),
- **extra_args,
- },
- )
- prompt = {"prompt_token_ids": [0]}
- outputs = runner.generate([prompt], sampling)
-
- assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
- mm = outputs[0].multimodal_output
- assert isinstance(mm, dict) and "audio" in mm
- audio = _flatten_audio(mm["audio"])
- assert 12_000 <= audio.size
- assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
- duration_s = audio.size / SAMPLE_RATE
- assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
-
-
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None:
- """SVS from precomputed note/lyric metadata.
-
- Exercises the SVS DiT + vocoder serving path while skipping the heavy
- upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in
- the merge gate without those non-pip dependencies.
- """
- base_dir, _svc_dir, preprocess_dir = soulx_weights
-
- if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file():
- pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}")
- if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file():
- pytest.skip(
- "SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). "
- "See `examples/offline_inference/text_to_speech/README.md`."
- )
-
- json.dump(
- {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1},
- open(base_dir / "config.json", "w"),
- indent=4,
- )
-
- with OmniRunner(
- str(base_dir),
- stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"),
- async_chunk=False,
- ) as runner:
- sampling = OmniDiffusionSamplingParams(
- num_inference_steps=4,
- guidance_scale=3.0,
- seed=42,
- extra_args={
- "language": "Mandarin",
- "control": "score",
- "prompt_metadata_path": str(PROMPT_METADATA),
- "target_metadata_path": str(TARGET_METADATA),
- "audio_path": str(PROMPT_AUDIO),
- "preprocess_weights_dir": str(preprocess_dir),
- },
- )
- prompt = {"prompt_token_ids": [0]}
- outputs = runner.generate([prompt], sampling)
-
- assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output"
- mm = outputs[0].multimodal_output
- assert isinstance(mm, dict) and "audio" in mm
- audio = _flatten_audio(mm["audio"])
- assert 12_000 <= audio.size
- assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4
- duration_s = audio.size / SAMPLE_RATE
- assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s"
diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py
index 00ecd09d9..759c6c500 100644
--- a/tests/e2e/offline_inference/test_stable_audio_expansion.py
+++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stable Audio offline e2e: real weights, FP8 + TeaCache (single job to save GPU).
@@ -15,7 +15,6 @@ import numpy as np
import pytest
import torch
-from tests.helpers import skip_if_gated_repo_inaccessible
from tests.helpers.assertions import assert_audio_valid
from tests.helpers.mark import hardware_test
from vllm_omni import Omni
@@ -25,8 +24,6 @@ from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
-_MODEL_REPO = "stabilityai/stable-audio-open-1.0"
-
_SAMPLE_RATE = 44100
_CLIP_DURATION_S = 2.0
@@ -80,12 +77,10 @@ def test_stable_audio_quantization_and_teacache() -> None:
CI should provide ``HF_TOKEN`` if the checkpoint is gated.
"""
- skip_if_gated_repo_inaccessible(_MODEL_REPO)
# ``model_class_name`` must be passed explicitly: the default-stage-cfg
# factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when
# deciding ``final_output_type`` (#2077), and at construction time the
- # auto-resolution from ``model_index.json`` has not run yet. AudioX's
- # offline test follows the same pattern.
+ # auto-resolution from ``model_index.json`` has not run yet.
m = Omni(
model="stabilityai/stable-audio-open-1.0",
model_class_name="StableAudioPipeline",
diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py
index 483396b23..3a1b60c94 100644
--- a/tests/e2e/offline_inference/test_step_audio2_expansion.py
+++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py
@@ -19,7 +19,7 @@ from vllm_omni.outputs import OmniRequestOutput
MODEL = "stepfun-ai/Step-Audio-2-mini"
DEPLOY_CONFIG = get_deploy_config_path("step_audio2_ci.yaml")
-TEST_PARAMS = [(MODEL, None, {"deploy_config": DEPLOY_CONFIG, "trust_remote_code": True})]
+TEST_PARAMS = [(MODEL, DEPLOY_CONFIG, {"trust_remote_code": True})]
pytestmark = [pytest.mark.slow, pytest.mark.tts]
diff --git a/tests/e2e/offline_inference/test_voxcpm2_tts.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py
index c10694a40..af55c51aa 100644
--- a/tests/e2e/offline_inference/test_voxcpm2_tts.py
+++ b/tests/e2e/offline_inference/test_voxcpm2_tts.py
@@ -39,8 +39,7 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor:
return audio
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None:
@@ -53,8 +52,7 @@ def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None:
assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None:
@@ -95,8 +93,7 @@ def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None:
assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s"
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxcpm2_prefill_decode_mixed_batch_003(omni_runner: OmniRunner) -> None:
diff --git a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py
index 2441c46dc..dcff41cda 100644
--- a/tests/e2e/offline_inference/test_voxtral_tts_expansion.py
+++ b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py
@@ -102,7 +102,7 @@ def test_voxtral_tts_offline_streaming():
async def _run():
async_omni = AsyncOmni(
model=MODEL,
- stage_configs_path=STAGE_CONFIG,
+ deploy_config=STAGE_CONFIG,
stage_init_timeout=300,
enforce_eager=True,
)
diff --git a/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py b/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py
new file mode 100644
index 000000000..184f0f7d2
--- /dev/null
+++ b/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E smoke test for Wan2.1 AutoRound MXFP4 inference."""
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+QUANTIZED_MODEL = os.environ.get(
+ "WAN21_AUTOROUND_MXFP4_MODEL",
+ "INCModel/Wan2.1-T2V-1.3B-Diffusers-MXFP4-AutoRound",
+)
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"xpu": "B60"})
+@pytest.mark.parametrize(
+ "omni_runner",
+ [(QUANTIZED_MODEL, None)],
+ indirect=True,
+)
+def test_wan21_autoround_mxfp4_generates_video(omni_runner, omni_runner_handler):
+ sampling_params = OmniDiffusionSamplingParams(
+ height=256,
+ width=256,
+ num_frames=5,
+ num_inference_steps=1,
+ guidance_scale=1.0,
+ generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+ )
+ response = omni_runner_handler.send_diffusion_request(
+ {
+ "prompt": "A red toy car driving on a table",
+ "sampling_params": sampling_params,
+ }
+ )
+
+ assert response.success, response.error_message
+ assert response.images is not None and len(response.images) == 1
+ frames = np.asarray(response.images[0])
+ assert frames.shape[1:4] == (5, 256, 256)
+ assert frames.std() > 0.01
diff --git a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py
index 044435cb5..4ac51a344 100644
--- a/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py
+++ b/tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E tests for Wan2.2 AutoRound W4A16 quantized inference.
These tests cover I2V (image-to-video) and T2V (text-to-video) generation
@@ -19,8 +19,8 @@ import pytest
import torch
from PIL import Image
-from tests.helpers.env import DeviceMemoryMonitor
from tests.helpers.mark import hardware_test
+from tests.helpers.monitor import DeviceMemoryMonitor
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
@@ -91,7 +91,7 @@ def _create_test_image(width: int = WIDTH, height: int = HEIGHT) -> Image.Image:
return Image.fromarray(arr)
-def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table, smooth motion") -> tuple:
+def _generate_i2v_video(offline_client, prompt: str = "A cat sitting on a table, smooth motion") -> tuple:
"""Generate one I2V video, return (frames, peak_memory_mb)."""
gc.collect()
current_omni_platform.empty_cache()
@@ -101,7 +101,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
monitor.start()
image = _create_test_image()
- response = omni_runner_handler.send_diffusion_request(
+ response = offline_client.send_diffusion_request(
{
"prompt": prompt,
"images": image,
@@ -112,7 +112,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
peak = monitor.peak_used_mb
monitor.stop()
- assert response.success, f"Request failed: {response.error_message}"
+ assert response.success, "Request failed"
assert response.images is not None and len(response.images) > 0, "Expected image output"
frames = response.images[0]
@@ -122,7 +122,7 @@ def _generate_i2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
return frames, peak
-def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a table") -> tuple:
+def _generate_t2v_video(offline_client, prompt: str = "A cat sitting on a table") -> tuple:
"""Generate one T2V video, return (frames, peak_memory_mb)."""
gc.collect()
current_omni_platform.empty_cache()
@@ -131,7 +131,7 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
- response = omni_runner_handler.send_diffusion_request(
+ response = offline_client.send_diffusion_request(
{
"prompt": prompt,
"sampling_params": _sampling_params_t2v(),
@@ -141,7 +141,7 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
peak = monitor.peak_used_mb
monitor.stop()
- assert response.success, f"Request failed: {response.error_message}"
+ assert response.success, "Request failed"
assert response.images is not None and len(response.images) > 0, "Expected image output"
frames = response.images[0]
@@ -158,9 +158,9 @@ def _generate_t2v_video(omni_runner_handler, prompt: str = "A cat sitting on a t
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True)
-def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler):
+def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, offline_client):
"""Load the W4A16 quantized Wan2.2 I2V model and verify it produces a valid video."""
- frames, _ = _generate_i2v_video(omni_runner_handler)
+ frames, _ = _generate_i2v_video(offline_client)
assert frames is not None, "Expected video frames output"
assert hasattr(frames, "shape"), "Expected frames to have a shape attribute"
@@ -182,9 +182,9 @@ def test_wan22_i2v_autoround_w4a16_generates_video(omni_runner, omni_runner_hand
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True)
-def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_handler):
+def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, offline_client):
"""Load the W4A16 quantized Wan2.2 T2V model and verify it produces a valid video."""
- frames, _ = _generate_t2v_video(omni_runner_handler)
+ frames, _ = _generate_t2v_video(offline_client)
assert frames is not None, "Expected video frames output"
assert hasattr(frames, "shape"), "Expected frames to have a shape attribute"
@@ -204,9 +204,9 @@ def test_wan22_t2v_autoround_w4a16_generates_video(omni_runner, omni_runner_hand
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", quant_i2v_params, indirect=True)
-def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
+def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, offline_client):
"""Measure peak GPU memory of W4A16 quantized I2V model."""
- frames, peak = _generate_i2v_video(omni_runner_handler)
+ frames, peak = _generate_i2v_video(offline_client)
assert frames is not None, "Expected video frames output"
_memory_results["quant_i2v"] = peak
@@ -220,9 +220,9 @@ def test_wan22_i2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", baseline_i2v_params, indirect=True)
-def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler):
+def test_wan22_i2v_autoround_w4a16_baseline_peak(omni_runner, offline_client):
"""Measure peak GPU memory of BF16 baseline I2V model."""
- frames, peak = _generate_i2v_video(omni_runner_handler)
+ frames, peak = _generate_i2v_video(offline_client)
assert frames is not None, "Expected video frames output"
_memory_results["baseline_i2v"] = peak
@@ -261,9 +261,9 @@ def test_wan22_i2v_autoround_w4a16_memory_savings():
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", quant_t2v_params, indirect=True)
-def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
+def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, offline_client):
"""Measure peak GPU memory of W4A16 quantized T2V model."""
- frames, peak = _generate_t2v_video(omni_runner_handler)
+ frames, peak = _generate_t2v_video(offline_client)
assert frames is not None, "Expected video frames output"
_memory_results["quant_t2v"] = peak
@@ -277,9 +277,9 @@ def test_wan22_t2v_autoround_w4a16_quant_peak(omni_runner, omni_runner_handler):
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_runner", baseline_t2v_params, indirect=True)
-def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, omni_runner_handler):
+def test_wan22_t2v_autoround_w4a16_baseline_peak(omni_runner, offline_client):
"""Measure peak GPU memory of BF16 baseline T2V model."""
- frames, peak = _generate_t2v_video(omni_runner_handler)
+ frames, peak = _generate_t2v_video(offline_client)
assert frames is not None, "Expected video frames output"
_memory_results["baseline_t2v"] = peak
diff --git a/tests/e2e/offline_inference/test_wan22_t2v.py b/tests/e2e/offline_inference/test_wan22_t2v.py
index f3d03f399..805fbd1fe 100644
--- a/tests/e2e/offline_inference/test_wan22_t2v.py
+++ b/tests/e2e/offline_inference/test_wan22_t2v.py
@@ -1,7 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -13,7 +16,7 @@ NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 1})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_text_to_video_001(omni_runner_handler: OmniRunnerHandler):
+def test_text_to_video_001(offline_client: OfflineOmniClient):
sampling = OmniDiffusionSamplingParams(
height=512,
width=512,
@@ -29,4 +32,4 @@ def test_text_to_video_001(omni_runner_handler: OmniRunnerHandler):
"negative_prompt": NEGATIVE_PROMPT,
"sampling_params": sampling,
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/offline_inference/test_z_image.py b/tests/e2e/offline_inference/test_z_image.py
index 1770ccd85..ede9ea1ba 100644
--- a/tests/e2e/offline_inference/test_z_image.py
+++ b/tests/e2e/offline_inference/test_z_image.py
@@ -1,21 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniRunnerHandler
+from tests.helpers.runtime import OfflineOmniClient
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
MODEL = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard."
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 1, "xpu": 2})
@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
-def test_zimage(omni_runner_handler: OmniRunnerHandler):
+def test_zimage(offline_client: OfflineOmniClient):
# high resolution may cause OOM on L4
sampling = OmniDiffusionSamplingParams(
height=256,
@@ -30,4 +30,4 @@ def test_zimage(omni_runner_handler: OmniRunnerHandler):
"prompt": PROMPT,
"sampling_params": sampling,
}
- omni_runner_handler.send_diffusion_request(request_config)
+ offline_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
index cb77c40e7..01afe689a 100644
--- a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
+++ b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests."""
from __future__ import annotations
@@ -11,21 +14,12 @@ from types import SimpleNamespace
import pytest
from huggingface_hub import snapshot_download
-from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from tests.helpers.runtime import OmniServerParams, get_model_prefix
+from tests.helpers.stage_config import get_deploy_config_path, get_deploy_duplex_max_sessions
MODEL = "openbmb/MiniCPM-o-4_5"
-DEPLOY_CONFIG = modify_stage_config(
- get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
- updates={
- "base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
- # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
- # is left automatic so duplex matches the simplex deploy profiles.
- "stages": {
- 1: {"max_model_len": 4096},
- },
- },
-)
+DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml"
+DEPLOY_CONFIG = get_deploy_config_path(DEPLOY_CONFIG_REL)
ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
@@ -46,6 +40,16 @@ SERVER_PARAMS = [
]
+def deploy_max_sessions() -> int:
+ """Concurrent duplex sessions the deploy config under test admits.
+
+ The admission probe has to expect the capacity the server is actually
+ started with. Hardcoding it silently drifts the moment the deploy config
+ changes its capacity, turning that change into an unrelated probe timeout.
+ """
+ return get_deploy_duplex_max_sessions(DEPLOY_CONFIG_REL)
+
+
def validated_wav(path: Path, expected_sha256: str) -> Path:
actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
if actual_sha256 != expected_sha256:
@@ -73,10 +77,11 @@ def validated_soft_interrupt_wav() -> Path:
return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256)
-def resolve_ref_audio(model_prefix: str) -> Path:
+def resolve_ref_audio() -> Path:
# ``MODEL`` may be a local checkpoint directory instead of a Hub repo id
# (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only
# fall back to the Hub cache when no such directory exists on disk.
+ model_prefix = get_model_prefix()
model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
if not model_root.is_dir():
model_root = Path(snapshot_download(MODEL, local_files_only=True))
@@ -86,6 +91,20 @@ def resolve_ref_audio(model_prefix: str) -> Path:
return ref_audio
+def duplex_camera_frames(*, seconds: int, cache_dir: Path) -> list[str]:
+ """Base64 JPEG frames of a moving synthetic scene, one per second.
+
+ Mirrors what the realtime web client captures (1 fps, no client-side
+ resize: the server normalizes at ``scale_resolution=448``).
+ """
+ from tests.e2e.online_serving.helpers.minicpmo_realtime_duplex_scenarios import _video_frames_from_file
+ from tests.helpers.media import generate_synthetic_video
+
+ video = generate_synthetic_video(448, 448, 30 * seconds, cache_dir=cache_dir)
+ frames, _ = _video_frames_from_file(Path(video["file_path"]))
+ return frames
+
+
def realtime_url(omni_server) -> str:
return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"
@@ -108,6 +127,8 @@ def demo_args(
output_audio_format="pcm16",
chunk_ms=200,
realtime_input=True,
+ input_video=None,
+ stack_frames=1,
first_turn_ms=1400,
turn_duration_ms=[],
first_turn_transcript="duplex CI speech",
diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
index 17e0cb2eb..7776f2a0f 100644
--- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
+++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Strict MiniCPM-o 4.5 Realtime duplex E2E scenario harness.
This script is intentionally scenario-based instead of a generic chat client.
@@ -34,8 +37,10 @@ import json
import sys
import time
import wave
+from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
+from types import SimpleNamespace
_REPO_ROOT = Path(__file__).resolve().parents[4]
if str(_REPO_ROOT) not in sys.path:
@@ -52,19 +57,100 @@ from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
PCM16_SAMPLE_RATE,
RealtimeEventCollector,
build_realtime_url,
+ duplex_unit_boundary_ms,
read_pcm16_wav,
summarize_session_request_metrics,
)
+from vllm_omni.experimental.fullduplex.client import ( # noqa: E402
+ reference_audio_data_url as _ref_audio_data_url,
+)
+from vllm_omni.experimental.fullduplex.video_stacking import ( # noqa: E402
+ concat_frames_b64,
+ unit_subframe_offsets,
+)
_url_with_model = build_realtime_url
_read_wav_pcm16 = read_pcm16_wav
+DemoArgs = argparse.Namespace | SimpleNamespace
-def _ref_audio_data_url(path: str | None) -> str | None:
- if path is None:
- return None
- ref_path = Path(path).expanduser()
- return "data:audio/wav;base64," + base64.b64encode(ref_path.read_bytes()).decode("ascii")
+def _video_frames_from_file(
+ path: Path,
+ *,
+ fps: float = 1.0,
+ jpeg_quality: int = 70,
+ stack_frames: int = 1,
+) -> tuple[list[str], list[str | None]]:
+ """Decode ``path`` into base64 JPEG frames at the omni duplex cadence.
+
+ One frame per second matches what the realtime web client captures, so the
+ stream stays representative of a live camera. ``stack_frames > 1`` also
+ samples each unit's interior sub-frames and tiles them into one composite,
+ the official way to raise visual refresh rate without changing the audio
+ cadence. Returns ``(base_frames, stacked_frames)`` with the tracks parallel.
+ """
+ import cv2
+
+ offsets = unit_subframe_offsets(stack_frames)
+ capture = cv2.VideoCapture(str(path))
+ if not capture.isOpened():
+ raise ValueError(f"cannot open video: {path}")
+
+ def encode(frame_bgr, index: int) -> str:
+ ok, jpeg = cv2.imencode(".jpg", frame_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
+ if not ok:
+ raise ValueError(f"cannot JPEG-encode frame {index} of {path}")
+ return base64.b64encode(jpeg.tobytes()).decode("ascii")
+
+ try:
+ source_fps = capture.get(cv2.CAP_PROP_FPS) or 0.0
+ stride = max(1, round(source_fps / fps)) if source_fps > 0 else 1
+ # Sub-frame k of a unit sits ``offset * stride`` frames past its base.
+ sub_strides = {max(1, round(offset * stride)): position for position, offset in enumerate(offsets)}
+ frames: list[str] = []
+ interiors: list[list[str]] = []
+ index = 0
+ while True:
+ ok, frame_bgr = capture.read()
+ if not ok:
+ break
+ position, within_unit = divmod(index, stride)
+ if within_unit == 0:
+ frames.append(encode(frame_bgr, index))
+ interiors.append([])
+ elif within_unit in sub_strides and position < len(interiors):
+ interiors[position].append(encode(frame_bgr, index))
+ index += 1
+ finally:
+ capture.release()
+ if not frames:
+ raise ValueError(f"video has no decodable frames: {path}")
+ stacked = [concat_frames_b64(interior) if interior else None for interior in interiors]
+ return frames, stacked
+
+
+def _resolve_video_frames(args: DemoArgs) -> tuple[list[str], list[str | None]]:
+ """Base64 JPEG frames for the duplex camera track, newest last.
+
+ Returns the base track plus the parallel stacked composites. Only a video
+ source has sub-frames to stack; an explicit frame list or a still image
+ stacks nothing.
+ """
+ frames = list(getattr(args, "video_frames_b64", None) or [])
+ if frames:
+ stacked = list(getattr(args, "video_stacked_frames_b64", None) or [])
+ stacked += [None] * (len(frames) - len(stacked))
+ return frames, stacked[: len(frames)]
+ input_video = getattr(args, "input_video", None)
+ if input_video:
+ return _video_frames_from_file(
+ Path(input_video).expanduser(),
+ stack_frames=max(1, int(getattr(args, "stack_frames", 1) or 1)),
+ )
+ frame_image = getattr(args, "frame_image", None)
+ if frame_image:
+ return [base64.b64encode(Path(frame_image).expanduser().read_bytes()).decode("ascii")], [None]
+ return [], []
@dataclass
@@ -374,17 +460,20 @@ class DemoState:
commit_indices = [
index for index, event in enumerate(self.events) if event.get("type") == "input_audio_buffer.committed"
]
- decision_indices = [
- index
- for index, event in enumerate(self.events)
- if event.get("type") == "response.created"
- or (
- event.get("type") == "response.listen"
- and isinstance(event.get("response"), dict)
- and isinstance(event["response"].get("metadata"), dict)
- and event["response"]["metadata"].get("model_listen") is True
- )
- ]
+ decision_indices: list[int] = []
+ for index, event in enumerate(self.events):
+ event_type = event.get("type")
+ if event_type == "response.created":
+ decision_indices.append(index)
+ continue
+ if event_type != "response.listen":
+ continue
+ response = event.get("response")
+ if not isinstance(response, dict):
+ continue
+ metadata = response.get("metadata")
+ if isinstance(metadata, dict) and metadata.get("model_listen") is True:
+ decision_indices.append(index)
first_input_index = self.first_index("input_audio_buffer.speech_started")
if first_input_index is None:
first_input_index = commit_indices[0] if commit_indices else None
@@ -501,7 +590,7 @@ class DemoState:
return stale
-def _session_update_event(args: argparse.Namespace) -> dict[str, object]:
+def _session_update_event(args: DemoArgs) -> dict[str, object]:
session_payload: dict[str, object] = {
"model": args.model,
"modalities": ["audio", "text"],
@@ -837,7 +926,8 @@ async def _send_pcm16(
hints: dict[str, object] | None = None,
first_chunk_hints: dict[str, object] | None = None,
on_model_unit_ready=None,
- frame_b64: str | None = None,
+ frames_b64: Sequence[str] | None = None,
+ stacked_frames_b64: Sequence[str | None] | None = None,
) -> None:
hints = hints or {}
first_chunk_hints = first_chunk_hints or {}
@@ -845,6 +935,8 @@ async def _send_pcm16(
model_unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
audio_ms = 0
next_model_unit_bytes = model_unit_bytes
+ frames = list(frames_b64 or [])
+ stacked = list(stacked_frames_b64 or [])
frames_sent = 0
for offset in range(0, len(pcm16), chunk_bytes):
chunk = pcm16[offset : offset + chunk_bytes]
@@ -853,9 +945,18 @@ async def _send_pcm16(
chunk_hints = dict(hints)
if offset == 0:
chunk_hints.update(first_chunk_hints)
- if frame_b64 is not None and audio_ms > frames_sent * 1000:
- # Omni duplex cadence: one camera frame per 1 s of audio.
- chunk_hints["video_frames"] = [frame_b64]
+ if frames and audio_ms >= duplex_unit_boundary_ms(frames_sent):
+ # Omni duplex cadence: one camera frame per model unit, riding the
+ # append that closes it, so the frame and the second of audio it was
+ # captured during enter the same unit. A video advances one frame per
+ # unit and holds its last frame when the audio outlives the clip; a
+ # still image is a one-element list and therefore repeats.
+ # A stacked composite of that unit's interior sub-frames rides the
+ # same append right after the base frame (wire maximum of 2).
+ index = min(frames_sent, len(frames) - 1)
+ composite = stacked[index] if index < len(stacked) else None
+ base = frames[index]
+ chunk_hints["video_frames"] = [base] if composite is None else [base, composite]
frames_sent += 1
await ws.send(
json.dumps(
@@ -936,7 +1037,8 @@ async def _send_clean_turn(
realtime_input: bool = False,
model_policy_settle_s: float = 2.0,
commit_input: bool = True,
- frame_b64: str | None = None,
+ frames_b64: Sequence[str] | None = None,
+ stacked_frames_b64: Sequence[str | None] | None = None,
) -> tuple[str | None, str]:
before_created = state.count("response.created")
before_model_listen = state.model_listen_count
@@ -958,7 +1060,8 @@ async def _send_clean_turn(
chunk_ms=chunk_ms,
realtime_delay=realtime_input,
hints={"transcript": transcript} if send_transcript_hint else {},
- frame_b64=frame_b64,
+ frames_b64=frames_b64,
+ stacked_frames_b64=stacked_frames_b64,
)
if commit_input:
state.input_commit_sent_at_s.append(time.monotonic())
@@ -1164,7 +1267,8 @@ async def _send_listen_only_overlap_pair(
model_policy_settle_s: float,
validation_mode: str,
silence_ms: int,
- frame_b64: str | None = None,
+ frames_b64: Sequence[str] | None = None,
+ stacked_frames_b64: Sequence[str | None] | None = None,
) -> tuple[list[str | None], list[str], bool]:
if not realtime_input:
raise ValueError("listen-only-overlap requires --realtime-input")
@@ -1195,7 +1299,8 @@ async def _send_listen_only_overlap_pair(
realtime_delay=True,
hints={"transcript": transcripts[0]} if send_transcript_hint else {},
on_model_unit_ready=continue_until_first_speak,
- frame_b64=frame_b64,
+ frames_b64=frames_b64,
+ stacked_frames_b64=stacked_frames_b64,
)
state.input_commit_sent_at_s.append(time.monotonic())
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
@@ -1234,7 +1339,8 @@ async def _send_listen_only_overlap_pair(
realtime_delay=True,
hints={"transcript": transcripts[1]} if send_transcript_hint else {},
on_model_unit_ready=record_model_unit_ready,
- frame_b64=frame_b64,
+ frames_b64=frames_b64,
+ stacked_frames_b64=stacked_frames_b64,
)
state.input_commit_sent_at_s.append(time.monotonic())
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
@@ -1298,10 +1404,8 @@ async def _send_listen_only_overlap_pair(
return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok
-async def run_demo(args: argparse.Namespace) -> dict[str, object]:
- demo_frame_b64: str | None = None
- if getattr(args, "frame_image", None):
- demo_frame_b64 = base64.b64encode(Path(args.frame_image).read_bytes()).decode("ascii")
+async def run_demo(args: DemoArgs) -> dict[str, object]:
+ demo_frames_b64, demo_stacked_b64 = _resolve_video_frames(args)
turn_input_paths = _turn_input_paths(
Path(args.input_wav),
list(getattr(args, "turn_input_wav", []) or []),
@@ -1351,6 +1455,9 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
try:
await ws.send(json.dumps(_session_update_event(args)))
await _wait_for(state, lambda: state.count("session.created") > 0, timeout_s=20, label="session.created")
+ start_barrier = getattr(args, "start_barrier", None)
+ if start_barrier is not None:
+ await start_barrier.wait()
turn_transcripts = _turn_transcripts(args.first_turn_transcript, turns=args.turns)
turn_specs = list(zip(turn_transcripts, turn_durations, strict=True))
@@ -1369,7 +1476,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
validation_mode=validation_mode,
silence_ms=max(0, args.silence_ms),
- frame_b64=demo_frame_b64,
+ frames_b64=demo_frames_b64,
+ stacked_frames_b64=demo_stacked_b64,
)
for turn_index in range(2, len(turn_specs)):
transcript, duration_ms = turn_specs[turn_index]
@@ -1387,7 +1495,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
realtime_input=realtime_input,
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
commit_input=not continuous_input,
- frame_b64=demo_frame_b64,
+ frames_b64=demo_frames_b64,
+ stacked_frames_b64=demo_stacked_b64,
)
turn_response_ids.append(response_id)
turn_outcomes.append(outcome)
@@ -1407,7 +1516,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
realtime_input=realtime_input,
model_policy_settle_s=max(0.0, args.model_policy_settle_ms / 1000),
commit_input=not continuous_input,
- frame_b64=demo_frame_b64,
+ frames_b64=demo_frames_b64,
+ stacked_frames_b64=demo_stacked_b64,
)
turn_response_ids.append(response_id)
turn_outcomes.append(outcome)
@@ -1582,6 +1692,8 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
"continuous_input": continuous_input,
"continuous_input_ok": continuous_input_ok,
"input_chunk_ms": args.chunk_ms,
+ "video_frame_count": len(demo_frames_b64),
+ "video_stacked_frame_count": sum(frame is not None for frame in demo_stacked_b64),
"model_policy_settle_ms": args.model_policy_settle_ms,
"turn_outcomes": turn_outcomes,
"turn_outcomes_ok": turn_outcomes_ok,
@@ -1639,6 +1751,19 @@ def parse_args() -> argparse.Namespace:
default=None,
help="Optional image file sent as an omni-duplex camera frame (one per 1 s of audio)",
)
+ parser.add_argument(
+ "--input-video",
+ default=None,
+ help="Optional video file streamed as omni-duplex camera frames (decoded at 1 fps, "
+ "overrides --frame-image; the last frame is held if the audio outlives the clip)",
+ )
+ parser.add_argument(
+ "--stack-frames",
+ type=int,
+ default=1,
+ choices=[1, 2],
+ help="Frames per append: 2 also resends the previous frame for motion context.",
+ )
parser.add_argument(
"--turn-input-wav",
action="append",
diff --git a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py
new file mode 100644
index 000000000..b623446ad
--- /dev/null
+++ b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py
@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import argparse
+import asyncio
+import base64
+import json
+import math
+import uuid
+import wave
+from collections.abc import Sequence
+from contextlib import asynccontextmanager, suppress
+from pathlib import Path
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import numpy as np
+from scipy.signal import resample_poly
+
+from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, wait_for, write_pcm16_wav
+
+INPUT_SAMPLE_RATE_HZ = 16_000
+OUTPUT_SAMPLE_RATE_HZ = 22_050
+FRAME_SAMPLES = 1_280
+FRAME_PERIOD_S = FRAME_SAMPLES / INPUT_SAMPLE_RATE_HZ
+DEFAULT_FUNCTION_TOOLS = [
+ {
+ "type": "function",
+ "name": "generate_random_number",
+ "description": "Generate a random integer between min and max (inclusive).",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "min": {"type": "integer", "description": "Minimum value (inclusive)"},
+ "max": {"type": "integer", "description": "Maximum value (inclusive)"},
+ },
+ "required": ["min", "max"],
+ },
+ }
+]
+DEFAULT_INSTRUCTIONS = "You are NVIDIA Voice Chat. Answer briefly. Start by greeting the user."
+DEFAULT_FUNCTION_INSTRUCTIONS = (
+ "You are NVIDIA Voice Chat. If the user's request matches an available tool, "
+ "you MUST call that tool instead of answering from your own knowledge. "
+ "Use only argument values spoken by the user and never invent missing values."
+)
+
+
+def _url(base_url: str, model: str, session_id: str) -> str:
+ parts = urlsplit(base_url)
+ query = dict(parse_qsl(parts.query, keep_blank_values=True))
+ query.update(duplex="1", model=model, autostart="0", session_id=session_id)
+ return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
+
+
+def _read_wav(path: Path, *, input_channel: int = 0) -> np.ndarray:
+ with wave.open(str(path), "rb") as wav_file:
+ channels = wav_file.getnchannels()
+ width = wav_file.getsampwidth()
+ source_rate = wav_file.getframerate()
+ raw = wav_file.readframes(wav_file.getnframes())
+ dtypes = {1: np.uint8, 2: np.dtype(" 1:
+ pcm = pcm.reshape(-1, channels)[:, input_channel]
+ if source_rate != INPUT_SAMPLE_RATE_HZ:
+ divisor = math.gcd(source_rate, INPUT_SAMPLE_RATE_HZ)
+ pcm = resample_poly(pcm, up=INPUT_SAMPLE_RATE_HZ // divisor, down=source_rate // divisor)
+ return np.ascontiguousarray(pcm, dtype=" int:
+ count = math.ceil(pcm.size / FRAME_SAMPLES)
+ if max_frames is not None:
+ count = min(count, max_frames)
+ for seq in range(count):
+ frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
+ frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size)).astype(" list[dict[str, object]]:
+ return [event for event in client.events.events if event.get("type") == event_type]
+
+
+async def _return_function_output_when_ready(
+ client: RealtimeDuplexClient,
+ *,
+ output: str,
+ timeout_s: float,
+) -> tuple[str, int]:
+ await wait_for(
+ lambda: bool(client.events.errors()) or client.events.count("response.function_call_arguments.done") > 0,
+ timeout_s=timeout_s,
+ label="function call to execute",
+ )
+ if client.events.errors():
+ raise AssertionError(f"function call failed before tool execution: {client.events.errors()}")
+ function_done = _events(client, "response.function_call_arguments.done")[-1]
+ call_id = function_done.get("call_id")
+ if not isinstance(call_id, str) or not call_id:
+ raise AssertionError(f"completed function call has no call_id: {function_done}")
+ event_count_before_output = len(client.events.events)
+ await client.send(
+ {
+ "type": "conversation.item.create",
+ "item": {
+ "type": "function_call_output",
+ "call_id": call_id,
+ "output": output,
+ },
+ }
+ )
+ return call_id, event_count_before_output
+
+
+def _write_events(path: Path, client: RealtimeDuplexClient) -> None:
+ path.write_text(
+ "".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
+ encoding="utf-8",
+ )
+
+
+@asynccontextmanager
+async def _managed_client(client: RealtimeDuplexClient, *, timeout_s: float):
+ async with client:
+ try:
+ yield
+ finally:
+ if client.events.count("session.created") and not client.events.count("session.closed"):
+ with suppress(Exception):
+ await client.close_session(timeout_s=min(timeout_s, 30.0))
+
+
+async def run(args: argparse.Namespace) -> dict[str, object]:
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+ if args.expect_function_call and args.instructions == DEFAULT_INSTRUCTIONS:
+ instructions = DEFAULT_FUNCTION_INSTRUCTIONS
+ else:
+ instructions = args.instructions
+ tools = DEFAULT_FUNCTION_TOOLS if args.expect_function_call else None
+ session_id = f"nemotron-voicechat-{uuid.uuid4().hex}"
+ client = RealtimeDuplexClient(_url(args.url, args.model, session_id))
+ async with _managed_client(client, timeout_s=args.timeout_s):
+ session_payload: dict[str, object] = {
+ "session_id": session_id,
+ "model": args.model,
+ "modalities": ["audio", "text"],
+ "input_audio_format": "pcm_f32le",
+ "output_audio_format": "pcm16",
+ "instructions": instructions,
+ "idle_timeout_s": args.timeout_s,
+ "turn_detection": None,
+ "extra_body": {"auto_response": True},
+ }
+ if tools is not None:
+ session_payload["tools"] = tools
+ await client.send({"type": "session.update", "session": session_payload})
+ await wait_for(
+ lambda: client.events.count("session.created") > 0 or bool(client.events.errors()),
+ timeout_s=args.timeout_s,
+ label="session.created",
+ )
+ if client.events.errors():
+ raise AssertionError(f"session setup failed: {client.events.errors()}")
+ created = _events(client, "session.created")[-1]
+ session = created.get("session")
+ capabilities = session.get("capabilities") if isinstance(session, dict) else None
+ expected = dict(
+ implementation_level="model_native_duplex",
+ chunk_period_ms=80,
+ supports_core_resumable_request=True,
+ supports_core_kv_lease=False,
+ supports_multi_session=False,
+ )
+ if not isinstance(capabilities, dict) or any(capabilities.get(key) != value for key, value in expected.items()):
+ raise AssertionError(f"unexpected capabilities: {capabilities}")
+
+ function_output_task = (
+ asyncio.create_task(
+ _return_function_output_when_ready(
+ client,
+ output=args.function_output,
+ timeout_s=args.timeout_s,
+ )
+ )
+ if args.function_output is not None
+ else None
+ )
+ pcm = _read_wav(Path(args.input_wav), input_channel=args.input_channel)
+ frame_count = await _stream(client, pcm, max_frames=args.max_frames, realtime=not args.no_realtime)
+ completed_responses_at_commit = client.events.count("response.done")
+ await client.send({"type": "input_audio_buffer.commit", "final": True})
+ await wait_for(
+ lambda: (
+ bool(client.events.errors())
+ or (
+ client.events.count("response.function_call_arguments.done") > 0
+ if args.expect_function_call
+ else (
+ client.events.count("response.audio.delta") >= args.minimum_audio_chunks
+ and (client.events.count("response.done") > completed_responses_at_commit)
+ )
+ )
+ ),
+ timeout_s=args.timeout_s,
+ label="model output",
+ )
+ await asyncio.sleep(args.drain_s)
+ if client.events.errors():
+ raise AssertionError(f"Realtime session emitted errors: {client.events.errors()}")
+ done_events = _events(client, "response.done")
+ if not args.expect_function_call:
+ response = done_events[-1].get("response") if done_events else None
+ status = response.get("status") if isinstance(response, dict) else None
+ if status != "completed":
+ raise AssertionError(f"response did not complete successfully: {done_events[-1:]}")
+
+ function_events = [
+ event for event in client.events.events if str(event.get("type", "")).startswith("response.function_call")
+ ]
+ function_items = [
+ event
+ for event in _events(client, "response.output_item.done")
+ if isinstance(event.get("item"), dict) and event["item"].get("type") == "function_call"
+ ]
+ if args.expect_function_call and not any(
+ event.get("type") == "response.function_call_arguments.done" for event in function_events
+ ):
+ raise AssertionError(f"no completed function call: {function_events}")
+ if args.expect_function_call:
+ matching_items = [
+ event for event in function_items if event["item"].get("name") == args.expected_function_name
+ ]
+ if not matching_items:
+ raise AssertionError(f"expected {args.expected_function_name!r}, got {function_items}")
+ function_item = matching_items[-1]["item"]
+ try:
+ function_arguments = json.loads(str(function_item.get("arguments", "")))
+ except json.JSONDecodeError as exc:
+ raise AssertionError(f"function arguments are not JSON: {function_item}") from exc
+ if args.expected_function_arguments is not None:
+ expected_arguments = json.loads(args.expected_function_arguments)
+ if function_arguments != expected_arguments:
+ raise AssertionError(
+ f"function arguments differ: expected={expected_arguments!r}, actual={function_arguments!r}"
+ )
+
+ if args.function_output is not None:
+ call_id = function_item.get("call_id")
+ if not isinstance(call_id, str) or not call_id:
+ raise AssertionError(f"function item has no call_id: {function_item}")
+ assert function_output_task is not None
+ returned_call_id, event_count_before_output = await function_output_task
+ assert returned_call_id == call_id
+
+ def tool_result_completed() -> bool:
+ later = client.events.events[event_count_before_output:]
+ transcript = "".join(
+ str(event.get("delta", ""))
+ for event in later
+ if event.get("type") == "response.audio_transcript.delta"
+ ).lower()
+ return (
+ any(event.get("type") == "response.audio.delta" for event in later)
+ and any(event.get("type") == "response.done" for event in later)
+ and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript)
+ )
+
+ await wait_for(
+ lambda: bool(client.events.errors()) or tool_result_completed(),
+ timeout_s=args.timeout_s,
+ label="completed response after function output",
+ )
+ if client.events.errors():
+ raise AssertionError(f"function output failed: {client.events.errors()}")
+ await asyncio.sleep(args.drain_s)
+
+ audio = client.events.audio_bytes()
+ audio_events = _events(client, "response.audio.delta")
+ rates = {event.get("sample_rate_hz") for event in audio_events}
+ if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}:
+ raise AssertionError(f"unexpected output sample rates: {rates}")
+ if not args.expect_function_call and args.minimum_audio_chunks and not audio:
+ raise AssertionError("model produced no audio")
+ expected_bytes = 2 * OUTPUT_SAMPLE_RATE_HZ * expected["chunk_period_ms"] // 1000
+ packet_sizes = [len(base64.b64decode(str(event.get("delta", "")), validate=True)) for event in audio_events]
+ if not args.expect_function_call and any(size != expected_bytes for size in packet_sizes):
+ raise AssertionError(f"audio deltas are not fixed 80 ms PCM16 packets: {packet_sizes}")
+ audio_pcm = np.frombuffer(audio, dtype=" argparse.Namespace:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--url", default="ws://127.0.0.1:8125/v1/realtime")
+ parser.add_argument("--model", required=True)
+ parser.add_argument("--input-wav", required=True)
+ parser.add_argument("--input-channel", type=int, default=0)
+ parser.add_argument("--output-dir", default="/tmp/nemotron-voicechat-duplex")
+ parser.add_argument("--instructions", default=DEFAULT_INSTRUCTIONS)
+ parser.add_argument("--max-frames", type=int)
+ parser.add_argument("--minimum-audio-chunks", type=int, default=1)
+ parser.add_argument("--minimum-audio-rms", type=float, default=1e-4)
+ parser.add_argument("--expect-function-call", action="store_true")
+ parser.add_argument("--expected-function-name", default="generate_random_number")
+ parser.add_argument("--expected-function-arguments")
+ parser.add_argument("--function-output")
+ parser.add_argument("--expected-post-tool-text")
+ parser.add_argument("--no-realtime", action="store_true")
+ parser.add_argument("--drain-s", type=float, default=2.0)
+ parser.add_argument("--timeout-s", type=float, default=600.0)
+ return parser.parse_args(argv)
+
+
+def main() -> None:
+ print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
index 331f0d20f..910a5a180 100644
--- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
+++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Concurrent MiniCPM-o Realtime duplex and resumable-session E2E driver."""
from __future__ import annotations
@@ -14,6 +17,7 @@ from types import SimpleNamespace
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import websockets
+from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed
SCRIPT_DIR = Path(__file__).resolve().parent
@@ -42,6 +46,67 @@ async def run_demo(args):
return await _scenario_module().run_demo(args)
+class _SynchronizedStartGate:
+ """One-shot, failure-propagating start gate for Python 3.10+."""
+
+ def __init__(self, parties: int, *, timeout_s: float):
+ if parties < 1:
+ raise ValueError("synchronized start parties must be positive")
+ if timeout_s <= 0:
+ raise ValueError("synchronized start timeout must be positive")
+ self._parties = parties
+ self._timeout_s = timeout_s
+ self._arrived = 0
+ self._released = False
+ self._failure: BaseException | None = None
+ self._event = asyncio.Event()
+ self._lock = asyncio.Lock()
+
+ async def wait(self) -> None:
+ async with self._lock:
+ if self._failure is not None:
+ raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure
+ if self._released:
+ return
+ self._arrived += 1
+ if self._arrived == self._parties:
+ self._released = True
+ self._event.set()
+ return
+
+ try:
+ await asyncio.wait_for(self._event.wait(), timeout=self._timeout_s)
+ except asyncio.TimeoutError as exc:
+ failure = TimeoutError(
+ f"synchronized start timed out after {self._timeout_s:.1f}s "
+ f"({self._arrived}/{self._parties} sessions ready)"
+ )
+ await self.abort(failure)
+ if self._released and self._failure is None:
+ return
+ raise failure from exc
+
+ if self._failure is not None:
+ raise RuntimeError(f"synchronized start aborted: {self._failure!r}") from self._failure
+
+ async def abort(self, failure: BaseException) -> None:
+ async with self._lock:
+ if self._released or self._failure is not None:
+ return
+ self._failure = failure
+ self._event.set()
+
+
+async def _run_demo_with_start_gate(args: SimpleNamespace):
+ start_gate = getattr(args, "start_barrier", None)
+ try:
+ return await run_demo(args)
+ except BaseException as exc:
+ if start_gate is not None:
+ await start_gate.abort(exc)
+ raise
+
+
def _with_resume_mode(url: str) -> str:
parts = urlsplit(url)
query = dict(parse_qsl(parts.query, keep_blank_values=True))
@@ -101,14 +166,24 @@ def _validate_semantic_isolation(
return True
if len(expected_tokens) != len(results):
return False
- for result, expected_token in zip(results, expected_tokens, strict=True):
+ normalized_expected_tokens = [token.strip().casefold() for token in expected_tokens]
+ if any(not token for token in normalized_expected_tokens):
+ return False
+ for result_index, (result, expected_token) in enumerate(zip(results, normalized_expected_tokens, strict=True)):
details = result.get("transcript_integrity")
transcripts = (
[str(item.get("transcript", "")) for item in details if isinstance(item, dict)]
if isinstance(details, list)
else []
)
- if expected_token not in "".join(transcripts):
+ joined_transcripts = "".join(transcripts).casefold()
+ if expected_token not in joined_transcripts:
+ return False
+ if any(
+ other_token in joined_transcripts
+ for other_index, other_token in enumerate(normalized_expected_tokens)
+ if other_index != result_index
+ ):
return False
return True
@@ -130,10 +205,14 @@ async def _receive_until(ws, event_type: str, *, timeout_s: float) -> tuple[dict
return await asyncio.wait_for(receive(), timeout=timeout_s)
+def _server_event_sequences(events: list[dict[str, object]]) -> list[int]:
+ return [sequence for event in events if isinstance(sequence := event.get("server_event_seq"), int)]
+
+
async def _open_admission_session(
args: argparse.Namespace,
session_id: str,
-) -> tuple[object, dict[str, object]]:
+) -> tuple[ClientConnection, dict[str, object]]:
url = _url_with_model(
args.url,
args.model,
@@ -159,7 +238,7 @@ async def _open_admission_session(
return ws, created
-async def _close_admission_session(ws, *, timeout_s: float) -> None:
+async def _close_admission_session(ws: ClientConnection, *, timeout_s: float) -> None:
await ws.send(json.dumps({"type": "session.close"}))
await _receive_until(ws, "session.closed", timeout_s=timeout_s)
await ws.close()
@@ -169,8 +248,8 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
if limit < 1:
raise ValueError("admission limit must be positive")
prefix = f"admission-{uuid.uuid4().hex}"
- accepted: list[tuple[object, dict[str, object]]] = []
- replacement: tuple[object, dict[str, object]] | None = None
+ accepted: list[tuple[ClientConnection, dict[str, object]]] = []
+ replacement: tuple[ClientConnection, dict[str, object]] | None = None
overflow_code = None
try:
for index in range(limit):
@@ -274,14 +353,7 @@ async def _resume_probe(
generation = created.get("attachment_generation")
if not isinstance(token, str) or not isinstance(incarnation, int):
raise RuntimeError("session.created omitted resumable credentials")
- last_seq = max(
- (
- event.get("server_event_seq", 0)
- for event in first_events
- if isinstance(event.get("server_event_seq"), int)
- ),
- default=0,
- )
+ last_seq = max(_server_event_sequences(first_events), default=0)
delay_s = args.expire_after_s if expect_expired else args.resume_after_ms / 1000
if delay_s > 0:
@@ -324,7 +396,7 @@ async def _resume_probe(
await second.send(json.dumps({"type": "session.close"}))
closed, close_events = await _receive_until(second, "session.closed", timeout_s=args.timeout_s)
- replay_sequences = [event["server_event_seq"] for event in replay if isinstance(event.get("server_event_seq"), int)]
+ replay_sequences = _server_event_sequences(replay)
return {
"ok": (
resumed.get("session_id") == session_id
@@ -386,14 +458,7 @@ async def _takeover_probe(
generation = created.get("attachment_generation")
if not isinstance(token, str) or not isinstance(incarnation, int) or not isinstance(generation, int):
raise RuntimeError("session.created omitted takeover credentials")
- last_seq = max(
- (
- event.get("server_event_seq", 0)
- for event in first_events
- if isinstance(event.get("server_event_seq"), int)
- ),
- default=0,
- )
+ last_seq = max(_server_event_sequences(first_events), default=0)
second = await websockets.connect(_with_resume_mode(url), max_size=64 * 1024 * 1024)
await second.send(
@@ -457,7 +522,9 @@ async def _takeover_probe(
await first.close()
-def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace:
+def _demo_args(
+ args: argparse.Namespace, index: int, start_barrier: _SynchronizedStartGate | None = None
+) -> SimpleNamespace:
validation_mode = "response-required" if args.response_required else "model-policy"
input_wav = args.session_input_wav[index] if args.session_input_wav else args.input_wav
return SimpleNamespace(
@@ -477,14 +544,16 @@ def _demo_args(args: argparse.Namespace, index: int) -> SimpleNamespace:
omit_transcript_hints=True,
validation_mode=validation_mode,
temperature=args.temperature,
- scenario="sequential",
+ scenario=getattr(args, "scenario", "sequential"),
+ silence_ms=getattr(args, "silence_ms", 500),
require_audio=args.response_required,
require_distinct_inputs=False,
- expect_empty_turn=[],
+ expect_empty_turn=list(getattr(args, "expect_empty_turn", []) or []),
short_ack_ms=350,
turns=args.turns,
timeout_s=args.timeout_s,
model_policy_settle_ms=args.model_policy_settle_ms,
+ start_barrier=start_barrier,
)
@@ -511,8 +580,16 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
)
lifecycle_result = await run_lifecycle_probes(args)
+ start_barrier = (
+ _SynchronizedStartGate(
+ args.sessions,
+ timeout_s=min(float(args.timeout_s), 30.0),
+ )
+ if getattr(args, "synchronized_start", False)
+ else None
+ )
session_results = await asyncio.gather(
- *(run_demo(_demo_args(args, index)) for index in range(args.sessions)),
+ *(_run_demo_with_start_gate(_demo_args(args, index, start_barrier)) for index in range(args.sessions)),
return_exceptions=True,
)
failures = [repr(result) for result in session_results if isinstance(result, BaseException)]
@@ -596,7 +673,14 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--turns", type=int, default=1)
parser.add_argument("--first-turn-ms", type=int, default=1400)
parser.add_argument("--turn-duration-ms", type=int, action="append", default=[])
+ parser.add_argument("--expect-empty-turn", type=int, action="append", default=[])
parser.add_argument("--response-required", action="store_true")
+ parser.add_argument(
+ "--scenario",
+ choices=["sequential", "listen-only-overlap"],
+ default="sequential",
+ )
+ parser.add_argument("--silence-ms", type=int, default=500)
parser.add_argument("--temperature", type=float, default=None)
parser.add_argument("--disconnect-session-index", type=int)
parser.add_argument("--takeover-session-index", type=int)
@@ -604,6 +688,11 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--expire-session-index", type=int)
parser.add_argument("--expire-after-s", type=float, default=6.0)
parser.add_argument("--verify-admission-limit", type=int)
+ parser.add_argument(
+ "--synchronized-start",
+ action="store_true",
+ help="Wait until every session is created before starting input.",
+ )
parser.add_argument("--model-policy-settle-ms", type=int, default=600)
parser.add_argument("--timeout-s", type=float, default=120.0)
args = parser.parse_args()
@@ -615,6 +704,13 @@ def parse_args() -> argparse.Namespace:
parser.error("provide exactly one --session-expected-token per session")
if args.session_expected_token and not args.session_input_wav:
parser.error("--session-expected-token requires --session-input-wav")
+ normalized_expected_tokens = [token.strip().casefold() for token in args.session_expected_token]
+ for token_index, token in enumerate(normalized_expected_tokens):
+ if any(
+ token in other_token or other_token in token
+ for other_token in normalized_expected_tokens[token_index + 1 :]
+ ):
+ parser.error("--session-expected-token values must not overlap after normalization")
return args
diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
new file mode 100644
index 000000000..4047a0901
--- /dev/null
+++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py
@@ -0,0 +1,85 @@
+import base64
+from pathlib import Path
+
+from vllm_omni.experimental.fullduplex.client import (
+ PCM16_BYTES_PER_SAMPLE,
+ PCM16_SAMPLE_RATE,
+ RealtimeDuplexClient,
+ RealtimeEventCollector,
+ build_realtime_url,
+ read_pcm16_wav,
+ wait_for,
+)
+
+
+def _events(events: list[dict[str, object]], response_id: str, kind: str) -> list[dict[str, object]]:
+ return [e for e in events if e.get("type") == kind and RealtimeEventCollector.response_id(e) == response_id]
+
+
+async def run_server_vad_interrupt(args) -> dict[str, object]:
+ initial = read_pcm16_wav(Path(args.input_wav))
+ interrupt = read_pcm16_wav(Path(args.interrupt_wav))[: PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * 47 // 10]
+ session_id = f"server-vad-hard-interrupt-{id(args)}"
+ client = RealtimeDuplexClient(build_realtime_url(args.url, args.model, autostart=False, session_id=session_id))
+
+ async def until(predicate, label: str) -> None:
+ await wait_for(predicate, timeout_s=args.timeout_s, label=label)
+
+ try:
+ await client.__aenter__()
+ await client.configure(
+ args.model,
+ ref_audio="data:audio/wav;base64," + base64.b64encode(Path(args.ref_audio).read_bytes()).decode(),
+ session_id=session_id,
+ turn_detection={"type": "server_vad", "interrupt_response": True},
+ timeout_s=args.timeout_s,
+ )
+ await client.stream_pcm16(initial + bytes(16_000 * 2 * 1600 // 1000), chunk_ms=200, realtime=True)
+ await client.commit()
+ await until(lambda: any(e.get("type") == "response.created" for e in client.events.events), "response.created")
+ target_id = RealtimeEventCollector.response_id(
+ next(e for e in client.events.events if e.get("type") == "response.created")
+ )
+ assert target_id
+ await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta")
+ cursor = len(client.events.events)
+ await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True)
+ await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done")
+ await client.commit()
+ await until(
+ lambda: any(
+ e.get("type") == "response.done" and RealtimeEventCollector.response_id(e) != target_id
+ for e in client.events.events[cursor:]
+ ),
+ "follow-up response.done",
+ )
+ events = client.events.events
+ finally:
+ try:
+ await client.close_session(timeout_s=args.timeout_s)
+ finally:
+ await client.__aexit__(None, None, None)
+
+ done = _events(events, target_id, "response.done")
+ terminal = done[0] if len(done) == 1 else None
+ trailing = events[events.index(terminal) + 1 :] if terminal is not None else events
+ stale = any(
+ e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing
+ )
+ followup_audio = any(
+ e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id
+ for e in events[cursor:]
+ )
+ ok = (
+ terminal is not None
+ and isinstance(terminal.get("response"), dict)
+ and terminal["response"].get("status") == "cancelled"
+ and terminal["response"].get("status_details", {}).get("reason") == "turn_detected"
+ and not stale
+ and sum(e.get("type") == "input_audio_buffer.speech_started" for e in events[cursor:]) == 1
+ and sum(e.get("type") == "input_audio_buffer.speech_stopped" for e in events[cursor:]) == 1
+ and any(e.get("type") == "input_audio_buffer.committed" for e in events[cursor:])
+ and followup_audio
+ and not any(e.get("type") == "error" for e in events)
+ )
+ return {"ok": ok}
diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
index d1b42cba4..3aff0cf8d 100644
--- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
+++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
@@ -252,6 +252,19 @@ def summarize_artifacts(
(summary["done_indices"][0] for summary in reversed(response_summaries) if summary.get("done_indices")),
None,
)
+ # A soft interrupt's follow-up response can drain past the commit (residual
+ # model unit), so its response.done lands after input_audio_buffer.committed.
+ # Anchor the "listen after last done, before commit" sandwich on the last
+ # turn that closed its floor strictly before the commit, not the
+ # globally-last done (which may sit after the commit and leave the interval empty).
+ last_done_before_commit_index = next(
+ (
+ summary["done_indices"][0]
+ for summary in reversed(response_summaries)
+ if summary.get("done_indices") and commit_index is not None and summary["done_indices"][0] < commit_index
+ ),
+ None,
+ )
listen_indices = [index for index, event in enumerate(events) if event.get("type") == "response.listen"]
effective_min_responses = min_responses if validation_mode == "response-required" else 1
enough_responses = len(response_summaries) >= effective_min_responses
@@ -279,9 +292,9 @@ def summarize_artifacts(
)
listen_after_last_done = last_done_index is not None and any(index > last_done_index for index in listen_indices)
listen_after_response_before_commit = (
- last_done_index is not None
+ last_done_before_commit_index is not None
and commit_index is not None
- and any(last_done_index < index < commit_index for index in listen_indices)
+ and any(last_done_before_commit_index < index < commit_index for index in listen_indices)
)
final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices)
transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries)
diff --git a/tests/e2e/online_serving/test_audiox_expansion.py b/tests/e2e/online_serving/test_audiox_expansion.py
deleted file mode 100644
index 612f397e0..000000000
--- a/tests/e2e/online_serving/test_audiox_expansion.py
+++ /dev/null
@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E online serving test for AudioX text-to-audio diffusion.
-
-Mirrors `tests/e2e/online_serving/test_sd3_expansion.py` for image diffusion: spin up
-`vllm-omni` with `--model-class-name AudioXPipeline` and validate via the standard
-`send_diffusion_request` helper (now audio-aware).
-"""
-
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
-
-# Tiny / random checkpoint usable in CI; override to a real bundle locally.
-AUDIOX_TEST_MODEL = os.environ.get("AUDIOX_TEST_MODEL", "zhangj1an/audiox_random")
-T2A_PROMPT = "A quiet living room with soft fabric rustle and gentle cat breathing."
-
-SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
-
-
-def _audiox_server_cases(model: str):
- return [
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=["--model-class-name", "AudioXPipeline"],
- ),
- id="t2a",
- marks=SINGLE_CARD_FEATURE_MARKS,
- ),
- ]
-
-
-@pytest.mark.slow
-@pytest.mark.diffusion
-@pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True)
-def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
- """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`.
-
- AudioX is registered in ``vllm_omni/model_extras`` (``AudioXPipeline``), so the
- ``audiox_task`` / ``seconds_*`` / ``sigma_*`` keys below flow through the standard
- ``apply_declared_extra_args`` server path into ``sampling_params.extra_args`` rather
- than the old top-level ``extra_args`` escape hatch.
- """
- request_config = {
- "model": omni_server.model,
- "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT),
- "extra_body": {
- "num_inference_steps": 4,
- "guidance_scale": 6.0,
- "seed": 42,
- "audiox_task": "t2a",
- "seconds_start": 0.0,
- "seconds_total": 2.0,
- "sigma_min": 0.03,
- "sigma_max": 1000.0,
- },
- }
- openai_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_bagel.py b/tests/e2e/online_serving/test_bagel.py
index b22762acd..84c2302cc 100644
--- a/tests/e2e/online_serving/test_bagel.py
+++ b/tests/e2e/online_serving/test_bagel.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end online serving test for Bagel text2img and img2img generation.
@@ -75,12 +75,11 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
]
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_bagel_text2img_online(omni_server, openai_client) -> None:
+def test_bagel_text2img_online(omni_server, online_client) -> None:
"""Test Bagel text2img via OpenAI-compatible chat completions API."""
request_config = {
"model": omni_server.model,
@@ -95,15 +94,14 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None:
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_bagel_img2img_online(omni_server, openai_client) -> None:
+def test_bagel_img2img_online(omni_server, online_client) -> None:
"""Test Bagel img2img via OpenAI-compatible chat completions API."""
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
buffer = BytesIO()
@@ -121,4 +119,4 @@ def test_bagel_img2img_online(omni_server, openai_client) -> None:
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py
index fa17791e5..e22614a6b 100644
--- a/tests/e2e/online_serving/test_bagel_expansion.py
+++ b/tests/e2e/online_serving/test_bagel_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode
and are supported by the following models:
@@ -19,10 +22,10 @@ assert_diffusion_response validates successful generation and the expected
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
@@ -183,7 +186,7 @@ def _get_diffusion_feature_cases(model: str):
)
def test_bagel(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""L4 diffusion feature coverage for Bagel on H100.
@@ -218,4 +221,4 @@ def test_bagel(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_bagel_multi_replicas.py b/tests/e2e/online_serving/test_bagel_multi_replicas.py
index cb7aeeb94..c360b63e8 100644
--- a/tests/e2e/online_serving/test_bagel_multi_replicas.py
+++ b/tests/e2e/online_serving/test_bagel_multi_replicas.py
@@ -1,6 +1,11 @@
-"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving.
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Weekly e2e coverage for BAGEL diffusion multi-replica serving.
This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs.
+It is covered by the weekly Diffusion · H100 · Multi-GPU lane
+(``.buildkite/cuda/test-weekly.yml``, ``mirror_hardwares: h100_4``).
"""
import os
@@ -17,7 +22,7 @@ MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml")
ROUTE_STRESS_REQUESTS = 6
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
test_params = [
OmniServerParams(
@@ -31,7 +36,7 @@ test_params = [
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_client) -> None:
+def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": dummy_messages_from_mix_data(
@@ -47,7 +52,7 @@ def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_
},
}
- responses = openai_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+ responses = online_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
assert len(responses) == ROUTE_STRESS_REQUESTS
assert all(response.success for response in responses)
diff --git a/tests/e2e/online_serving/test_boogu_image.py b/tests/e2e/online_serving/test_boogu_image.py
index 5a5494d7d..440e031f1 100644
--- a/tests/e2e/online_serving/test_boogu_image.py
+++ b/tests/e2e/online_serving/test_boogu_image.py
@@ -16,7 +16,7 @@ parity vs. the Diffusers baseline is covered separately (see the support plan, s
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image.py \
- -m "full_model and diffusion and H100" --run-level=full_model
+ -m "slow and diffusion and H100"
"""
import os
@@ -28,7 +28,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Base"
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
diff --git a/tests/e2e/online_serving/test_boogu_image_edit.py b/tests/e2e/online_serving/test_boogu_image_edit.py
index 45f999d26..118b46335 100644
--- a/tests/e2e/online_serving/test_boogu_image_edit.py
+++ b/tests/e2e/online_serving/test_boogu_image_edit.py
@@ -1,15 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
-Online serving tests for Boogu-Image-0.1-Edit (image-to-image via chat completions).
+Online serving tests for Boogu-Image-0.1-Edit and Edit-Turbo.
-The Edit checkpoint shares the ``BooguImagePipeline`` class with the Base
+Both Edit checkpoints share the ``BooguImagePipeline`` class with the Base
text-to-image checkpoint; the editing (TI2I) path activates when the request
-carries a reference image. Text guidance rides on ``guidance_scale`` (upstream
-default 4.0), and image guidance rides on ``guidance_scale_2`` (upstream default
-1.0 = off; a value > 1 enables the double-guidance path). Only a single
-reference image is supported.
+carries a reference image. The regular Edit checkpoint uses text guidance via
+``guidance_scale`` and optional image guidance via ``guidance_scale_2``.
+Edit-Turbo is guidance-distilled and uses four steps with guidance scale 1.0.
+Only a single reference image is supported.
- ``test_single_image_to_image_001``: one reference image, text guidance only.
- ``test_double_guidance_001``: one reference image, text + image guidance.
@@ -19,7 +19,7 @@ reference image is supported.
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_edit.py \
- -m "full_model and diffusion and H100" --run-level=full_model
+ -m "slow and diffusion and H100"
"""
import base64
@@ -35,11 +35,15 @@ from tests.helpers.runtime import (
dummy_messages_from_mix_data,
)
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
EDIT_PROMPT = "Change the style to a colored pencil drawing."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+EDIT_MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
+
def _get_diffusion_feature_cases(model: str):
"""Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
@@ -54,13 +58,44 @@ def _get_diffusion_feature_cases(model: str):
]
+def _get_single_edit_cases():
+ """Return the regular and distilled Edit server/request configurations."""
+ return [
+ pytest.param(
+ OmniServerParams(model=EDIT_MODEL),
+ {
+ "num_inference_steps": 2,
+ "guidance_scale": 5.0,
+ },
+ id="edit",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=EDIT_TURBO_MODEL,
+ server_args=["--revision", EDIT_TURBO_REVISION],
+ ),
+ {
+ "num_inference_steps": 4,
+ "guidance_scale": 1.0,
+ },
+ id="edit-turbo-hotfix-1k",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ ]
+
+
@pytest.mark.parametrize(
- "omni_server",
- _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
- indirect=True,
+ ("omni_server", "sampling_params"),
+ _get_single_edit_cases(),
+ indirect=["omni_server"],
)
-def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Single-reference text-guided edit smoke for ``Boogu/Boogu-Image-0.1-Edit``."""
+def test_single_image_to_image_001(
+ omni_server: OmniServer,
+ sampling_params: dict,
+ openai_client: OpenAIClientHandler,
+):
+ """Single-reference text-guided edit smoke for regular and Turbo Edit."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -69,9 +104,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
"model": omni_server.model,
"messages": messages,
"extra_body": {
- "num_inference_steps": 2,
- # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
- "guidance_scale": 5.0,
+ **sampling_params,
"seed": 42,
},
}
@@ -81,7 +114,7 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+ _get_diffusion_feature_cases(EDIT_MODEL),
indirect=True,
)
def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
@@ -105,15 +138,20 @@ def test_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClien
openai_client.send_diffusion_request(request_config)
-def _edits_request_config(model: str, *, guidance_scale_2: float | None = None) -> dict:
+def _edits_request_config(
+ model: str,
+ *,
+ num_inference_steps: int = 2,
+ guidance_scale: float = 5.0,
+ guidance_scale_2: float | None = None,
+) -> dict:
"""Multipart ``/v1/images/edits`` config: one synthetic reference + smoke-depth params."""
image_bytes = base64.b64decode(generate_synthetic_image(512, 512)["base64"])
data = {
"model": model,
"prompt": EDIT_PROMPT,
- "num_inference_steps": 2,
- # Boogu text guidance (upstream default 4.0); > 1 enables CFG.
- "guidance_scale": 5.0,
+ "num_inference_steps": num_inference_steps,
+ "guidance_scale": guidance_scale,
"seed": 42,
}
if guidance_scale_2 is not None:
@@ -133,19 +171,25 @@ def _assert_edits_response_has_image(responses) -> None:
@pytest.mark.parametrize(
- "omni_server",
- _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
- indirect=True,
+ ("omni_server", "sampling_params"),
+ _get_single_edit_cases(),
+ indirect=["omni_server"],
)
-def test_images_edits_endpoint_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Single-reference text-guided edit via the OpenAI-compatible ``/v1/images/edits`` endpoint."""
- responses = openai_client.send_images_edits_http_request(_edits_request_config(omni_server.model))
+def test_images_edits_endpoint_001(
+ omni_server: OmniServer,
+ sampling_params: dict,
+ openai_client: OpenAIClientHandler,
+):
+ """Single-reference regular and Turbo Edit via ``/v1/images/edits``."""
+ responses = openai_client.send_images_edits_http_request(
+ _edits_request_config(omni_server.model, **sampling_params)
+ )
_assert_edits_response_has_image(responses)
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Boogu/Boogu-Image-0.1-Edit"),
+ _get_diffusion_feature_cases(EDIT_MODEL),
indirect=True,
)
def test_images_edits_endpoint_double_guidance_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
diff --git a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py
index efc968eb6..df7650b65 100644
--- a/tests/e2e/online_serving/test_boogu_image_edit_expansion.py
+++ b/tests/e2e/online_serving/test_boogu_image_edit_expansion.py
@@ -1,8 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
-L4 expansion coverage for ``Boogu/Boogu-Image-0.1-Edit`` (image editing / TI2I).
+L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I).
The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
text-guided editing and the text+image double-guidance path. This expansion
@@ -16,16 +16,17 @@ module adds the two edit paths smoke does not exercise:
``height`` / ``width``: asserts the output resolution follows the reference
(upstream ``align_res``, on by default for a single reference).
-The Edit checkpoint shares ``BooguImagePipeline`` with the Base checkpoint; the
-TI2I path activates when the request carries a reference image. Only a single
-reference image is supported. Cases stay smoke-depth (``num_inference_steps=2``).
-CPU offload, Cache-DiT, and multi-GPU parallelism are not yet supported for this
-model and are intentionally omitted.
+Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
+the TI2I path activates when the request carries a reference image. Only a
+single reference image is supported. The regular Edit cases stay smoke-depth
+(``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
+configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
+supported for these models and are intentionally omitted.
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_edit_expansion.py \
- -m "full_model and diffusion and H100" --run-level=full_model
+ -m "slow and diffusion and H100"
"""
import os
@@ -38,9 +39,11 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Edit"
+EDIT_TURBO_MODEL = "Boogu/Boogu-Image-0.1-Edit-Turbo"
+EDIT_TURBO_REVISION = "hotfix-1k-20260708"
EDIT_PROMPT = "Change the style to a colored pencil drawing."
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
@@ -57,6 +60,33 @@ def _get_diffusion_feature_cases(model: str):
]
+def _get_align_res_cases():
+ """Return regular and Turbo Edit configurations for align-res coverage."""
+ return [
+ pytest.param(
+ OmniServerParams(model=MODEL),
+ {
+ "num_inference_steps": 2,
+ "guidance_scale": 5.0,
+ },
+ id="edit",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=EDIT_TURBO_MODEL,
+ server_args=["--revision", EDIT_TURBO_REVISION],
+ ),
+ {
+ "num_inference_steps": 4,
+ "guidance_scale": 1.0,
+ },
+ id="edit-turbo-hotfix-1k",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ ]
+
+
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
"""Image-only guidance: text CFG off (1.0), image CFG on (2.0) -> image-only CFG branch."""
@@ -76,8 +106,16 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
openai_client.send_diffusion_request(request_config)
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+@pytest.mark.parametrize(
+ ("omni_server", "sampling_params"),
+ _get_align_res_cases(),
+ indirect=["omni_server"],
+)
+def test_align_res_output_size(
+ omni_server: OmniServer,
+ sampling_params: dict,
+ openai_client: OpenAIClientHandler,
+) -> None:
"""Output resolution follows the reference image (upstream ``align_res``).
No ``height`` / ``width`` is requested. The reference is 512x384 — both
@@ -92,8 +130,7 @@ def test_align_res_output_size(omni_server: OmniServer, openai_client: OpenAICli
"model": omni_server.model,
"messages": messages,
"extra_body": {
- "num_inference_steps": 2,
- "guidance_scale": 5.0,
+ **sampling_params,
"seed": 42,
},
}
diff --git a/tests/e2e/online_serving/test_boogu_image_expansion.py b/tests/e2e/online_serving/test_boogu_image_expansion.py
index ef6d3a62e..a56e5a338 100644
--- a/tests/e2e/online_serving/test_boogu_image_expansion.py
+++ b/tests/e2e/online_serving/test_boogu_image_expansion.py
@@ -30,7 +30,7 @@ omitted (add rows here when those features land).
From ``tests/``::
pytest -s -v e2e/online_serving/test_boogu_image_expansion.py \
- -m "full_model and diffusion and H100" --run-level=full_model
+ -m "slow and diffusion and H100"
"""
import os
@@ -42,7 +42,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Boogu/Boogu-Image-0.1-Base"
T2I_PROMPT = "A mountain lake at sunset, photorealistic, cinematic lighting"
diff --git a/tests/e2e/online_serving/test_cosmos3.py b/tests/e2e/online_serving/test_cosmos3.py
index 09cb2365b..f28e2ae26 100644
--- a/tests/e2e/online_serving/test_cosmos3.py
+++ b/tests/e2e/online_serving/test_cosmos3.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Online serving L2 smoke tests for ``nvidia/Cosmos3-Nano``.
@@ -8,6 +8,7 @@ broader image/video similarity checks stay in ``tests/e2e/accuracy``.
"""
import os
+from typing import Any
import pytest
@@ -16,7 +17,7 @@ from tests.helpers.media import generate_synthetic_image, generate_synthetic_vid
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
)
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -56,9 +57,9 @@ def _get_diffusion_feature_cases(model: str):
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
"""Default Cosmos3 T2I smoke through ``/v1/images/generations``."""
- responses = openai_client.send_images_generations_http_request(
+ responses = online_client.send_images_generations_http_request(
{
"json": {
"model": omni_server.model,
@@ -85,12 +86,12 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_video_generation_modes_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
"""Default Cosmos3 T2V, I2V, and V2V smoke through the async ``/v1/videos`` endpoint."""
image_reference = f"data:image/jpeg;base64,{generate_synthetic_image(256, 256, seed=42)['base64']}"
video_reference = f"data:video/mp4;base64,{generate_synthetic_video(256, 256, 5)['base64']}"
- video_cases = [
+ video_cases: list[dict[str, Any]] = [
{
"form_data": {
**VIDEO_SHAPE,
@@ -118,8 +119,8 @@ def test_video_generation_modes_001(omni_server: OmniServer, openai_client: Open
]
for case in video_cases:
- request_config = {
+ request_config: dict[str, Any] = {
"model": omni_server.model,
**case,
}
- openai_client.send_video_diffusion_request(request_config)
+ online_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
index 4b21f956e..544afc7bd 100644
--- a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
+++ b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for CosyVoice3 TTS model with voice cloning.
@@ -14,18 +14,26 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.platforms import current_omni_platform
-pytestmark = [pytest.mark.slow, pytest.mark.tts]
+if current_omni_platform.is_rocm():
+ os.environ.setdefault("COSYVOICE3_TRT", "0")
+
+pytestmark = [
+ pytest.mark.slow,
+ pytest.mark.tts,
+ pytest.mark.core_model,
+]
MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
# Official CosyVoice zero-shot prompt audio and its transcript. Vendored under
# tests/assets/ so the server does not depend on raw.githubusercontent.com being
# reachable at request time (same rationale as issue #3263 for Qwen3-TTS).
-REF_AUDIO_URL = load_test_audio_data_url("cosyvoice3/zero_shot_prompt.wav")
+REF_AUDIO_URL = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True)
REF_TEXT = "希望你以后能够做的比我还好呦。"
@@ -42,6 +50,12 @@ def get_prompt(prompt_type="zh"):
return prompts.get(prompt_type, prompts["zh"])
+def _set_rocm_request_timeout(request_config: dict) -> None:
+ if current_omni_platform.is_rocm():
+ # Cold AITER MHA builds on loaded MI300 agents can exceed 20 minutes.
+ request_config["timeout"] = 1800.0
+
+
tts_server_params = [
pytest.param(
OmniServerParams(
@@ -65,10 +79,9 @@ tts_async_chunk_server_params = [
]
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_zh_001(omni_server, openai_client) -> None:
+def test_voice_clone_zh_001(omni_server, online_client) -> None:
"""
Test voice cloning TTS with Chinese text via OpenAI API.
Deploy Setting: default yaml
@@ -85,13 +98,13 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ _set_rocm_request_timeout(request_config)
+ online_client.send_audio_speech_request(request_config)
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
-def test_voice_clone_zh_002(omni_server, openai_client) -> None:
+def test_voice_clone_zh_002(omni_server, online_client) -> None:
"""
Test voice cloning TTS with Chinese text via async_chunk streaming.
Deploy Setting: cosyvoice3.yaml with default ``async_chunk: true``
@@ -109,13 +122,13 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ _set_rocm_request_timeout(request_config)
+ online_client.send_audio_speech_request(request_config)
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644")
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_en_001(omni_server, openai_client) -> None:
+def test_voice_clone_en_001(omni_server, online_client) -> None:
"""
Test voice cloning TTS with English text via OpenAI API.
Deploy Setting: default yaml
@@ -132,4 +145,5 @@ def test_voice_clone_en_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ _set_rocm_request_timeout(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_covo_audio_expansion.py b/tests/e2e/online_serving/test_covo_audio_expansion.py
index 33771d647..6a387269b 100644
--- a/tests/e2e/online_serving/test_covo_audio_expansion.py
+++ b/tests/e2e/online_serving/test_covo_audio_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Covo-Audio-Chat model.
"""
@@ -37,7 +37,7 @@ test_params = [
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_to_audio_001(omni_server, openai_client) -> None:
+def test_audio_to_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text + audio
Output Modal: audio
@@ -58,4 +58,4 @@ def test_audio_to_audio_001(omni_server, openai_client) -> None:
"stream": False,
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py
index 1f16a584b..c4cd40ff8 100644
--- a/tests/e2e/online_serving/test_dynin_omni_expansion.py
+++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Example online tests for Dynin-Omni model.
"""
@@ -64,7 +64,7 @@ def _build_i2i_messages(prompt: str) -> list[dict]:
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-def test_send_i2i_request_001(omni_server, openai_client) -> None:
+def test_send_i2i_request_001(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _build_i2i_messages(I2I_PROMPT),
@@ -73,27 +73,27 @@ def test_send_i2i_request_001(omni_server, openai_client) -> None:
"sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-def test_send_t2i_request_001(omni_server, openai_client) -> None:
+def test_send_t2i_request_001(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _build_t2i_messages(T2I_PROMPT),
"modalities": ["image"],
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-def test_send_t2s_request_001(omni_server, openai_client) -> None:
+def test_send_t2s_request_001(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _build_t2s_messages(T2S_PROMPT),
"modalities": ["audio"],
"audio_ref_text": T2S_PROMPT,
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
diff --git a/tests/e2e/online_serving/test_ernie_image_expansion.py b/tests/e2e/online_serving/test_ernie_image_expansion.py
index 848fec4e6..b06267b39 100644
--- a/tests/e2e/online_serving/test_ernie_image_expansion.py
+++ b/tests/e2e/online_serving/test_ernie_image_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
End-to-end test for ErnieImage text-to-image generation via online serving.
@@ -11,7 +11,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -45,7 +45,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_ernie_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_ernie_image(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test each diffusion feature with ERNIE-Image (text-to-image), one feature per case."""
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
request_config = {
@@ -60,4 +60,4 @@ def test_ernie_image(omni_server: OmniServer, openai_client: OpenAIClientHandler
"seed": 42,
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py
index 2659cb927..d7c20e832 100644
--- a/tests/e2e/online_serving/test_flux2_expansion.py
+++ b/tests/e2e/online_serving/test_flux2_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Tests for Flux2 Klein; currently Dev is implemented separately,
but ideally these models will fold together in the future.
@@ -15,7 +18,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -113,7 +116,7 @@ def _get_diffusion_feature_cases(model: str):
),
indirect=True,
)
-def test_flux2_klein(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux2_klein(omni_server: OmniServer, online_client: OnlineOmniClient):
messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
request_config = {
"model": omni_server.model,
@@ -128,4 +131,4 @@ def test_flux2_klein(omni_server: OmniServer, openai_client: OpenAIClientHandler
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py
index 19a14ce39..932b3cfa5 100644
--- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py
+++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
@@ -19,7 +22,7 @@ image resolution.
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -116,7 +119,7 @@ def _get_flux_2_dev_feature_cases(model: str):
)
def test_flux_2_dev(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Validate FLUX.2-dev online serving with CPU offload."""
@@ -135,7 +138,7 @@ def test_flux_2_dev(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
@@ -145,14 +148,14 @@ def test_flux_2_dev(
)
def test_flux_2_dev_batched_chat_completions(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Validate batched chat completions for diffusion models."""
messages = [
[{"role": "user", "content": PROMPT}],
[{"role": "user", "content": "A sunset over the ocean."}],
]
- responses = openai_client.send_batched_chat_completions_http_request(
+ responses = online_client.send_batched_chat_completions_http_request(
{
"json": {
"model": omni_server.model,
@@ -167,5 +170,7 @@ def test_flux_2_dev_batched_chat_completions(
assert responses and len(responses) == 1
resp = responses[0]
assert resp.success
- choices = resp.json_body["choices"]
+ body = resp.json_body
+ assert isinstance(body, dict)
+ choices = body["choices"]
assert len(choices) == len(messages)
diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py
index ede5ff56b..79460c631 100644
--- a/tests/e2e/online_serving/test_flux_kontext_expansion.py
+++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode
and are supported by the FluxKontext model.
@@ -7,7 +10,7 @@ import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -32,7 +35,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux_kontext_text_to_image(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test text-to-image generation with FluxKontext in regular end-user scenarios."""
messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard")
@@ -47,7 +50,7 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
@@ -55,7 +58,7 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux_kontext_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test image editing with FluxKontext in regular end-user scenarios."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -74,7 +77,7 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
@@ -82,7 +85,7 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test image editing with FluxKontext without negative prompt."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -99,7 +102,7 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
@@ -107,7 +110,7 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux_kontext_high_resolution(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test high-resolution generation with FluxKontext."""
messages = dummy_messages_from_mix_data(content_text="A beautiful landscape with mountains and a lake")
@@ -122,7 +125,7 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
@@ -130,7 +133,7 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_flux_kontext_multiple_outputs(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_flux_kontext_multiple_outputs(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test generating multiple outputs with FluxKontext."""
messages = dummy_messages_from_mix_data(content_text="A photo of a cat sitting on a laptop keyboard")
@@ -146,4 +149,4 @@ def test_flux_kontext_multiple_outputs(omni_server: OmniServer, openai_client: O
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_glm_image_expansion.py b/tests/e2e/online_serving/test_glm_image_expansion.py
index d2a08f633..290272364 100644
--- a/tests/e2e/online_serving/test_glm_image_expansion.py
+++ b/tests/e2e/online_serving/test_glm_image_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode
and are supported by the following model:
@@ -25,12 +28,12 @@ from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image")
@@ -152,7 +155,7 @@ MODES = ["t2i", "i2i"]
def test_glm_image(
omni_server: OmniServer,
mode: str,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""GLM-Image diffusion feature coverage for t2i and i2i on two H100s."""
@@ -180,4 +183,4 @@ def test_glm_image(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_glm_tts_expansion.py b/tests/e2e/online_serving/test_glm_tts_expansion.py
index c558e0c40..89a98f6fb 100644
--- a/tests/e2e/online_serving/test_glm_tts_expansion.py
+++ b/tests/e2e/online_serving/test_glm_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E expansion tests for GLM-TTS online serving (nightly CI).
@@ -14,7 +14,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -24,7 +24,7 @@ MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS")
REF_TEXT = "他当时还跟线下其他的站姐吵架,然后,打架进局子了。"
DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml")
-REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav")
+REF_AUDIO_URL = get_asset_path("glm_tts/jiayan_zh.wav", as_data_url=True)
SYNC_EXTRA_ARGS = [
"--trust-remote-code",
@@ -64,7 +64,7 @@ tts_async_chunk_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
-def test_voice_clone_zh_sync(omni_server, openai_client) -> None:
+def test_voice_clone_zh_sync(omni_server, online_client) -> None:
"""
Test voice cloning TTS with Chinese text.
Deploy Setting: glm_tts.yaml with ``--no-async-chunk`` (sync two-stage)
@@ -80,12 +80,12 @@ def test_voice_clone_zh_sync(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
-def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None:
+def test_voice_clone_zh_async_chunk(omni_server, online_client) -> None:
"""
Test voice cloning TTS with Chinese text via async_chunk streaming.
Deploy Setting: glm_tts.yaml default ``async_chunk: true``
@@ -102,12 +102,12 @@ def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True)
-def test_models_endpoint(omni_server, openai_client) -> None:
+def test_models_endpoint(omni_server, online_client) -> None:
"""Test the /v1/models endpoint returns loaded model."""
- models = openai_client.client.models.list()
+ models = online_client.client.models.list()
assert len(models.data) > 0
diff --git a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py
index a362d4902..dea3b3c9a 100644
--- a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py
+++ b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
End-to-end diffusion coverage for HiDream-I1-Full in online serving mode.
@@ -12,7 +15,7 @@ image resolution.
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -49,7 +52,7 @@ def _get_hidream_i1_image_feature_cases(model: str):
@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4662")
def test_hidream_i1_image(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Validate HiDream-I1-Full online serving with CPU offload."""
@@ -67,4 +70,4 @@ def test_hidream_i1_image(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py
index 8927eb039..e549c6bfc 100644
--- a/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py
+++ b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py
@@ -1,17 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online expansion tests for higgs-audio v2 against /v1/audio/speech.
v1 scope is plain text -> 24 kHz speech plus shallow voice clone via
ref_audio + ref_text (inline) or voice= (after POST /v1/audio/voices).
-The model-aware request validator
-(vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2::validate)
-rejects multi-speaker tags, language overrides, task_type, and bare
-voice= for names that do not match an uploaded speaker — this suite
-exercises both the happy path (plain text in, audio bytes out) and the
-validator rejections, so a regression that loosens the schema will fail
-loudly.
+Model-aware validator rejections live in L1
+``tests/entrypoints/openai_api/test_tts_adapter.py``
+(``test_higgs_audio_v2_validate_*``).
"""
from __future__ import annotations
@@ -27,7 +23,7 @@ os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -63,7 +59,7 @@ _MIN_AUDIO_BYTES = 20_000
# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py
# for the asset rationale — keeping a single shared reference clip across TTS
# tests avoids duplicating WAVs in the repo.
-_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -72,9 +68,9 @@ class TestHiggsAudioV2OnlineHappyPath:
"""Plain-text -> audio happy paths against the live HTTP server."""
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_wav(self, omni_server, openai_client) -> None:
+ def test_plain_text_wav(self, omni_server, online_client) -> None:
"""Single non-streaming WAV request — covers the canonical TTS happy path."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
@@ -86,7 +82,7 @@ class TestHiggsAudioV2OnlineHappyPath:
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+ def test_plain_text_pcm_streaming(self, omni_server, online_client) -> None:
"""Streaming PCM via the shared-memory connector's codec_streaming path.
higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``,
@@ -102,7 +98,7 @@ class TestHiggsAudioV2OnlineHappyPath:
content correctness is separately guarded by the full_model transcript
check.
"""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Streaming the quick brown fox over the lazy dog.",
@@ -116,9 +112,9 @@ class TestHiggsAudioV2OnlineHappyPath:
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+ def test_plain_text_with_max_new_tokens(self, omni_server, online_client) -> None:
"""max_new_tokens is one of the few extra fields the higgs validator accepts."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Innovation distinguishes between a leader and a follower.",
@@ -131,9 +127,9 @@ class TestHiggsAudioV2OnlineHappyPath:
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+ def test_concurrent_plain_text(self, omni_server, online_client) -> None:
"""Three concurrent non-streaming requests — guards the per-slot audio state."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "It was the night before my birthday.",
@@ -146,104 +142,6 @@ class TestHiggsAudioV2OnlineHappyPath:
)
-# ---------------------------------------------------------------------------
-# Validator rejections served over HTTP. Each case targets one out-of-scope
-# field; the validator returns 4xx via the OpenAI error path. We keep the
-# error-message substring loose so a phrasing tweak in the validator does not
-# break CI, but tight enough to catch a regression that silently accepts the
-# field.
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-class TestHiggsAudioV2OnlineValidatorRejections:
- """Out-of-scope fields must come back as 4xx with a higgs-named message."""
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None:
- """Voice clone needs the transcript too — half-supplied is a 4xx."""
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": "Hello world.",
- "ref_audio": _REF_AUDIO_URL,
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "ref_text",
- }
- )
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None:
- """Symmetric guard: ref_text alone is not enough — must come with ref_audio."""
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": "Hello world.",
- "ref_text": "some transcript",
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "ref_audio",
- }
- )
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_task_type(self, omni_server, openai_client) -> None:
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": "Hello world.",
- "task_type": "Base",
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "task_type",
- }
- )
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_language_override(self, omni_server, openai_client) -> None:
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": "Hello world.",
- "language": "Chinese",
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "language",
- }
- )
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None:
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": "[SPEAKER0] hi",
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "multi-speaker",
- }
- )
-
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_rejects_empty_input(self, omni_server, openai_client) -> None:
- openai_client.send_audio_speech_request(
- {
- "model": omni_server.model,
- "input": " ",
- "response_format": "wav",
- "timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "status_code": (400, 422),
- "err_message": "empty",
- }
- )
-
-
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestHiggsAudioV2OnlineVoiceClone:
"""Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice.
@@ -256,9 +154,9 @@ class TestHiggsAudioV2OnlineVoiceClone:
"""
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_basic(self, omni_server, openai_client) -> None:
+ def test_voice_clone_basic(self, omni_server, online_client) -> None:
"""Single non-streaming WAV request driven by the qwen3_tts ref clip."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
@@ -272,9 +170,9 @@ class TestHiggsAudioV2OnlineVoiceClone:
)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None:
+ def test_voice_clone_pcm_streaming(self, omni_server, online_client) -> None:
"""Voice clone over the streaming PCM path."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Innovation distinguishes a leader from a follower.",
diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py
index d26e6bff4..e9ba68909 100644
--- a/tests/e2e/online_serving/test_higgs_audio_v3.py
+++ b/tests/e2e/online_serving/test_higgs_audio_v3.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""End-to-end online tests for higgs-audio v3 against /v1/audio/speech.
Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out
@@ -21,7 +21,7 @@ os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -51,7 +51,7 @@ _MIN_AUDIO_BYTES = 20_000
# Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech)
# vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS
# tests rather than duplicating asset bytes.
-_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -59,13 +59,12 @@ _REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know w
class TestHiggsAudioV3OnlineHappyPath:
"""Plain-text -> audio happy paths against the live HTTP server."""
- @pytest.mark.core_model
- @pytest.mark.advanced_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_wav(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_plain_text_wav(self, omni_server, online_client) -> None:
"""Single non-streaming WAV request - canonical TTS happy path."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
@@ -79,31 +78,29 @@ class TestHiggsAudioV3OnlineHappyPath:
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_plain_text_with_max_new_tokens(self, omni_server, online_client) -> None:
"""``max_new_tokens`` is one of the few extra fields the v3 validator accepts."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Innovation distinguishes between a leader and a follower.",
"stream": False,
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
- "max_new_tokens": 32,
+ "max_new_tokens": 1024,
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
)
- @pytest.mark.skip(reason="issue#4411")
- @pytest.mark.core_model
- @pytest.mark.advanced_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_concurrent_plain_text(self, omni_server, online_client) -> None:
"""Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "it was the night before my birthday",
@@ -114,11 +111,10 @@ class TestHiggsAudioV3OnlineHappyPath:
request_num=3,
)
- @pytest.mark.skip(reason="issue#4411")
- @pytest.mark.advanced_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_plain_text_pcm_streaming(self, omni_server, online_client) -> None:
"""Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path.
The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks
@@ -134,7 +130,7 @@ class TestHiggsAudioV3OnlineHappyPath:
~0.16 dB margin over measured single-request output (1.16 dB on L4), so
flake risk is too high there.
"""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The quick brown fox jumps over the lazy dog.",
@@ -146,11 +142,10 @@ class TestHiggsAudioV3OnlineHappyPath:
}
)
- @pytest.mark.skip(reason="issue#4411")
- @pytest.mark.advanced_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_concurrent_pcm_streaming(self, omni_server, online_client) -> None:
"""Three concurrent streaming requests - guards per-request frame cursors
in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR.
@@ -160,7 +155,7 @@ class TestHiggsAudioV3OnlineHappyPath:
over the measured worst-of-3 on L4 (-0.42 dB) while still well above the
catastrophic-failure region.
"""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "She sells seashells by the seashore.",
@@ -184,17 +179,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
tokens (pause, sfx) go inline. SFX tokens must pair with their written
onomatopoeia.
- These tests check the *serving surface* - the validator accepts the
- payload, the engine produces audio, and the WAV has non-trivial size.
- They do not assert audio quality (out of scope for CI).
+ These tests check the serving surface and spoken text. They do not assert
+ whether the requested control changes audio quality.
"""
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_inline_emotion_and_expressive(self, omni_server, online_client) -> None:
"""Delivery tokens (emotion + expressive_high) at the start of input."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."),
@@ -202,15 +196,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
+ "transcript_expected_text": "Wait, that was actually hilarious.",
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_inline_style_whispering(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_inline_style_whispering(self, omni_server, online_client) -> None:
"""Style token at the start - ``whispering``."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "<|style:whispering|>It is just between you and me, alright?",
@@ -218,15 +213,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
+ "transcript_expected_text": "It is just between you and me, alright?",
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_inline_prosody_speed_and_pitch(self, omni_server, online_client) -> None:
"""Two prosody tokens at the start - slow speed plus low pitch."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": (
@@ -236,15 +232,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
+ "transcript_expected_text": "The radar shows a storm approaching from the east.",
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_inline_pause_mid_text(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_inline_pause_mid_text(self, omni_server, online_client) -> None:
"""Positional pause token placed inline between two clauses."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hold on a moment <|prosody:pause|> let me think about it.",
@@ -252,15 +249,16 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
+ "transcript_expected_text": "Hold on a moment. Let me think about it.",
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_inline_sfx_with_onomatopoeia(self, omni_server, online_client) -> None:
"""SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``)."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": (
@@ -271,6 +269,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
"response_format": "wav",
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
"min_audio_bytes": _MIN_AUDIO_BYTES,
+ "transcript_expected_text": "I cannot believe that just happened. I am still recovering from it.",
}
)
@@ -279,13 +278,12 @@ class TestHiggsAudioV3OnlineInlineControlTokens:
class TestHiggsAudioV3OnlineVoiceClone:
"""Voice clone via the two payload shapes the model serves."""
- @pytest.mark.core_model
- @pytest.mark.advanced_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_voice_clone_ref_audio_ref_text(self, omni_server, online_client) -> None:
"""Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
@@ -297,17 +295,17 @@ class TestHiggsAudioV3OnlineVoiceClone:
}
)
- @pytest.mark.core_model
+ @pytest.mark.slow
@pytest.mark.tts
- @hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
+ @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+ def test_voice_clone_references_alias(self, omni_server, online_client) -> None:
"""BosonAI cookbook payload: ``references=[{audio_path, text}]``.
``normalize_references_alias`` (``protocol/audio.py``) is supposed to
translate the cookbook field into ``ref_audio`` / ``ref_text``; without
it the request would silently fall through to zero-shot synthesis.
"""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello world.",
diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
index a448335c6..0bc3f449f 100644
--- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
+++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py
@@ -1,20 +1,32 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Tests of common diffusion feature combinations in online serving mode
for HunyuanVideo-1.5-T2V (480p).
Coverage (H100, since model cannot fit L4):
-- CacheDiT + Layerwise CPU offloading (1 GPU)
-- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs)
+- CPU offloading (1 GPU) — ``core_model`` + ``advanced_model``
+- CacheDiT + Layerwise CPU offloading (1 GPU) — ``full_model``
+- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) — ``full_model``
+
+HunyuanVideo-1.5 is a high-priority model, so only the most basic single-card deployment
+row runs on every PR (L2) and on merge (L3). The heavyweight feature combinations stay
+nightly-only (L4), together with the video similarity suites in
+``tests/e2e/accuracy/hunyuanvideo15_{t2v,i2v}/``.
+
+From ``tests/``::
+
+ pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "core_model and diffusion" --run-level=core_model
+ pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model and diffusion" --run-level=advanced_model
+ pytest -s -v e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and diffusion" --run-level=full_model
"""
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
-pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+pytestmark = [pytest.mark.diffusion]
PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion."
NEGATIVE_PROMPT = "low quality, blurry, distorted"
@@ -24,9 +36,14 @@ MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+# Only the most basic single-card deployment row is cheap enough for PR (L2) / merge (L3)
+# CI; the CacheDiT / parallel feature combinations run nightly (L4).
+_CORE_MARKS = [pytest.mark.core_model, pytest.mark.advanced_model]
+_FULL_MARKS = [pytest.mark.full_model]
+
def _get_diffusion_feature_cases(model: str):
- """Return L4 diffusion feature cases for HunyuanVideo-1.5.
+ """Return diffusion feature cases for HunyuanVideo-1.5.
Designed for 2x H100 environment per issue #1832.
"""
@@ -40,7 +57,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="single_card_cpu_offload",
- marks=SINGLE_CARD_MARKS,
+ marks=SINGLE_CARD_MARKS + _CORE_MARKS,
),
# (1 GPU) CacheDiT + Layerwise CPU offloading
pytest.param(
@@ -53,7 +70,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="single_card_cachedit_layerwise",
- marks=SINGLE_CARD_MARKS,
+ marks=SINGLE_CARD_MARKS + _FULL_MARKS,
),
# (2 GPUs) CacheDiT + TP=2 + VAE patch parallel=2
pytest.param(
@@ -70,7 +87,7 @@ def _get_diffusion_feature_cases(model: str):
],
),
id="parallel_cachedit_tp2_vae2",
- marks=PARALLEL_MARKS,
+ marks=PARALLEL_MARKS + _FULL_MARKS,
),
]
@@ -82,9 +99,9 @@ def _get_diffusion_feature_cases(model: str):
)
def test_hunyuan_video_15_t2v(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
- """L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
+ """Diffusion feature coverage for HunyuanVideo-1.5-T2V on H100."""
form_data = {
"prompt": PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
@@ -101,4 +118,4 @@ def test_hunyuan_video_15_t2v(
"form_data": form_data,
}
- openai_client.send_video_diffusion_request(request_config)
+ online_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_indextts2_expansion.py b/tests/e2e/online_serving/test_indextts2_expansion.py
index d971ef166..da9bb6dd8 100644
--- a/tests/e2e/online_serving/test_indextts2_expansion.py
+++ b/tests/e2e/online_serving/test_indextts2_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint.
Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV.
@@ -15,7 +15,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -23,7 +23,7 @@ pytestmark = [pytest.mark.slow, pytest.mark.tts]
MODEL = "IndexTeam/IndexTTS-2"
-REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav")
+REF_AUDIO_URL = get_asset_path("indextts2/ref_audio.wav", as_data_url=True)
tts_server_params = [
pytest.param(
@@ -39,7 +39,7 @@ tts_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_basic_english(omni_server, openai_client) -> None:
+def test_basic_english(omni_server, online_client) -> None:
"""
Basic English TTS: text + ref_audio, no emotion.
Deploy Setting: default yaml
@@ -55,12 +55,12 @@ def test_basic_english(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"min_audio_bytes": 1024,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_streaming(omni_server, openai_client) -> None:
+def test_streaming(omni_server, online_client) -> None:
"""
stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the
OpenAI speech endpoint should still accept the streaming request path and
@@ -80,4 +80,4 @@ def test_streaming(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"min_audio_bytes": 1024,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_krea2_expansion.py b/tests/e2e/online_serving/test_krea2_expansion.py
index f76a567f4..eb1d6d078 100644
--- a/tests/e2e/online_serving/test_krea2_expansion.py
+++ b/tests/e2e/online_serving/test_krea2_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline.
@@ -29,7 +29,7 @@ import requests
from PIL import Image
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -76,7 +76,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_krea2(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Serve Krea 2 and validate the OpenAI-compatible image-generation route.
guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG,
@@ -96,7 +96,7 @@ def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler):
"seed": 42,
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
def _basic_lora_payload(lora: dict | None = None) -> dict:
diff --git a/tests/e2e/online_serving/test_lance.py b/tests/e2e/online_serving/test_lance.py
index 9c1a40be9..54e782e10 100644
--- a/tests/e2e/online_serving/test_lance.py
+++ b/tests/e2e/online_serving/test_lance.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""End-to-end online serving test for Lance (single-stage).
Verifies that the Lance pipeline can serve text-to-image and image-edit
@@ -100,7 +100,7 @@ def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_lance_text2img_online(omni_server, openai_client) -> None:
+def test_lance_text2img_online(omni_server, online_client) -> None:
"""Lance text2img via the OpenAI-compatible chat completions API."""
request_config = {
"model": omni_server.model,
@@ -115,7 +115,7 @@ def test_lance_text2img_online(omni_server, openai_client) -> None:
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.core_model
@@ -123,7 +123,7 @@ def test_lance_text2img_online(omni_server, openai_client) -> None:
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_lance_img2img_online(omni_server, openai_client) -> None:
+def test_lance_img2img_online(omni_server, online_client) -> None:
"""Lance image_edit via the OpenAI-compatible chat completions API."""
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
buffer = BytesIO()
@@ -141,4 +141,4 @@ def test_lance_img2img_online(omni_server, openai_client) -> None:
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py
index 777880732..7e4d16568 100644
--- a/tests/e2e/online_serving/test_lingbot_video_moe.py
+++ b/tests/e2e/online_serving/test_lingbot_video_moe.py
@@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH
)
-@pytest.mark.full_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True)
def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
diff --git a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py
index 7c7e0788c..af6bbac50 100644
--- a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py
+++ b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Recommended tests of diffusion features that are available in online serving mode
and are supported by the following model:
@@ -16,7 +19,7 @@ import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_longcat_image_edit(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test the recommended feature combinations for LongCat-Image-Edit."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -84,4 +87,4 @@ def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClient
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_longcat_image_expansion.py b/tests/e2e/online_serving/test_longcat_image_expansion.py
index 5ed7bb156..5616708c1 100644
--- a/tests/e2e/online_serving/test_longcat_image_expansion.py
+++ b/tests/e2e/online_serving/test_longcat_image_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Recommended tests of diffusion features that are available in online serving mode
and are supported by the following model:
@@ -15,7 +18,7 @@ This validates:
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
@@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(MODEL),
indirect=True,
)
-def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_longcat_image(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test the recommended feature combinations for LongCat-Image."""
messages = dummy_messages_from_mix_data(content_text=TEXT_TO_IMAGE_PROMPT)
@@ -83,4 +86,4 @@ def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandl
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_ltx.py b/tests/e2e/online_serving/test_ltx.py
new file mode 100644
index 000000000..e4d8f51a8
--- /dev/null
+++ b/tests/e2e/online_serving/test_ltx.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Online-serving smoke tests for the public LTX-2 pipeline entries."""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+
+DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
+BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
+PROMPT = "A serene lake at sunset with mountains in the background."
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+def _server(
+ model: str,
+ model_class_name: str,
+ *extra_args: str,
+) -> OmniServerParams:
+ return OmniServerParams(
+ model=model,
+ server_args=[
+ "--enforce-eager",
+ *extra_args,
+ "--model-class-name",
+ model_class_name,
+ ],
+ )
+
+
+def _cases():
+ return [
+ pytest.param(
+ _server(BASE_MODEL, "LTX2Pipeline"),
+ 3,
+ id="one_stage",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ _server(DISTILLED_MODEL, "LTX2DistilledPipeline"),
+ 8,
+ id="distilled",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ _server(
+ DISTILLED_MODEL,
+ "LTX2DistilledPipeline",
+ "--use-hsdp",
+ "--hsdp-shard-size",
+ "2",
+ ),
+ 8,
+ id="distilled_hsdp2",
+ marks=HSDP_MARKS,
+ ),
+ pytest.param(
+ _server(
+ BASE_MODEL,
+ "LTX2TwoStagePipeline",
+ "--enable-layerwise-offload",
+ ),
+ 3,
+ id="two_stage",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ ]
+
+
+@pytest.mark.parametrize(
+ ("omni_server", "num_inference_steps"),
+ _cases(),
+ indirect=["omni_server"],
+)
+def test_ltx_pipeline_entries(
+ omni_server: OmniServer,
+ num_inference_steps: int,
+ openai_client: OpenAIClientHandler,
+ subtests,
+):
+ for task in ("t2v", "i2v"):
+ with subtests.test(task=task):
+ form_data = {
+ "prompt": PROMPT,
+ "model": omni_server.model,
+ "height": 128,
+ "width": 128,
+ "num_frames": 9,
+ "fps": 8,
+ "num_inference_steps": num_inference_steps,
+ "seed": 42,
+ }
+ request_config = {
+ "model": omni_server.model,
+ "form_data": form_data,
+ }
+ if task == "i2v":
+ request_config["image_reference"] = (
+ f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+ )
+
+ openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_ltx25.py b/tests/e2e/online_serving/test_ltx25.py
new file mode 100644
index 000000000..4975abcc0
--- /dev/null
+++ b/tests/e2e/online_serving/test_ltx25.py
@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Online serving smokes for all four canonical LTX-2.5 pipelines."""
+
+import os
+
+import pytest
+
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+DEFAULT_MODEL = "Lightricks/LTX-2.5-Diffusers"
+DEFAULT_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
+MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL)
+MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "")
+PROMPT = "A red fox walks through a snowy forest while the camera remains fixed."
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+@pytest.fixture(scope="module", autouse=True)
+def require_ltx25_model_access() -> None:
+ if not os.path.isdir(MODEL):
+ skip_if_gated_repo_inaccessible(
+ MODEL,
+ revision=MODEL_REVISION or None,
+ filename="model_index.json",
+ )
+
+
+def _server(model_class_name: str) -> OmniServerParams:
+ return OmniServerParams(
+ model=MODEL,
+ server_args=[
+ *(["--revision", MODEL_REVISION] if MODEL_REVISION else []),
+ "--model-class-name",
+ model_class_name,
+ "--enforce-eager",
+ "--enable-layerwise-offload",
+ "--diffusion-attention-backend",
+ "CUDNN_ATTN",
+ ],
+ )
+
+
+def _cases():
+ return [
+ pytest.param(_server("LTX2Pipeline"), 30, id="full_one_stage", marks=SINGLE_CARD_MARKS),
+ pytest.param(_server("LTX2TwoStagePipeline"), 30, id="full_two_stage", marks=SINGLE_CARD_MARKS),
+ pytest.param(_server("LTX2DistilledOneStagePipeline"), 8, id="distilled_one_stage", marks=SINGLE_CARD_MARKS),
+ pytest.param(_server("LTX2DistilledTwoStagePipeline"), 8, id="distilled_two_stage", marks=SINGLE_CARD_MARKS),
+ ]
+
+
+@pytest.mark.parametrize(("omni_server", "num_inference_steps"), _cases(), indirect=["omni_server"])
+def test_ltx25_pipeline_entries(
+ omni_server: OmniServer,
+ num_inference_steps: int,
+ openai_client: OpenAIClientHandler,
+ subtests,
+) -> None:
+ """Generate T2V and first-frame I2V through the synchronous video API."""
+ for task in ("t2v", "i2v"):
+ with subtests.test(task=task):
+ request_config = {
+ "model": omni_server.model,
+ "form_data": {
+ "model": omni_server.model,
+ "prompt": PROMPT,
+ "height": 256,
+ "width": 256,
+ "num_frames": 9,
+ "fps": 24,
+ "num_inference_steps": num_inference_steps,
+ "seed": 42,
+ },
+ }
+ if task == "i2v":
+ request_config["image_reference"] = (
+ f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+ )
+
+ openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py
deleted file mode 100644
index 4d0eb6c90..000000000
--- a/tests/e2e/online_serving/test_ltx2_expansion.py
+++ /dev/null
@@ -1,91 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Expansion tests for LTX-2 two-stage pipelines.
-
-Coverage:
-- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video)
-"""
-
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
-PROMPT = "A serene lake at sunset with mountains in the background."
-PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
-
-
-def _cases():
- cases = []
-
- cases.append(
- pytest.param(
- OmniServerParams(
- model=MODEL,
- server_args=[
- *HSDP_ARGS,
- "--model-class-name",
- "LTX2DistilledPipeline",
- ],
- ),
- False,
- id="t2v_hsdp",
- marks=PARALLEL_MARKS,
- )
- )
-
- cases.append(
- pytest.param(
- OmniServerParams(
- model=MODEL,
- server_args=[
- *HSDP_ARGS,
- "--model-class-name",
- "LTX2DistilledPipeline",
- ],
- ),
- True,
- id="i2v_hsdp",
- marks=PARALLEL_MARKS,
- )
- )
-
- return cases
-
-
-@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
-def test_ltx2_two_stage_hsdp(
- omni_server: OmniServer,
- is_i2v: bool,
- openai_client: OpenAIClientHandler,
-):
- # Keep CI small while exercising the fixed 8+3-step distilled recipe.
- # Height and width describe the final output; Stage 1 runs at half size.
- form_data = {
- "prompt": PROMPT,
- "model": omni_server.model,
- "height": 128,
- "width": 128,
- "num_frames": 9,
- "fps": 8,
- "num_inference_steps": 8,
- "seed": 42,
- }
-
- request_config = {
- "model": omni_server.model,
- "form_data": form_data,
- }
-
- if is_i2v:
- request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
-
- openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_mimo_audio_expansion.py b/tests/e2e/online_serving/test_mimo_audio_expansion.py
index e5722ddb8..51f4d7ed9 100644
--- a/tests/e2e/online_serving/test_mimo_audio_expansion.py
+++ b/tests/e2e/online_serving/test_mimo_audio_expansion.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
-E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100).
+E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100 · Single-GPU).
"""
import os
@@ -77,7 +77,7 @@ def get_max_batch_size(size_type="few"):
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(omni_server, openai_client) -> None:
+def test_text_to_text_001(omni_server, online_client) -> None:
"""
Test text input processing and text-only output generation via OpenAI API.
Deploy Setting: default yaml
@@ -95,13 +95,13 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.skip(reason="CI failed 8571")
-def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
+def test_audio_to_text_audio_001(omni_server, online_client) -> None:
"""
Test audio and text input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -126,4 +126,4 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
},
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
diff --git a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py
index 8ab327022..97a634906 100644
--- a/tests/e2e/online_serving/test_ming_flash_omni_expansion.py
+++ b/tests/e2e/online_serving/test_ming_flash_omni_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E online serving tests for Ming-flash-omni-2.0 model (Thinker stage).
Tests multimodal understanding via OpenAI-compatible API.
@@ -20,7 +20,7 @@ from tests.helpers.stage_config import get_deploy_config_path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-pytestmark = [pytest.mark.omni, pytest.mark.full_model]
+pytestmark = [pytest.mark.omni, pytest.mark.slow]
_SKIP_NEED_4_H100_NOT_CI = pytest.mark.skip(
reason="Requires 4x H100 GPUs; skipped in CI for now.",
@@ -70,7 +70,7 @@ def get_max_batch_size(size_type="few"):
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(omni_server, openai_client) -> None:
+def test_text_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: text
Output Modal: text
@@ -90,13 +90,13 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_stream_001(omni_server, openai_client) -> None:
+def test_text_to_text_stream_001(omni_server, online_client) -> None:
"""
Input Modal: text
Output Modal: text
@@ -116,13 +116,13 @@ def test_text_to_text_stream_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_image_to_text_001(omni_server, openai_client) -> None:
+def test_image_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: image + text
Output Modal: text
@@ -143,13 +143,13 @@ def test_image_to_text_001(omni_server, openai_client) -> None:
"modalities": ["text"],
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_to_text_001(omni_server, openai_client) -> None:
+def test_audio_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: audio + text
Output Modal: text
@@ -170,13 +170,13 @@ def test_audio_to_text_001(omni_server, openai_client) -> None:
"modalities": ["text"],
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_video_to_text_001(omni_server, openai_client) -> None:
+def test_video_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: video + text
Output Modal: text
@@ -197,13 +197,13 @@ def test_video_to_text_001(omni_server, openai_client) -> None:
"modalities": ["text"],
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@_SKIP_NEED_4_H100_NOT_CI
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_mix_to_text_001(omni_server, openai_client) -> None:
+def test_mix_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: text + audio + image + video
Output Modal: text
@@ -228,4 +228,4 @@ def test_mix_to_text_001(omni_server, openai_client) -> None:
"modalities": ["text"],
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
diff --git a/tests/e2e/online_serving/test_ming_tts_expansion.py b/tests/e2e/online_serving/test_ming_tts_expansion.py
index d4fbb7026..4352a8966 100644
--- a/tests/e2e/online_serving/test_ming_tts_expansion.py
+++ b/tests/e2e/online_serving/test_ming_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E expansion tests for Ming-omni-tts online serving (nightly CI).
@@ -18,7 +18,6 @@ from tests.helpers.stage_config import get_deploy_config_path
pytestmark = [
pytest.mark.slow,
pytest.mark.tts,
- pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"),
]
MODEL = "inclusionAI/Ming-omni-tts-0.5B"
@@ -57,7 +56,7 @@ def get_prompt(prompt_type="zh"):
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True)
-def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None:
+def test_text_to_audio_non_streaming_001(omni_server, online_client) -> None:
"""
Deploy Setting: ming_tts.yaml with --no-async-chunk
Input Modal: text
@@ -72,12 +71,12 @@ def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None:
"response_format": "wav",
"timeout": 300.0,
}
- openai_client.send_audio_speech_request(request_config, request_num=2)
+ online_client.send_audio_speech_request(request_config, request_num=2)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True)
-def test_text_to_audio_streaming_001(omni_server, openai_client) -> None:
+def test_text_to_audio_streaming_001(omni_server, online_client) -> None:
"""
Deploy Setting: ming_tts.yaml (async_chunk=true)
Input Modal: text + voice
@@ -94,4 +93,4 @@ def test_text_to_audio_streaming_001(omni_server, openai_client) -> None:
"response_format": "wav",
"timeout": 300.0,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py
index 6953ef291..cb8995206 100644
--- a/tests/e2e/online_serving/test_minicpmo_4_5.py
+++ b/tests/e2e/online_serving/test_minicpmo_4_5.py
@@ -1,12 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
Talker, and Code2Wav stages.
"""
+import json
import os
+import tempfile
+from copy import deepcopy
+from pathlib import Path
import pytest
+from vllm.logger import DEFAULT_LOGGING_CONFIG
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -19,10 +27,36 @@ _MODEL = "openbmb/MiniCPM-o-4_5"
_CI_DEPLOY = modify_stage_config(
get_deploy_config_path("minicpmo_4_5.yaml"),
updates={
- "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}}
+ "stages": {
+ 0: {"default_sampling_params.max_tokens": 64},
+ 1: {
+ "default_sampling_params.max_tokens": 1024,
+ # Content-consistency assertions must not depend on a random
+ # codec trajectory: different valid TTS samples can append an
+ # audible tail even when Thinker text is identical.
+ "default_sampling_params.temperature": 0.0,
+ },
+ },
},
)
+_PROMPT_LOG_TEMP_DIR = tempfile.TemporaryDirectory(prefix="minicpmo45-prompt-selection-")
+_PROMPT_SELECTION_LOG = Path(_PROMPT_LOG_TEMP_DIR.name) / "prompt-selection.log"
+_PROMPT_LOGGING_CONFIG_PATH = Path(_PROMPT_LOG_TEMP_DIR.name) / "logging.json"
+_PROMPT_LOGGING_CONFIG = deepcopy(DEFAULT_LOGGING_CONFIG)
+_PROMPT_LOGGING_CONFIG["handlers"]["minicpmo_prompt_selection"] = {
+ "class": "logging.FileHandler",
+ "filename": str(_PROMPT_SELECTION_LOG),
+ "formatter": "vllm",
+ "level": "DEBUG",
+}
+_PROMPT_LOGGING_CONFIG["loggers"]["vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav"] = {
+ "handlers": ["minicpmo_prompt_selection"],
+ "level": "DEBUG",
+ "propagate": True,
+}
+_PROMPT_LOGGING_CONFIG_PATH.write_text(json.dumps(_PROMPT_LOGGING_CONFIG), encoding="utf-8")
+
test_params = [
pytest.param(
OmniServerParams(
@@ -30,6 +64,10 @@ test_params = [
stage_config_path=_CI_DEPLOY,
use_stage_cli=False,
server_args=["--trust-remote-code", "--async-chunk"],
+ env_dict={
+ "VLLM_CONFIGURE_LOGGING": "1",
+ "VLLM_LOGGING_CONFIG_PATH": str(_PROMPT_LOGGING_CONFIG_PATH),
+ },
),
id="async_chunk",
),
@@ -65,6 +103,27 @@ def get_max_batch_size(size_type="few"):
return batch_sizes.get(size_type, 5)
+# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning.
+_TTS_EXTRA_BODY = {
+ "chat_template_kwargs": {
+ "use_tts_template": True,
+ "enable_thinking": False,
+ }
+}
+
+
+def _prompt_selection_log_offset() -> int:
+ return _PROMPT_SELECTION_LOG.stat().st_size if _PROMPT_SELECTION_LOG.exists() else 0
+
+
+def _read_prompt_selection_log(offset: int) -> str:
+ if not _PROMPT_SELECTION_LOG.exists():
+ return ""
+ with open(_PROMPT_SELECTION_LOG, "rb") as log_file:
+ log_file.seek(offset)
+ return log_file.read().decode("utf-8")
+
+
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@@ -111,11 +170,57 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
"messages": messages,
"stream": True,
"key_words": {"audio": ["Beijing"]},
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config)
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_with_reference_audio(omni_server, openai_client) -> None:
+ request_config = {
+ "model": omni_server.model,
+ "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()),
+ "stream": True,
+ "modalities": ["text", "audio"],
+ "extra_body": {
+ **_TTS_EXTRA_BODY,
+ "ref_audio": f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}",
+ },
+ }
+
+ log_offset = _prompt_selection_log_offset()
+ responses = openai_client.send_omni_request(request_config)
+ prompt_log = _read_prompt_selection_log(log_offset)
+
+ assert responses[0].success
+ assert responses[0].audio_bytes
+ assert "prompt_cache_id=runtime-ref-" in prompt_log
+ assert "minicpmo45_ref_" in prompt_log
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_with_default_reference(omni_server, openai_client) -> None:
+ request_config = {
+ "model": omni_server.model,
+ "messages": dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt()),
+ "stream": True,
+ "modalities": ["text", "audio"],
+ "extra_body": _TTS_EXTRA_BODY,
+ }
+
+ responses = openai_client.send_omni_request(request_config)
+
+ assert responses[0].success
+ assert responses[0].audio_bytes
+
+
@pytest.mark.full_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@@ -141,12 +246,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
"stream": True,
"key_words": {"text": ["Beijing"]},
"modalities": ["text", "audio"],
- "extra_body": {
- "chat_template_kwargs": {
- "use_tts_template": True,
- "enable_thinking": False,
- }
- },
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -175,6 +275,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None:
"model": omni_server.model,
"messages": messages,
"stream": True,
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -203,6 +304,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
"model": omni_server.model,
"messages": messages,
"stream": True,
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -236,6 +338,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"model": omni_server.model,
"messages": messages,
"stream": True,
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
index d63811022..bbdf85c09 100644
--- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
+++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""CI coverage for the MiniCPM-o 4.5 native-duplex Realtime API."""
from __future__ import annotations
@@ -13,6 +16,7 @@ import websockets
from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
SERVER_PARAMS,
demo_args,
+ duplex_camera_frames,
multi_session_args,
realtime_url,
resolve_ref_audio,
@@ -27,10 +31,16 @@ from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import
)
from tests.helpers.mark import hardware_test
from vllm_omni.experimental.fullduplex.client import build_realtime_url
+from vllm_omni.experimental.fullduplex.video_stacking import concat_frames_b64
pytestmark = pytest.mark.omni
+def _assert_positive_int(value: object) -> None:
+ assert isinstance(value, int)
+ assert value > 0
+
+
def _assert_request_metrics(metrics: object, *, expected_count: int) -> None:
assert isinstance(metrics, list)
assert len(metrics) == expected_count
@@ -107,8 +117,8 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[
@pytest.mark.core_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
-def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
- ref_audio = resolve_ref_audio(model_prefix)
+def test_duplex_websocket_protocol_smoke(omni_server) -> None:
+ ref_audio = resolve_ref_audio()
events = asyncio.run(
_run_protocol_smoke(
url=realtime_url(omni_server),
@@ -125,11 +135,11 @@ def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
-def test_duplex_single_session_response_required(omni_server, model_prefix: str, tmp_path: Path) -> None:
+def test_duplex_single_session_response_required(omni_server, tmp_path: Path) -> None:
args = demo_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
- ref_audio=resolve_ref_audio(model_prefix),
+ ref_audio=resolve_ref_audio(),
output_dir=tmp_path / "single_session",
)
args.turns = 2
@@ -139,7 +149,7 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
args.turn_duration_ms = [args.first_turn_ms] * args.turns
result = asyncio.run(run_demo(args))
assert result["ok"] is True
- assert result["audio_delta_count"] > 0
+ _assert_positive_int(result["audio_delta_count"])
assert result["done_count"] == 2
assert result["error_count"] == 0
assert result["all_audio_responses_have_transcript"] is True
@@ -151,13 +161,51 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
-def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str, tmp_path: Path) -> None:
+def test_duplex_single_session_video_input(omni_server, tmp_path: Path) -> None:
+ """Audio plus a 1 fps camera track, the omni-duplex video contract.
+
+ Frames ride the audio appends, so the response contract is unchanged: what
+ this covers is that interleaved vision input keeps the turn intact instead
+ of stalling or erroring out mid-segment.
+ """
+ args = demo_args(
+ omni_server=omni_server,
+ input_wav=validated_input_wav(),
+ ref_audio=resolve_ref_audio(),
+ output_dir=tmp_path / "video_input",
+ )
+ args.turns = 2
+ args.turn_duration_ms = [args.first_turn_ms] * args.turns
+ frames = duplex_camera_frames(seconds=4, cache_dir=tmp_path / "camera")
+ args.video_frames_b64 = frames
+ # Each unit also carries a composite of its interior sub-frames, so the
+ # append exercises the official two-image frame_list that carries motion.
+ args.video_stacked_frames_b64 = [concat_frames_b64([frames[index]] * 2) for index in range(len(frames))]
+
+ result = asyncio.run(run_demo(args))
+
+ assert result["ok"] is True
+ assert result["video_frame_count"] == 4
+ assert result["video_stacked_frame_count"] == 4
+ _assert_positive_int(result["audio_delta_count"])
+ assert result["done_count"] == 2
+ assert result["error_count"] == 0
+ assert result["all_audio_responses_have_transcript"] is True
+ assert result["transcript_delta_done_ok"] is True
+ _assert_request_metrics(result["request_metrics"], expected_count=2)
+ _assert_session_metrics(result["session_metrics"], expected_count=2)
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_two_sessions_resume_and_takeover(omni_server, tmp_path: Path) -> None:
result = asyncio.run(
run_multi_session(
multi_session_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
- ref_audio=resolve_ref_audio(model_prefix),
+ ref_audio=resolve_ref_audio(),
output_dir=tmp_path / "multi_session",
response_required=True,
)
diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
index 68204fe28..5b35c4b53 100644
--- a/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+++ b/tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Nightly lifecycle coverage for the MiniCPM-o 4.5 native-duplex API."""
from __future__ import annotations
@@ -12,6 +15,7 @@ import pytest
from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
SERVER_PARAMS,
SOFT_INTERRUPT_SHA256,
+ deploy_max_sessions,
multi_session_args,
realtime_url,
resolve_ref_audio,
@@ -21,6 +25,7 @@ from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
run_lifecycle_probes,
)
+from tests.e2e.online_serving.run_minicpmo_realtime_duplex_server_vad import run_server_vad_interrupt
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import (
run_soft_interrupt,
)
@@ -31,11 +36,11 @@ pytestmark = [pytest.mark.full_model, pytest.mark.omni]
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
-def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_path: Path) -> None:
+def test_duplex_admission_and_expiry_reaper(omni_server, tmp_path: Path) -> None:
args = multi_session_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
- ref_audio=resolve_ref_audio(model_prefix),
+ ref_audio=resolve_ref_audio(),
output_dir=tmp_path / "admission_expiry",
response_required=False,
)
@@ -43,7 +48,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_
args.disconnect_session_index = None
args.takeover_session_index = None
args.expire_session_index = 0
- args.verify_admission_limit = 2
+ args.verify_admission_limit = deploy_max_sessions()
result = asyncio.run(run_lifecycle_probes(args))
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
@@ -59,7 +64,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_
# so it never reaches a mid-stream decision point.
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
-def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None:
+def test_duplex_soft_interrupt(omni_server, tmp_path: Path) -> None:
input_wav = validated_soft_interrupt_wav()
result = asyncio.run(
run_soft_interrupt(
@@ -67,7 +72,7 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -
url=realtime_url(omni_server),
model=omni_server.model,
input_wav=str(input_wav),
- ref_audio=str(resolve_ref_audio(model_prefix)),
+ ref_audio=str(resolve_ref_audio()),
output_dir=str(tmp_path / "soft_interrupt"),
summary_output=None,
chunk_ms=200,
@@ -88,3 +93,22 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -
assert result["response_lifecycle_ok"] is True
assert result["response_audio_contract_ok"] is True
assert result["followup_response_transcript_ok"] is True
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
+def test_duplex_server_vad_hard_interrupt(omni_server) -> None:
+ result = asyncio.run(
+ run_server_vad_interrupt(
+ SimpleNamespace(
+ url=realtime_url(omni_server),
+ model=omni_server.model,
+ input_wav=str(validated_input_wav()),
+ interrupt_wav=str(validated_soft_interrupt_wav()),
+ ref_audio=str(resolve_ref_audio()),
+ chunk_ms=200,
+ timeout_s=180.0,
+ )
+ )
+ )
+ assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py
index e84985e3e..e7fa50c16 100644
--- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py
+++ b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py
@@ -1,7 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""E2E online expansion tests for MiniCPM-o 4.5.
These cover modality combinations, separate Code2Wav behavior, sequential
-request isolation, longer audio output, and Chinese speech.
+request isolation, longer audio output, long-form generation, and Chinese speech.
Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and
``async_chunk`` exercises ``--async-chunk``.
@@ -75,6 +78,15 @@ def get_max_batch_size(size_type="few"):
return batch_sizes.get(size_type, 5)
+# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning.
+_TTS_EXTRA_BODY = {
+ "chat_template_kwargs": {
+ "use_tts_template": True,
+ "enable_thinking": False,
+ }
+}
+
+
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_video_to_text_001(omni_server, openai_client) -> None:
@@ -127,6 +139,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None:
"model": omni_server.model,
"messages": messages_1,
"stream": True,
+ "extra_body": _TTS_EXTRA_BODY,
},
request_num=1,
)
@@ -138,6 +151,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None:
"messages": messages_2,
"stream": True,
"key_words": {"text": ["Beijing"]},
+ "extra_body": _TTS_EXTRA_BODY,
},
request_num=1,
)
@@ -164,20 +178,36 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
"messages": messages,
"modalities": ["text", "audio"],
"stream": True,
- # Delimit the assistant answer with the TTS template so the Talker
- # does not spend its codec-token budget speaking hidden reasoning.
- "extra_body": {
- "chat_template_kwargs": {
- "use_tts_template": True,
- "enable_thinking": False,
- }
- },
+ "extra_body": _TTS_EXTRA_BODY,
"key_words": {"audio": ["Beijing"]},
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_long_form_001(omni_server, openai_client) -> None:
+ """
+ Input Modal: text only (long-form generation prompt).
+ Output Modal: text, audio (default ``modalities``);
+ Input Setting: stream=True
+ Datasets: single request
+ """
+ messages = dummy_messages_from_mix_data(
+ system_prompt=get_system_prompt(),
+ content_text="帮我讲一个100字的故事.",
+ )
+
+ request_config = {
+ "model": omni_server.model,
+ "messages": messages,
+ "stream": True,
+ "extra_body": _TTS_EXTRA_BODY,
+ }
+ openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+
+
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_chinese_text_to_audio(omni_server, openai_client) -> None:
@@ -193,5 +223,6 @@ def test_chinese_text_to_audio(omni_server, openai_client) -> None:
"messages": messages,
"stream": True,
"key_words": {"text": ["北京"]},
+ "extra_body": _TTS_EXTRA_BODY,
}
openai_client.send_omni_request(request_config)
diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
index b2e40c58f..9569078d2 100644
--- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
+++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
import base64
import importlib.util
@@ -15,6 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
DRIVER_DIR = Path(__file__).resolve().parent
HELPERS_DIR = DRIVER_DIR / "helpers"
DEMO_PATH = HELPERS_DIR / "minicpmo_realtime_duplex_scenarios.py"
+EXAMPLE_DEMO_PATH = DRIVER_DIR.parents[2] / "examples/online_serving/minicpmo/realtime_duplex_demo.py"
MULTI_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_multi_session.py"
PAIR_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_demo_pair.py"
SOFT_INTERRUPT_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_soft_interrupt.py"
@@ -47,6 +51,18 @@ def _load_pair_demo_module():
return module
+def _load_example_demo_module():
+ spec = importlib.util.spec_from_file_location(
+ "minicpmo_realtime_duplex_example_demo_test",
+ EXAMPLE_DEMO_PATH,
+ )
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = module
+ spec.loader.exec_module(module)
+ return module
+
+
def _load_soft_interrupt_demo_module():
spec = importlib.util.spec_from_file_location(
"minicpmo_realtime_duplex_soft_interrupt_test",
@@ -524,6 +540,111 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp
assert summary["response_summaries"][1]["transcript"] == "一加一等于二。"
+def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_path):
+ """A follow-up response can drain a residual model unit past the final commit.
+
+ Sequence: ``r1.done -> listen -> r2.created -> ... -> commit -> r2.done -> listen``.
+
+ The globally-last ``response.done`` (r2) lands *after* the commit, so anchoring
+ ``listen_after_response_before_commit`` on it leaves the ``(last_done, commit)``
+ interval empty and fails the contract even though the turn yielded the floor
+ correctly. The check must anchor on the last done strictly before the commit
+ (r1.done), with the listen at 3.1 s satisfying the sandwich.
+ """
+ demo = _load_soft_interrupt_demo_module()
+ output = tmp_path / "soft_interrupt_residual_drain"
+ output.mkdir()
+ first_response_id = "resp-first"
+ second_response_id = "resp-second"
+ events = [
+ {"type": "response.listen", "_client_received_at_s": 1.0},
+ {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0},
+ {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0},
+ {
+ "type": "response.audio.delta",
+ "response_id": first_response_id,
+ "delta": "AAAA",
+ "_client_received_at_s": 2.1,
+ },
+ {
+ "type": "response.audio_transcript.delta",
+ "response_id": first_response_id,
+ "delta": "中国古代四大发明",
+ "_client_received_at_s": 2.1,
+ },
+ {
+ "type": "response.audio.delta",
+ "response_id": first_response_id,
+ "delta": "AAAA",
+ "_client_received_at_s": 2.6,
+ },
+ {
+ "type": "response.audio_transcript.delta",
+ "response_id": first_response_id,
+ "delta": "是造纸术。",
+ "_client_received_at_s": 2.6,
+ },
+ {"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 3.0},
+ {"type": "response.listen", "_client_received_at_s": 3.1},
+ {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.5},
+ {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 3.5},
+ {
+ "type": "response.audio.delta",
+ "response_id": second_response_id,
+ "delta": "AAAA",
+ "_client_received_at_s": 3.6,
+ },
+ {
+ "type": "response.audio_transcript.delta",
+ "response_id": second_response_id,
+ "delta": "一加一等于",
+ "_client_received_at_s": 3.6,
+ },
+ {
+ "type": "response.audio.delta",
+ "response_id": second_response_id,
+ "delta": "AAAA",
+ "_client_received_at_s": 3.7,
+ },
+ {
+ "type": "response.audio_transcript.delta",
+ "response_id": second_response_id,
+ "delta": "二。",
+ "_client_received_at_s": 3.7,
+ },
+ # The final commit lands while the follow-up response is still draining.
+ {"type": "input_audio_buffer.committed", "_client_received_at_s": 4.0},
+ {"type": "response.done", "response_id": second_response_id, "_client_received_at_s": 4.5},
+ {"type": "response.listen", "_client_received_at_s": 4.6},
+ ]
+ (output / "events.jsonl").write_text(
+ "".join(demo.json.dumps(event) + "\n" for event in events),
+ encoding="utf-8",
+ )
+ (output / "result.json").write_text(
+ demo.json.dumps({"ok": True, "response_ids": [first_response_id, second_response_id]}),
+ encoding="utf-8",
+ )
+
+ summary = demo.summarize_artifacts(
+ output_dir=output,
+ validation_mode="response-required",
+ min_responses=2,
+ min_audio_deltas_per_response=2,
+ expect_followup_response_substring="一加一等于二",
+ )
+
+ assert summary["ok"] is True
+ # Anchored on r1.done (the last done before the commit), not the late r2.done.
+ assert summary["listen_after_response_before_commit"] is True
+ # listen_after_last_done still keys off the globally-last done, hence the final listen.
+ assert summary["listen_after_last_done"] is True
+ assert summary["final_listen_after_commit"] is True
+ assert summary["second_response_before_final_commit"] is True
+ assert summary["followup_response_transcript_ok"] is True
+ assert summary["response_summaries"][1]["transcript"] == "一加一等于二。"
+
+
def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp_path):
demo = _load_soft_interrupt_demo_module()
output = tmp_path / "model_policy"
@@ -758,6 +879,34 @@ def test_realtime_duplex_multi_session_rejects_cross_session_response_identity()
)
+def test_realtime_duplex_multi_session_rejects_overlapping_expected_tokens(monkeypatch, capsys):
+ demo = _load_multi_demo_module()
+ monkeypatch.setattr(
+ demo.sys,
+ "argv",
+ [
+ "multi.py",
+ "--sessions",
+ "2",
+ "--input-wav",
+ "fallback.wav",
+ "--session-input-wav",
+ "a.wav",
+ "--session-input-wav",
+ "b.wav",
+ "--session-expected-token",
+ "cat",
+ "--session-expected-token",
+ "cater",
+ ],
+ )
+
+ with pytest.raises(SystemExit):
+ demo.parse_args()
+
+ assert "must not overlap after normalization" in capsys.readouterr().err
+
+
def test_realtime_duplex_multi_session_reads_nested_terminal_identity():
demo = _load_multi_demo_module()
@@ -768,6 +917,67 @@ def test_realtime_duplex_multi_session_reads_nested_terminal_identity():
)
+def test_realtime_duplex_synchronized_start_gate_releases_all_participants():
+ demo = _load_multi_demo_module()
+
+ async def exercise():
+ gate = demo._SynchronizedStartGate(3, timeout_s=0.5)
+ released = []
+
+ async def participant(index):
+ await gate.wait()
+ released.append(index)
+
+ await asyncio.wait_for(
+ asyncio.gather(*(participant(index) for index in range(3))),
+ timeout=1.0,
+ )
+ return released
+
+ assert set(asyncio.run(exercise())) == {0, 1, 2}
+
+
+def test_realtime_duplex_synchronized_start_gate_aborts_waiters_on_failure(monkeypatch):
+ demo = _load_multi_demo_module()
+
+ async def fake_run_demo(args):
+ if args.fail:
+ raise RuntimeError("session creation failed")
+ await args.start_barrier.wait()
+ return {"ok": True}
+
+ monkeypatch.setattr(demo, "run_demo", fake_run_demo)
+
+ async def exercise():
+ gate = demo._SynchronizedStartGate(2, timeout_s=0.5)
+ return await asyncio.wait_for(
+ asyncio.gather(
+ demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=True)),
+ demo._run_demo_with_start_gate(SimpleNamespace(start_barrier=gate, fail=False)),
+ return_exceptions=True,
+ ),
+ timeout=1.0,
+ )
+
+ failed, aborted = asyncio.run(exercise())
+ assert isinstance(failed, RuntimeError)
+ assert str(failed) == "session creation failed"
+ assert isinstance(aborted, RuntimeError)
+ assert "synchronized start aborted" in str(aborted)
+ assert "session creation failed" in str(aborted)
+
+
+def test_realtime_duplex_synchronized_start_gate_has_bounded_wait():
+ demo = _load_multi_demo_module()
+
+ async def exercise():
+ gate = demo._SynchronizedStartGate(2, timeout_s=0.01)
+ await gate.wait()
+
+ with pytest.raises(TimeoutError, match=r"1/2 sessions ready"):
+ asyncio.run(exercise())
+
+
def test_realtime_duplex_demo_resolves_distinct_turn_inputs():
demo = _load_demo_module()
primary = Path("first.wav")
@@ -1721,6 +1931,241 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak():
assert len(ws.messages) == 10
+def _sent_video_frames(messages: list[dict[str, object]]) -> list[list[str]]:
+ frames: list[list[str]] = []
+ for message in messages:
+ raw = message.get("video_frames")
+ if not isinstance(raw, list):
+ continue
+ frames.append([frame for frame in raw if isinstance(frame, str)])
+ return frames
+
+
+def _decode_jpeg_size(frame_b64: str) -> tuple[int, int]:
+ import io
+
+ from PIL import Image
+
+ return Image.open(io.BytesIO(base64.b64decode(frame_b64))).size
+
+
+class _RecordingWebSocket:
+ def __init__(self, demo):
+ self._demo = demo
+ self.messages: list[dict[str, object]] = []
+
+ async def send(self, payload):
+ self.messages.append(self._demo.json.loads(payload))
+
+
+def _send_seconds_of_audio(demo, *, seconds: int, frames: list[str], stacked: list[str | None] | None = None):
+ ws = _RecordingWebSocket(demo)
+ asyncio.run(
+ demo._send_pcm16(
+ ws,
+ b"\x01\x00" * (demo.PCM16_SAMPLE_RATE * seconds),
+ chunk_ms=200,
+ realtime_delay=False,
+ frames_b64=frames,
+ stacked_frames_b64=stacked,
+ )
+ )
+ return ws.messages
+
+
+def test_realtime_duplex_demo_streams_one_video_frame_per_model_unit():
+ demo = _load_demo_module()
+
+ messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2", "f3"])
+
+ # Frame k rides the append that closes model unit k, so Stage0 can bind it
+ # to that unit's audio. 3 s of audio closes units at 1030 ms and 2030 ms and
+ # leaves a residual too short for a third, so only two frames go out.
+ assert _sent_video_frames(messages) == [["f0"], ["f1"]]
+ frame_indices = [index for index, message in enumerate(messages) if "video_frames" in message]
+ assert frame_indices == [5, 10]
+ # The carrying appends are exactly the ones that reach a unit boundary.
+ assert [messages[index]["audio_end_ms"] for index in frame_indices] == [1200, 2200]
+
+
+def test_realtime_duplex_demo_never_sends_a_frame_before_its_unit_can_close():
+ """Regression: frame 0 used to ride the very first 200 ms append.
+
+ Stage0 cannot close a unit until 1030 ms of audio has arrived, and it does
+ not carry frames across appends, so a frame sent that early was silently
+ dropped and every later frame ended up one unit ahead of its audio.
+ """
+ demo = _load_demo_module()
+
+ messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2"])
+
+ carried = [message for message in messages if "video_frames" in message]
+ assert carried, "expected at least one frame to be sent"
+ for message in carried:
+ assert message["audio_end_ms"] >= demo.duplex_unit_boundary_ms(0)
+
+
+def test_realtime_duplex_demo_holds_the_last_video_frame_when_audio_outlives_the_clip():
+ demo = _load_demo_module()
+
+ # 4 s of audio closes three units (1030/2030/3030 ms); the two-frame clip
+ # holds its last frame for the third.
+ assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=4, frames=["a", "b"])) == [["a"], ["b"], ["b"]]
+ # A still image is a one-element clip and therefore repeats every unit.
+ assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=3, frames=["still"])) == [["still"], ["still"]]
+
+
+def test_realtime_duplex_demo_sends_each_units_stacked_composite_next_to_its_base_frame():
+ demo = _load_demo_module()
+
+ messages = _send_seconds_of_audio(
+ demo,
+ seconds=4,
+ frames=["f0", "f1", "f2"],
+ stacked=["s0", "s1", None],
+ )
+
+ # Official pairing is frame_list=[base, composite of that unit's interior],
+ # so the composite belongs to the same unit as the base beside it -- never
+ # to the previous one. A unit without interior sub-frames sends base alone.
+ assert _sent_video_frames(messages) == [["f0", "s0"], ["f1", "s1"], ["f2"]]
+
+
+def test_minicpmo_duplex_camera_fixture_returns_flat_base_frame_track(tmp_path, monkeypatch):
+ from tests.e2e.online_serving.helpers import minicpmo_4_5_duplex as fixtures
+ from tests.e2e.online_serving.helpers import minicpmo_realtime_duplex_scenarios as demo
+ from tests.helpers import media
+
+ video_path = tmp_path / "camera.mp4"
+ monkeypatch.setattr(
+ media,
+ "generate_synthetic_video",
+ lambda *args, **kwargs: {"file_path": str(video_path)},
+ )
+ monkeypatch.setattr(
+ demo,
+ "_video_frames_from_file",
+ lambda path: (["frame-0", "frame-1"], [None, None]),
+ )
+
+ frames = fixtures.duplex_camera_frames(seconds=2, cache_dir=tmp_path)
+
+ assert frames == ["frame-0", "frame-1"]
+
+
+def test_realtime_duplex_demo_decodes_a_video_file_into_one_frame_per_second(tmp_path):
+ pytest.importorskip("cv2")
+ pytest.importorskip("imageio")
+ from tests.helpers.media import generate_synthetic_video
+
+ demo = _load_demo_module()
+ # 30 fps synthetic clip: 90 frames is 3 s of camera input.
+ video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path)
+
+ frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"]))
+
+ assert len(frames) == 3
+ assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in frames)
+ assert len(set(frames)) == 3
+ # stack_frames defaults to 1, so there is nothing to stack.
+ assert stacked == [None, None, None]
+
+
+def test_realtime_duplex_demo_stacks_a_units_interior_subframes_into_one_composite(tmp_path):
+ pytest.importorskip("cv2")
+ pytest.importorskip("imageio")
+ from tests.helpers.media import generate_synthetic_video
+
+ demo = _load_demo_module()
+ video = generate_synthetic_video(64, 64, 90, cache_dir=tmp_path)
+
+ frames, stacked = demo._resolve_video_frames(SimpleNamespace(input_video=video["file_path"], stack_frames=5))
+
+ # stack_frames=5 raises the visual refresh rate to 5 fps, but the four
+ # interior sub-frames of each second collapse into a single composite, so
+ # the wire still carries 2 images per unit and the audio cadence is
+ # untouched.
+ assert len(frames) == 3
+ assert len(stacked) == 3
+ assert all(frame is not None for frame in stacked)
+ assert all(base64.b64decode(frame, validate=True)[:3] == b"\xff\xd8\xff" for frame in stacked)
+ composite = _decode_jpeg_size(stacked[0])
+ assert composite != _decode_jpeg_size(frames[0])
+
+
+def test_example_demo_uses_external_wav_duration_and_skips_video_soundtrack(tmp_path, monkeypatch):
+ demo = _load_example_demo_module()
+ wav = tmp_path / "external.wav"
+ demo.write_pcm16_wav(wav, b"\x00\x00" * (demo.PCM16_SAMPLE_RATE * 2), sample_rate_hz=demo.PCM16_SAMPLE_RATE)
+ soundtrack_calls: list[object] = []
+
+ def fake_soundtrack(*args, **kwargs):
+ soundtrack_calls.append(1)
+ raise AssertionError("video soundtrack must not be loaded when --input-wav is set")
+
+ def fake_frames(video_path, work_dir, *, fps, max_side, duration_s, stack_frames):
+ assert video_path == Path("/tmp/silent.mp4")
+ assert fps == 1.0
+ assert duration_s == pytest.approx(2.0)
+ return ["f0", "f1"], [None, None]
+
+ monkeypatch.setattr(demo, "_extract_video_soundtrack", fake_soundtrack)
+ monkeypatch.setattr(demo, "_extract_video_frames", fake_frames)
+
+ wav_out, frames, stacked = demo._resolve_duplex_av_inputs(
+ input_wav=str(wav),
+ input_video="/tmp/silent.mp4",
+ work_dir=tmp_path / "video_input",
+ fps=1.0,
+ max_side=0,
+ )
+
+ assert wav_out == str(wav)
+ assert frames == ["f0", "f1"]
+ assert stacked == [None, None]
+ assert soundtrack_calls == []
+
+
+def test_example_demo_extracts_video_soundtrack_only_when_wav_is_absent(tmp_path, monkeypatch):
+ demo = _load_example_demo_module()
+ extracted = tmp_path / "from_video.wav"
+ demo.write_pcm16_wav(extracted, b"\x00\x00" * demo.PCM16_SAMPLE_RATE, sample_rate_hz=demo.PCM16_SAMPLE_RATE)
+
+ def fake_extract(video_path, work_dir, *, fps, max_side, stack_frames):
+ assert video_path == Path("/tmp/clip.mp4")
+ return extracted, ["v0"], ["s0"]
+
+ monkeypatch.setattr(demo, "_extract_video_input", fake_extract)
+
+ wav_out, frames, stacked = demo._resolve_duplex_av_inputs(
+ input_wav=None,
+ input_video="/tmp/clip.mp4",
+ work_dir=tmp_path / "video_input",
+ fps=1.0,
+ max_side=0,
+ )
+
+ assert wav_out == str(extracted)
+ assert frames == ["v0"]
+ assert stacked == ["s0"]
+
+
+def test_realtime_duplex_demo_prefers_a_video_over_a_still_frame_image(tmp_path):
+ demo = _load_demo_module()
+ still = tmp_path / "frame.jpg"
+ still.write_bytes(b"\xff\xd8\xff-still")
+
+ # A still has no interior sub-frames, so it never stacks.
+ assert demo._resolve_video_frames(SimpleNamespace(frame_image=str(still))) == (
+ [base64.b64encode(still.read_bytes()).decode("ascii")],
+ [None],
+ )
+ assert demo._resolve_video_frames(
+ SimpleNamespace(frame_image=str(still), video_frames_b64=["preset"]),
+ ) == (["preset"], [None])
+ assert demo._resolve_video_frames(SimpleNamespace()) == ([], [])
+
+
def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch):
demo = _load_demo_module()
state = demo.DemoState()
diff --git a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
new file mode 100644
index 000000000..148520ff0
--- /dev/null
+++ b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MiniMax-H3 DLO + DP2 online-serving smoke test."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import io
+import json
+import os
+
+import av
+import pytest
+import requests
+
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = os.environ.get("VLLM_TEST_MINIMAX_H3_MODEL", "MiniMaxAI/MiniMax-H3")
+WIDTH = 1344
+HEIGHT = 768
+FPS = 24
+NUM_INFERENCE_STEPS = 4
+REQUEST_TIMEOUT_SECONDS = 1800
+H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+SERVER_ARGS = [
+ "--trust-remote-code",
+ "--task-type",
+ "fl2va",
+ "--num-gpus",
+ "2",
+ "--tensor-parallel-size",
+ "1",
+ "--data-parallel-size",
+ "2",
+ "--request-batch-max-wait-ms",
+ "500",
+ "--usp",
+ "1",
+ "--ring",
+ "1",
+ "--text-encoder-tp-size",
+ "1",
+ "--vae-patch-parallel-size",
+ "1",
+ "--vae-parallel-mode",
+ "tile",
+ "--vae-use-tiling",
+ "--enable-distributed-layerwise-offload",
+]
+
+
+def _assert_audio_stream_present(video: bytes) -> None:
+ """Assert that the generated MP4 contains decodable audio samples."""
+ with av.open(io.BytesIO(video)) as container:
+ audio_streams = [stream for stream in container.streams if stream.type == "audio"]
+ assert audio_streams, "MiniMax-H3 MP4 has no audio stream"
+ audio_frame = next(container.decode(audio=0), None)
+ assert audio_frame is not None and audio_frame.samples > 0, "MiniMax-H3 MP4 audio stream is empty"
+
+
+def _run_t2va_request(client: OpenAIClientHandler, seed: int) -> bytes:
+ """Submit one synchronous T2VA request and return its MP4 body."""
+ request_data = {
+ "model": MODEL,
+ "prompt": "In a snowy blue-purple forest, a traveler walks past a sleeping giant; footsteps crunch in the snow while the creature softly breathes.",
+ "width": str(WIDTH),
+ "height": str(HEIGHT),
+ "fps": str(FPS),
+ "num_inference_steps": str(NUM_INFERENCE_STEPS),
+ "flow_shift": "12",
+ "seed": str(seed),
+ "extra_params": json.dumps(
+ {
+ "task": "t2va",
+ "duration": 4.0,
+ "aspect_ratio": "16:9",
+ "audio_flow_shift": 3.0,
+ },
+ separators=(",", ":"),
+ ),
+ }
+ response = requests.post(
+ f"{client.base_url.rstrip('/')}/v1/videos/sync",
+ data=request_data,
+ headers={"Accept": "video/mp4"},
+ timeout=REQUEST_TIMEOUT_SECONDS,
+ )
+ response.raise_for_status()
+ assert response.headers.get("content-type", "").startswith("video/mp4")
+ assert response.content, "MiniMax-H3 returned an empty video body"
+ return response.content
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@pytest.mark.parametrize(
+ "omni_server",
+ [
+ pytest.param(
+ OmniServerParams(
+ model=MODEL,
+ server_args=SERVER_ARGS,
+ stage_init_timeout=1800,
+ init_timeout=1800,
+ ),
+ id="minimax_h3_dlo_dp2_t2va",
+ marks=H100_TWO_CARD_MARKS,
+ )
+ ],
+ indirect=True,
+)
+def test_minimax_h3_dlo_dp2_t2va(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+ """Validate one complete DLO all-gather DP2 wave with two T2VA jobs."""
+ with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+ futures = [executor.submit(_run_t2va_request, openai_client, seed) for seed in (1101, 1102)]
+ videos = [future.result() for future in futures]
+
+ for video in videos:
+ assert_video_valid(video, width=WIDTH, height=HEIGHT, fps=FPS)
+ _assert_audio_stream_present(video)
diff --git a/tests/e2e/online_serving/test_minimax_music3.py b/tests/e2e/online_serving/test_minimax_music3.py
new file mode 100644
index 000000000..3d8718e70
--- /dev/null
+++ b/tests/e2e/online_serving/test_minimax_music3.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for MiniMax Music 3 text-to-music.
+
+This is a music model on the speech endpoint, so the request shape differs
+from every other TTS model here: ``input`` carries the lyrics, ``instructions``
+carries the caption that decides genre and arrangement, and both are required.
+
+Weekly rather than per-PR: the checkpoint is 27 GB across five components and a
+single request occupies two decode rows, because classifier-free guidance is
+mandatory for this model. Rejection cases live in
+``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+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
+
+MODEL = "MiniMaxAI/MiniMax-Music3"
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0
+
+# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The
+# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room
+# for the model ending a song before the cap.
+_FRAMES_10S = 250
+_FRAMES_30S = 750
+_MIN_BYTES_10S = 500_000
+_MIN_BYTES_30S = 1_500_000
+
+LYRICS = (
+ "[Verse]\nWalking down the empty street at midnight\n"
+ "Streetlights flicker like a broken dream\n"
+ "[Chorus]\nAnd I keep on walking\nTill the morning finds me"
+)
+
+CAPTION = (
+ "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, "
+ "soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. "
+ "Intimate bedroom production, gentle tape saturation, no bright cymbals."
+)
+
+
+tts_server_params = [
+ pytest.param(
+ OmniServerParams(
+ model=MODEL,
+ stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
+ server_args=["--trust-remote-code", "--disable-log-stats"],
+ ),
+ id="minimax_music3",
+ )
+]
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_001(omni_server, openai_client) -> None:
+ """Baseline: lyrics plus caption in, one ten-second song out.
+
+ Deploy Setting: minimax_music3.yaml
+ Input Modal: text (lyrics) + text (caption)
+ Output Modal: audio, 32 kHz stereo
+ Input Setting: stream=False, seed pinned, 250 frames
+ """
+ request_config = {
+ "model": omni_server.model,
+ "input": LYRICS,
+ "instructions": CAPTION,
+ "seed": 1,
+ "max_new_tokens": _FRAMES_10S,
+ "stream": False,
+ "response_format": "wav",
+ "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+ "min_audio_bytes": _MIN_BYTES_10S,
+ }
+ openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.slow
+@pytest.mark.tts
+@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1})
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_music_multi_window_002(omni_server, openai_client) -> None:
+ """A song long enough to span several acoustic windows.
+
+ The acoustic stage decodes overlapping 200-frame windows on a 100-frame
+ hop and crops the overlap back off, so a request that fits in one window
+ never exercises the stitching. 750 frames covers seven of them.
+
+ Deploy Setting: minimax_music3.yaml
+ Input Setting: stream=False, seed pinned, 750 frames
+ """
+ request_config = {
+ "model": omni_server.model,
+ "input": LYRICS,
+ "instructions": CAPTION,
+ "seed": 1,
+ "max_new_tokens": _FRAMES_30S,
+ "stream": False,
+ "response_format": "wav",
+ "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+ "min_audio_bytes": _MIN_BYTES_30S,
+ }
+ openai_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py
index 2d048fe52..598326a1e 100644
--- a/tests/e2e/online_serving/test_moss_tts_expansion.py
+++ b/tests/e2e/online_serving/test_moss_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E online tests for the full MOSS-TTS family via /v1/audio/speech.
@@ -30,10 +30,14 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
+_SKIP_ISSUE_6417 = pytest.mark.skip(
+ reason="https://github.com/vllm-project/vllm-omni/issues/6417",
+)
+
pytestmark = [
pytest.mark.slow,
pytest.mark.tts,
- pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"),
+ _SKIP_ISSUE_6417,
]
MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
@@ -92,7 +96,7 @@ tts_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None:
"""
Realtime voice_clone via /v1/audio/speech, non-streaming.
Deploy Setting: moss_tts_realtime.yaml
@@ -114,12 +118,12 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_002_streaming(omni_server, online_client, ref_audio_data_url) -> None:
"""
Realtime voice_clone via /v1/audio/speech, streaming PCM.
Deploy Setting: moss_tts_realtime.yaml
@@ -143,12 +147,12 @@ def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_
"min_hnr_db": -5.0,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_003_chinese(omni_server, online_client, ref_audio_data_url) -> None:
"""
Realtime voice_clone via /v1/audio/speech, Chinese input.
Deploy Setting: moss_tts_realtime.yaml
@@ -168,4 +172,4 @@ def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_ur
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py
index 3468a9656..3ea8c91ea 100644
--- a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py
+++ b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for MOSS-TTS-Nano model with text input and audio output.
@@ -86,7 +86,7 @@ tts_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_001(omni_server, online_client, ref_audio_data_url) -> None:
"""
Test voice_clone mode (no ref_text) via /v1/audio/speech.
Deploy Setting: default yaml
@@ -103,12 +103,12 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No
"ref_audio": ref_audio_data_url,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_002(omni_server, online_client, ref_audio_data_url) -> None:
"""
Test streaming voice_clone via /v1/audio/speech.
Deploy Setting: default yaml
@@ -135,12 +135,12 @@ def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> No
"min_hnr_db": -5.0,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_003(omni_server, online_client, ref_audio_data_url) -> None:
"""
Test Chinese voice_clone via /v1/audio/speech.
Deploy Setting: default yaml
@@ -157,12 +157,12 @@ def test_text_to_audio_003(omni_server, openai_client, ref_audio_data_url) -> No
"ref_audio": ref_audio_data_url,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audio_data_url) -> None:
+def test_text_to_audio_004_ref_text_ignored(omni_server, online_client, ref_audio_data_url) -> None:
"""
Sending ``ref_text`` alongside ``ref_audio`` is accepted but ignored.
@@ -180,4 +180,4 @@ def test_text_to_audio_004_ref_text_ignored(omni_server, openai_client, ref_audi
"ref_text": REF_AUDIO_TRANSCRIPT,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py
new file mode 100644
index 000000000..acdffd20d
--- /dev/null
+++ b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py
@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import asyncio
+import os
+from pathlib import Path
+
+import pytest
+from huggingface_hub import snapshot_download
+
+from tests.e2e.online_serving.nemotron_voicechat_realtime_duplex import parse_args, run
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, get_model_prefix
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_MODEL", "nvidia/NVIDIA-NemotronLabs-VoiceChat-11B")
+DEPLOY_CONFIG = get_deploy_config_path("nemotron_labs_voicechat_duplex.yaml")
+TOKENIZER = os.environ.get("VLLM_TEST_NEMOTRON_VOICECHAT_LLM_PATH")
+
+# The client gives up after CLIENT_TIMEOUT_S. The server-side stage-input
+# deadline must fire first, otherwise a stalled stream surfaces as a bare
+# client TimeoutError with no server-side error at all. See #6816.
+CLIENT_TIMEOUT_S = 300
+SERVER_INPUT_WAIT_TIMEOUT_S = 240
+
+_SERVER_ENV: dict[str, str] = {"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S": str(SERVER_INPUT_WAIT_TIMEOUT_S)}
+if TOKENIZER:
+ _SERVER_ENV["NEMOTRON_VOICECHAT_LLM_PATH"] = TOKENIZER
+
+pytestmark = [pytest.mark.full_model, pytest.mark.omni]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize(
+ "omni_server",
+ [
+ pytest.param(
+ OmniServerParams(
+ model=MODEL,
+ stage_config_path=DEPLOY_CONFIG,
+ env_dict=_SERVER_ENV,
+ ),
+ id="native-duplex",
+ )
+ ],
+ indirect=True,
+)
+def test_native_duplex_turn_taking_streams_model_audio(omni_server, tmp_path: Path) -> None:
+ model_prefix = get_model_prefix()
+ root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
+ if not root.is_dir():
+ root = Path(snapshot_download(MODEL, local_files_only=True))
+ args = parse_args(
+ (
+ f"--url ws://{omni_server.host}:{omni_server.port}/v1/realtime --model {omni_server.model} "
+ f"--input-wav {root / 'turn_taking.wav'} --input-channel 0 --max-frames 190 "
+ f"--minimum-audio-chunks 48 --minimum-audio-rms 0.001 --no-realtime --timeout-s {CLIENT_TIMEOUT_S} "
+ f"--output-dir {tmp_path / 'native_duplex'}"
+ ).split()
+ )
+ result = asyncio.run(run(args))
+ assert result["ok"] is True
+ assert result["input_frames"] == 190
+ assert result["event_counts"]["response.speak"] > 0
+ assert result["audio_bytes"] >= result["input_frames"] * 1764 * 2 // 4
diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py
index 381079320..3e0bc53b5 100644
--- a/tests/e2e/online_serving/test_nextstep_expansion.py
+++ b/tests/e2e/online_serving/test_nextstep_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Online serving E2E for NextStep-1.1 text-to-image (tensor parallel).
"""
@@ -10,7 +13,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -66,7 +69,7 @@ def _get_diffusion_feature_cases(model: str):
_get_diffusion_feature_cases(model=os.environ.get("VLLM_TEST_NEXTSTEP_MODEL", _DEFAULT_MODEL)),
indirect=True,
)
-def test_nextstep_11(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_nextstep_11(omni_server: OmniServer, online_client: OnlineOmniClient):
messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
request_config = {
"model": omni_server.model,
@@ -82,4 +85,4 @@ def test_nextstep_11(omni_server: OmniServer, openai_client: OpenAIClientHandler
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py
index ab35b74ea..3a042082e 100644
--- a/tests/e2e/online_serving/test_omnivoice_expansion.py
+++ b/tests/e2e/online_serving/test_omnivoice_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint.
@@ -14,9 +14,10 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME
try:
from transformers import HiggsAudioV2TokenizerModel # noqa: F401
@@ -46,7 +47,7 @@ TEST_PARAMS = [
_DEFAULT_MIN_AUDIO_BYTES = 5000
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -62,7 +63,7 @@ class TestOmniVoiceTTS:
"""E2E tests for OmniVoice TTS model."""
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_speech_auto_voice(self, omni_server, openai_client) -> None:
+ def test_speech_auto_voice(self, omni_server, online_client) -> None:
"""Test auto voice TTS generation (text only, no reference audio)."""
request_config = {
"model": omni_server.model,
@@ -71,7 +72,7 @@ class TestOmniVoiceTTS:
"timeout": 180.0,
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
@@ -79,7 +80,7 @@ class TestOmniVoiceSeed:
"""E2E tests for OmniVoice Seed Params."""
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_speech_auto_voice_seed_deterministic(self, omni_server, openai_client) -> None:
+ def test_speech_auto_voice_seed_deterministic(self, omni_server, online_client) -> None:
cfg = {
"model": omni_server.model,
"input": get_prompt("text"),
@@ -88,12 +89,12 @@ class TestOmniVoiceSeed:
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
- r1 = openai_client.send_audio_speech_request(cfg)[0]
- r2 = openai_client.send_audio_speech_request(cfg)[0]
+ r1 = online_client.send_audio_speech_request(cfg)[0]
+ r2 = online_client.send_audio_speech_request(cfg)[0]
assert r1.audio_bytes == r2.audio_bytes
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_speech_auto_voice_seed_non_deterministic(self, omni_server, openai_client) -> None:
+ def test_speech_auto_voice_seed_non_deterministic(self, omni_server, online_client) -> None:
cfg = {
"model": omni_server.model,
"input": get_prompt("text"),
@@ -104,8 +105,8 @@ class TestOmniVoiceSeed:
cfg1 = {**cfg, "seed": 42}
cfg2 = {**cfg, "seed": 43}
- r1 = openai_client.send_audio_speech_request(cfg1)[0]
- r2 = openai_client.send_audio_speech_request(cfg2)[0]
+ r1 = online_client.send_audio_speech_request(cfg1)[0]
+ r2 = online_client.send_audio_speech_request(cfg2)[0]
assert r1.audio_bytes != r2.audio_bytes
@@ -115,7 +116,7 @@ class TestOmniVoiceVoiceCloning:
"""E2E tests for OmniVoice voice cloning functionality."""
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_ref_audio_only(self, omni_server, openai_client) -> None:
+ def test_voice_clone_ref_audio_only(self, omni_server, online_client) -> None:
"""Test voice cloning with ref_audio only (x_vector mode)."""
request_config = {
"model": omni_server.model,
@@ -125,10 +126,10 @@ class TestOmniVoiceVoiceCloning:
"timeout": 180.0,
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
- def test_voice_clone_ref_audio_and_text(self, omni_server, openai_client) -> None:
+ def test_voice_clone_ref_audio_and_text(self, omni_server, online_client) -> None:
"""Test voice cloning with ref_audio and ref_text (in-context mode)."""
request_config = {
"model": omni_server.model,
@@ -139,4 +140,108 @@ class TestOmniVoiceVoiceCloning:
"timeout": 180.0,
"min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+@pytest.mark.parametrize("voice", [_DEFAULT_VOICE_NAME, {"id": _DEFAULT_VOICE_NAME}])
+def test_speech_with_voice_param_accepted(omni_server, online_client, voice) -> None:
+ """The default voice should be accepted as both a plain string and a
+ VoiceID dict, even when the model has no uploaded speakers."""
+ request_config = {
+ "model": omni_server.model,
+ "input": get_prompt("text"),
+ "voice": voice,
+ "response_format": "wav",
+ "timeout": 180.0,
+ "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+ }
+ online_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_voice_request_before_and_after_clone_registration(omni_server, online_client) -> None:
+ """The default voice should always be listed and usable, even after
+ a different cloned voice is registered."""
+ speech_config = {
+ "model": omni_server.model,
+ "input": get_prompt("text"),
+ "voice": _DEFAULT_VOICE_NAME,
+ }
+
+ # 1) "default" should appear in the voices list before any uploads
+ voices = online_client.send_audio_voices_list_http_request()[0]
+ assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
+
+ # 2) Speech with the default voice should succeed
+ online_client.send_audio_speech_request(speech_config)
+
+ # 3) Register a cloned voice under a different name
+ register_config = {
+ "data": {"name": "bar", "consent": "test-consent"},
+ "files": {
+ "audio_sample": (
+ "clone_2.wav",
+ get_asset_path("qwen3_tts/clone_2.wav").read_bytes(),
+ "audio/wav",
+ ),
+ },
+ }
+ online_client.send_audio_voices_create_http_request(register_config)
+
+ # 4) "default" should still be listed alongside the new voice
+ voices = online_client.send_audio_voices_list_http_request()[0]
+ assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
+ assert "bar" in voices.json_body["voices"]
+
+ # 5) Speech with the default voice should still succeed
+ online_client.send_audio_speech_request(speech_config)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_registered_default_voice_overrides_placeholder(omni_server, online_client) -> None:
+ """When a real voice is uploaded under the default name, it should be
+ used as a cloned voice; after deletion, the placeholder behavior resumes."""
+ speech_config = {
+ "model": omni_server.model,
+ "input": get_prompt("text"),
+ "voice": _DEFAULT_VOICE_NAME,
+ }
+
+ # 1) Placeholder default works before any registration
+ online_client.send_audio_speech_request(speech_config)
+
+ # 2) Register a real voice under the default name
+ register_config = {
+ "data": {"name": _DEFAULT_VOICE_NAME, "consent": "test-consent"},
+ "files": {
+ "audio_sample": (
+ "clone_2.wav",
+ get_asset_path("qwen3_tts/clone_2.wav").read_bytes(),
+ "audio/wav",
+ ),
+ },
+ }
+ online_client.send_audio_voices_create_http_request(register_config)
+
+ # 3) After creating the voice with the default name, it's visible as an uploaded_voice
+ voices = online_client.send_audio_voices_list_http_request()[0]
+ assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
+ uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]]
+ assert _DEFAULT_VOICE_NAME in uploaded_names
+ online_client.send_audio_speech_request(speech_config)
+
+ # 4) Delete the uploaded default voice
+ online_client.send_audio_voices_delete_http_request(
+ {"name": _DEFAULT_VOICE_NAME},
+ )
+
+ # 5) Ensure the placeholder behavior is consistent, but it's not an uploaded voice now
+ voices = online_client.send_audio_voices_list_http_request()[0]
+ assert _DEFAULT_VOICE_NAME in voices.json_body["voices"]
+ uploaded_names = [v["name"] for v in voices.json_body["uploaded_voices"]]
+ assert _DEFAULT_VOICE_NAME not in uploaded_names
+ online_client.send_audio_speech_request(speech_config)
diff --git a/tests/e2e/online_serving/test_pi0_expansion.py b/tests/e2e/online_serving/test_pi0_expansion.py
new file mode 100644
index 000000000..3a9c3b177
--- /dev/null
+++ b/tests/e2e/online_serving/test_pi0_expansion.py
@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving.
+
+Boots ``vllm serve --omni --deploy-config pi0.yaml`` and drives the real OpenPI
+websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses
+(handshake metadata → send observation → receive action chunk). Mirrors
+``tests/e2e/online_serving/test_dreamzero_expansion.py``. Needs a GPU + a
+pi0_base checkpoint; skipped in CI unless explicitly run.
+
+The in-process LeRobot parity oracle lives separately in
+``tests/diffusion/models/pi0/test_pi0_parity.py``.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import (
+ OmniServerParams,
+ get_open_port,
+ pi0_openpi_require_dependencies,
+ pi0_openpi_run_policy_session,
+ pi0_openpi_validate_session_result,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "lerobot/pi0_base"
+
+
+test_params = [
+ OmniServerParams(
+ model=MODEL,
+ port=8092,
+ server_args=[
+ "--deploy-config",
+ "vllm_omni/deploy/pi0.yaml",
+ "--served-model-name",
+ "pi0",
+ "--enforce-eager",
+ "--disable-log-stats",
+ ],
+ env_dict={
+ "ATTENTION_BACKEND": "torch",
+ "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+ "VLLM_DISABLE_COMPILE_CACHE": "1",
+ "MASTER_PORT": str(get_open_port()),
+ },
+ )
+]
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_pi0_openpi_online(omni_server):
+ try:
+ pi0_openpi_require_dependencies()
+ except ModuleNotFoundError as exc:
+ pytest.skip(str(exc))
+
+ result = pi0_openpi_run_policy_session(
+ host=omni_server.host,
+ port=omni_server.port,
+ prompt="pick up the red block and place it in the bin",
+ session_id="pi0-online-e2e",
+ num_steps=2,
+ )
+
+ # Asserts every returned chunk is [50, 32] + finite, and the handshake
+ # metadata matches pi0.yaml's policy_server_config.
+ pi0_openpi_validate_session_result(result)
+
+ metadata = result["metadata"]
+ assert tuple(metadata["image_resolution"]) == (224, 224)
+ assert metadata["needs_wrist_camera"] is True
+ assert metadata["needs_session_id"] is False
+ assert metadata["action_space"] == "joint_position"
diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
index bbf9ca2b6..1e5e59aed 100644
--- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
+++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
E2E Online tests for Qwen2.5-Omni model with video input and audio output.
"""
@@ -59,7 +62,7 @@ def get_max_batch_size(size_type="few"):
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+def test_mix_to_text_audio_001(omni_server, online_client) -> None:
"""
Test multi-modal input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -90,13 +93,13 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
}
# Test single completion
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@pytest.mark.slow
@pytest.mark.omni
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(omni_server, openai_client) -> None:
+def test_text_to_text_001(omni_server, online_client) -> None:
"""
Test text input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -114,4 +117,4 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py
index a54d2fe8a..f6ede8944 100644
--- a/tests/e2e/online_serving/test_qwen3_omni.py
+++ b/tests/e2e/online_serving/test_qwen3_omni.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
E2E Online tests for Qwen3-Omni model with video input and audio output.
"""
@@ -82,7 +85,7 @@ def get_max_batch_size(size_type="few"):
@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+def test_mix_to_text_audio_001(omni_server, online_client) -> None:
"""
Test multi-modal input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -113,7 +116,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
}
# Test single completion
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@pytest.mark.advanced_model
@@ -122,7 +125,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(omni_server, openai_client) -> None:
+def test_text_to_text_001(omni_server, online_client) -> None:
"""
Test text input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -140,18 +143,27 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
-def _run_prefix_cache_check(openai_client, request_config: dict):
+def _run_prefix_cache_check(online_client, request_config: dict):
"""Make two requests given a request config, and validate that:
1. The second request actually had cached tokens
2. The number of cached tokens is divisible by the block size used in
test_params, because currently upstream vLLM does not cache partial
blocks.
+ 3. The first request also reports prompt token details; details must be
+ present whenever --enable-prompt-tokens-details is on, even when zero
+ tokens are cached. (An exact cached_tokens == 0 check would be
+ order-dependent here: the module-scoped server may have already cached
+ the shared system-prompt blocks from earlier tests.)
+
+ Returns the (first, second) responses for additional per-test assertions.
"""
- openai_client.send_omni_request(request_config, request_num=1)[0]
- cached_response = openai_client.send_omni_request(request_config, request_num=1)[0]
+ first_response = online_client.send_omni_request(request_config, request_num=1)[0]
+ cached_response = online_client.send_omni_request(request_config, request_num=1)[0]
+
+ assert first_response.cached_tokens is not None
# Ensure that we have a prefix cache hit on the second request and that only the last
# partial block is uncached (since currently we don't cache partial blocks).
@@ -163,13 +175,15 @@ def _run_prefix_cache_check(openai_client, request_config: dict):
assert num_cached_tokens % BLOCK_SIZE == 0
assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens
+ return first_response, cached_response
+
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
+def test_thinker_prefix_caching_text_output(omni_server, online_client) -> None:
"""
Test thinker prefix caching by sending identical requests with an image (i.e.,
a large shared prefix) and verifying that the second request uses cached tokens
@@ -195,7 +209,10 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
"stream": False,
"modalities": ["text"],
}
- _run_prefix_cache_check(openai_client, request_config)
+ first_response, _ = _run_prefix_cache_check(online_client, request_config)
+
+ assert first_response.multimodal_tokens is not None
+ assert first_response.multimodal_tokens.get("image", 0) > 0
@pytest.mark.advanced_model
@@ -203,7 +220,7 @@ def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None:
+def test_thinker_prefix_caching_audio_output(omni_server, online_client) -> None:
"""
Verify that thinker prefix caching does not hang when the request
produces audio output (text + audio modalities). Sends two identical
@@ -225,7 +242,7 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
},
}
- _run_prefix_cache_check(openai_client, request_config)
+ _run_prefix_cache_check(online_client, request_config)
@pytest.mark.advanced_model
@@ -233,13 +250,13 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
+def test_completions_rejected_for_thinker_talker(omni_server, online_client) -> None:
"""Ensure Thinker-talker models reject /v1/completions; we do this because the
thinker-talker handoff implementations currently use ChatML <|im_start|> and
<|im_end|> markers to segment the input sequence; when we don't have them,
the talker does not get any embeddings, which breaks the server.
"""
- responses = openai_client.send_completions_http_request(
+ responses = online_client.send_completions_http_request(
{
"json": {
"model": omni_server.model,
diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
index 1ba94c011..10c5dd8cb 100644
--- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py
+++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Qwen3-Omni model.
"""
@@ -7,7 +7,6 @@ E2E Online tests for Qwen3-Omni model.
import os
import pytest
-from openai import BadRequestError
from tests.helpers.mark import hardware_test
from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -89,7 +88,7 @@ test_params = [
),
]
-test_token_params = [
+reduced_token_params = [
pytest.param(
OmniServerParams(
model=model,
@@ -97,7 +96,7 @@ test_token_params = [
use_stage_cli=True,
server_args=["--async-chunk"],
),
- id="batch_token_64",
+ id="batch_token_2048",
)
]
@@ -138,30 +137,8 @@ def get_max_batch_size(size_type="few"):
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client) -> None:
- """
- Input Modal: text
- Output Modal: audio
- Input Setting: stream=True
- Datasets: single request
- """
- messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "modalities": ["audio"],
- "stream": True,
- "key_words": {"text": ["beijing"]},
- }
-
- openai_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_to_text_audio_001(omni_server, openai_client) -> None:
+@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
+def test_text_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text
Output Modal: text, audio
@@ -176,12 +153,12 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None:
"key_words": {"text": ["beijing"]},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_video_to_text_001(omni_server, openai_client) -> None:
+def test_text_video_to_text_001(omni_server, online_client) -> None:
"""
Input Modal: long synthetic video (120s @ 30fps, LONG_VIDEO_FRAMES frames)
Output Modal: text
@@ -200,40 +177,17 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None:
"key_words": {"video": VIDEO_KEY},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
- """
- Input Modal: text, audio
- Output Modal: text, audio
- Input Setting: stream=False
- Datasets: single request
- """
- audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
- messages = dummy_messages_from_mix_data(
- audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
- )
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "key_words": {"audio": AUDIO_KEY},
- }
-
- openai_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None:
+@pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
+def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
"""
Input Modal: text, long-duration audio (~LONG_AUDIO_DURATION_SEC s WAV)
Output Modal: text, audio
Input Setting: stream=False
- Datasets: single request
+ Datasets: few requests
"""
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(LONG_AUDIO_DURATION_SEC, 1)['base64']}"
messages = dummy_messages_from_mix_data(
@@ -248,41 +202,17 @@ def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None:
"key_words": {"audio": AUDIO_KEY},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
- """
- Input Modal: text, image
- Output Modal: text, audio
- Input Setting: stream=False
- Datasets: single request
- """
- image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
-
- messages = dummy_messages_from_mix_data(
- image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image")
- )
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "key_words": {"image": IMAGE_KEY},
- }
-
- openai_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_large_image_to_text_audio_001(omni_server, openai_client) -> None:
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, high-resolution image (1080p-class JPEG)
Output Modal: text, audio
Input Setting: stream=False
- Datasets: single request
+ Datasets: few requests
"""
image_data_url = (
f"data:image/jpeg;base64,{generate_synthetic_image(LARGE_IMAGE_WIDTH, LARGE_IMAGE_HEIGHT)['base64']}"
@@ -300,12 +230,12 @@ def test_large_image_to_text_audio_001(omni_server, openai_client) -> None:
"key_words": {"image": IMAGE_KEY},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
"""
Input Modal: text, video
Output Modal: text, audio
@@ -325,42 +255,12 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
"key_words": {"video": VIDEO_KEY},
}
- openai_client.send_omni_request(request_config)
-
-
-@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
-def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
- """
- Input Modal: text, audio, image, video
- Output Modal: text, audio
- Input Setting: stream=True
- Datasets: few requests
- """
- video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
- image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
- audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
- messages = dummy_messages_from_mix_data(
- system_prompt=get_system_prompt(),
- video_data_url=video_data_url,
- image_data_url=image_data_url,
- audio_data_url=audio_data_url,
- content_text=get_prompt("mix"),
- )
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "stream": True,
- "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
- }
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_in_video_001(omni_server, openai_client) -> None:
+def test_audio_in_video_001(omni_server, online_client) -> None:
"""
Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
Output Modal: text, audio
@@ -381,12 +281,12 @@ def test_audio_in_video_001(omni_server, openai_client) -> None:
"use_audio_in_video": True,
"key_words": {"video": VIDEO_KEY},
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_in_video_002(omni_server, openai_client) -> None:
+def test_audio_in_video_002(omni_server, online_client) -> None:
"""
Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
Output Modal: text, audio
@@ -408,12 +308,12 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
"key_words": {"video": VIDEO_KEY},
}
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
+def test_audio_in_video_default_loader_sampling_regression(omni_server, online_client) -> None:
"""
Regression: ``use_audio_in_video`` with default video loader sampling (no
``media_io_kwargs.video.fps``) after a prior request that pinned fps.
@@ -438,7 +338,7 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c
}
# Prime mm-cache with explicit loader fps (historical failure mode).
- openai_client.send_omni_request(
+ online_client.send_omni_request(
{
**base_config,
"extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
@@ -446,17 +346,37 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_c
)
# Default loader sampling: must succeed without explicit media_io fps.
- openai_client.send_omni_request(dict(base_config))
+ online_client.send_omni_request(dict(base_config))
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_one_word_prompt_001(omni_server, openai_client) -> None:
- """
+def test_one_word_prompt_001(omni_server, online_client) -> None:
+ """Catastrophic-regression gate on one-word pronunciation.
+
+ Whether a one-word answer is pronounced intelligibly is sampled, not guaranteed: the
+ talker runs at temperature 0.9 (the upstream default, see qwen3_omni_moe.yaml stage 1),
+ so asserting it per request gates on a random variable, which is what the old 10-retry
+ loop was compensating for.
+
+ Sends a fixed 40 requests and requires at least 29 to transcribe as the expected
+ word. 85.6% is a predeclared minimum acceptable success probability, taken from a
+ measured HF reference run (77/90); vLLM-Omni itself measured 92.2%. Assuming 40
+ independent Bernoulli trials, >=29/40 has a nominal rejection probability of 0.9%
+ at p=.856, with 92.9% power at p=.60 and 55.9% at p=.70.
+
+ Those are conditional on p, not a measured false-failure rate for this CI: 77/90 is
+ one point estimate (Wilson 76.8%-91.4%), and at the low end of that interval the gate
+ would reject about 20% of runs. Pinning down a real rate would take repeated
+ report-only runs on the target hardware.
+
+ A drop to 70% is caught only about half the time, so this is a smoke gate for gross
+ pipeline breakage, not a quality tracker -- resolving moderate drift needs ~80+
+ samples and belongs in an accuracy benchmark. See #5395.
+
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
- Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
@@ -474,47 +394,43 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None:
"transcript_language": "en",
}
- # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/helpers/assertions.py).
- _similarity_assert_msg = "The audio content is not same as the text"
- _max_retries = 10
- for attempt in range(_max_retries):
- try:
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
- break
- except AssertionError as e:
- if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
- raise
- print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
+ # The judge runs per request so mismatches can be counted, which also transcribes
+ # concurrently where the previous request_num=5 path judged serially -- worth watching
+ # on tighter cards, since each judge loads whisper on the same device as the talker and
+ # code2wav stages.
+ online_client.send_omni_request(
+ request_config, request_num=40, min_successes=29, max_concurrency=get_max_batch_size()
+ )
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_001(omni_server, openai_client) -> None:
+def test_speaker_001(omni_server, online_client) -> None:
"""
- Input Modal: text only (one-word answer constraint).
- Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
+ Input Modal: Chinese text.
+ Output Modal: Chinese text and audio using an uppercase preset speaker name.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
- content_text=get_prompt("text"),
+ content_text=get_prompt("text_chinese"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
- "speaker": "Chelsie",
- "key_words": {"text": ["beijing"]},
+ "speaker": "CHELSIE",
+ "key_words": {"text": ["北京"]},
}
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_002(omni_server, openai_client) -> None:
+def test_speaker_002(omni_server, online_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
@@ -539,7 +455,7 @@ def test_speaker_002(omni_server, openai_client) -> None:
_max_retries = 10
for attempt in range(_max_retries):
try:
- openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ online_client.send_omni_request(request_config, request_num=get_max_batch_size())
break
except AssertionError as e:
if _gender_assert_substr not in str(e) or attempt == _max_retries - 1:
@@ -549,61 +465,12 @@ def test_speaker_002(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_speaker_003(omni_server, openai_client) -> None:
- """
- Input Modal: text only (one-word answer constraint).
- Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
- Input Setting: stream=True
- Datasets: single request
- """
- messages = dummy_messages_from_mix_data(
- system_prompt=get_system_prompt(),
- content_text=get_prompt("text"),
- )
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "stream": True,
- "speaker": "CHELSIE",
- "key_words": {"text": ["beijing"]},
- }
-
- openai_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_language_001(omni_server, openai_client) -> None:
- """
- Input Modal: text only (one-word answer constraint).
- Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
- Input Setting: stream=True
- Datasets: single request
- """
- messages = dummy_messages_from_mix_data(
- system_prompt=get_system_prompt(),
- content_text=get_prompt("text_chinese"),
- )
-
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "stream": True,
- "key_words": {"text": ["北京"]},
- }
-
- openai_client.send_omni_request(request_config)
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
+def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
"""
Input Modal: text only (long-form generation prompt).
Output Modal: text, audio (default ``modalities``);
Input Setting: stream=True
- Datasets: single request
+ Datasets: few requests
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
@@ -611,25 +478,7 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
)
request_config = {"model": omni_server.model, "messages": messages, "stream": True}
- responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+ responses = online_client.send_omni_request(request_config, request_num=get_max_batch_size())
text = responses[0].text_content if responses else ""
word_count = len(text.split())
assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
-
-
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
-def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
- """Ensure invalid audio format is rejected with 400 before streaming starts."""
- messages = dummy_messages_from_mix_data(
- system_prompt=get_system_prompt(),
- content_text=get_prompt(),
- )
- with pytest.raises(BadRequestError, match="audio format"):
- openai_client.client.chat.completions.create(
- model=omni_server.model,
- messages=messages,
- modalities=["text", "audio"],
- audio={"voice": "alloy", "format": "aac"},
- stream=True,
- )
diff --git a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
index f269836f3..1a27a5449 100644
--- a/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
+++ b/tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Core-model CI guard for Qwen3-Omni multi-replica stage-pool routing on 4 GPUs.
"""
@@ -90,7 +93,7 @@ def _assert_audio_outputs(responses: list[OmniResponse], *, expect_text: bool) -
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -> None:
+def test_text_only_batch_uses_multi_replica_talker(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _text_messages(),
@@ -98,7 +101,7 @@ def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -
"modalities": ["text"],
}
- responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+ responses = online_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
_assert_batch_size(responses, ROUTE_STRESS_REQUESTS)
_assert_text_outputs(responses)
@@ -107,7 +110,7 @@ def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, openai_client) -> None:
+def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _text_messages(),
@@ -115,7 +118,7 @@ def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, open
"modalities": ["audio"],
}
- responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+ responses = online_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
_assert_batch_size(responses, ROUTE_STRESS_REQUESTS)
_assert_audio_outputs(responses, expect_text=False)
@@ -124,13 +127,13 @@ def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, open
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=4)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, openai_client) -> None:
+def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, online_client) -> None:
request_config = {
"model": omni_server.model,
"messages": _mixed_messages(),
"stream": True,
}
- responses = openai_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS)
+ responses = online_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS)
_assert_batch_size(responses, MIXED_MODAL_REQUESTS)
_assert_audio_outputs(responses, expect_text=True)
diff --git a/tests/e2e/online_serving/test_qwen3_tts_base.py b/tests/e2e/online_serving/test_qwen3_tts_base.py
index 60caf8f82..dd25520f9 100644
--- a/tests/e2e/online_serving/test_qwen3_tts_base.py
+++ b/tests/e2e/online_serving/test_qwen3_tts_base.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Qwen3-TTS model with text input and audio output.
@@ -14,7 +14,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -24,7 +24,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0
# Vendored under tests/assets/qwen3_tts/clone_2.wav so the server does not need
# to fetch the reference audio over HTTPS at request time (CI runners in some
# regions cannot reach the original Aliyun OSS host; see issue #3263).
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -55,11 +55,10 @@ tts_server_params = [
@pytest.mark.advanced_model
-@pytest.mark.core_model
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client) -> None:
+def test_text_to_audio_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -80,14 +79,14 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))
+ online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))
@pytest.mark.advanced_model
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002(omni_server, openai_client) -> None:
+def test_text_to_audio_002(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -108,4 +107,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py
index 0edaa4668..e126877db 100644
--- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py
+++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Qwen3-TTS model with text input and audio output.
@@ -12,7 +12,7 @@ import os
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
@@ -23,7 +23,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
# See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale.
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
@@ -41,7 +41,7 @@ def get_max_batch_size(size_type="few"):
return batch_sizes.get(size_type, 5)
-tts_server_params = [
+tts_async_chunk_server_params = [
pytest.param(
OmniServerParams(
model=MODEL,
@@ -50,9 +50,12 @@ tts_server_params = [
),
id="async_chunk",
),
- # Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone
- # flips the deploy yaml's bool and the pipeline dispatches to the
- # end-to-end codec processor. No variant yaml / pipeline needed.
+]
+
+# Synchronous (no async-chunk) variant — ``--no-async-chunk`` alone
+# flips the deploy yaml's bool and the pipeline dispatches to the
+# end-to-end codec processor. No variant yaml / pipeline needed.
+tts_no_async_chunk_server_params = [
pytest.param(
OmniServerParams(
model=MODEL,
@@ -63,10 +66,16 @@ tts_server_params = [
),
]
+DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0
+
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
+@pytest.mark.parametrize(
+ "omni_server",
+ tts_async_chunk_server_params + tts_no_async_chunk_server_params,
+ indirect=True,
+)
+def test_voice_clone_streaming_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -87,12 +96,16 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))
+ online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_response_format_001(omni_server, openai_client) -> None:
+@pytest.mark.parametrize(
+ "omni_server",
+ tts_async_chunk_server_params + tts_no_async_chunk_server_params,
+ indirect=True,
+)
+def test_response_format_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -111,4 +124,31 @@ def test_response_format_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
+def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, online_client) -> None:
+ """Regression for #5049: x-vector-only then ICL on the same ``ref_audio``.
+
+ Pre-fix, readiness was mode-agnostic: an x-vector request marked the
+ artifact ready (embedding only, no ``ref_code``), so the following ICL
+ request took the artifact-only path and killed EngineCore. After the fix
+ (readiness keyed by ``(artifact_key, x_vector_only)``), both modes succeed
+ and the server stays up. Only runs on the ``async_chunk`` server param.
+ """
+ base_request = {
+ "model": omni_server.model,
+ "input": get_prompt(),
+ "stream": False,
+ "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
+ "response_format": "wav",
+ "task_type": "Base",
+ "ref_audio": REF_AUDIO_URL,
+ "min_audio_bytes": 1,
+ }
+
+ online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})
+ online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT})
+ online_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})
diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py
index 7ebe08e1a..6192f8410 100644
--- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Qwen3-TTS model with text input and audio output.
@@ -20,7 +20,6 @@ from tests.helpers.stage_config import (
get_deploy_config_stage,
modify_stage_config,
)
-from vllm_omni.platforms import current_omni_platform
MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
@@ -92,13 +91,9 @@ def test_default_cuda_graph_startup(omni_server) -> None:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.tts
-@pytest.mark.skipif(
- current_omni_platform.is_npu(),
- reason="#issue 5479",
-)
@hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client) -> None:
+def test_text_to_audio_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -116,19 +111,15 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
"voice": "vivian",
}
- openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
+ online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.tts
-@pytest.mark.skipif(
- current_omni_platform.is_npu(),
- reason="#issue 5479",
-)
@hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002(omni_server, openai_client) -> None:
+def test_text_to_audio_002(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -147,4 +138,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None:
"voice": "vivian",
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py
index 9c941cad8..07808af27 100644
--- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py
+++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for Qwen3-TTS model with text input and audio output.
@@ -62,7 +62,7 @@ tts_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_001(omni_server, openai_client) -> None:
+def test_voice_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -84,7 +84,7 @@ def test_voice_001(omni_server, openai_client) -> None:
# Retry once on assertion failures from transcript similarity / gender checks (flaky ASR or estimators).
for attempt in range(2):
try:
- openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
+ online_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
break
except AssertionError:
if attempt == 0:
@@ -94,7 +94,7 @@ def test_voice_001(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_002(omni_server, openai_client) -> None:
+def test_voice_002(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -112,12 +112,12 @@ def test_voice_002(omni_server, openai_client) -> None:
"voice": "Serena",
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_003(omni_server, openai_client) -> None:
+def test_voice_003(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -135,12 +135,12 @@ def test_voice_003(omni_server, openai_client) -> None:
"voice": "SERENA",
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_language_001(omni_server, openai_client) -> None:
+def test_language_001(omni_server, online_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
@@ -158,4 +158,4 @@ def test_language_001(omni_server, openai_client) -> None:
"voice": "vivian",
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py b/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py
deleted file mode 100644
index e6c396932..000000000
--- a/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py
+++ /dev/null
@@ -1,88 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""E2E regression for Qwen3-TTS ref_audio artifact mode collision (#5049).
-
-Reproduces the issue's exact sequence on a live server: an
-``x_vector_only_mode=true`` request followed by an ICL
-(``x_vector_only_mode=false`` + ``ref_text``) request using the **same**
-``ref_audio``.
-
-Before the fix, the ref_audio artifact readiness was mode-agnostic: the
-x-vector-only request (which caches a speaker embedding but no ``ref_code``)
-marked the artifact "ready", so the ICL request took the artifact-only serving
-path (``ref_audio`` stripped before dispatch), found no ``ref_code`` and no
-audio to recompute from, and raised in ``build_prompt_embeds`` ->
-``EngineDeadError`` — taking the whole server down rather than failing the one
-request.
-
-After the fix (readiness keyed by ``(artifact_key, x_vector_only)``), the ICL
-request no longer matches the x-vector artifact, so it sends ``ref_audio`` and
-the worker computes ``ref_code``. Both requests succeed and the server stays up.
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-import pytest
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
-from tests.helpers.runtime import OmniServerParams
-from tests.helpers.stage_config import get_deploy_config_path
-
-MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
-DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0
-
-# Vendored under tests/assets/qwen3_tts/clone_2.wav (see test_qwen3_tts_base.py).
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
-REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-INPUT_TEXT = "The weather is nice today, perfect for a walk in the park."
-
-tts_server_params = [
- pytest.param(
- OmniServerParams(
- model=MODEL,
- stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
- server_args=["--trust-remote-code"],
- ),
- id="async_chunk",
- )
-]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, openai_client) -> None:
- """Regression for #5049: x-vector-only then ICL on the same ref_audio.
-
- Both requests must return audio and the engine must stay alive (pre-fix the
- second request killed EngineCore).
- """
- base_request = {
- "model": omni_server.model,
- "input": INPUT_TEXT,
- "stream": False,
- "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S,
- "response_format": "wav",
- "task_type": "Base",
- "ref_audio": REF_AUDIO_URL,
- # Assert non-empty audio came back (not just HTTP 200).
- "min_audio_bytes": 1,
- }
-
- # 1) x-vector-only request: caches a speaker embedding (no ref_code) and, on
- # completion, marks the ref_audio artifact "ready".
- openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})
-
- # 2) ICL request with the SAME ref_audio. Pre-fix this reused the x-vector
- # artifact via the artifact-only path, hit the missing ref_code, and
- # killed EngineCore (#5049). It must now succeed with audio.
- openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT})
-
- # 3) The server must still be serving after the ICL request (pre-fix this
- # would fail with EngineDeadError / connection refused).
- openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True})
diff --git a/tests/e2e/online_serving/test_qwen_image.py b/tests/e2e/online_serving/test_qwen_image.py
index c40cfbda8..aa037d9b8 100644
--- a/tests/e2e/online_serving/test_qwen_image.py
+++ b/tests/e2e/online_serving/test_qwen_image.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Online serving tests for ``Qwen/Qwen-Image`` (text-to-image).
@@ -24,7 +24,7 @@ import os
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -61,7 +61,7 @@ def _get_default_case(model: str):
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
"""Default Qwen-Image T2I smoke (single ``default`` server config)."""
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
request_config = {
@@ -76,14 +76,14 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH
"seed": 42,
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
-def test_batch_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_batch_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
"""Concurrent T2I: one ``request_config`` dict per prompt (``send_diffusion_request`` list mode)."""
request_config = [
{
@@ -100,4 +100,4 @@ def test_batch_001(omni_server: OmniServer, openai_client: OpenAIClientHandler)
}
for prompt in TEST_PROMPTS
]
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen_image_edit.py b/tests/e2e/online_serving/test_qwen_image_edit.py
index 738a2e50c..dc1f60521 100644
--- a/tests/e2e/online_serving/test_qwen_image_edit.py
+++ b/tests/e2e/online_serving/test_qwen_image_edit.py
@@ -1,21 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
-Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
+Online serving tests for ``Qwen/Qwen-Image-Edit-2511`` (image-to-image via chat completions).
-- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
- resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
- scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
+- ``test_single_image_to_image_001``: one reference image, fixed 512×512.
+- ``test_multi_images_to_image_001``: two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``slow`` — mixed input resolutions; ``extra_body`` uses
+ per-output ``width``/``height`` lists; the test client sends one scalar-size request
+ per list index in parallel and merges images (see ``OnlineOmniClient.send_diffusion_request``).
-``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row per model.
+``_get_diffusion_feature_cases`` registers a single ``default`` ``OmniServerParams`` row.
From ``tests/``::
- pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "core_model and diffusion" --run-level=core_model
- pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "advanced_model and diffusion" --run-level=advanced_model
+ pytest -s -v e2e/online_serving/test_qwen_image_edit.py -m "slow and diffusion" --run-level=full_model
"""
import pytest
@@ -25,10 +24,11 @@ from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
+MODEL = "Qwen/Qwen-Image-Edit-2511"
EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
MULTI_EDIT_PROMPT = (
"Transform the first image into a Dadaism collage art. "
@@ -39,7 +39,7 @@ NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-def _get_diffusion_feature_cases(model: str):
+def _get_diffusion_feature_cases(model: str = MODEL):
"""Return one ``default`` ``OmniServerParams`` row for ``model`` (no extra ``server_args``)."""
return [
pytest.param(
@@ -52,15 +52,15 @@ def _get_diffusion_feature_cases(model: str):
]
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
+ _get_diffusion_feature_cases(),
indirect=True,
)
-def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
+def test_single_image_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient):
+ """Single-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511`` (CFG when negative prompt + ``true_cfg_scale`` > 1)."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
@@ -79,18 +79,17 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
+ _get_diffusion_feature_cases(),
indirect=True,
)
-def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_multi_images_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
@@ -111,17 +110,17 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
+ _get_diffusion_feature_cases(),
indirect=True,
)
-def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_different_sizes_001(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Multi-reference edit with distinct input resolutions; per-output ``width``/``height`` as lists (client splits into concurrent scalar-size calls)."""
image_data_url_list = [
f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}",
@@ -144,4 +143,4 @@ def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClien
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
index b3d6f8b19..47bc2809b 100644
--- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py
@@ -1,20 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode
-and are supported by the following models:
-- Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2511: single image input and two image inputs
+and are supported by Qwen-Image-Edit-2511 (single image input and two image inputs).
"""
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
-EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
-SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
+MODEL = "Qwen/Qwen-Image-Edit-2511"
+SINGLE_EDIT_PROMPT = "Restyle this image into a Vincent van Gogh style impressionist painting."
MULTI_EDIT_PROMPT = (
"Transform the first image into a Dadaism collage art. "
"Transform the second image into a Vincent van Gogh style painting. "
@@ -25,9 +26,7 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-# This test file targets two models, so I write a helper function.
-# If a similar test only involves one model, one can just define a global list variable.
-def _get_diffusion_feature_cases(model: str):
+def _get_diffusion_feature_cases(model: str = MODEL):
return [
pytest.param(
OmniServerParams(
@@ -123,38 +122,10 @@ def _get_diffusion_feature_cases(model: str):
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"),
+ _get_diffusion_feature_cases(),
indirect=True,
)
-def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios."""
- image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
-
- messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
-
- # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "extra_body": {
- "height": 512,
- "width": 512,
- "num_inference_steps": 2,
- "negative_prompt": NEGATIVE_PROMPT,
- "true_cfg_scale": 4.0,
- "seed": 42,
- },
- }
-
- openai_client.send_diffusion_request(request_config)
-
-
-@pytest.mark.parametrize(
- "omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
- indirect=True,
-)
-def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test Qwen-Image-Edit-2511 with a single image input.
Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
@@ -166,7 +137,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
"""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
- messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
+ messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT)
request_config = {
"model": omni_server.model,
@@ -181,15 +152,15 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
"omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
+ _get_diffusion_feature_cases(),
indirect=True,
)
-def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test Qwen-Image-Edit-2511 with two image inputs."""
image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -212,4 +183,4 @@ def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client:
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py
index c900ca4ce..890b077d2 100644
--- a/tests/e2e/online_serving/test_qwen_image_expansion.py
+++ b/tests/e2e/online_serving/test_qwen_image_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode
and are supported by the following text-to-image models:
@@ -5,15 +8,14 @@ and are supported by the following text-to-image models:
- Qwen-Image-2512
One feature per test case, matching the Test Plan in PR #1682 (Qwen-Image-Edit).
-Supported features for Qwen-Image series: TeaCache, Cache-DiT, Ulysses-SP, Ring-Attention,
-CFG-Parallel, Tensor-Parallel, VAE-Patch-Parallel, CPU offload (layerwise).
-See docs/user_guide/diffusion_acceleration.md.
+Nightly covers each feature once, alternating models (5× Qwen-Image, 4× Qwen-Image-2512).
+Ulysses stays on Qwen-Image. See docs/user_guide/diffusion_acceleration.md.
"""
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
@@ -22,118 +24,116 @@ NEGATIVE_PROMPT = "blurry, low quality"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+MODEL_IMAGE = "Qwen/Qwen-Image"
+MODEL_2512 = "Qwen/Qwen-Image-2512"
-def _get_diffusion_feature_cases(model: str):
- return [
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=["--enable-cpu-offload"],
- ),
- id="cpu_offload",
- marks=SINGLE_CARD_FEATURE_MARKS,
+# One server per feature. Alternate models; keep feature pytest ids unchanged.
+FEATURE_CASES = [
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_IMAGE,
+ server_args=["--enable-cpu-offload"],
),
- pytest.param(
- OmniServerParams(model=model, server_args=["--step-execution"]),
- id="step_execution",
- marks=SINGLE_CARD_FEATURE_MARKS,
+ id="cpu_offload",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(model=MODEL_2512, server_args=["--step-execution"]),
+ id="step_execution",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(model=MODEL_IMAGE, server_args=["--cache-backend", "tea_cache"]),
+ id="cache_tea_cache",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_2512,
+ server_args=[
+ "--cache-backend",
+ "cache_dit",
+ "--enable-layerwise-offload",
+ ],
),
- pytest.param(
- OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
- id="cache_tea_cache",
- marks=SINGLE_CARD_FEATURE_MARKS,
+ id="layerwise_offload",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_IMAGE,
+ server_args=[
+ "--cache-backend",
+ "cache_dit",
+ "--ulysses-degree",
+ "2",
+ ],
),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--cache-backend",
- "cache_dit",
- "--enable-layerwise-offload",
- ],
- ),
- id="layerwise_offload",
- marks=SINGLE_CARD_FEATURE_MARKS,
+ id="ulysses_2",
+ marks=PARALLEL_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_2512,
+ server_args=[
+ "--cache-backend",
+ "cache_dit",
+ "--ring",
+ "2",
+ ],
),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--cache-backend",
- "cache_dit",
- "--ulysses-degree",
- "2",
- ],
- ),
- id="ulysses_2",
- marks=PARALLEL_FEATURE_MARKS,
+ id="ring_2",
+ marks=PARALLEL_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_IMAGE,
+ server_args=[
+ "--cache-backend",
+ "tea_cache",
+ "--cfg-parallel-size",
+ "2",
+ ],
),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--cache-backend",
- "cache_dit",
- "--ring",
- "2",
- ],
- ),
- id="ring_2",
- marks=PARALLEL_FEATURE_MARKS,
+ id="cfg_parallel_2",
+ marks=PARALLEL_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_2512,
+ server_args=[
+ "--cache-backend",
+ "cache_dit",
+ "--tensor-parallel-size",
+ "2",
+ "--vae-patch-parallel-size",
+ "2",
+ "--vae-use-tiling",
+ "--diffusion-quantization-config",
+ '{"method":"fp8"}',
+ ],
),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--cache-backend",
- "tea_cache",
- "--cfg-parallel-size",
- "2",
- ],
- ),
- id="cfg_parallel_2",
- marks=PARALLEL_FEATURE_MARKS,
+ id="vae_patch_parallel_2",
+ marks=PARALLEL_FEATURE_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=MODEL_IMAGE,
+ server_args=[
+ "--use-hsdp",
+ "--hsdp-shard-size",
+ "2",
+ ],
),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--cache-backend",
- "cache_dit",
- "--tensor-parallel-size",
- "2",
- "--vae-patch-parallel-size",
- "2",
- "--vae-use-tiling",
- "--diffusion-quantization-config",
- '{"method":"fp8"}',
- ],
- ),
- id="vae_patch_parallel_2",
- marks=PARALLEL_FEATURE_MARKS,
- ),
- pytest.param(
- OmniServerParams(
- model=model,
- server_args=[
- "--use-hsdp",
- "--hsdp-shard-size",
- "2",
- ],
- ),
- id="parallel_hsdp",
- marks=PARALLEL_FEATURE_MARKS,
- ),
- ]
+ id="parallel_hsdp",
+ marks=PARALLEL_FEATURE_MARKS,
+ ),
+]
-@pytest.mark.parametrize(
- "omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image"),
- indirect=True,
-)
-def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Test each diffusion feature with Qwen-Image (text-to-image), one feature per case."""
+@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
+def test_qwen_image(omni_server: OmniServer, online_client: OnlineOmniClient):
+ """One diffusion feature per case; model is chosen in FEATURE_CASES."""
messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
request_config = {
"model": omni_server.model,
@@ -147,27 +147,4 @@ def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler)
"seed": 42,
},
}
- openai_client.send_diffusion_request(request_config)
-
-
-@pytest.mark.parametrize(
- "omni_server",
- _get_diffusion_feature_cases("Qwen/Qwen-Image-2512"),
- indirect=True,
-)
-def test_qwen_image_2512(omni_server: OmniServer, openai_client: OpenAIClientHandler):
- """Test each diffusion feature with Qwen-Image-2512 (text-to-image), one feature per case."""
- messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
- request_config = {
- "model": omni_server.model,
- "messages": messages,
- "extra_body": {
- "height": 512,
- "width": 512,
- "num_inference_steps": 2,
- "negative_prompt": NEGATIVE_PROMPT,
- "true_cfg_scale": 4.0,
- "seed": 42,
- },
- }
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen_image_layered.py b/tests/e2e/online_serving/test_qwen_image_layered.py
index b29609351..298f4ca9e 100644
--- a/tests/e2e/online_serving/test_qwen_image_layered.py
+++ b/tests/e2e/online_serving/test_qwen_image_layered.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Online serving tests for ``Qwen/Qwen-Image-Layered`` (layer decomposition).
@@ -18,7 +18,7 @@ from tests.helpers.media import generate_synthetic_image
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -44,11 +44,10 @@ def _get_diffusion_feature_cases(model: str):
]
-@pytest.mark.core_model
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_single_image_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Default Qwen-Image-Layered smoke (single ``default`` server config)."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -65,4 +64,4 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
index 315fba6ef..bc36c9fa8 100644
--- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
+++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Comprehensive tests of diffusion features that are available in online serving mode
and are supported by Qwen-Image-Layered model.
@@ -12,13 +12,18 @@ Kept cases (maximum feature coverage, no cpu-offload):
Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
"""
+import base64
+import io
+from typing import Any
+
import pytest
+from PIL import Image
from tests.helpers.mark import hardware_marks
-from tests.helpers.media import decode_b64_image, generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Qwen/Qwen-Image-Layered"
EDIT_PROMPT = "Decompose this image into layers."
@@ -82,7 +87,7 @@ FEATURE_CASES = [
@pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True)
-def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_feature(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test feature combinations with Qwen-Image-Layered."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -99,7 +104,7 @@ def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler):
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
# ── Guard: Issue #1969 – layered output must return `layers` images ─────────
@@ -123,19 +128,19 @@ LAYERS_GUARD_CASES = [
]
-def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config: dict):
+def _collect_image_url_items(online_client: OnlineOmniClient, request_config: dict):
"""Send a request and return all image_url content items from the API response.
Handles both object-style (Pydantic) and dict-style (raw JSON) content items,
because the OpenAI SDK may return either depending on how the server serializes
the multimodal content list.
"""
- chat_completion = openai_client.client.chat.completions.create(
+ chat_completion = online_client.client.chat.completions.create(
model=request_config["model"],
messages=request_config["messages"],
extra_body=request_config.get("extra_body"),
)
- image_items = []
+ image_items: list[Any] = []
for choice in chat_completion.choices:
content = getattr(choice.message, "content", None)
assert content is not None, "API response content is None"
@@ -164,7 +169,7 @@ def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config:
)
def test_layered_output_image_count(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
expected_layers: int,
):
"""Guard for https://github.com/vllm-project/vllm-omni/issues/1969
@@ -189,7 +194,7 @@ def test_layered_output_image_count(
},
}
- image_items = _collect_image_url_items(openai_client, request_config)
+ image_items = _collect_image_url_items(online_client, request_config)
# Core assertion: the API must NOT collapse multiple layers into 1 image
assert len(image_items) != 1 or expected_layers == 1, (
@@ -209,7 +214,8 @@ def test_layered_output_image_count(
url = item.image_url.url
assert url.startswith("data:image"), f"image_url item {i} is not a data URI: {url[:80]}"
b64 = url.split(",", 1)[1]
- img = decode_b64_image(b64)
+ img = Image.open(io.BytesIO(base64.b64decode(b64)))
+ img.load()
assert img is not None, f"Failed to decode image at index {i}"
@@ -231,7 +237,7 @@ PROMPT_CASES = [
@pytest.mark.parametrize("omni_server", PROMPT_CASES, indirect=True)
-def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_empty_prompt(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test feature combinations with Qwen-Image-Layered."""
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
@@ -248,4 +254,4 @@ def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandle
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_sana_video_expansion.py b/tests/e2e/online_serving/test_sana_video_expansion.py
new file mode 100644
index 000000000..31e0f8552
--- /dev/null
+++ b/tests/e2e/online_serving/test_sana_video_expansion.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""SANA-Video 2B online smoke tests for the native and Diffusers backends."""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL_480P = os.environ.get(
+ "SANA_VIDEO_480P_MODEL",
+ "Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+)
+MODEL_720P = os.environ.get(
+ "SANA_VIDEO_720P_MODEL",
+ "Efficient-Large-Model/SANA-Video_2B_720p_diffusers",
+)
+PROMPT = "A cat walking on grass toward the camera. motion score: 30."
+NEGATIVE_PROMPT = "blurry, low quality, temporal artifacts"
+
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _t2v_backend_cases():
+ cases = []
+ for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)):
+ cases.extend(
+ [
+ pytest.param(
+ OmniServerParams(model=model, server_args=["--model-class-name", "SanaVideoPipeline"]),
+ id=f"native-{variant}",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=model,
+ server_args=[
+ "--diffusion-load-format",
+ "diffusers",
+ "--diffusion-attention-backend",
+ "TORCH_SDPA",
+ ],
+ ),
+ id=f"diffusers-adapter-{variant}",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ ]
+ )
+ return cases
+
+
+@pytest.mark.parametrize("omni_server", _t2v_backend_cases(), indirect=True)
+def test_sana_video_t2v_backends(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+ request_config = {
+ "model": omni_server.model,
+ "form_data": {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ "height": 192,
+ "width": 320,
+ "num_frames": 9,
+ "fps": 8,
+ "num_inference_steps": 2,
+ "guidance_scale": 4.0,
+ "seed": 42,
+ },
+ }
+ openai_client.send_video_diffusion_request(request_config)
+
+
+def _i2v_cases():
+ cases = []
+ for variant, model in (("480p", MODEL_480P), ("720p", MODEL_720P)):
+ cases.extend(
+ [
+ pytest.param(
+ OmniServerParams(model=model, server_args=["--model-class-name", "SanaImageToVideoPipeline"]),
+ id=f"native-{variant}",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ OmniServerParams(
+ model=model,
+ server_args=[
+ "--model-class-name",
+ "SanaImageToVideoPipeline",
+ "--diffusion-load-format",
+ "diffusers",
+ "--diffusion-attention-backend",
+ "TORCH_SDPA",
+ ],
+ ),
+ id=f"diffusers-adapter-{variant}",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ ]
+ )
+ return cases
+
+
+@pytest.mark.parametrize("omni_server", _i2v_cases(), indirect=True)
+def test_sana_video_i2v_variants(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+ request_config = {
+ "model": omni_server.model,
+ "form_data": {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ "height": 192,
+ "width": 320,
+ "num_frames": 9,
+ "fps": 8,
+ "num_inference_steps": 2,
+ "guidance_scale": 4.0,
+ "seed": 42,
+ },
+ "image_reference": (f"data:image/jpeg;base64,{generate_synthetic_image(320, 192, seed=42)['base64']}"),
+ }
+ openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_sana_wm.py b/tests/e2e/online_serving/test_sana_wm.py
new file mode 100644
index 000000000..883e4c367
--- /dev/null
+++ b/tests/e2e/online_serving/test_sana_wm.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for SANA-WM image-to-video (first-frame I2V via ``/v1/videos``).
+
+Boots the server straight from the model id — the Stage-1 repo's
+``model_index.json`` names ``SanaWmPipeline``, so the class resolves on its own
+— submits one async ``/v1/videos`` job with a first-frame reference image, and
+asserts the job completes and returns video bytes.
+
+From ``tests/``::
+
+ pytest -s -v e2e/online_serving/test_sana_wm.py -m "advanced_model and diffusion" --run-level=advanced_model
+"""
+
+import base64
+import json
+import os
+from io import BytesIO
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from vllm_omni.diffusion.models.sana_wm import (
+ SANA_WM_MODEL_ID,
+ SANA_WM_OUTPUT_HEIGHT,
+ SANA_WM_OUTPUT_WIDTH,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = SANA_WM_MODEL_ID
+PROMPT = "A slow forward camera move through a quiet city street."
+NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
+
+# Smoke knobs: a handful of frames / steps so the job finishes quickly.
+SMOKE_NUM_FRAMES = 9
+SMOKE_NUM_INFERENCE_STEPS = 2
+# First-frame I2V needs a camera trajectory; "w-" is the forward-move action DSL.
+# Passed via the ``sana_wm`` form field (a JSON string parsed server-side); explicit
+# intrinsics avoid the optional Pi3X camera-calibration dependency.
+SANA_WM_PARAMS = {
+ "action": f"w-{SMOKE_NUM_FRAMES - 1}",
+ "translation_speed": 0.055,
+ "rotation_speed_deg": 1.2,
+ "intrinsics": {
+ "fx": SANA_WM_OUTPUT_WIDTH / 2,
+ "fy": SANA_WM_OUTPUT_WIDTH / 2,
+ "cx": SANA_WM_OUTPUT_WIDTH / 2,
+ "cy": SANA_WM_OUTPUT_HEIGHT / 2,
+ },
+}
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _first_frame_data_url() -> str:
+ """A solid-color first frame encoded as a PNG ``data:`` URL for I2V."""
+ from PIL import Image
+
+ image = Image.new("RGB", (SANA_WM_OUTPUT_WIDTH, SANA_WM_OUTPUT_HEIGHT), (96, 128, 160))
+ buffer = BytesIO()
+ image.save(buffer, format="PNG")
+ encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+ return f"data:image/png;base64,{encoded}"
+
+
+def _get_diffusion_feature_cases(model: str):
+ """Single default server row."""
+ return [
+ pytest.param(
+ OmniServerParams(model=model),
+ id="default",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_image_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+ """Default SANA-WM I2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
+ request_config = {
+ "model": omni_server.model,
+ "image_reference": _first_frame_data_url(),
+ "form_data": {
+ "prompt": PROMPT,
+ "negative_prompt": NEGATIVE_PROMPT,
+ # The camera payload rides inside extra_params; there is no
+ # top-level sana_wm form field on the video endpoints, so sending
+ # one is silently dropped and the request fails in the
+ # preprocessor for missing camera control.
+ "extra_params": json.dumps({"sana_wm": SANA_WM_PARAMS}),
+ "height": SANA_WM_OUTPUT_HEIGHT,
+ "width": SANA_WM_OUTPUT_WIDTH,
+ "num_frames": SMOKE_NUM_FRAMES,
+ "fps": 16,
+ "num_inference_steps": SMOKE_NUM_INFERENCE_STEPS,
+ "guidance_scale": 5.0,
+ "seed": 42,
+ },
+ }
+ openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py
index f881a212b..9793d8a4b 100644
--- a/tests/e2e/online_serving/test_sd3_expansion.py
+++ b/tests/e2e/online_serving/test_sd3_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Tests for Stable Diffusion 3.5 medium model.
"""
@@ -8,7 +11,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -51,10 +54,7 @@ def _get_diffusion_feature_cases(model: str):
"2",
],
),
- marks=[
- *FOUR_CARD_FEATURE_MARKS,
- pytest.mark.skip(reason="#3432"),
- ],
+ marks=FOUR_CARD_FEATURE_MARKS,
),
]
@@ -66,7 +66,7 @@ def _get_diffusion_feature_cases(model: str):
),
indirect=True,
)
-def test_sd3_medium(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_sd3_medium(omni_server: OmniServer, online_client: OnlineOmniClient):
messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
request_config = {
"model": omni_server.model,
@@ -81,4 +81,4 @@ def test_sd3_medium(omni_server: OmniServer, openai_client: OpenAIClientHandler)
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_sdxl.py b/tests/e2e/online_serving/test_sdxl.py
index 4d7fc34d5..052d9e57e 100644
--- a/tests/e2e/online_serving/test_sdxl.py
+++ b/tests/e2e/online_serving/test_sdxl.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Tests for Stable Diffusion XL Base 1.0 model.
@@ -10,7 +13,7 @@ from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import (
OmniServer,
OmniServerParams,
- OpenAIClientHandler,
+ OnlineOmniClient,
dummy_messages_from_mix_data,
)
@@ -41,7 +44,7 @@ def _get_diffusion_feature_cases(model: str):
),
indirect=True,
)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_text_to_image_001(omni_server: OmniServer, online_client: OnlineOmniClient):
messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
request_config = {
"model": omni_server.model,
@@ -56,4 +59,4 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_stable_audio_online_expansion.py b/tests/e2e/online_serving/test_stable_audio_online_expansion.py
new file mode 100644
index 000000000..fbd6bdbfd
--- /dev/null
+++ b/tests/e2e/online_serving/test_stable_audio_online_expansion.py
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving test for Stable Audio Open text-to-audio diffusion.
+
+Stable Audio Open is served through the OpenAI-compatible
+`POST /v1/audio/generate` endpoint (JSON in, binary WAV out).
+"""
+
+import os
+
+import pytest
+
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+# Stable Audio Open is gated; override to a local bundle locally if needed.
+STABLE_AUDIO_TEST_MODEL = os.environ.get("STABLE_AUDIO_TEST_MODEL", "stabilityai/stable-audio-open-1.0")
+T2A_PROMPT = "A piano playing a gentle melody with soft room ambience."
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
+
+
+def _stable_audio_server_cases(model: str):
+ return [
+ pytest.param(
+ OmniServerParams(model=model),
+ id="t2a",
+ marks=SINGLE_CARD_FEATURE_MARKS,
+ ),
+ ]
+
+
+@pytest.fixture
+def _require_stable_audio_access() -> None:
+ """Skip cleanly (before the server boots) if the gated checkpoint is inaccessible."""
+ skip_if_gated_repo_inaccessible(STABLE_AUDIO_TEST_MODEL)
+
+
+@pytest.mark.slow
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _stable_audio_server_cases(STABLE_AUDIO_TEST_MODEL), indirect=True)
+def test_stable_audio_t2a_online(
+ _require_stable_audio_access: None,
+ omni_server: OmniServer,
+ openai_client: OpenAIClientHandler,
+) -> None:
+ """Stable Audio Open text-to-audio: `/v1/audio/generate` returns a non-empty WAV.
+
+ Uses tiny steps / short duration to keep CI light, matching the offline smoke
+ test in `tests/e2e/offline_inference/test_stable_audio_expansion.py`.
+ """
+ request_config = {
+ "json": {
+ "model": omni_server.model,
+ "input": T2A_PROMPT,
+ "audio_length": 2.0,
+ "num_inference_steps": 4,
+ "guidance_scale": 7.0,
+ "negative_prompt": "Low quality.",
+ "seed": 42,
+ "response_format": "wav",
+ },
+ "timeout": 300,
+ }
+ responses = openai_client.send_audio_generate_http_request(request_config)
+ assert responses, "no response from /v1/audio/generate"
+ resp = responses[0]
+ assert resp.success, f"audio generate failed: {resp.status_code} {resp.error_message}"
+ assert resp.status_code == 200
diff --git a/tests/e2e/online_serving/test_step_audio2_expansion.py b/tests/e2e/online_serving/test_step_audio2_expansion.py
index 9c8280b59..5542d81df 100644
--- a/tests/e2e/online_serving/test_step_audio2_expansion.py
+++ b/tests/e2e/online_serving/test_step_audio2_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E expansion tests for Step-Audio2 online serving (weekly CI).
@@ -94,22 +94,22 @@ def _build_audio_to_text_request_config(omni_server, audio_base64: str) -> dict:
@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_single_audio_to_text_request(omni_server, openai_client) -> None:
+def test_single_audio_to_text_request(omni_server, online_client) -> None:
"""Test a single audio-to-text request via the chat completions API."""
request_config = _build_audio_to_text_request_config(
omni_server,
_synthetic_audio_base64(duration_sec=2),
)
- openai_client.send_omni_request(request_config)
+ online_client.send_omni_request(request_config)
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("request_num", [2, 4])
-def test_concurrent_audio_to_text_requests(omni_server, openai_client, request_num: int) -> None:
+def test_concurrent_audio_to_text_requests(omni_server, online_client, request_num: int) -> None:
"""Test concurrent audio-to-text requests."""
request_config = _build_audio_to_text_request_config(
omni_server,
_synthetic_audio_base64(duration_sec=2),
)
- responses = openai_client.send_omni_request(request_config, request_num=request_num)
+ responses = online_client.send_omni_request(request_config, request_num=request_num)
assert len(responses) == request_num
diff --git a/tests/e2e/online_serving/test_voxcpm2_tts.py b/tests/e2e/online_serving/test_voxcpm2_tts.py
index 19ddc0401..0d0364138 100644
--- a/tests/e2e/online_serving/test_voxcpm2_tts.py
+++ b/tests/e2e/online_serving/test_voxcpm2_tts.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online tests for VoxCPM2 with text input and audio output.
@@ -46,12 +46,11 @@ tts_server_params = [
]
-@pytest.mark.advanced_model
-@pytest.mark.core_model
+@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_server, openai_client) -> None:
+def test_text_to_audio_001(omni_server, online_client) -> None:
"""
Test zero-shot text-to-audio via OpenAI API.
Deploy Setting: voxcpm2.yaml
@@ -69,14 +68,14 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
"voice": "default",
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
+ online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.tts
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_text_to_audio_002(omni_server, openai_client) -> None:
+def test_text_to_audio_002(omni_server, online_client) -> None:
"""
Test zero-shot streaming text-to-audio via OpenAI API.
Deploy Setting: voxcpm2.yaml
@@ -95,4 +94,4 @@ def test_text_to_audio_002(omni_server, openai_client) -> None:
"voice": "default",
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
index 996585a1d..d630c8627 100644
--- a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
+++ b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E Online expansion tests for VoxCPM2 (voice clone, PCM, concurrency).
@@ -12,11 +12,11 @@ import os
import pytest
from tests.helpers.mark import hardware_test
-from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -25,7 +25,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0
_MIN_AUDIO_BYTES = 40_000
MAX_CONCURRENT = 4
-REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
def get_prompt(prompt_type="text"):
@@ -49,7 +49,7 @@ tts_server_params = [
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
+def test_voice_clone_streaming_001(omni_server, online_client) -> None:
"""
Test voice-clone text-to-audio via OpenAI API.
Deploy Setting: voxcpm2.yaml
@@ -69,12 +69,12 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
"ref_audio": REF_AUDIO_URL,
"min_audio_bytes": _MIN_AUDIO_BYTES,
}
- openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
+ online_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT)
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_response_format_001(omni_server, openai_client) -> None:
+def test_response_format_001(omni_server, online_client) -> None:
"""
Test voice-clone non-stream PCM output via OpenAI API.
Deploy Setting: voxcpm2.yaml
@@ -94,4 +94,4 @@ def test_response_format_001(omni_server, openai_client) -> None:
"min_audio_bytes": _MIN_AUDIO_BYTES,
"min_hnr_db": -2.0,
}
- openai_client.send_audio_speech_request(request_config)
+ online_client.send_audio_speech_request(request_config)
diff --git a/tests/e2e/online_serving/test_voxtral_tts_expansion.py b/tests/e2e/online_serving/test_voxtral_tts_expansion.py
index dc8db9d73..278e563e8 100644
--- a/tests/e2e/online_serving/test_voxtral_tts_expansion.py
+++ b/tests/e2e/online_serving/test_voxtral_tts_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E tests for Voxtral TTS online serving.
@@ -31,9 +31,9 @@ pytestmark = [
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
-def test_speech_english_basic(omni_server, openai_client) -> None:
+def test_speech_english_basic(omni_server, online_client) -> None:
"""Test basic English TTS generation."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "how are you",
@@ -47,9 +47,9 @@ def test_speech_english_basic(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
-def test_speech_english_streaming(omni_server, openai_client) -> None:
+def test_speech_english_streaming(omni_server, online_client) -> None:
"""Test basic streaming English TTS generation (PCM via streaming API)."""
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Hello, how are you?",
@@ -65,11 +65,11 @@ def test_speech_english_streaming(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_speech_different_voices(omni_server, openai_client) -> None:
+def test_speech_different_voices(omni_server, online_client) -> None:
"""Test TTS with different voice presets."""
voices = ["casual_female", "neutral_male"]
for voice in voices:
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Testing voice selection.",
@@ -82,11 +82,11 @@ def test_speech_different_voices(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_speech_speed(omni_server, openai_client) -> None:
+def test_speech_speed(omni_server, online_client) -> None:
"""Request with speed parameters."""
speeds = [0.5, 1, 1.5, 2, 2.5]
for speed in speeds:
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The boy was there when the sun rose.",
@@ -100,7 +100,7 @@ def test_speech_speed(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_speech_instructions(omni_server, openai_client) -> None:
+def test_speech_instructions(omni_server, online_client) -> None:
"""Request with instructions parameters."""
instructions = [
"Speak formally",
@@ -109,7 +109,7 @@ def test_speech_instructions(omni_server, openai_client) -> None:
"Speak with a chirpy happy voice",
]
for instruction in instructions:
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "The boy was there when the sun rose.",
@@ -134,11 +134,11 @@ def test_speech_instructions(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_speech_response_formats(omni_server, openai_client) -> None:
+def test_speech_response_formats(omni_server, online_client) -> None:
"""Test TTS with different response formats."""
response_formats = ["wav", "mp3"]
for response_format in response_formats:
- openai_client.send_audio_speech_request(
+ online_client.send_audio_speech_request(
{
"model": omni_server.model,
"input": "Testing various response formats.",
@@ -151,7 +151,7 @@ def test_speech_response_formats(omni_server, openai_client) -> None:
@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_speech_batches(omni_server, openai_client) -> None:
+def test_speech_batches(omni_server, online_client) -> None:
"""Test TTS batches."""
items = [
{"input": "The birch canoe slid on the smooth planks."},
@@ -161,7 +161,7 @@ def test_speech_batches(omni_server, openai_client) -> None:
{"input": "Rice is often served in round bowls."},
]
- openai_client.send_audio_speech_batch_http_request(
+ online_client.send_audio_speech_batch_http_request(
{
"json": {
"model": omni_server.model,
diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py
index 85eee128a..322cde3cd 100644
--- a/tests/e2e/online_serving/test_wan22_expansion.py
+++ b/tests/e2e/online_serving/test_wan22_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Comprehensive tests of diffusion features that are available in online serving mode.
@@ -19,7 +22,7 @@ import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
@@ -140,7 +143,7 @@ def _get_wan22_feature_cases():
)
def test_wan22_diffusion_features(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
model_path = omni_server.model
is_i2v_or_ti2v = any(kw in model_path for kw in ["I2V", "TI2V"])
@@ -176,4 +179,4 @@ def test_wan22_diffusion_features(
if is_i2v_or_ti2v:
request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
- openai_client.send_video_diffusion_request(request_config)
+ online_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_wan22_t2v.py b/tests/e2e/online_serving/test_wan22_t2v.py
index e6fe38086..f3f21ae7d 100644
--- a/tests/e2e/online_serving/test_wan22_t2v.py
+++ b/tests/e2e/online_serving/test_wan22_t2v.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Online serving smoke for ``Wan-AI/Wan2.2-T2V-A14B-Diffusers`` (text-to-video via ``/v1/videos``).
@@ -18,7 +18,7 @@ import os
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -44,7 +44,7 @@ def _get_diffusion_feature_cases(model: str):
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+def test_text_to_video_001(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
"""Default Wan2.2 T2V smoke: async ``/v1/videos`` job completes and returns video bytes."""
request_config = {
"model": omni_server.model,
@@ -60,4 +60,4 @@ def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientH
"seed": 42,
},
}
- openai_client.send_video_diffusion_request(request_config)
+ online_client.send_video_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py
index d58bb9e32..cf24e34b0 100644
--- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py
+++ b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Comprehensive e2e tests of diffusion features for Wan2.1-VACE in online serving mode.
@@ -24,9 +24,9 @@ Coverage:
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
PROMPT = "A cat walking slowly across a sunlit garden path"
@@ -150,9 +150,9 @@ def _get_vace_feature_cases():
_get_vace_feature_cases(),
indirect=True,
)
-def test_wan_2_1_vace(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+def test_wan_2_1_vace(omni_server: OmniServer, online_client: OnlineOmniClient):
"""Test VACE T2V generation with all supported diffusion acceleration features."""
- openai_client.send_video_diffusion_request(
+ online_client.send_video_diffusion_request(
{
"model": MODEL,
"form_data": {
diff --git a/tests/e2e/online_serving/test_z_image.py b/tests/e2e/online_serving/test_z_image.py
index 2ef78ce0f..7e3af1132 100644
--- a/tests/e2e/online_serving/test_z_image.py
+++ b/tests/e2e/online_serving/test_z_image.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Tests of common diffusion feature combinations in online serving mode
for Z-Image.
@@ -11,7 +14,7 @@ exercise Z-Image's supported parallel feature combinations:
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
MODEL = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard."
@@ -32,7 +35,7 @@ def _get_diffusion_feature_cases():
]
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -41,7 +44,7 @@ def _get_diffusion_feature_cases():
)
def test_basic_001(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Exercise supported Z-Image diffusion features in minimal CI cases."""
messages = dummy_messages_from_mix_data(content_text=PROMPT)
@@ -57,10 +60,10 @@ def test_basic_001(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
-@pytest.mark.advanced_model
+@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
@@ -69,7 +72,7 @@ def test_basic_001(
)
def test_different_sizes_001(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Exercise supported Z-Image diffusion features in minimal CI cases."""
messages = dummy_messages_from_mix_data(content_text=PROMPT)
@@ -85,4 +88,4 @@ def test_different_sizes_001(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py
index ac27267e5..c37a947f4 100644
--- a/tests/e2e/online_serving/test_zimage_expansion.py
+++ b/tests/e2e/online_serving/test_zimage_expansion.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Tests of common diffusion feature combinations in online serving mode
for Z-Image.
@@ -14,7 +17,7 @@ exercise Z-Image's supported feature combinations:
import pytest
from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
@@ -110,7 +113,7 @@ def _get_diffusion_feature_cases():
)
def test_zimage(
omni_server: OmniServer,
- openai_client: OpenAIClientHandler,
+ online_client: OnlineOmniClient,
):
"""Exercise supported Z-Image diffusion features in minimal CI cases."""
request_config = {
@@ -125,4 +128,4 @@ def test_zimage(
},
}
- openai_client.send_diffusion_request(request_config)
+ online_client.send_diffusion_request(request_config)
diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py
index 055866bd5..fb8231a75 100644
--- a/tests/engine/test_arg_utils.py
+++ b/tests/engine/test_arg_utils.py
@@ -4,11 +4,12 @@ invariant to the specific attributes of vLLM config except in cases where we
explicitly patch values that differ from vLLM.
"""
-import argparse
import inspect
import json
import os
import shutil
+import sys
+import types
from types import SimpleNamespace
from unittest.mock import Mock
@@ -20,8 +21,9 @@ from vllm.engine.arg_utils import EngineArgs
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,6 +34,11 @@ def test_sync_config_is_omni():
assert isinstance(cfg, OmniModelConfig)
+def test_session_mode_reaches_omni_model_config():
+ assert OmniEngineArgs().create_model_config().session_mode == "turn"
+ assert OmniEngineArgs(session_mode="duplex").create_model_config().session_mode == "duplex"
+
+
def test_default_stage_id_is_concrete_int():
"""Ensure `stage_id` stays safe for downstream arithmetic/indexing."""
engine_args = OmniEngineArgs()
@@ -44,7 +51,102 @@ def test_default_stage_id_is_concrete_int():
assert cfg.stage_id == 0
-def test_multimodal_kwarg_overrides(mocker):
+def test_full_payload_capability_reaches_omni_model_config(monkeypatch):
+ """Keep the pipeline capability intact through the engine adapter."""
+ captured: dict[str, object] = {}
+ baseline_config = Mock()
+ worker_module = types.ModuleType("test_connector_worker")
+
+ class ConnectorRunner(OmniConnectorModelRunnerMixin):
+ pass
+
+ class ConnectorWorker:
+ model_runner_cls = ConnectorRunner
+
+ worker_module.ConnectorWorker = ConnectorWorker
+ monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+ monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None)
+ monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config)
+
+ def capture_omni_config(_cls, model_config, **omni_kwargs):
+ assert model_config is baseline_config
+ captured.update(omni_kwargs)
+ return baseline_config
+
+ monkeypatch.setattr(
+ OmniModelConfig,
+ "from_vllm_model_config",
+ classmethod(capture_omni_config),
+ )
+
+ OmniEngineArgs(
+ stage_id=1,
+ requires_full_payload_input=True,
+ worker_cls="test_connector_worker.ConnectorWorker",
+ ).create_model_config()
+
+ assert captured["requires_full_payload_input"] is True
+
+
+def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch):
+ worker_module = types.ModuleType("test_async_scheduler_worker")
+
+ class WorkerWithoutConnector:
+ model_runner_cls = object
+
+ worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+ monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+ args = OmniEngineArgs(
+ stage_id=1,
+ async_chunk=True,
+ worker_cls="test_async_scheduler_worker.WorkerWithoutConnector",
+ )
+
+ assert args.requires_full_payload_input is False
+
+
+def test_full_payload_capability_requires_selected_worker_connector(monkeypatch):
+ worker_module = types.ModuleType("test_worker_without_connector")
+
+ class WorkerWithoutConnector:
+ model_runner_cls = object
+
+ worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+ monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
+ with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+ OmniEngineArgs(
+ stage_id=1,
+ requires_full_payload_input=True,
+ worker_cls="test_worker_without_connector.WorkerWithoutConnector",
+ )
+
+
+def test_full_payload_capability_validates_platform_selected_worker(monkeypatch):
+ worker_module = types.ModuleType("test_platform_worker_without_connector")
+
+ class WorkerWithoutConnector:
+ model_runner_cls = object
+
+ worker_module.WorkerWithoutConnector = WorkerWithoutConnector
+ monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+ monkeypatch.setattr(
+ current_omni_platform,
+ "get_omni_ar_worker_cls",
+ lambda: "test_platform_worker_without_connector.WorkerWithoutConnector",
+ )
+
+ with pytest.raises(ValueError, match="does not provide an Omni connector model runner"):
+ OmniEngineArgs(
+ stage_id=1,
+ requires_full_payload_input=True,
+ worker_type="ar",
+ )
+
+
+def test_multimodal_kwarg_overrides(monkeypatch):
"""Ensure that overrides in the multimodal config are preserved."""
sig = inspect.signature(OmniEngineArgs)
default_mm_cache = sig.parameters["mm_processor_cache_gb"].default
@@ -58,8 +160,8 @@ def test_multimodal_kwarg_overrides(mocker):
assert self.mm_processor_cache_gb == override_val
return fake_model_config
- mocker.patch.object(EngineArgs, "create_model_config", _fake_parent_create_model_config)
- mocker.patch.object(OmniModelConfig, "from_vllm_model_config", side_effect=lambda model_config, **_: model_config)
+ monkeypatch.setattr(EngineArgs, "create_model_config", _fake_parent_create_model_config)
+ monkeypatch.setattr(OmniModelConfig, "from_vllm_model_config", lambda model_config, **_: model_config)
cfg = OmniEngineArgs(
model="Qwen/Qwen2-VL-2B-Instruct",
@@ -149,19 +251,6 @@ def test_qwen3_tts_startup_task_type_is_validated():
)
-def test_from_cli_args_picks_up_stage_configs_path():
- """from_cli_args should pick up stage_configs_path from namespace."""
- ns = argparse.Namespace(
- model="facebook/opt-125m",
- stage_configs_path="/some/path.yaml",
- custom_pipeline_args=None,
- )
-
- args = OmniEngineArgs.from_cli_args(ns)
- assert args.stage_configs_path == "/some/path.yaml"
- assert args.custom_pipeline_args is None
-
-
def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
"""Ensure Code2Wav mirrors stage max_model_len into nested HF overrides.
@@ -197,6 +286,28 @@ def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
}
+def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, monkeypatch, mocker):
+ """A successful remote tokenizer download must not enter the error path."""
+ tokenizer_dir = tmp_path / "CosyVoice-BlankEN"
+ tokenizer_dir.mkdir()
+ baseline_config = Mock()
+ warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning")
+
+ monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *args, **kwargs: str(tmp_path))
+ monkeypatch.setattr(OmniEngineArgs, "_patch_empty_hf_config", lambda *args, **kwargs: None)
+ monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config)
+ monkeypatch.setattr(
+ OmniModelConfig,
+ "from_vllm_model_config",
+ classmethod(lambda cls, model_config, **omni_kwargs: model_config),
+ )
+
+ args = OmniEngineArgs(model="remote/cosyvoice3", model_arch="CosyVoice3Model")
+ assert args.create_model_config() is baseline_config
+ assert args.tokenizer == str(tokenizer_dir)
+ warning.assert_not_called()
+
+
def test_patch_missing_local_hf_config(tmp_path):
"""Native model bundles may only contain a checkpoints/ directory."""
(tmp_path / "checkpoints").mkdir()
@@ -418,130 +529,6 @@ def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
-def test_stage_configs_path_field():
- """OmniEngineArgs with stage_configs_path should construct without error."""
- args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
- assert args.stage_configs_path == "/some/path.yaml"
-
-
-def test_strip_single_engine_args():
- """_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
- kwargs = {
- # Parent EngineArgs fields — stripped unless explicitly allowlisted
- "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
- "tensor_parallel_size": 4,
- "gpu_memory_utilization": 0.9,
- "model": "some/model",
- # Parent field that should be kept (allowlisted)
- "worker_extension_cls": "some.Extension",
- # OmniEngineArgs-only / non-engine fields — should pass through
- "stage_configs_path": "/path/to/yaml",
- "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
- "mode": "text-to-image",
- "lora_path": "/some/lora",
- }
-
- filtered = AsyncOmniEngine._strip_single_engine_args(kwargs)
-
- # Stripped — parent EngineArgs fields
- assert "compilation_config" not in filtered
- assert filtered["tensor_parallel_size"] == 4
- assert "gpu_memory_utilization" not in filtered
- assert "model" not in filtered
-
- # Stripped — orchestrator-level OmniEngineArgs field
- assert "stage_configs_path" not in filtered
-
- # Kept
- assert filtered["worker_extension_cls"] == "some.Extension"
- assert filtered["custom_pipeline_args"] == {"pipeline_class": "my.Pipeline"}
- assert filtered["mode"] == "text-to-image"
- assert filtered["lora_path"] == "/some/lora"
-
-
-def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
- """model is always in kwargs (callers set it via from_cli_args/asdict),
- so it should not cause the override warning by itself or appear in it."""
- mock_warn = mocker.patch("vllm_omni.engine.async_omni_engine.logger.warning")
-
- # Typical caller kwargs: model is always present, no other parent
- # EngineArgs fields are explicitly overridden.
- AsyncOmniEngine._strip_single_engine_args(
- {
- "model": "some/model",
- "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
- }
- )
- mock_warn.assert_not_called()
-
- # When there *are* genuinely surprising overrides alongside model,
- # the warning should mention them but not model. Keep-listed fields such as
- # tensor_parallel_size are intentionally passed through and should not warn.
- AsyncOmniEngine._strip_single_engine_args(
- {
- "model": "some/model",
- "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
- "tensor_parallel_size": 4,
- "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
- }
- )
- mock_warn.assert_called_once()
- warned_args = mock_warn.call_args[0][-1] # the formatted arg list
- assert "compilation_config" in warned_args
- assert "tensor_parallel_size" not in warned_args
- assert "model" not in warned_args
-
-
-def test_new_format_deploy_yaml_does_not_strip_interleave_cli_overrides(tmp_path):
- """New-format deploy YAMLs must not strip interleave_mm_strings / media_io_kwargs.
-
- Pytest historically passed ``--stage-configs-path`` for deploy YAMLs; stripping
- those CLI flags dropped MiniCPM interleaved AV packing and cut Daily-Omni
- accuracy from ~78% to ~70%.
- """
- from vllm_omni.entrypoints.utils import is_new_format_deploy_config
-
- yaml_path = tmp_path / "minicpm_deploy.yaml"
- yaml_path.write_text(
- """
-pipeline: minicpmo_4_5
-async_chunk: true
-stages:
- - stage_id: 0
- devices: "0"
-""".strip(),
- encoding="utf-8",
- )
- path = str(yaml_path)
- assert is_new_format_deploy_config(path)
-
- kwargs = {
- "stage_configs_path": path,
- "trust_remote_code": True,
- "interleave_mm_strings": True,
- "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
- "gpu_memory_utilization": 0.55,
- }
- # Mirror AsyncOmniEngine._resolve_stage_configs strip decision.
- if kwargs.get("stage_configs_path") is not None and not is_new_format_deploy_config(kwargs["stage_configs_path"]):
- base = AsyncOmniEngine._strip_single_engine_args(dict(kwargs))
- else:
- base = dict(kwargs)
- assert base.get("interleave_mm_strings") is True
- assert base.get("media_io_kwargs") == {"video": {"fps": 1, "num_frames": 128}}
- # Parent EngineArgs fields also remain available for deploy merge.
- assert base.get("gpu_memory_utilization") == 0.55
-
- legacy = tmp_path / "legacy.yaml"
- legacy.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
- assert not is_new_format_deploy_config(str(legacy))
- legacy_kwargs = dict(kwargs, stage_configs_path=str(legacy))
- stripped = AsyncOmniEngine._strip_single_engine_args(legacy_kwargs)
- assert "interleave_mm_strings" not in stripped
- assert "media_io_kwargs" not in stripped
- assert "gpu_memory_utilization" not in stripped
-
-
# For https://github.com/vllm-project/vllm-omni/issues/3293
def test_tensor_parallel_size_none_is_handled():
"""Ensure the tensor parallel size of None isn't forwarded."""
diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py
index 40f89df84..56b25704d 100644
--- a/tests/engine/test_async_omni_engine_abort.py
+++ b/tests/engine/test_async_omni_engine_abort.py
@@ -68,7 +68,7 @@ async def test_abort():
with ExitStack() as after:
engine = AsyncOmni(
model=model,
- stage_configs_path=stage_config,
+ deploy_config=stage_config,
)
after.callback(engine.shutdown)
diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py
new file mode 100644
index 000000000..86276ec66
--- /dev/null
+++ b/tests/engine/test_async_omni_engine_abort_ack.py
@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for acknowledged abort on AsyncOmniEngine."""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+from types import SimpleNamespace
+
+import pytest
+from pytest_mock import MockerFixture
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage
+from vllm_omni.engine.rpc_result_router import CorrelatedRpcClient
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine:
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = SimpleNamespace(sync_q=request_q)
+ engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q)
+ engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
+ return engine
+
+
+def test_abort_result_message_correlation_key():
+ msg = AbortResultMessage(rpc_id="abc", success=True)
+ assert msg.rpc_correlation_key == ("abort", "abc")
+
+
+def test_sync_abort_is_fire_and_forget():
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+
+ engine.abort(["req-1"])
+
+ msg = request_q.get_nowait()
+ assert isinstance(msg, AbortRequestMessage)
+ assert msg.request_ids == ["req-1"]
+ assert msg.rpc_id is None
+ assert rpc_q.empty()
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_waits_for_ack(mocker: MockerFixture):
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+ return_value=SimpleNamespace(hex="abort-rpc-1"),
+ )
+
+ async def _run() -> None:
+ task = asyncio.create_task(engine.abort_async(["req-a", "req-b"], timeout=1))
+ msg = await asyncio.to_thread(request_q.get, True, 1)
+ assert isinstance(msg, AbortRequestMessage)
+ assert msg.request_ids == ["req-a", "req-b"]
+ assert msg.rpc_id == "abort-rpc-1"
+ assert not task.done()
+
+ rpc_q.put(AbortResultMessage(rpc_id="abort-rpc-1", success=True))
+ result = await task
+ assert result == []
+
+ try:
+ asyncio.run(_run())
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_raises_on_orchestrator_error(mocker: MockerFixture):
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+ return_value=SimpleNamespace(hex="abort-rpc-err"),
+ )
+
+ async def _run() -> None:
+ task = asyncio.create_task(engine.abort_async(["req-x"], timeout=1))
+ await asyncio.to_thread(request_q.get, True, 1)
+ rpc_q.put(
+ AbortResultMessage(
+ rpc_id="abort-rpc-err",
+ success=False,
+ error="stage abort failed",
+ )
+ )
+ with pytest.raises(RuntimeError, match="stage abort failed"):
+ await task
+
+ try:
+ asyncio.run(_run())
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_times_out_without_result(mocker: MockerFixture):
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+ return_value=SimpleNamespace(hex="abort-rpc-timeout"),
+ )
+
+ async def _run() -> None:
+ with pytest.raises(TimeoutError, match="abort timed out"):
+ await engine.abort_async(["req-timeout"], timeout=0.05)
+
+ try:
+ asyncio.run(_run())
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture):
+ class SignallingQueue(queue.Queue):
+ def __init__(self) -> None:
+ super().__init__(maxsize=1)
+ self.put_attempted = threading.Event()
+
+ def put(self, item, block=True, timeout=None):
+ self.put_attempted.set()
+ return super().put(item, block=block, timeout=timeout)
+
+ request_q = SignallingQueue()
+ request_q.put("queue-is-full")
+ request_q.put_attempted.clear()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.uuid.uuid4",
+ return_value=SimpleNamespace(hex="blocked-abort"),
+ )
+
+ async def _run() -> None:
+ task = asyncio.create_task(engine.abort_async(["req-blocked"], timeout=1))
+ assert await asyncio.to_thread(request_q.put_attempted.wait, 1)
+ assert not task.done()
+
+ assert request_q.get(timeout=1) == "queue-is-full"
+ msg = await asyncio.to_thread(request_q.get, True, 1)
+ assert isinstance(msg, AbortRequestMessage)
+ assert msg.rpc_id == "blocked-abort"
+ rpc_q.put(AbortResultMessage(rpc_id="blocked-abort", success=True))
+ await task
+
+ try:
+ asyncio.run(_run())
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_is_noop_after_shutdown_starts():
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = True
+ engine.request_queue = None
+ engine._correlated_rpc_client = None
+
+ asyncio.run(engine.abort_async(["req-late"]))
+
+
+def test_abort_async_is_noop_for_empty_request_ids():
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = None
+ engine._correlated_rpc_client = None
+
+ asyncio.run(engine.abort_async([]))
+
+
+def test_abort_async_tolerates_request_queue_close_during_shutdown(
+ mocker: MockerFixture,
+):
+ class SyntheticSyncQueueShutDownError(Exception):
+ pass
+
+ rpc_q: queue.Queue = queue.Queue()
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+
+ class ClosingRequestQueue:
+ def put(self, item) -> None:
+ del item
+ engine._shutdown_called = True
+ raise SyntheticSyncQueueShutDownError
+
+ request_q = ClosingRequestQueue()
+ engine.request_queue = SimpleNamespace(sync_q=request_q)
+ engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ SyntheticSyncQueueShutDownError,
+ )
+
+ try:
+ asyncio.run(engine.abort_async(["req-submit-race"], timeout=1))
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_surfaces_request_queue_close_while_engine_is_live(
+ mocker: MockerFixture,
+):
+ class SyntheticSyncQueueShutDownError(Exception):
+ pass
+
+ class ClosedRequestQueue:
+ def put(self, item) -> None:
+ del item
+ raise SyntheticSyncQueueShutDownError
+
+ request_q = ClosedRequestQueue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = SimpleNamespace(sync_q=request_q)
+ engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q)
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ SyntheticSyncQueueShutDownError,
+ )
+
+ try:
+ with pytest.raises(SyntheticSyncQueueShutDownError):
+ asyncio.run(engine.abort_async(["req-live-submit"], timeout=1))
+ finally:
+ engine._correlated_rpc_client.close()
+
+
+def test_abort_async_tolerates_rpc_router_close_during_shutdown():
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+
+ async def _run() -> None:
+ task = asyncio.create_task(engine.abort_async(["req-ack-race"], timeout=1))
+ await asyncio.to_thread(request_q.get, True, 1)
+ engine._shutdown_called = True
+ engine._correlated_rpc_client.close()
+ await task
+
+ asyncio.run(_run())
+
+
+def test_abort_async_surfaces_rpc_router_close_while_engine_is_live():
+ request_q: queue.Queue = queue.Queue()
+ rpc_q: queue.Queue = queue.Queue()
+ engine = _make_engine(request_q, rpc_q)
+
+ async def _run() -> None:
+ task = asyncio.create_task(engine.abort_async(["req-live"], timeout=1))
+ await asyncio.to_thread(request_q.get, True, 1)
+ engine._correlated_rpc_client.close()
+ with pytest.raises(RuntimeError, match="RPC result router closed"):
+ await task
+
+ asyncio.run(_run())
diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py
index 740064173..7ad111e25 100644
--- a/tests/engine/test_async_omni_engine_input.py
+++ b/tests/engine/test_async_omni_engine_input.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import pytest
import torch
from pytest_mock import MockerFixture
@@ -14,6 +17,104 @@ from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+class _SyntheticSyncQueueShutDownError(Exception):
+ pass
+
+
+def test_abort_is_noop_for_empty_request_ids(mocker: MockerFixture):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = mocker.Mock()
+
+ engine.abort([])
+
+ engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_abort_is_noop_after_shutdown_starts(mocker: MockerFixture):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = True
+ engine.request_queue = mocker.Mock()
+
+ engine.abort(["req-1"])
+
+ engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_abort_tolerates_queue_close_race_during_shutdown(mocker: MockerFixture):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = mocker.Mock()
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ _SyntheticSyncQueueShutDownError,
+ )
+
+ def close_queue_during_abort(*args, **kwargs):
+ del args, kwargs
+ engine._shutdown_called = True
+ raise _SyntheticSyncQueueShutDownError
+
+ engine.request_queue.sync_q.put.side_effect = close_queue_during_abort
+
+ engine.abort(["req-1"])
+
+
+def test_abort_tolerates_legacy_janus_close_error_during_shutdown(mocker: MockerFixture):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = mocker.Mock()
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ None,
+ )
+
+ def close_queue_during_abort(*args, **kwargs):
+ del args, kwargs
+ engine._shutdown_called = True
+ raise RuntimeError("Operation on the closed queue is forbidden")
+
+ engine.request_queue.sync_q.put.side_effect = close_queue_during_abort
+
+ engine.abort(["req-1"])
+
+
+def test_abort_surfaces_unrelated_legacy_runtime_error_during_shutdown(
+ mocker: MockerFixture,
+):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = mocker.Mock()
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ None,
+ )
+
+ def fail_during_abort(*args, **kwargs):
+ del args, kwargs
+ engine._shutdown_called = True
+ raise RuntimeError("unrelated queue failure")
+
+ engine.request_queue.sync_q.put.side_effect = fail_during_abort
+
+ with pytest.raises(RuntimeError, match="unrelated queue failure"):
+ engine.abort(["req-1"])
+
+
+def test_abort_surfaces_unexpected_closed_queue(mocker: MockerFixture):
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = False
+ engine.request_queue = mocker.Mock()
+ mocker.patch(
+ "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN",
+ _SyntheticSyncQueueShutDownError,
+ )
+ engine.request_queue.sync_q.put.side_effect = _SyntheticSyncQueueShutDownError
+
+ with pytest.raises(_SyntheticSyncQueueShutDownError):
+ engine.abort(["req-1"])
+
+
def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest:
return EngineCoreRequest(
request_id=request_id,
@@ -67,6 +168,39 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
output_processor.add_request.assert_not_called()
+def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
+ 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 = ("speech",)
+
+ input_processor = mocker.Mock()
+ input_processor.process_inputs.return_value = _make_engine_core_request()
+ engine.input_processor = input_processor
+ engine.output_processors = [mocker.Mock()]
+
+ seen_global_ids: list[list[str]] = []
+
+ def transform(prompt, _sampling_params):
+ seen_global_ids.append(prompt["additional_information"]["global_request_id"])
+ return prompt
+
+ engine.prompt_transform_func = transform
+ prompt = {"prompt_token_ids": [1, 1, 1]}
+
+ message = engine._build_add_request_message(
+ request_id="req-1",
+ prompt=prompt,
+ sampling_params_list=[params],
+ final_stage_id=0,
+ arrival_time=0.0,
+ )
+
+ assert seen_global_ids == [["req-1"]]
+ assert message.original_prompt["additional_information"]["global_request_id"] == ["req-1"]
+
+
def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture):
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
@@ -371,3 +505,106 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er
)
assert stage_pool.get_bound_replica_id("req-error") is None
+
+
+def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture):
+ """A guided request is all-or-nothing.
+
+ Admitting the parent before its companion is built leaves an orphan: a
+ model whose guidance is mandatory never completes the pair, so the request
+ holds scheduler and KV capacity for the scheduler's whole hold budget and
+ then produces no audio. The build has to raise before anything is enqueued.
+ """
+ engine = object.__new__(AsyncOmniEngine)
+ params = SamplingParams(max_tokens=8)
+ engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin"))
+ engine.supported_tasks = ("generate",)
+ engine.input_processor = mocker.Mock()
+ engine.request_queue = mocker.Mock()
+
+ with pytest.raises(ValueError, match="null twin"):
+ engine._build_cfg_companions(
+ parent_id="req",
+ original_prompt={"prompt": "positive"},
+ stage0_params=params,
+ sampling_params_list=[params],
+ )
+
+ engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture):
+ """The same holds when expansion succeeds but input processing does not."""
+ engine = object.__new__(AsyncOmniEngine)
+ params = SamplingParams(max_tokens=8)
+ engine.prompt_expand_func = lambda *_args: [
+ ExpandedPrompt(
+ prompt={"prompt": "negative"},
+ role="cfg_text",
+ request_id_suffix="__cfg_text",
+ )
+ ]
+ engine.supported_tasks = ("generate",)
+ engine.input_processor = mocker.Mock()
+ engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt")
+ engine.request_queue = mocker.Mock()
+
+ with pytest.raises(RuntimeError, match="tokenizer rejected"):
+ engine._build_cfg_companions(
+ parent_id="req",
+ original_prompt={"prompt": "positive"},
+ stage0_params=params,
+ sampling_params_list=[params],
+ )
+
+ engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture):
+ """Construction and admission are separate steps."""
+ engine = object.__new__(AsyncOmniEngine)
+ params = SamplingParams(max_tokens=8)
+ engine.prompt_expand_func = lambda *_args: [
+ ExpandedPrompt(
+ prompt={"prompt": "negative"},
+ role="cfg_text",
+ request_id_suffix="__cfg_text",
+ )
+ ]
+ engine.supported_tasks = ("generate",)
+ engine.input_processor = mocker.Mock()
+ engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
+ engine.request_queue = mocker.Mock()
+
+ companions = engine._build_cfg_companions(
+ parent_id="req",
+ original_prompt={"prompt": "positive"},
+ stage0_params=params,
+ sampling_params_list=[params],
+ )
+
+ assert len(companions) == 1
+ assert companions[0].companion_id == "req__cfg_text"
+ assert companions[0].parent_id == "req"
+ engine.request_queue.sync_q.put.assert_not_called()
+
+
+def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture):
+ """An unguided request expands to no companions and admits normally."""
+ engine = object.__new__(AsyncOmniEngine)
+ params = SamplingParams(max_tokens=8)
+ engine.prompt_expand_func = lambda *_args: []
+ engine.supported_tasks = ("generate",)
+ engine.input_processor = mocker.Mock()
+ engine.request_queue = mocker.Mock()
+
+ assert (
+ engine._build_cfg_companions(
+ parent_id="req",
+ original_prompt={"prompt": "positive"},
+ stage0_params=params,
+ sampling_params_list=[params],
+ )
+ == []
+ )
+ engine.request_queue.sync_q.put.assert_not_called()
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 27aee8706..c1e56b7e1 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -1,6 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import concurrent.futures
import contextlib
import importlib
+import json
import os
import time
import types
@@ -15,6 +19,7 @@ from vllm_omni.engine.stage_init_utils import (
ReplicaInitPlan,
build_stage0_input_processor,
compute_replica_layout,
+ split_devices_for_replicas,
)
from vllm_omni.engine.stage_runtime import StageRuntime
@@ -161,6 +166,16 @@ def _make_diffusion_plan(
)
+def _make_stage_runtime() -> StageRuntime:
+ return StageRuntime(
+ stage_configs=[],
+ model="dummy-model",
+ config_path="dummy-config",
+ stage_init_timeout=1,
+ async_chunk=False,
+ )
+
+
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
@@ -180,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
engine.model = "dummy-model"
engine.config_path = "dummy-config"
engine.single_stage_mode = False
- engine.diffusion_batch_size = 1
engine.async_chunk = False
engine.tokenizer = None
engine._single_stage_id_filter = None
@@ -221,6 +235,75 @@ def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
assert replica_devices_map == {0: ["0,1", "2,3"]}
+@pytest.mark.parametrize("num_replicas", [1, 2, 4])
+def test_split_devices_for_replicas_returns_one_slot_per_replica_when_devices_unset(num_replicas):
+ """``devices`` is optional, and the result is indexed by replica id.
+
+ A stage that declares ``num_replicas`` without ``devices`` lets every
+ replica inherit the launcher's CUDA_VISIBLE_DEVICES, so each one still
+ needs its own (empty) slot.
+ """
+ assert split_devices_for_replicas(None, num_replicas, 1, 0) == [None] * num_replicas
+
+
+def test_compute_replica_layout_covers_every_replica_when_devices_unset():
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={},
+ runtime=types.SimpleNamespace(num_replicas=3),
+ )
+
+ replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+
+ assert replicas_per_stage == [3]
+ assert replica_devices_map == {0: [None, None, None]}
+
+
+def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatch):
+ """Regression: a multi-replica stage with no ``devices`` must still plan.
+
+ ``_build_logical_stage_init_plans`` indexes ``replica_devices_map`` by
+ replica id, so a short list raised ``IndexError`` for every replica after
+ the first.
+ """
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={},
+ runtime=types.SimpleNamespace(num_replicas=3),
+ )
+ runtime = StageRuntime(
+ stage_configs=[stage_cfg],
+ model="dummy-model",
+ config_path="dummy-config",
+ stage_init_timeout=1,
+ async_chunk=False,
+ )
+
+ monkeypatch.setattr(
+ runtime_mod,
+ "extract_legacy_stage_metadata",
+ lambda cfg: _make_llm_metadata(cfg.stage_id),
+ )
+ monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {})
+ monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+ monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {})
+ monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_args, **_kwargs: (types.SimpleNamespace(), object))
+
+ replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+ stage_plans = runtime._build_logical_stage_init_plans(
+ omni_transfer_config=None,
+ replicas_per_stage=replicas_per_stage,
+ replica_devices_map=replica_devices_map,
+ )
+
+ assert [replica.replica_id for replica in stage_plans[0].replicas] == [0, 1, 2]
+ assert [replica.stage_cfg.runtime.devices for replica in stage_plans[0].replicas] == [None, None, None]
+
+
def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
shared = types.SimpleNamespace(name="shared")
extra = types.SimpleNamespace(name="extra")
@@ -233,6 +316,48 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
assert cleanup_clients == [shared, extra]
+def test_initialize_local_diffusion_replica_scopes_runtime_env(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+ from vllm_omni.platforms import current_omni_platform
+
+ runtime = _make_stage_runtime()
+ plan = _make_diffusion_plan(0, stage_id=0).replicas[0]
+
+ runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"
+ device_env_var = current_omni_platform.device_control_env_var
+ runtime._init_visible_devices_baseline = "0,1"
+ plan.metadata.runtime_cfg = {
+ "devices": "0",
+ "env": {runtime_env_var: "stage-value"},
+ }
+ monkeypatch.delenv(runtime_env_var, raising=False)
+ monkeypatch.setenv(device_env_var, "0,1")
+
+ captured: dict[str, str | None] = {}
+ monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
+
+ def _capture_launch_diffusion_stage_replica(**_kwargs):
+ captured["runtime_env"] = os.environ.get(runtime_env_var)
+ captured["device_env"] = os.environ.get(device_env_var)
+ raise RuntimeError("stop after capturing launch environment")
+
+ monkeypatch.setattr(
+ runtime_mod,
+ "launch_diffusion_stage_replica",
+ _capture_launch_diffusion_stage_replica,
+ )
+
+ with pytest.raises(RuntimeError, match="stop after capturing launch environment"):
+ runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
+
+ assert captured == {
+ "runtime_env": "stage-value",
+ "device_env": "0",
+ }
+ assert runtime_env_var not in os.environ
+ assert os.environ[device_env_var] == "0,1"
+
+
def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
@@ -243,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -271,16 +395,23 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
os.environ[env_var] = old_env
-def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
+@pytest.mark.parametrize(
+ ("num_stages", "expected_inline"),
+ [(1, True), (2, False)],
+)
+def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(
+ monkeypatch,
+ num_stages,
+ expected_inline,
+):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
runtime = StageRuntime(
- stage_configs=[types.SimpleNamespace()],
+ stage_configs=[types.SimpleNamespace()] * num_stages,
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=4,
async_chunk=False,
)
@@ -293,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
def _capture_launch_diffusion_stage_replica(**kwargs):
captured["stage_id"] = kwargs["metadata"].stage_id
captured["stage_init_timeout"] = kwargs["stage_init_timeout"]
- captured["batch_size"] = kwargs["batch_size"]
captured["use_inline"] = kwargs["use_inline"]
captured["omni_master_server"] = kwargs["omni_master_server"]
return types.SimpleNamespace(), StageReplicaResources()
@@ -305,12 +435,238 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
assert captured == {
"stage_id": 0,
"stage_init_timeout": 302,
- "batch_size": 4,
- "use_inline": True,
+ "use_inline": expected_inline,
"omni_master_server": None,
}
+def test_initialize_local_diffusion_replica_uses_explicit_inline_opt_in(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+ from vllm_omni.engine.stage_engine_startup import StageReplicaResources
+
+ runtime = _make_stage_runtime()
+ plan = _make_diffusion_plan(0, stage_id=1).replicas[0]
+ plan.stage_cfg.engine_args["inline_diffusion"] = True
+ captured: dict[str, object] = {}
+
+ monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
+ monkeypatch.setattr(
+ runtime_mod,
+ "launch_diffusion_stage_replica",
+ lambda **kwargs: (
+ captured.update(use_inline=kwargs["use_inline"]),
+ (types.SimpleNamespace(), StageReplicaResources()),
+ )[1],
+ )
+
+ runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
+
+ assert captured["use_inline"] is True
+
+
+def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+
+ runtime = _make_stage_runtime()
+ plan = _make_llm_plan(0, stage_id=0, vllm_config=object()).replicas[0]
+ plan.engine_args_dict = {}
+
+ runtime_env_var = "VLLM_OMNI_TEST_STAGE_RUNTIME_ENV"
+ runtime._init_visible_devices_baseline = "0"
+ plan.metadata.runtime_cfg = {
+ "devices": "0",
+ "env": {runtime_env_var: "stage-value"},
+ }
+ monkeypatch.setenv(runtime_env_var, "parent-value")
+ # Pin the logical->physical device mapping so the assertion below is
+ # deterministic regardless of the ambient CUDA_VISIBLE_DEVICES (which on a
+ # rebase host can be e.g. "3,4", mapping logical device 0 -> physical 3).
+ runtime._init_visible_devices_baseline = "0"
+
+ captured: dict[str, str | None] = {}
+
+ @contextlib.contextmanager
+ def _capture_launch_stage_replica(*, stage_visible_devices, **_kwargs):
+ captured["runtime_env"] = os.environ.get(runtime_env_var)
+ captured["stage_visible_devices"] = stage_visible_devices
+ yield None
+
+ monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
+ monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
+
+ with pytest.raises(RuntimeError, match="launcher returned no resources"):
+ runtime._initialize_local_llm_replica(plan, stage_init_timeout=1)
+
+ assert captured == {
+ "runtime_env": "stage-value",
+ "stage_visible_devices": "0",
+ }
+ assert os.environ[runtime_env_var] == "parent-value"
+
+
+def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatch):
+ import vllm_omni.diffusion.stage_diffusion_client as client_mod
+ import vllm_omni.engine.stage_init_utils as init_mod
+
+ od_config = types.SimpleNamespace(max_num_seqs=4)
+ captured: dict[str, object] = {}
+ monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+
+ def _capture_client(model, config, metadata, stage_init_timeout, use_inline):
+ captured.update(
+ model=model,
+ config=config,
+ metadata=metadata,
+ stage_init_timeout=stage_init_timeout,
+ use_inline=use_inline,
+ )
+ return object()
+
+ monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client)
+ metadata = _make_diffusion_metadata(0)
+
+ init_mod.initialize_diffusion_stage(
+ 0,
+ "dummy-model",
+ types.SimpleNamespace(),
+ metadata,
+ stage_init_timeout=12,
+ use_inline=True,
+ )
+
+ assert od_config.max_num_seqs == 4
+ assert captured == {
+ "model": "dummy-model",
+ "config": od_config,
+ "metadata": metadata,
+ "stage_init_timeout": 12,
+ "use_inline": True,
+ }
+
+
+def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch):
+ import vllm_omni.diffusion.stage_diffusion_client as client_mod
+ import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
+ import vllm_omni.engine.stage_engine_startup as startup_mod
+
+ od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
+ monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
+ monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
+ monkeypatch.setattr(
+ startup_mod,
+ "register_stage_with_omni_master",
+ lambda **kwargs: types.SimpleNamespace(
+ handshake_address="tcp://127.0.0.1:26001",
+ input_address="tcp://127.0.0.1:26002",
+ output_address="tcp://127.0.0.1:26003",
+ ),
+ )
+
+ omni_master_server = types.SimpleNamespace(
+ address="127.0.0.1",
+ port=25000,
+ release_route_port_reservations=lambda *args, **kwargs: None,
+ )
+ proc_manager = types.SimpleNamespace(
+ addresses=types.SimpleNamespace(
+ inputs=["tcp://127.0.0.1:26002"],
+ outputs=["tcp://127.0.0.1:26003"],
+ )
+ )
+ monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
+ sentinel_client = object()
+ monkeypatch.setattr(
+ client_mod.StageDiffusionClient,
+ "from_addresses",
+ lambda metadata, **kwargs: sentinel_client,
+ )
+
+ result, resources = startup_mod.launch_diffusion_stage_replica(
+ model="dummy-model",
+ stage_config=types.SimpleNamespace(),
+ metadata=types.SimpleNamespace(stage_id=0),
+ stage_init_timeout=12,
+ use_inline=False,
+ omni_master_server=omni_master_server,
+ )
+
+ assert result is sentinel_client
+ assert od_config.max_num_seqs == 4
+ assert resources.manager is proc_manager
+
+
+def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
+ import vllm_omni.diffusion.stage_diffusion_client as client_mod
+ import vllm_omni.engine.stage_init_utils as init_mod
+
+ od_config = types.SimpleNamespace(max_num_seqs=8, step_execution=True)
+ monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
+ monkeypatch.setattr(client_mod, "create_diffusion_client", lambda *args: object())
+ metadata = _make_diffusion_metadata(0)
+
+ init_mod.initialize_diffusion_stage(
+ 0,
+ "dummy-model",
+ types.SimpleNamespace(),
+ metadata,
+ stage_init_timeout=12,
+ use_inline=True,
+ )
+
+ assert od_config.max_num_seqs == 8
+
+
+def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(monkeypatch):
+ import vllm_omni.diffusion.stage_diffusion_client as client_mod
+ import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
+ import vllm_omni.engine.stage_engine_startup as startup_mod
+
+ od_config = types.SimpleNamespace(
+ max_num_seqs=8,
+ step_execution=True,
+ parallel_config=types.SimpleNamespace(world_size=1),
+ )
+ monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
+ monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
+ monkeypatch.setattr(
+ startup_mod,
+ "register_stage_with_omni_master",
+ lambda **kwargs: types.SimpleNamespace(
+ handshake_address="tcp://127.0.0.1:26001",
+ input_address="tcp://127.0.0.1:26002",
+ output_address="tcp://127.0.0.1:26003",
+ ),
+ )
+ omni_master_server = types.SimpleNamespace(
+ address="127.0.0.1",
+ port=25000,
+ release_route_port_reservations=lambda *args, **kwargs: None,
+ )
+ proc_manager = types.SimpleNamespace(
+ addresses=types.SimpleNamespace(
+ inputs=["tcp://127.0.0.1:26002"],
+ outputs=["tcp://127.0.0.1:26003"],
+ )
+ )
+ monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager)
+ monkeypatch.setattr(
+ client_mod.StageDiffusionClient,
+ "from_addresses",
+ lambda metadata, **kwargs: object(),
+ )
+
+ startup_mod.launch_diffusion_stage_replica(
+ model="dummy-model",
+ stage_config=types.SimpleNamespace(),
+ metadata=types.SimpleNamespace(stage_id=0),
+ stage_init_timeout=12,
+ use_inline=False,
+ omni_master_server=omni_master_server,
+ )
+
+ assert od_config.max_num_seqs == 8
+
+
def test_stage_runtime_initializes_stage_pools(monkeypatch):
import vllm_omni.engine.stage_runtime as runtime_mod
@@ -319,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -385,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -409,11 +763,12 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
- monkeypatch.setattr(
- runtime_mod.StageEngineCoreClientBase,
- "make_async_mp_client",
- lambda **kwargs: captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client,
- )
+
+ def _make_async_mp_client(**kwargs):
+ captured["client_log_stats"] = kwargs["log_stats"]
+ return stage_client
+
+ monkeypatch.setattr(runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", _make_async_mp_client)
assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
assert captured["log_stats"] is True
@@ -428,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -481,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -522,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -560,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys():
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
- diffusion_batch_size=1,
async_chunk=False,
)
plan = _make_llm_plan(
@@ -586,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -623,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
@@ -635,7 +984,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
def _initialize_stage_replicas(_stage_plans, _stage_init_timeout):
exc = RuntimeError("replica launch failed")
- exc._initialized_clients_by_stage = {0: [None, initialized_client]}
+ setattr(exc, "_initialized_clients_by_stage", {0: [None, initialized_client]})
raise exc
monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas)
@@ -662,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
- diffusion_batch_size=1,
async_chunk=False,
)
stage_init_timeout = 302
@@ -732,6 +1080,8 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p
default_sampling_params={},
)
+ (tmp_path / "llm").mkdir()
+
engine_args = build_engine_args_dict(
stage_cfg,
str(tmp_path),
@@ -793,6 +1143,9 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
default_sampling_params={},
)
+ (tmp_path / "llm").mkdir()
+ (tmp_path / "tokenizer").mkdir()
+
engine_args = build_engine_args_dict(
stage_cfg,
str(tmp_path),
@@ -803,6 +1156,445 @@ def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer")
+def test_build_engine_args_rejects_missing_subdir_in_local_model_dir(tmp_path):
+ """A local model directory without the stage subfolder must fail closed.
+
+ Forwarding the joined path as ``EngineArgs.model`` sends a non-directory
+ string to HuggingFace, which reports it as a malformed repo id.
+ """
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+ (tmp_path / "tokenizer").mkdir()
+
+ with pytest.raises(RuntimeError, match="language_model"):
+ build_engine_args_dict(stage_cfg, str(tmp_path))
+
+
+# A real snapshot subfolder always carries artifacts; empty directories would
+# assert the broken exists-means-complete predicate this suite regresses.
+_SUBDIR_ARTIFACT = {"language_model": "model.safetensors", "tokenizer": "tokenizer.json"}
+
+
+def _make_snapshot(root, subdirs):
+ for subdir in subdirs:
+ folder = root / subdir
+ folder.mkdir(parents=True, exist_ok=True)
+ (folder / _SUBDIR_ARTIFACT.get(subdir, "data.bin")).write_text("x")
+ return str(root)
+
+
+def test_build_engine_args_pulls_stage_subdirs_missing_from_cached_snapshot(monkeypatch, tmp_path):
+ """Regression for issue #6638.
+
+ ``snapshot_download(local_files_only=True)`` hands back a snapshot root as
+ soon as the repo is partly cached, even when this stage's subfolders were
+ never materialized. Stage init must fetch exactly those subfolders instead
+ of joining onto a path that exists nowhere.
+ """
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ _make_snapshot(snapshot, ["tokenizer"])
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if kwargs.get("local_files_only"):
+ return str(snapshot)
+ _make_snapshot(snapshot, ["language_model"])
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+ assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer")
+ assert os.path.isdir(engine_args["model"])
+ # Only the stage's own subfolders are pulled, not the whole repo.
+ assert calls[-1]["allow_patterns"] == ["language_model/*", "tokenizer/*"]
+
+
+def test_build_engine_args_skips_hub_call_when_cached_snapshot_is_complete(monkeypatch, tmp_path):
+ """A warm cache stays offline-friendly: no outgoing Hub request."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ _make_snapshot(snapshot, ["language_model", "tokenizer"])
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ assert kwargs.get("local_files_only"), "warm cache must not reach the Hub"
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_a_partial_subdir(monkeypatch, tmp_path):
+ """A subfolder holding only config.json is an interrupted download.
+
+ The warm-cache early return must not accept it: converting the Hub ID into
+ that local path leaves vLLM's loader without any fallback for the missing
+ weights.
+ """
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ partial = snapshot / "language_model"
+ partial.mkdir(parents=True)
+ (partial / "config.json").write_text("{}")
+ _make_snapshot(snapshot, ["tokenizer"])
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if kwargs.get("local_files_only"):
+ return str(snapshot)
+ _make_snapshot(snapshot, ["language_model"])
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert any(not call.get("local_files_only") for call in calls), "a partial subdir must fall through to the Hub"
+ assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_when_index_lists_missing_shards(monkeypatch, tmp_path):
+ """An index plus a subset of its shards is an interrupted download.
+
+ The index downloads early; accepting it as proof of weights converts the
+ Hub ID into a local path vLLM cannot fetch the remaining shards for.
+ """
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ partial = snapshot / "language_model"
+ partial.mkdir(parents=True)
+ (partial / "model.safetensors.index.json").write_text(
+ json.dumps(
+ {
+ "weight_map": {
+ "a.weight": "model-00001-of-00002.safetensors",
+ "b.weight": "model-00002-of-00002.safetensors",
+ }
+ }
+ )
+ )
+ (partial / "model-00001-of-00002.safetensors").write_text("x")
+ _make_snapshot(snapshot, ["tokenizer"])
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if not kwargs.get("local_files_only"):
+ (partial / "model-00002-of-00002.safetensors").write_text("x")
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert any(not call.get("local_files_only") for call in calls), "missing shards must fall through to the Hub"
+ assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+
+
+def test_build_engine_args_redownloads_shards_without_their_index(monkeypatch, tmp_path):
+ """Shard-named weights always ship an index; one without it is partial."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ partial = snapshot / "language_model"
+ partial.mkdir(parents=True)
+ (partial / "model-00001-of-00004.safetensors").write_text("x")
+ _make_snapshot(snapshot, ["tokenizer"])
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if not kwargs.get("local_files_only"):
+ shards = {f"w{i}.weight": f"model-0000{i}-of-00004.safetensors" for i in range(1, 5)}
+ (partial / "model.safetensors.index.json").write_text(json.dumps({"weight_map": shards}))
+ for shard in shards.values():
+ (partial / shard).write_text("x")
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert any(not call.get("local_files_only") for call in calls)
+
+
+def test_build_engine_args_redownloads_a_tokenizer_folder_without_vocabulary(monkeypatch, tmp_path):
+ """Templates and configs download first; alone they are not a tokenizer."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ _make_snapshot(snapshot, ["language_model"])
+ partial = snapshot / "tokenizer"
+ partial.mkdir(parents=True)
+ (partial / "chat_template.jinja").write_text("{{ messages }}")
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if not kwargs.get("local_files_only"):
+ _make_snapshot(snapshot, ["tokenizer"])
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert any(not call.get("local_files_only") for call in calls)
+ assert engine_args["tokenizer"] == os.path.join(str(snapshot), "tokenizer")
+
+
+def test_build_engine_args_forwards_revision_and_download_dir(monkeypatch, tmp_path):
+ """revision/download_dir must shape snapshot selection (they cannot be
+ corrected downstream once the repo ID is a local path)."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "pinned"
+ _make_snapshot(snapshot, ["language_model", "tokenizer"])
+ seen = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ seen.append(kwargs)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={
+ "model_subdir": "language_model",
+ "tokenizer_subdir": "tokenizer",
+ "revision": "v2.0",
+ "download_dir": str(tmp_path / "cache"),
+ },
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert engine_args["model"] == os.path.join(str(snapshot), "language_model")
+ assert seen, "snapshot resolution must have been consulted"
+ for call in seen:
+ assert call.get("revision") == "v2.0"
+ assert call.get("cache_dir") == str(tmp_path / "cache")
+
+
+def test_build_engine_args_resolves_tokenizer_revision_separately(monkeypatch, tmp_path):
+ """A tokenizer pinned to a different revision resolves against its own snapshot."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ model_snapshot = tmp_path / "snapshots" / "model-rev"
+ _make_snapshot(model_snapshot, ["language_model"])
+ tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev"
+ _make_snapshot(tokenizer_snapshot, ["tokenizer"])
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ if kwargs.get("revision") == "tok-rev":
+ return str(tokenizer_snapshot)
+ return str(model_snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={
+ "model_subdir": "language_model",
+ "tokenizer_subdir": "tokenizer",
+ "revision": "model-rev",
+ "tokenizer_revision": "tok-rev",
+ },
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model")
+ assert engine_args["tokenizer"] == os.path.join(str(tokenizer_snapshot), "tokenizer")
+
+
+def test_build_engine_args_resolves_root_tokenizer_revision_separately(monkeypatch, tmp_path):
+ """An empty tokenizer_subdir targets the snapshot root and still honors
+ its own revision."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ model_snapshot = tmp_path / "snapshots" / "model-rev"
+ _make_snapshot(model_snapshot, ["language_model"])
+ # An empty tokenizer_subdir means the tokenizer IS the snapshot root, so
+ # the vocabulary artifact must sit at the root. Building it under
+ # ``tokenizer/`` returned a root that holds no loadable tokenizer, which
+ # the root-tokenizer validation now correctly rejects.
+ tokenizer_snapshot = tmp_path / "snapshots" / "tok-rev"
+ tokenizer_snapshot.mkdir(parents=True, exist_ok=True)
+ (tokenizer_snapshot / _SUBDIR_ARTIFACT["tokenizer"]).write_text("x")
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ if kwargs.get("revision") == "tok-rev":
+ return str(tokenizer_snapshot)
+ return str(model_snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={
+ "model_subdir": "language_model",
+ "tokenizer_subdir": "",
+ "revision": "model-rev",
+ "tokenizer_revision": "tok-rev",
+ },
+ default_sampling_params={},
+ )
+
+ engine_args = build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+ assert engine_args["model"] == os.path.join(str(model_snapshot), "language_model")
+ assert engine_args["tokenizer"] == str(tokenizer_snapshot)
+
+
+def test_build_engine_args_fails_closed_when_subdir_cannot_be_downloaded(monkeypatch, tmp_path):
+ """An undownloadable subfolder raises here instead of reaching HuggingFace."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ snapshot = tmp_path / "snapshots" / "deadbeef"
+ _make_snapshot(snapshot, ["tokenizer"])
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ if kwargs.get("local_files_only"):
+ return str(snapshot)
+ raise OSError("offline")
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ with pytest.raises(RuntimeError, match="language_model"):
+ build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+
+def test_build_engine_args_fails_closed_on_cold_cache_instead_of_joining_repo_id(monkeypatch):
+ """With nothing cached, the stage must not join a subdir onto the repo id."""
+ from huggingface_hub import HfApi
+
+ from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ raise OSError("offline")
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ stage_cfg = types.SimpleNamespace(
+ stage_id=0,
+ stage_type="llm",
+ engine_args={"model_subdir": "language_model", "tokenizer_subdir": "tokenizer"},
+ default_sampling_params={},
+ )
+
+ with pytest.raises(RuntimeError, match="nothing is cached locally"):
+ build_engine_args_dict(stage_cfg, "MiniMaxAI/MiniMax-Music3")
+
+
+def test_model_path_resolver_is_generic_and_model_owned(tmp_path):
+ from vllm_omni.engine.stage_init_utils import _resolve_model_path
+
+ engine_args = {
+ "model_path_resolver": ("vllm_omni.model_executor.models.minimax_h3.checkpoint.resolve_minimax_h3_model_root"),
+ "revision": None,
+ "task_type": "ref2va",
+ }
+
+ resolved = _resolve_model_path(str(tmp_path), engine_args)
+
+ assert resolved == str(tmp_path / "Ref2VA" / "text_encoder")
+ assert "model_path_resolver" not in engine_args
+
+
def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
import vllm_omni.engine.stage_init_utils as init_mod
diff --git a/tests/engine/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py
index 4720cd8ea..df0d7b2ed 100644
--- a/tests/engine/test_data_entry_keys.py
+++ b/tests/engine/test_data_entry_keys.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for data_entry_keys."""
import msgspec
@@ -71,6 +74,25 @@ class TestOmniPayloadStruct:
for top, sub in original.items():
assert sorted(d[top].keys()) == sorted(sub.keys())
+ def test_text_conditioning_uses_shared_payload_fields(self):
+ hidden = torch.zeros(4, 5120)
+ token_role_ids = torch.tensor([[1], [1], [0], [0]])
+
+ payload = to_dict(
+ to_struct(
+ {
+ "hidden_states": {"output": hidden},
+ "meta": {"token_role_ids": token_role_ids},
+ }
+ )
+ )
+
+ assert torch.equal(payload["hidden_states"]["output"], hidden)
+ assert torch.equal(payload["meta"]["token_role_ids"], token_role_ids)
+ for legacy_key in ("encoder_hidden_states", "token_tags"):
+ with pytest.raises(msgspec.ValidationError, match="unknown field"):
+ to_struct({legacy_key: torch.zeros(1)})
+
def test_to_dict_drops_unset_fields(self):
s = OmniPayloadStruct(meta=MetaStruct(left_context_size=10))
d = to_dict(s)
@@ -372,12 +394,12 @@ class TestSerializeDeserializePayload:
assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
def test_tensor_dtype_preserved(self):
- # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
- for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
+ for dtype in [torch.float16, torch.bfloat16, torch.float32, torch.int64, torch.int32, torch.bool]:
original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
wire = serialize_payload(original)
restored = deserialize_payload(wire)
assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
+ assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
def test_tensor_shape_preserved(self):
t = torch.randn(3, 4, 5)
@@ -394,3 +416,12 @@ class TestSerializeDeserializePayload:
original: OmniPayload = {"meta": {"finished": None}}
wire = serialize_payload(original)
assert wire is None
+
+ def test_minimax_h3_prepared_reference_video_descriptor_round_trip(self):
+ descriptor = '{"artifact_dir":"/tmp/h3","videos":[]}'
+
+ wire = serialize_payload({"meta": {"minimax_h3_prepared_reference_videos": descriptor}})
+ assert wire is not None
+ restored = deserialize_payload(wire)
+
+ assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor
diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py
index 6eae2e698..d04ed2e70 100644
--- a/tests/engine/test_duplex_import_boundary.py
+++ b/tests/engine/test_duplex_import_boundary.py
@@ -3,6 +3,7 @@
from __future__ import annotations
+import os
import subprocess
import sys
from pathlib import Path
@@ -14,9 +15,18 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
def _assert_isolated_import_succeeds(script: str) -> None:
+ # These tests validate Python import/package boundaries and do not use a
+ # GPU. Other tests in the combined entrypoints/engine suite may temporarily
+ # configure only one of the CUDA/ROCm visibility variables. Do not let an
+ # unrelated, inconsistent device-selection environment prevent the child
+ # interpreter from reaching the boundary assertion.
+ env = os.environ.copy()
+ env.pop("CUDA_VISIBLE_DEVICES", None)
+ env.pop("HIP_VISIBLE_DEVICES", None)
result = subprocess.run(
[sys.executable, "-c", script],
cwd=REPO_ROOT,
+ env=env,
capture_output=True,
text=True,
timeout=60,
@@ -25,6 +35,14 @@ def _assert_isolated_import_succeeds(script: str) -> None:
assert result.returncode == 0, result.stderr
+def test_isolated_import_ignores_conflicting_gpu_visibility(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Import-boundary checks must not depend on GPU state leaked by prior tests."""
+ monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0")
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
+
+ _assert_isolated_import_succeeds("import vllm_omni.outputs")
+
+
def test_stable_engine_imports_do_not_load_experimental_duplex() -> None:
_assert_isolated_import_succeeds("""
import sys
diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py
index 1b8bf2adb..5f0416377 100644
--- a/tests/engine/test_membership_controller.py
+++ b/tests/engine/test_membership_controller.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -30,6 +30,14 @@ class FakeHub:
self.closed = True
+async def _wait_until(predicate, timeout=1):
+ async def _poll():
+ while not predicate():
+ await asyncio.sleep(0)
+
+ await asyncio.wait_for(_poll(), timeout=timeout)
+
+
class FakePool:
def __init__(self, stage_id: int):
self.stage_id = stage_id
@@ -39,6 +47,7 @@ class FakePool:
self.invalidated = []
self.hub = None
self.lb = None
+ self.replica_ids = {}
def attach_hub(self, hub):
self.hub = hub
@@ -49,7 +58,9 @@ class FakePool:
def add_client(self, input_addr, client, *, replica_id=None):
self.added.append((input_addr, client, replica_id))
self.clients.append(client)
- return len(self.clients) - 1 if replica_id is None else replica_id
+ resolved_replica_id = len(self.clients) - 1 if replica_id is None else replica_id
+ self.replica_ids[input_addr] = resolved_replica_id
+ return resolved_replica_id
def invalidate_addr(self, input_addr):
self.invalidated.append(input_addr)
@@ -59,6 +70,12 @@ class FakePool:
self.removed.append(input_addr)
return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
+ def get_client_by_addr(self, input_addr):
+ return SimpleNamespace() if input_addr in self.replica_ids else None
+
+ def get_replica_id_by_addr(self, input_addr):
+ return self.replica_ids.get(input_addr)
+
def _snapshot(*replicas):
return SimpleNamespace(replicas=list(replicas))
@@ -159,6 +176,29 @@ async def test_do_register_offloads_remote_factory(monkeypatch):
assert pool.added[0] == ("tcp://stage-0-replica-1", pool.clients[0], 1)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("failure", ["missing-address", "add-client"])
+async def test_do_register_shuts_down_rejected_client(monkeypatch, failure):
+ shutdown_calls = 0
+
+ class Client:
+ client_addresses = {} if failure == "missing-address" else {"input_address": "tcp://rejected"}
+
+ def shutdown(self):
+ nonlocal shutdown_calls
+ shutdown_calls += 1
+
+ pool = FakePool(stage_id=0)
+ if failure == "add-client":
+ monkeypatch.setattr(pool, "add_client", lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("reject")))
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=lambda *_args: Client())
+
+ with pytest.raises(RuntimeError):
+ await controller._do_register(0, 1)
+
+ assert shutdown_calls == 1
+
+
@pytest.mark.asyncio
async def test_unregister_then_register_restores_coordinator_replica_slot(monkeypatch):
class Client:
@@ -178,6 +218,16 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
return replacement
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=factory)
+ removed_replicas = []
+
+ async def cleanup(_request_ids):
+ return None
+
+ controller.install_unregister_handlers(
+ output_queue=asyncio.Queue(),
+ cleanup_callback=cleanup,
+ replica_removed_callback=lambda stage_id, replica_id: removed_replicas.append((stage_id, replica_id)),
+ )
await controller.handle_unregister(0, "tcp://old-0")
await controller.handle_unregister(0, "tcp://old-1")
@@ -188,3 +238,368 @@ async def test_unregister_then_register_restores_coordinator_replica_slot(monkey
assert pool.get_replica_id_by_addr("tcp://new-1") == 1
assert pool.get_replica_id_by_addr("tcp://old-1") is None
assert pool.available_replica_ids() == [1]
+ assert removed_replicas == [(0, 0), (0, 1)]
+
+
+@pytest.mark.asyncio
+async def test_handle_register_deduplicates_inflight_and_attached_replica(monkeypatch):
+ pool = FakePool(stage_id=0)
+ factory_started = threading.Event()
+ release_factory = threading.Event()
+ factory_calls = []
+
+ def _factory(stage_id, replica_id):
+ factory_calls.append((stage_id, replica_id))
+ factory_started.set()
+ assert release_factory.wait(timeout=1)
+ return SimpleNamespace(
+ client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
+ shutdown=lambda: None,
+ )
+
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+
+ await asyncio.gather(
+ controller.handle_register(0, 2),
+ controller.handle_register(0, 2),
+ )
+ assert await asyncio.to_thread(factory_started.wait, 1)
+
+ await controller.handle_register(0, 2)
+ release_factory.set()
+ await controller.drain_tasks(timeout=1)
+
+ await controller.handle_register(0, 2)
+ await controller.drain_tasks(timeout=1)
+
+ assert factory_calls == [(0, 2)]
+ assert len(pool.added) == 1
+
+
+@pytest.mark.asyncio
+async def test_handle_register_allows_retry_after_factory_failure(monkeypatch):
+ pool = FakePool(stage_id=0)
+ factory_calls = 0
+
+ def _factory(stage_id, replica_id):
+ nonlocal factory_calls
+ factory_calls += 1
+ if factory_calls == 1:
+ raise RuntimeError("injected registration failure")
+ return SimpleNamespace(
+ client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
+ shutdown=lambda: None,
+ )
+
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+
+ await controller.handle_register(0, 3)
+ await controller.drain_tasks(timeout=1)
+ await controller.handle_register(0, 3)
+ await controller.drain_tasks(timeout=1)
+
+ assert factory_calls == 2
+ assert len(pool.added) == 1
+
+
+@pytest.mark.asyncio
+async def test_unregister_allows_same_replica_to_register_again(monkeypatch):
+ pool = FakePool(stage_id=0)
+ factory_calls = 0
+
+ def _factory(stage_id, replica_id):
+ nonlocal factory_calls
+ factory_calls += 1
+ return SimpleNamespace(
+ client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
+ shutdown=lambda: None,
+ )
+
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+
+ await controller.handle_register(0, 4)
+ await controller.drain_tasks(timeout=1)
+ await controller.handle_unregister(0, "tcp://stage-0-replica-4")
+ await controller.handle_register(0, 4)
+ await controller.drain_tasks(timeout=1)
+
+ assert factory_calls == 2
+ assert len(pool.added) == 2
+
+
+@pytest.mark.asyncio
+async def test_handle_register_ignores_registration_after_shutdown(monkeypatch):
+ pool = FakePool(stage_id=0)
+ factory_calls = 0
+
+ def _factory(stage_id, replica_id):
+ nonlocal factory_calls
+ factory_calls += 1
+ return SimpleNamespace(
+ client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
+ shutdown=lambda: None,
+ )
+
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+ controller.shutdown()
+
+ await controller.handle_register(0, 5)
+ await controller.drain_tasks(timeout=1)
+
+ assert factory_calls == 0
+ assert pool.added == []
+
+
+@pytest.mark.asyncio
+async def test_shutdown_discards_client_from_inflight_registration(monkeypatch):
+ pool = FakePool(stage_id=0)
+ factory_started = threading.Event()
+ release_factory = threading.Event()
+ client_shutdown = threading.Event()
+
+ def _factory(stage_id, replica_id):
+ factory_started.set()
+ assert release_factory.wait(timeout=1)
+ return SimpleNamespace(
+ client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"},
+ shutdown=client_shutdown.set,
+ )
+
+ controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
+
+ await controller.handle_register(0, 6)
+ assert await asyncio.to_thread(factory_started.wait, 1)
+ controller.shutdown()
+ release_factory.set()
+ await controller.drain_tasks(timeout=1)
+
+ assert client_shutdown.is_set()
+ assert pool.added == []
+
+
+@pytest.mark.asyncio
+async def test_watcher_reconciles_replica_that_disappears_during_registration(monkeypatch):
+ input_addr = "tcp://stage-0-replica-7"
+ pool = StagePool(0, [])
+ hub = FakeHub(
+ snapshots=[
+ _snapshot(_replica(0, input_addr)),
+ _snapshot(),
+ ]
+ )
+ factory_started = threading.Event()
+ release_factory = threading.Event()
+ client_shutdown = threading.Event()
+
+ def _factory(stage_id, replica_id):
+ assert (stage_id, replica_id) == (0, 7)
+ factory_started.set()
+ assert release_factory.wait(timeout=1)
+ return SimpleNamespace(
+ client_addresses={"input_address": input_addr},
+ shutdown=client_shutdown.set,
+ )
+
+ controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
+ controller.WATCH_INTERVAL_S = 0
+ original_handle_unregister = controller.handle_unregister
+ first_unregister = asyncio.Event()
+ second_unregister = asyncio.Event()
+ unregister_calls = 0
+
+ async def _track_unregister(stage_id, addr, *args, **kwargs):
+ nonlocal unregister_calls
+ await original_handle_unregister(stage_id, addr, *args, **kwargs)
+ unregister_calls += 1
+ if unregister_calls == 1:
+ first_unregister.set()
+ elif unregister_calls == 2:
+ second_unregister.set()
+
+ controller.handle_unregister = _track_unregister # type: ignore[method-assign]
+
+ await controller.handle_register(0, 7)
+ assert await asyncio.to_thread(factory_started.wait, 1)
+ watcher = controller.start()
+
+ try:
+ await asyncio.wait_for(first_unregister.wait(), timeout=1)
+ assert pool.live_num_replicas == 0
+
+ release_factory.set()
+ await asyncio.wait_for(second_unregister.wait(), timeout=1)
+
+ assert client_shutdown.is_set()
+ assert pool.live_num_replicas == 0
+ assert controller._attached_remote_replicas == set()
+ assert controller._remote_replica_keys_by_input_addr == {}
+ finally:
+ release_factory.set()
+ controller.shutdown()
+ try:
+ await watcher
+ except asyncio.CancelledError:
+ pass
+ await controller.drain_tasks(timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_attached_replica_survives_until_hub_observes_it_up(monkeypatch):
+ class EmptyHub:
+ def __init__(self):
+ self.polls = 0
+
+ def get_replica_list(self):
+ self.polls += 1
+ return SimpleNamespace(replicas=[], timestamp=0.0)
+
+ def close(self):
+ pass
+
+ input_addr = "tcp://stage-0-replica-8"
+ pool = StagePool(0, [])
+ shutdown_calls = []
+ hub = EmptyHub()
+
+ def _factory(stage_id, replica_id):
+ return SimpleNamespace(
+ client_addresses={"input_address": input_addr},
+ shutdown=lambda: shutdown_calls.append((stage_id, replica_id)),
+ )
+
+ controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
+ controller.WATCH_INTERVAL_S = 0.001
+
+ await controller.handle_register(0, 8)
+ await controller.drain_tasks(timeout=1)
+ assert pool.live_num_replicas == 1
+
+ watcher = controller.start()
+ try:
+ await _wait_until(lambda: hub.polls >= 3)
+
+ assert shutdown_calls == []
+ assert pool.live_num_replicas == 1
+ finally:
+ controller.shutdown()
+ try:
+ await watcher
+ except asyncio.CancelledError:
+ pass
+ await controller.drain_tasks(timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_up_observed_before_attach_does_not_qualify_attachment(monkeypatch):
+ import vllm_omni.engine.membership_controller as membership_mod
+
+ input_addr = "tcp://stage-0-replica-10"
+ pool = StagePool(0, [])
+ hub = FakeHub(
+ snapshots=[
+ _snapshot(_replica(0, input_addr)),
+ _snapshot(),
+ ]
+ )
+ shutdown_calls = []
+ sleep_gates = asyncio.Queue()
+
+ async def _controlled_sleep(_delay):
+ gate = asyncio.Event()
+ await sleep_gates.put(gate)
+ await gate.wait()
+
+ monkeypatch.setattr(membership_mod.asyncio, "sleep", _controlled_sleep)
+
+ def _factory(stage_id, replica_id):
+ return SimpleNamespace(
+ client_addresses={"input_address": input_addr},
+ shutdown=lambda: shutdown_calls.append((stage_id, replica_id)),
+ )
+
+ controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
+ watcher = controller.start()
+
+ try:
+ first_sleep = await asyncio.wait_for(sleep_gates.get(), timeout=1)
+ await controller.handle_register(0, 10)
+ await controller.drain_tasks(timeout=1)
+ assert pool.live_num_replicas == 1
+
+ first_sleep.set()
+ await asyncio.wait_for(sleep_gates.get(), timeout=1)
+ await controller.drain_tasks(timeout=1)
+
+ assert shutdown_calls == []
+ assert pool.live_num_replicas == 1
+ finally:
+ controller.shutdown()
+ try:
+ await watcher
+ except asyncio.CancelledError:
+ pass
+ await controller.drain_tasks(timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_reattached_address_requires_a_fresh_up_observation(monkeypatch):
+ class MutableHub:
+ def __init__(self, input_addr):
+ self.input_addr = input_addr
+ self.up = True
+ self.polls = 0
+
+ def get_replica_list(self):
+ self.polls += 1
+ replicas = [_replica(0, self.input_addr)] if self.up else []
+ return _snapshot(*replicas)
+
+ def close(self):
+ pass
+
+ class Client:
+ def __init__(self, input_addr):
+ self.client_addresses = {"input_address": input_addr}
+ self.shutdown_calls = 0
+
+ def shutdown(self):
+ self.shutdown_calls += 1
+
+ input_addr = "tcp://stage-0-replica-9"
+ pool = StagePool(0, [])
+ hub = MutableHub(input_addr)
+ clients = []
+
+ def _factory(_stage_id, _replica_id):
+ client = Client(input_addr)
+ clients.append(client)
+ return client
+
+ controller = _controller(monkeypatch, pool, hub, remote_replica_factory=_factory)
+ controller.WATCH_INTERVAL_S = 0.001
+
+ await controller.handle_register(0, 9)
+ await controller.drain_tasks(timeout=1)
+ watcher = controller.start()
+
+ try:
+ await _wait_until(lambda: hub.polls >= 2)
+ hub.up = False
+ await _wait_until(lambda: clients[0].shutdown_calls == 1)
+ assert pool.live_num_replicas == 0
+
+ await controller.handle_register(0, 9)
+ await controller.drain_tasks(timeout=1)
+ polls_after_reattach = hub.polls
+ await _wait_until(lambda: hub.polls >= polls_after_reattach + 3)
+
+ assert len(clients) == 2
+ assert clients[1].shutdown_calls == 0
+ assert pool.live_num_replicas == 1
+ finally:
+ controller.shutdown()
+ try:
+ await watcher
+ except asyncio.CancelledError:
+ pass
+ await controller.drain_tasks(timeout=1)
diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py
index fa25bba9e..eb1bdfd41 100644
--- a/tests/engine/test_omni_request_output.py
+++ b/tests/engine/test_omni_request_output.py
@@ -270,3 +270,56 @@ class TestMsgpackRoundTrip:
# Content merged from legacy nested key
assert result.prompt == "legacy prompt"
assert result.prompt_token_ids == [1, 2]
+
+
+class TestRequestOutputFieldParity:
+ """``OmniRequestOutput`` must expose everything ``RequestOutput`` sets.
+
+ ``RequestOutput`` is a plain class, so the dataclass-generated ``__init__``
+ replaces its ``__init__`` and the inherited attributes only exist because
+ they are redeclared as fields. Any vLLM field that is not redeclared simply
+ does not exist on an omni output, and vLLM's own serving code reads those
+ attributes unconditionally — ``/v1/completions`` returned HTTP 500 with
+ ``AttributeError: 'OmniRequestOutput' object has no attribute
+ 'ec_transfer_params'`` for exactly that reason. This test fails on the next
+ vLLM bump that adds a field, instead of leaving it for an endpoint to hit.
+ """
+
+ def test_carries_every_attribute_request_output_sets(self) -> None:
+ reference = _make_text_request_output()
+ omni = OmniRequestOutput(request_id="req-parity")
+
+ missing = sorted(name for name in vars(reference) if not name.startswith("_") and not hasattr(omni, name))
+ assert not missing, (
+ f"OmniRequestOutput is missing RequestOutput attributes {missing}; "
+ "redeclare them as dataclass fields in vllm_omni/outputs/__init__.py"
+ )
+
+ def test_transfer_param_fields_default_to_none(self) -> None:
+ omni = OmniRequestOutput(request_id="req-defaults")
+ assert omni.kv_transfer_params is None
+ assert omni.ec_transfer_params is None
+ assert omni.num_cached_tokens is None
+ assert omni.num_cache_creation_tokens 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.
+
+ ``from_stage_output`` only carries the names listed in
+ ``_REQUEST_OUTPUT_CONTENT_ATTRS``. With the fields declared but absent
+ from that list the endpoint stops raising ``AttributeError`` while
+ silently dropping external-connector metadata and cache-creation
+ accounting, which is harder to notice than the crash it replaced.
+ """
+ source = _make_text_request_output()
+ source.num_cache_creation_tokens = 7
+ source.ec_transfer_params = {"connector": "shm", "handle": "abc"}
+ source.num_cached_tokens = 3
+ source.kv_transfer_params = {"role": "producer"}
+
+ omni = OmniRequestOutput.from_stage_output(source, stage_id=1)
+
+ assert omni.num_cache_creation_tokens == 7
+ assert omni.ec_transfer_params == {"connector": "shm", "handle": "abc"}
+ assert omni.num_cached_tokens == 3
+ assert omni.kv_transfer_params == {"role": "producer"}
diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py
index 3d63943a7..9dcf066cd 100644
--- a/tests/engine/test_orchestrator.py
+++ b/tests/engine/test_orchestrator.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import asyncio
@@ -14,11 +17,13 @@ import janus
import pytest
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs, FinishReason
from vllm.v1.engine.exceptions import EngineDeadError
from vllm.v1.metrics.stats import IterationStats
from vllm_omni.engine.messages import (
AbortRequestMessage,
+ AbortResultMessage,
AddCompanionRequestMessage,
CollectiveRPCRequestMessage,
CollectiveRPCResultMessage,
@@ -30,6 +35,7 @@ from vllm_omni.engine.messages import (
from vllm_omni.engine.orchestrator import (
Orchestrator,
OrchestratorRequestState,
+ StreamingSegmentState,
_build_terminal_empty_output,
)
from vllm_omni.engine.stage_pool import StagePool
@@ -102,6 +108,13 @@ class OrchestratorFixture:
result_future: concurrent.futures.Future[None]
+@dataclass
+class FakePromptRequest:
+ request_id: str
+ prompt_token_ids: list[int]
+ resumable: bool = True
+
+
class FakeStageClient:
def __init__(
self,
@@ -244,8 +257,32 @@ class FakeOutputProcessor:
)
def abort_requests(self, request_ids, internal: bool = False):
- self.abort_calls.append(request_ids)
- return request_ids
+ aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal)
+ return aborted_ids
+
+ def abort_requests_collecting_outputs(self, request_ids, *, internal: bool = False, commit_state: bool = True):
+ """Mirror MultimodalOutputProcessor collecting for AR abort-prefix tests."""
+ del internal, commit_state
+ ids = list(request_ids)
+ self.abort_calls.append(ids)
+ outputs: list[RequestOutput] = []
+ for rid in ids:
+ seeded = next(
+ (ro for ro in self.request_outputs if getattr(ro, "request_id", None) == rid),
+ None,
+ )
+ token_ids = list(seeded.outputs[0].token_ids) if seeded is not None and seeded.outputs else [1, 2]
+ # Intentionally stamp an internal-looking id on the RequestOutput so
+ # StagePool must re-key by the orchestrator id it passed in.
+ outputs.append(
+ _build_request_output(
+ f"engine-internal-{rid}",
+ token_ids=token_ids,
+ finished=True,
+ finish_reason="abort",
+ )
+ )
+ return ids, outputs
def update_scheduler_stats(self, _scheduler_stats) -> None:
return None
@@ -287,6 +324,20 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
+def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
+ return EngineCoreOutputs(
+ outputs=[
+ EngineCoreOutput(
+ request_id=request_id,
+ new_token_ids=[],
+ finish_reason=FinishReason.STOP,
+ )
+ ],
+ timestamp=1.0,
+ finished_requests={request_id},
+ )
+
+
def _build_request_output(
request_id: str,
*,
@@ -294,6 +345,7 @@ def _build_request_output(
prompt_token_ids: list[int] | None = None,
finished: bool = True,
text: str = "test",
+ finish_reason: str | None = None,
) -> RequestOutput:
completion = CompletionOutput(
index=0,
@@ -301,7 +353,7 @@ def _build_request_output(
token_ids=list(token_ids or [1, 2]),
cumulative_logprob=0.0,
logprobs=None,
- finish_reason="stop" if finished else None,
+ finish_reason=(finish_reason if finish_reason is not None else ("stop" if finished else None)),
stop_reason=None,
)
return RequestOutput(
@@ -760,6 +812,148 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
await _shutdown_orchestrator(orchestrator_fixture)
+@pytest.mark.asyncio
+async def test_async_chunk_raw_terminal_without_processed_output_finishes_request(orchestrator_factory) -> None:
+ stage0 = FakeStageClient(stage_type="llm", final_output=False)
+ stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
+ orchestrator_fixture = orchestrator_factory(
+ [stage0, stage1],
+ output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
+ async_chunk=True,
+ )
+ request = FakePromptRequest(
+ request_id="req-stream-terminal",
+ prompt_token_ids=[1, 2, 3, 4],
+ )
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id=request.request_id,
+ prompt=request,
+ original_prompt={"prompt": "stream audio"},
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ )
+
+ await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+ stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request.request_id))
+
+ output_msg = await _get_output_message(orchestrator_fixture)
+
+ assert output_msg.request_id == request.request_id
+ assert output_msg.stage_id == 1
+ assert output_msg.finished is True
+ assert output_msg.engine_outputs.finished is True
+ assert output_msg.engine_outputs.outputs[0].multimodal_output["audio"].numel() == 0
+ await _wait_for(lambda: request.request_id not in orchestrator_fixture.orchestrator.request_states)
+ finally:
+ await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_async_chunk_data_then_raw_terminal_finishes_once(orchestrator_factory) -> None:
+ request_id = "req-stream-data-terminal"
+ stage0 = FakeStageClient(stage_type="llm", final_output=False)
+ stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
+ final_data = _build_request_output(
+ request_id,
+ token_ids=[20, 21],
+ finished=False,
+ text="last audio chunk",
+ )
+ orchestrator_fixture = orchestrator_factory(
+ [stage0, stage1],
+ output_processors=[
+ FakeOutputProcessor(),
+ FakeOutputProcessor(request_outputs=[final_data]),
+ ],
+ async_chunk=True,
+ )
+ request = FakePromptRequest(
+ request_id=request_id,
+ prompt_token_ids=[1, 2, 3, 4],
+ )
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id=request_id,
+ prompt=request,
+ original_prompt={"prompt": "stream audio"},
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ )
+
+ await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+ stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id))
+
+ data_msg = await _get_output_message(orchestrator_fixture)
+ terminal_msg = await _get_output_message(orchestrator_fixture)
+
+ assert data_msg.engine_outputs is final_data
+ assert data_msg.finished is False
+ assert terminal_msg.request_id == request_id
+ assert terminal_msg.stage_id == 1
+ assert terminal_msg.finished is True
+ assert terminal_msg.engine_outputs.finished is True
+ await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states)
+ await asyncio.sleep(0.05)
+ with pytest.raises(queue.Empty):
+ orchestrator_fixture.output_sync_q.get_nowait()
+ finally:
+ await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_async_chunk_processed_terminal_and_raw_terminal_finishes_once(orchestrator_factory) -> None:
+ request_id = "req-stream-processed-terminal"
+ stage0 = FakeStageClient(stage_type="llm", final_output=False)
+ stage1 = FakeStageClient(stage_type="llm", final_output=True, final_output_type="audio")
+ processed_terminal = _build_request_output(
+ request_id,
+ token_ids=[20, 21],
+ finished=True,
+ text="last audio chunk",
+ )
+ orchestrator_fixture = orchestrator_factory(
+ [stage0, stage1],
+ output_processors=[
+ FakeOutputProcessor(),
+ FakeOutputProcessor(request_outputs=[processed_terminal]),
+ ],
+ async_chunk=True,
+ )
+ request = FakePromptRequest(
+ request_id=request_id,
+ prompt_token_ids=[1, 2, 3, 4],
+ )
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id=request_id,
+ prompt=request,
+ original_prompt={"prompt": "stream audio"},
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ )
+
+ await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+ stage1.push_engine_core_outputs(_terminal_engine_core_outputs(request_id))
+
+ terminal_msg = await _get_output_message(orchestrator_fixture)
+
+ assert terminal_msg.engine_outputs is processed_terminal
+ assert terminal_msg.finished is True
+ await _wait_for(lambda: request_id not in orchestrator_fixture.orchestrator.request_states)
+ await asyncio.sleep(0.05)
+ with pytest.raises(queue.Empty):
+ orchestrator_fixture.output_sync_q.get_nowait()
+ finally:
+ await _shutdown_orchestrator(orchestrator_fixture)
+
+
@pytest.mark.asyncio
async def test_run_shutdown(orchestrator_factory) -> None:
stages = [
@@ -805,10 +999,125 @@ async def test_run_abort(orchestrator_factory) -> None:
assert stages[0].abort_calls == [["req-abort"]]
assert stages[1].abort_calls == []
assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
+ # Fire-and-forget abort must not emit an RPC result.
+ assert orchestrator_fixture.queues[2].sync_q.empty()
finally:
await _shutdown_orchestrator(orchestrator_fixture)
+@pytest.mark.asyncio
+async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> None:
+ stages = [
+ FakeStageClient(stage_type="llm", final_output=False),
+ FakeStageClient(
+ stage_type="llm",
+ final_output=True,
+ next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
+ ),
+ ]
+ processors = [
+ FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[1], finished=True)]),
+ FakeOutputProcessor(request_outputs=[_build_request_output("req-ack", token_ids=[2], finished=True)]),
+ ]
+ orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
+ request = SimpleNamespace(request_id="req-ack", prompt_token_ids=[1, 2, 3])
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id="req-ack",
+ prompt=request,
+ original_prompt={"prompt": "cancel with ack"},
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ )
+ await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+ # Abort outputs come from the final AR stage's output processor, and
+ # StagePool.abort_requests only collects for requests with a live
+ # replica binding. Forward off stage-0 first so stage-1 is bound.
+ stages[0].push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+ await _wait_for(lambda: len(stages[1].add_request_calls) == 1)
+
+ orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["req-ack"], rpc_id="abort-1"))
+ result = await _get_rpc_message(orchestrator_fixture)
+ assert isinstance(result, AbortResultMessage)
+ assert result.rpc_id == "abort-1"
+ assert result.success is True
+ assert result.error is None
+ assert result.rpc_correlation_key == ("abort", "abort-1")
+ assert stages[0].abort_calls == [["req-ack"]]
+ assert stages[1].abort_calls == [["req-ack"]]
+ assert "req-ack" not in orchestrator_fixture.orchestrator.request_states
+ assert result.abort_outputs is not None
+ assert len(result.abort_outputs) == 1
+ abort_msg = result.abort_outputs[0]
+ assert abort_msg.request_id == "req-ack"
+ assert abort_msg.finished is True
+ assert abort_msg.stage_id == 1
+ assert list(abort_msg.engine_outputs.outputs[0].token_ids) == [2]
+ assert abort_msg.engine_outputs.outputs[0].finish_reason == "abort"
+ finally:
+ await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_abort_marks_only_last_final_stage_output_finished() -> None:
+ """A request with two final AR stages must keep earlier abort outputs unfinished."""
+ stages = [
+ FakeStageClient(stage_type="llm", final_output=True),
+ FakeStageClient(stage_type="llm", final_output=True),
+ ]
+ processors = [
+ FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[1], finished=True)]),
+ FakeOutputProcessor(request_outputs=[_build_request_output("req-two", token_ids=[2], finished=True)]),
+ ]
+ pools = _build_stage_pools([[stages[0]], [stages[1]]], output_processors=processors)
+ pools[0]._request_bindings["req-two"] = 0
+ pools[1]._request_bindings["req-two"] = 0
+
+ orchestrator = Orchestrator(
+ request_async_queue=asyncio.Queue(),
+ output_async_queue=asyncio.Queue(),
+ rpc_async_queue=asyncio.Queue(),
+ stage_pools=pools,
+ )
+ orchestrator.request_states["req-two"] = OrchestratorRequestState(
+ request_id="req-two",
+ final_stage_id=1,
+ final_output_stage_ids={0, 1},
+ )
+
+ outputs = await orchestrator._abort_request_ids(["req-two"])
+ assert [(msg.stage_id, msg.finished) for msg in outputs] == [(0, False), (1, True)]
+ assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [1]
+ assert list(outputs[1].engine_outputs.outputs[0].token_ids) == [2]
+
+
+@pytest.mark.asyncio
+async def test_handle_abort_emits_error_result_on_failure() -> None:
+ rpc_queue: asyncio.Queue = asyncio.Queue()
+ orchestrator = Orchestrator(
+ request_async_queue=asyncio.Queue(),
+ output_async_queue=asyncio.Queue(),
+ rpc_async_queue=rpc_queue,
+ stage_pools=[],
+ )
+
+ async def boom(_request_ids, *, abort=False):
+ del _request_ids, abort
+ raise RuntimeError("cleanup exploded")
+
+ orchestrator._cleanup_request_ids = boom # type: ignore[method-assign]
+
+ await orchestrator._handle_abort(AbortRequestMessage(request_ids=["req-fail"], rpc_id="abort-err"))
+
+ result = rpc_queue.get_nowait()
+ assert isinstance(result, AbortResultMessage)
+ assert result.rpc_id == "abort-err"
+ assert result.success is False
+ assert "cleanup exploded" in (result.error or "")
+
+
def _duplex_open_message(
session_id: str,
*,
@@ -1951,7 +2260,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None:
final_stage_id=0,
)
req_state.streaming.enabled = True
- req_state.streaming.segment_finished = True
+ req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
req_state.stage_submit_ts[0] = time.time()
orchestrator.request_states = {"req-stream": req_state}
orchestrator.stage_pools = [pool]
@@ -1984,7 +2293,12 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
)
output = SimpleNamespace(
request_id="req-stream",
- outputs=[SimpleNamespace(cumulative_token_ids=list(range(11)))],
+ outputs=[
+ SimpleNamespace(
+ cumulative_token_ids=list(range(11)),
+ finish_reason=FinishReason.LENGTH,
+ )
+ ],
)
metrics = pool.build_stage_metrics(
@@ -1996,6 +2310,29 @@ def test_stage_pool_metrics_use_resumable_segment_token_count() -> None:
assert metrics.num_tokens_out == 3
assert metrics.output_unit_count == 3
+ assert metrics.finish_reason == "length"
+
+
+def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None:
+ stage = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+ pool = StagePool(
+ 1,
+ [stage],
+ output_processor=FakeOutputProcessor(),
+ stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+ )
+ output = SimpleNamespace(request_id="req-image", _custom_output={"image": "non-empty"})
+ pool.record_output_timestamps([output], output_ts=132.0)
+
+ metrics = pool.build_stage_metrics(
+ [output],
+ submit_ts=130.0,
+ request_timestamp=120.0,
+ replica_id=0,
+ )
+
+ assert metrics.serving_time_to_first_output_ms == pytest.approx(12000.0)
+ assert metrics.image_time_to_first_output_ms == pytest.approx(2000.0)
@pytest.mark.asyncio
diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py
index 780c5a17b..90db55edb 100644
--- a/tests/engine/test_orchestrator_error_handling.py
+++ b/tests/engine/test_orchestrator_error_handling.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for error propagation paths within the Orchestrator.
Covers:
@@ -16,7 +16,9 @@ from types import SimpleNamespace
import janus
import pytest
+from vllm.sampling_params import SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.serial_utils import MsgpackEncoder
from vllm_omni.engine.messages import (
AddCompanionRequestMessage,
@@ -25,7 +27,11 @@ from vllm_omni.engine.messages import (
ShutdownRequestMessage,
StageSubmissionMessage,
)
-from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.orchestrator import (
+ Orchestrator,
+ OrchestratorRequestState,
+ cleanup_request_artifact_dirs,
+)
from vllm_omni.engine.stage_pool import StageUnavailableError
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -42,12 +48,22 @@ from .test_orchestrator import (
_wait_for,
)
+
+def test_request_artifact_cleanup_is_idempotent(tmp_path):
+ artifact_dir = tmp_path / "request-artifacts"
+ artifact_dir.mkdir()
+ (artifact_dir / "prepared.mp4").touch()
+
+ cleanup_request_artifact_dirs([str(artifact_dir)])
+ cleanup_request_artifact_dirs([str(artifact_dir)])
+
+ assert not artifact_dir.exists()
+
+
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _sampling_params(max_tokens: int = 4):
- from vllm.sampling_params import SamplingParams
-
return SamplingParams(max_tokens=max_tokens)
@@ -188,6 +204,45 @@ class _UnavailableDiffusionStageClient(FakeStageClient):
raise StageUnavailableError(f"stage {self.stage_id} has no live replica")
+class _OverflowOnEncodeStageClient(FakeStageClient):
+ """Msgpack-encodes the request's sampling params on dispatch, as the real
+ StageEngineCoreClient does. We use this to make sure overflow sampling params
+ do not break the orchestrator."""
+
+ async def add_request_async(self, request, *args, **kwargs) -> None:
+ MsgpackEncoder().encode(request.sampling_params)
+ await super().add_request_async(request, *args, **kwargs) # unreachable
+
+
+@pytest.mark.asyncio
+async def test_encode_overflow_should_fail_request_not_loop(orchestrator_factory) -> None:
+ """Regression test for over-range seed handling."""
+ stage0 = _OverflowOnEncodeStageClient(stage_type="llm", final_output=True)
+ orchestrator_fixture = orchestrator_factory([stage0])
+
+ # seed > 2**63-1 cannot be msgpack-serialized -> the malicious input under test.
+ overflow_params = SamplingParams(seed=2**70)
+
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id="req-overflow",
+ prompt=SimpleNamespace(
+ request_id="req-overflow",
+ prompt_token_ids=[1, 2],
+ sampling_params=overflow_params,
+ ),
+ original_prompt={"prompt": "hi"},
+ sampling_params_list=[overflow_params],
+ final_stage_id=0,
+ )
+
+ error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-overflow")
+ assert error_msg.fatal is False
+ assert error_msg.status_code == 400
+ assert orchestrator_fixture.thread.is_alive()
+ assert "req-overflow" not in orchestrator_fixture.orchestrator.request_states
+
+
@pytest.mark.asyncio
async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None:
"""Multi-replica stage: one replica dying must fail only the requests bound
@@ -411,6 +466,51 @@ async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orch
orchestrator_fixture.thread.join(timeout=5)
+@pytest.mark.asyncio
+async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_request(
+ orchestrator_factory,
+) -> None:
+ """R1.3 of #4855: an async-chunk request whose stage-0 prompt carries no
+ token ids cannot prewarm anything, so it must fail with a client error
+ instead of stalling.
+
+ Non-fatal is the point: ``fatal=True`` reads as "the engine died" and the
+ entrypoints act on it (``OmniLLM.generate`` raises out of its batch loop,
+ the serving layer calls ``terminate_if_errored``). One client's embeds-only
+ prompt must not reach other requests.
+ """
+ stage0 = FakeStageClient(stage_type="llm", final_output=False)
+ stage1 = FakeStageClient(stage_type="llm", final_output=True)
+ orchestrator_fixture = orchestrator_factory([stage0, stage1], async_chunk=True)
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id="req-embeds",
+ # No prompt_token_ids: what an embeds-only prompt looks like here.
+ prompt=SimpleNamespace(request_id="req-embeds"),
+ original_prompt={"prompt": "embeds only"},
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ )
+
+ error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-embeds")
+ assert error_msg.fatal is False
+ assert error_msg.status_code == 400
+ assert error_msg.error_type == "BadRequestError"
+ assert error_msg.stage_id == 0
+ assert "prompt_token_ids" in error_msg.error
+
+ # Nothing was prewarmed downstream, and the request is gone.
+ assert stage1.add_request_calls == []
+ await _wait_for(lambda: "req-embeds" not in orchestrator_fixture.orchestrator.request_states)
+ assert orchestrator_fixture.thread.is_alive()
+ finally:
+ if orchestrator_fixture.thread.is_alive():
+ orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+ orchestrator_fixture.thread.join(timeout=5)
+
+
# ───────── Direct unit tests for the fault-isolation helpers ─────────
@@ -835,6 +935,36 @@ async def test_diffusion_client_error_output_propagates_status_code(orchestrator
orchestrator_fixture.thread.join(timeout=5)
+@pytest.mark.asyncio
+async def test_duplex_input_processor_failure_is_request_scoped(orchestrator_factory, monkeypatch) -> None:
+ class FailingStage(FakeStageClient):
+ def process_engine_inputs(self, *_args, **_kwargs):
+ raise ValueError("No latent or hidden_states found in thinker output")
+
+ fixture = orchestrator_factory(
+ [FakeStageClient(stage_type="llm"), FailingStage(stage_type="llm", final_output=True)]
+ )
+ state = OrchestratorRequestState(
+ request_id="bad",
+ prompt=SimpleNamespace(request_id="bad", prompt_token_ids=[1]),
+ sampling_params_list=[_sampling_params(), _sampling_params()],
+ final_stage_id=1,
+ duplex_identity=SimpleNamespace(),
+ )
+
+ async def no_cleanup(*_args, **_kwargs):
+ pass
+
+ monkeypatch.setattr(fixture.orchestrator, "_cleanup_request_ids", no_cleanup)
+ try:
+ await fixture.orchestrator._forward_to_next_stage_unguarded("bad", 0, _build_request_output("raw"), state)
+ error = await _wait_for_error_message(fixture, request_id="bad")
+ assert error.fatal is False
+ finally:
+ fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+ fixture.thread.join(timeout=5)
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize(
("status_code", "error_type"),
diff --git a/tests/engine/test_orchestrator_event_driven.py b/tests/engine/test_orchestrator_event_driven.py
new file mode 100644
index 000000000..284dde55d
--- /dev/null
+++ b/tests/engine/test_orchestrator_event_driven.py
@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Event-driven orchestration loop (``VLLM_OMNI_EVENT_DRIVEN_ORCH=1``) tests.
+
+Parity suite: re-runs the legacy orchestration scenarios from
+``test_orchestrator.py`` / ``test_orchestrator_error_handling.py`` with the
+event-driven loop selected, so both loops are held to the same behavior. Plus
+event-driven-specific coverage: reader reconcile on client swap, the blocking
+final-output drain, and flag parsing.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import janus
+import pytest
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.messages import ShutdownRequestMessage
+from vllm_omni.engine.orchestrator import _event_driven_orch_enabled
+
+from . import test_orchestrator as legacy
+from . import test_orchestrator_error_handling as legacy_errors
+from .test_orchestrator import (
+ FakeOutputProcessor,
+ FakeStageClient,
+ OrchestratorFixture,
+ _build_harness,
+ _build_request_output,
+ _engine_core_outputs,
+ _enqueue_add_request,
+ _get_output_message,
+ _sampling_params,
+ _shutdown_orchestrator,
+ _wait_for,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def orchestrator_factory(monkeypatch):
+ """Flag-setting clone of the legacy harness fixture.
+
+ Sets ``VLLM_OMNI_EVENT_DRIVEN_ORCH=1`` before any Orchestrator is
+ constructed and asserts the flag actually took effect, so the parity tests
+ cannot silently exercise the legacy poll loop.
+ """
+ monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", "1")
+ fixtures: list[OrchestratorFixture] = []
+
+ def _factory(*args, **kwargs) -> OrchestratorFixture:
+ fixture = _build_harness(*args, **kwargs)
+ assert fixture.orchestrator._event_driven_orch is True
+ fixtures.append(fixture)
+ return fixture
+
+ yield _factory
+
+ for fixture in fixtures:
+ if fixture.thread.is_alive():
+ fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+ fixture.thread.join(timeout=5)
+ for q in fixture.queues:
+ q.close()
+
+
+# ---------------------------------------------------------------------------
+# Parity: the legacy scenario matrix, re-run through the event-driven loop
+# ---------------------------------------------------------------------------
+
+_PARITY_TESTS = [
+ legacy.test_run_two_stage_llm,
+ legacy.test_run_single_stage_diffusion,
+ legacy.test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks,
+ legacy.test_run_llm_to_diffusion,
+ legacy.test_run_async_chunk,
+ legacy.test_run_shutdown,
+ legacy.test_run_abort,
+ legacy.test_multi_replica_round_robin_distribution,
+ legacy.test_multi_replica_abort_broadcasts_to_all_replicas,
+ legacy.test_multi_replica_shutdown_all_replicas,
+ legacy.test_multi_replica_cfg_companion_inherits_parent_affinity,
+ # Stats plumbing. The scheduler-stats cases matter most here: a batch with
+ # no request outputs still carries SchedulerStats on throttled ticks, and a
+ # reader that drops every output-less batch silently stops reporting
+ # KV/queue gauges under the event-driven loop.
+ legacy.test_orchestrator_records_iteration_stats_without_scheduler_stats,
+ legacy.test_orchestrator_records_scheduler_stats_without_outputs,
+ legacy.test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch,
+ legacy.test_orchestrator_does_not_build_iteration_stats_without_stat_logger,
+ # Per-replica fault isolation (#4285): a dead replica must be evicted and
+ # the server kept up, on both the LLM reader path and the diffusion poller.
+ legacy_errors.test_engine_dead_error_evicts_replica_and_keeps_running,
+ legacy_errors.test_engine_dead_error_fails_only_dead_replica_requests,
+ legacy_errors.test_forward_to_dead_downstream_stage_fails_request_not_server,
+ legacy_errors.test_add_request_to_dead_stage_fails_request_not_server,
+ legacy_errors.test_diffusion_replica_death_on_poll_keeps_server,
+ legacy_errors.test_diffusion_error_output_routed_as_finished,
+ legacy_errors.test_diffusion_client_error_output_propagates_status_code,
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("legacy_test", _PARITY_TESTS, ids=lambda f: f.__name__)
+async def test_event_driven_parity(legacy_test, orchestrator_factory) -> None:
+ await legacy_test(orchestrator_factory)
+
+
+# ---------------------------------------------------------------------------
+# Event-driven-specific behavior
+# ---------------------------------------------------------------------------
+
+
+def test_flag_parsing(monkeypatch) -> None:
+ monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)
+ assert _event_driven_orch_enabled() is False
+ for value in ("1", "true", "True", "YES", "on"):
+ monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value)
+ assert _event_driven_orch_enabled() is True
+ for value in ("0", "false", "off", ""):
+ monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", value)
+ assert _event_driven_orch_enabled() is False
+
+
+def test_default_is_legacy_loop(monkeypatch) -> None:
+ """Without the env flag, the harness runs the legacy poll loop."""
+ monkeypatch.delenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", raising=False)
+ fixture = _build_harness([FakeStageClient(stage_type="llm", final_output=True)])
+ try:
+ assert fixture.orchestrator._event_driven_orch is False
+ finally:
+ fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+ fixture.thread.join(timeout=5)
+ for q in fixture.queues:
+ q.close()
+
+
+@pytest.mark.asyncio
+async def test_reader_reconcile_picks_up_swapped_client(orchestrator_factory) -> None:
+ """Outputs from a replica whose client object was replaced still flow.
+
+ The event-driven loop binds one reader task per client object; the
+ periodic reconcile must respawn the reader when ``pool.clients[replica]``
+ is swapped (replica replacement), otherwise the new client's outputs
+ would never be drained.
+ """
+ stage0 = FakeStageClient(stage_type="llm", final_output=True)
+ processor = FakeOutputProcessor(request_outputs=[_build_request_output("req-swap", token_ids=[3], finished=True)])
+ orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor])
+ request = SimpleNamespace(request_id="req-swap", prompt_token_ids=[1, 2])
+
+ try:
+ await _enqueue_add_request(
+ orchestrator_fixture,
+ request_id="req-swap",
+ prompt=request,
+ original_prompt={"prompt": "swap"},
+ sampling_params_list=[_sampling_params()],
+ final_stage_id=0,
+ )
+ await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+ # Swap in a fresh client for the same replica slot; keep the pool's
+ # other wiring intact. The reconcile tick (0.5 s) must respawn the
+ # reader bound to the new client object.
+ pool = orchestrator_fixture.orchestrator.stage_pools[0]
+ replacement = FakeStageClient(stage_type="llm", final_output=True)
+ replacement.stage_id = stage0.stage_id
+ replacement.replica_id = stage0.replica_id
+ pool.clients[0] = replacement
+
+ replacement.push_engine_core_outputs(_engine_core_outputs("swapped-raw", 1.0))
+
+ output_msg = await _get_output_message(orchestrator_fixture, timeout=5.0)
+ assert output_msg.request_id == "req-swap"
+ assert output_msg.finished is True
+ finally:
+ await _shutdown_orchestrator(orchestrator_fixture)
+
+
+# ---------------------------------------------------------------------------
+# Blocking final-output drain (AsyncOmniEngine.get_output_blocking_async)
+# ---------------------------------------------------------------------------
+
+
+def _drain_engine(alive: bool = True) -> AsyncOmniEngine:
+ engine = object.__new__(AsyncOmniEngine)
+ engine.output_queue = janus.Queue()
+ engine.orchestrator_thread = SimpleNamespace(is_alive=lambda: alive)
+ return engine
+
+
+def _drain_cleanup(engine: AsyncOmniEngine) -> None:
+ if engine._output_drain_executor is not None:
+ engine._output_drain_executor.shutdown(wait=False)
+ engine._output_drain_executor = None
+ engine.output_queue.close()
+
+
+@pytest.mark.asyncio
+async def test_blocking_drain_returns_queued_message() -> None:
+ engine = _drain_engine()
+ try:
+ engine.output_queue.sync_q.put_nowait("msg-1")
+ assert await engine.get_output_blocking_async(timeout=1.0) == "msg-1"
+ finally:
+ _drain_cleanup(engine)
+
+
+@pytest.mark.asyncio
+async def test_blocking_drain_wakes_on_late_message() -> None:
+ """A message put after the wait starts wakes the drain, no polling."""
+ engine = _drain_engine()
+ try:
+
+ async def _delayed_put() -> None:
+ await asyncio.sleep(0.05)
+ engine.output_queue.sync_q.put_nowait("late-msg")
+
+ put_task = asyncio.create_task(_delayed_put())
+ msg = await engine.get_output_blocking_async(timeout=5.0)
+ await put_task
+ assert msg == "late-msg"
+ finally:
+ _drain_cleanup(engine)
+
+
+@pytest.mark.asyncio
+async def test_blocking_drain_timeout_returns_none_when_alive() -> None:
+ engine = _drain_engine(alive=True)
+ try:
+ assert await engine.get_output_blocking_async(timeout=0.05) is None
+ finally:
+ _drain_cleanup(engine)
+
+
+@pytest.mark.asyncio
+async def test_blocking_drain_raises_when_orchestrator_dead() -> None:
+ engine = _drain_engine(alive=False)
+ try:
+ with pytest.raises(RuntimeError, match="Orchestrator died"):
+ await engine.get_output_blocking_async(timeout=0.05)
+ finally:
+ _drain_cleanup(engine)
diff --git a/tests/engine/test_orchestrator_segment_stage_keying.py b/tests/engine/test_orchestrator_segment_stage_keying.py
new file mode 100644
index 000000000..02bb9a338
--- /dev/null
+++ b/tests/engine/test_orchestrator_segment_stage_keying.py
@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Streaming segment-boundary state must be keyed by stage.
+
+``Orchestrator._orchestration_loop`` polls every stage into the same
+``OrchestratorRequestState``. While the segment-boundary fields lived in a single
+flat slot per request, whichever stage polled most recently overwrote the others,
+and the per-stage consumers then read back a different stage's boundary.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import (
+ Orchestrator,
+ OrchestratorRequestState,
+ StreamingSegmentState,
+)
+from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity
+from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+THINKER_STAGE = 0
+TALKER_STAGE = 1
+
+
+@dataclass
+class _StageMetrics:
+ pipeline_timings: dict[str, float] = field(default_factory=dict)
+
+
+@dataclass
+class _StageOutput:
+ """Stand-in for a processed stage output, with the fields the router reads."""
+
+ request_id: str
+ finished: bool = False
+ error: object | None = None
+
+
+class _RecordingPool:
+ """Stage pool that records the requests it was asked to build metrics for."""
+
+ def __init__(self) -> None:
+ self.built_for: list[str] = []
+
+ def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics:
+ self.built_for.extend(output.request_id for output in outputs)
+ return _StageMetrics()
+
+
+def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState:
+ req_state = OrchestratorRequestState(
+ request_id=request_id,
+ sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)],
+ final_stage_id=TALKER_STAGE,
+ )
+ req_state.streaming.enabled = True
+ req_state.duplex_identity = DuplexRequestIdentity(
+ session_id="session-1",
+ fence=DuplexFence(session_id="session-1"),
+ )
+ return req_state
+
+
+def test_duplex_output_context_reads_only_its_own_stage_segment() -> None:
+ req_state = _duplex_req_state()
+ req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(
+ finished=True,
+ token_ids=[11, 22],
+ output_metadata={"stage": "thinker"},
+ )
+ # The talker polls after the thinker and is still mid-segment.
+ req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False)
+
+ thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE)
+ talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE)
+
+ assert thinker is not None
+ assert talker is not None
+
+ # These are the assertions that fail on the flat field: the talker is written
+ # second, so the thinker loses its boundary to whichever stage polled last.
+ assert thinker.segment_finished is True
+ assert thinker.segment_token_ids == (11, 22)
+ assert thinker.segment_output_metadata == {"stage": "thinker"}
+
+ assert talker.segment_finished is False
+ assert talker.segment_token_ids == ()
+ assert talker.segment_output_metadata == {}
+
+
+@pytest.mark.asyncio
+async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None:
+ req_state = _duplex_req_state()
+ req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True)
+
+ talker_pool = _RecordingPool()
+ orchestrator = object.__new__(Orchestrator)
+ orchestrator.request_states = {req_state.request_id: req_state}
+ orchestrator.stage_pools = [_RecordingPool(), talker_pool]
+
+ routed: list[_StageMetrics | None] = []
+
+ async def record_route(
+ _stage: int,
+ _replica: int,
+ _output: _StageOutput,
+ _state: OrchestratorRequestState,
+ metrics: _StageMetrics | None,
+ ) -> None:
+ routed.append(metrics)
+
+ orchestrator._route_output = record_route
+
+ # The talker's own output is neither finished nor at a segment boundary, so
+ # it must not be treated as one just because the thinker reached one.
+ await orchestrator._handle_processed_outputs(
+ TALKER_STAGE,
+ 0,
+ [_StageOutput(request_id=req_state.request_id)],
+ )
+
+ assert talker_pool.built_for == []
+ assert routed == [None]
diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py
index cf347e49f..52cee4657 100644
--- a/tests/engine/test_orchestrator_stage_input_bridge.py
+++ b/tests/engine/test_orchestrator_stage_input_bridge.py
@@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams
from vllm_omni.engine.orchestrator import (
Orchestrator,
OrchestratorRequestState,
+ StreamingSegmentState,
_OrchestratorDuplexStagePort,
)
from vllm_omni.engine.stage_pool import StagePool
@@ -160,6 +161,26 @@ def _duplex_stage_port_submission():
return port, stage_pools, request_states, prewarm, submission
+def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None:
+ port, _, request_states, _, submission = _duplex_stage_port_submission()
+ request_state = request_states[submission.context.request_id]
+ bridge_state = request_state.streaming.bridge_states["duplex"]
+ bridge_state["model_turn_id"] = 0
+
+ next_context = DuplexStageRequestContext(
+ request_id=submission.context.request_id,
+ session_id=submission.context.session_id,
+ fence=DuplexFence(submission.context.session_id, turn_id=1),
+ stage_id=submission.context.stage_id,
+ final_stage_id=submission.context.final_stage_id,
+ config_generation=submission.context.config_generation,
+ sampling_params=submission.context.sampling_params,
+ )
+ port.ensure_request(next_context)
+
+ assert bridge_state["model_turn_id"] == 1
+
+
def _request_output(request_id: str) -> RequestOutput:
completion = CompletionOutput(
index=0,
@@ -256,12 +277,13 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
duplex_identity=SimpleNamespace(),
)
- await orchestrator._prewarm_async_chunk_stages(
+ prewarmed = await orchestrator._prewarm_async_chunk_stages(
"req-prewarm",
SimpleNamespace(prompt_token_ids=[1, 2], resumable=True),
req_state,
)
+ assert prewarmed is True
assert stage1.submitted == []
assert len(stage2.submitted) == 1
assert 1 not in req_state.stage_submit_ts
@@ -277,6 +299,7 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
@pytest.mark.asyncio
async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
+ prewarm.return_value = True
result = await port.submit(submission)
@@ -285,6 +308,30 @@ async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
+@pytest.mark.asyncio
+async def test_duplex_submit_bails_out_when_prewarm_failed_the_request() -> None:
+ """A failed prewarm has already aborted the request and popped its state.
+
+ Returning a success result here would hand the control plane a replica for a
+ request that no longer exists, and the trailing bookkeeping would re-register
+ a running counter the cleanup just released -- a leak that never decrements.
+ """
+ port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
+ counter = MagicMock()
+ port._running_counter = counter
+ prewarm.return_value = False
+ request_state = request_states["req-duplex"]
+
+ with pytest.raises(RuntimeError, match="prewarm failed"):
+ await port.submit(submission)
+
+ prewarm.assert_awaited_once()
+ assert request_state.duplex_stage_fences == {}
+ assert request_state.stage_submit_ts == {}
+ assert request_state.running_counter_registered is False
+ counter.increment.assert_not_called()
+
+
@pytest.mark.asyncio
async def test_async_route_forwards_to_outgoing_only_stage() -> None:
orchestrator = object.__new__(Orchestrator)
@@ -332,7 +379,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None:
final_output_stage_ids={0},
)
req_state.streaming.enabled = True
- req_state.streaming.segment_finished = True
+ req_state.streaming.segments[0] = StreamingSegmentState(finished=True)
output = SimpleNamespace(
request_id=req_state.request_id,
finished=True,
diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py
index 23ac601b3..62ed8334b 100644
--- a/tests/engine/test_output_processor.py
+++ b/tests/engine/test_output_processor.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
from types import SimpleNamespace
@@ -581,18 +581,18 @@ def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
def test_no_detokenizer_make_request_output_with_routed_experts():
"""make_request_output accepts the routed_experts arg that the multimodal
- output channel (_process_mm_only_outputs) passes positionally, and attaches
- it to the completion output.
+ output channel (_process_mm_only_outputs) passes, and attaches it to the
+ completion output.
- Regression: the call site passes 6 positional args
- (..., kv_transfer_params, routed_experts) but the override previously took
- only 5, raising TypeError on every generation-stage output.
+ routed_experts is omni-specific keyword-only: the 6th positional now
+ matches upstream's ec_transfer_params so that super().process_outputs()
+ calls do not misroute it.
"""
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
routed_experts = np.zeros((2, 3), dtype=np.int32)
# Mirror the exact call shape of _process_mm_only_outputs.
- result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts)
+ result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts)
assert result is not None
assert not isinstance(result, PoolingRequestOutput)
assert result.outputs[0].routed_experts is routed_experts
@@ -663,3 +663,99 @@ def test_mm_only_outputs_update_iteration_stats():
assert finished.finish_reason == FinishReason.STOP
assert finished.num_prompt_tokens == state.prompt_len
assert finished.num_generation_tokens == 2
+
+
+def test_ec_transfer_params_survive_both_construction_paths():
+ """v0.28 contract: encoder-cache transfer metadata must reach
+ RequestOutput.ec_transfer_params through BOTH omni construction paths —
+ the upstream-delegating one (logprobs processor present) and the
+ no-detokenizer generation-stage one (direct RequestOutput build).
+ Regression: the override accepted the argument but dropped it."""
+ ec = {"remote_handle": "enc-cache-1"}
+
+ state = _make_state(RequestOutputKind.CUMULATIVE)
+ out = state.make_request_output([1], None, None, None, {"kv": 1}, ec)
+ assert out is not None and out.ec_transfer_params == ec
+
+ kwargs = dict(_DEFAULT_STATE_KWARGS)
+ kwargs.update(logprobs_processor=None, detokenizer=None)
+ gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
+ completion = gen_state._new_completion_output([1], None, None)
+ direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec)
+ assert direct.ec_transfer_params == ec
+
+
+def test_num_cache_creation_tokens_reaches_direct_request_output():
+ """v0.28 contract: prefix-cache creation usage must reach
+ RequestOutput.num_cache_creation_tokens through the no-detokenizer
+ direct build, which previously passed only num_cached_tokens."""
+ kwargs = dict(_DEFAULT_STATE_KWARGS)
+ kwargs.update(logprobs_processor=None, detokenizer=None)
+ gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
+ gen_state.num_cached_tokens = 3
+ gen_state.num_cache_creation_tokens = 2
+
+ completion = gen_state._new_completion_output([1], None, None)
+ direct = gen_state._new_request_output("r", [completion], False, None, None)
+
+ assert direct.num_cached_tokens == 3
+ assert direct.num_cache_creation_tokens == 2
+
+
+def _abort_processor_with_parent(child_ids: list[str]):
+ processor = object.__new__(MultimodalOutputProcessor)
+ processor.request_states = {}
+ processor.external_req_ids = {}
+ processor.parent_requests = {}
+ processor._native_text_metrics_by_request = {}
+ processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None)
+ parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids))
+ processor.parent_requests["parent"] = parent
+ for child_id in child_ids:
+ req_state = SimpleNamespace(
+ parent_req=parent,
+ lora_name=None,
+ output_kind=RequestOutputKind.CUMULATIVE,
+ detokenizer=SimpleNamespace(output_token_ids=[7, 8]),
+ queue=None,
+ external_req_id="parent",
+ make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)),
+ )
+ processor.request_states[child_id] = req_state
+ return processor, parent
+
+
+def test_abort_last_parallel_child_drops_parent_request():
+ processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"])
+
+ aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True)
+ assert aborted == ["0_parent"]
+ assert "parent" in processor.parent_requests
+ assert parent.child_requests == {"1_parent"}
+
+ aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True)
+ assert aborted == ["1_parent"]
+ assert "parent" not in processor.parent_requests
+ assert not parent.child_requests
+
+
+def test_abort_snapshot_leaves_state_until_commit():
+ processor, parent = _abort_processor_with_parent(["0_parent"])
+ processor.external_req_ids["parent"] = ["0_parent"]
+
+ aborted, outputs = processor.abort_requests_collecting_outputs(
+ ["parent"],
+ internal=False,
+ commit_state=False,
+ )
+ assert aborted == ["0_parent"]
+ assert outputs
+ assert "0_parent" in processor.request_states
+ assert processor.external_req_ids["parent"] == ["0_parent"]
+ assert "parent" in processor.parent_requests
+ assert parent.child_requests == {"0_parent"}
+
+ processor.commit_aborted_request_state(["parent"], internal=False)
+ assert "0_parent" not in processor.request_states
+ assert "parent" not in processor.parent_requests
+ assert "parent" not in processor.external_req_ids
diff --git a/tests/engine/test_rpc_result_router.py b/tests/engine/test_rpc_result_router.py
index ec983392b..1b3159617 100644
--- a/tests/engine/test_rpc_result_router.py
+++ b/tests/engine/test_rpc_result_router.py
@@ -4,6 +4,7 @@ from typing import Literal
import pytest
from vllm_omni.engine.messages import (
+ AbortResultMessage,
CollectiveRPCResultMessage,
EngineQueueMessage,
ErrorMessage,
@@ -93,6 +94,19 @@ def test_rpc_result_router_does_not_broadcast_uncorrelated_nonfatal_errors():
router.close()
+def test_rpc_result_router_routes_abort_results_by_correlation_id():
+ source: queue.Queue = queue.Queue()
+ router = RpcResultRouter(source)
+ waiter = router.register(("abort", "abort-42"))
+
+ source.put(AbortResultMessage(rpc_id="abort-42", success=True))
+
+ result = waiter.get(timeout=1)
+ assert isinstance(result, AbortResultMessage)
+ assert result.rpc_id == "abort-42"
+ router.close()
+
+
def test_rpc_result_router_close_unblocks_waiters_and_stops_consumer():
source: queue.Queue = queue.Queue()
router = RpcResultRouter(source)
diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py
index bde8ba490..cacda1f60 100644
--- a/tests/engine/test_single_stage_mode.py
+++ b/tests/engine/test_single_stage_mode.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Unit tests for AsyncOmniEngine single-stage mode and OmniMasterServer."""
from __future__ import annotations
@@ -428,7 +431,7 @@ class TestSingleStageModeDetection:
assert engine.single_stage_mode is False
assert engine._single_stage_id_filter is None
- def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
+ def test_deploy_config_loads_duplex_runtime_config(self, mocker: MockerFixture):
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
get_pipeline_config = mocker.patch(
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
@@ -441,7 +444,7 @@ class TestSingleStageModeDetection:
engine = self._make_engine_no_thread(
mocker,
- stage_configs_path="/fake/duplex.yaml",
+ deploy_config="/fake/duplex.yaml",
)
get_pipeline_config.assert_called_once_with(
@@ -572,7 +575,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=stage_id_filter,
omni_master_address="127.0.0.1",
@@ -804,7 +806,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
)
mocker.patch.object(runtime_mod, "prepare_engine_environment")
@@ -857,7 +858,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -925,7 +925,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -947,7 +946,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -999,7 +997,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1065,7 +1062,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1117,7 +1113,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1138,7 +1133,8 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
- mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
+ od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
+ mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mock_register = mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
return_value=StageRegistrationResponse(
@@ -1174,6 +1170,7 @@ class TestSingleStageReplicaInitialization:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
+ assert od_config.max_num_seqs == 4
mock_register.assert_called_once_with(
omni_master_address="127.0.0.1",
omni_master_port=25000,
@@ -1183,7 +1180,7 @@ class TestSingleStageReplicaInitialization:
)
mock_manager.assert_called_once_with(
model="fake-model",
- od_config="diffusion-config",
+ od_config=od_config,
stage_init_timeout=60,
handshake_address="tcp://127.0.0.1:26001",
addresses=mocker.ANY,
@@ -1198,10 +1195,50 @@ class TestSingleStageReplicaInitialization:
request_address="tcp://127.0.0.1:26002",
response_address="tcp://127.0.0.1:26003",
proc_manager=mocker.ANY,
- batch_size=4,
)
assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
+ def test_initialize_local_custom_diffusion_replica_uses_inline_client(self, mocker: MockerFixture):
+ import vllm_omni.engine.stage_runtime as runtime_mod
+ from vllm_omni.platforms import current_omni_platform
+
+ runtime = DistStageRuntime(
+ stage_configs=[],
+ model="fake-model",
+ config_path="/fake/stages.yaml",
+ stage_init_timeout=60,
+ async_chunk=False,
+ single_stage_id_filter=None,
+ omni_master_address="127.0.0.1",
+ omni_master_port=26000,
+ )
+ runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
+ runtime._coordinator_runtime = None
+ plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0]
+ plan.stage_cfg.engine_args = {"custom_pipeline_args": {"pipeline_class": "test.CustomPipeline"}}
+ sentinel_client = SimpleNamespace()
+ mock_launch = mocker.patch.object(
+ runtime_mod,
+ "launch_diffusion_stage_replica",
+ return_value=(sentinel_client, StageReplicaResources()),
+ )
+ mocker.patch.object(runtime_mod, "inject_kv_stage_info")
+
+ device_env_var = current_omni_platform.device_control_env_var
+ prev_device_env = os.environ.get(device_env_var)
+ os.environ[device_env_var] = "0"
+ runtime._init_visible_devices_baseline = "0"
+ try:
+ result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60)
+ finally:
+ if prev_device_env is None:
+ os.environ.pop(device_env_var, None)
+ else:
+ os.environ[device_env_var] = prev_device_env
+
+ assert result is sentinel_client
+ assert mock_launch.call_args.kwargs["use_inline"] is True
+
def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture):
import vllm_omni.engine.stage_runtime as runtime_mod
from vllm_omni.platforms import current_omni_platform
@@ -1211,7 +1248,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
- diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1230,7 +1266,8 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
- mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config")
+ od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
+ mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
return_value=StageRegistrationResponse(
@@ -1371,7 +1408,7 @@ class TestConnectRemoteEngineCoresCoordinator:
omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2)
omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2)
mock_wait.assert_called_once()
- _, _, core_engines, parallel_config, *_ = mock_wait.call_args.args
+ _, core_engines, parallel_config, *_ = mock_wait.call_args.args
assert core_engines[0].local is False
assert parallel_config.data_parallel_size_local == 0
diff --git a/tests/engine/test_stage_device_layout.py b/tests/engine/test_stage_device_layout.py
new file mode 100644
index 000000000..0473dc341
--- /dev/null
+++ b/tests/engine/test_stage_device_layout.py
@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Regression tests for issue #5003: a per-stage world size that its assigned
+``devices`` cannot satisfy must fail early in ``build_vllm_config`` with a clear
+message, rather than surfacing as an opaque worker-side ``local rank ... out of
+bounds`` assertion.
+
+Root cause: a top-level ``--tensor-parallel-size`` is broadcast to every stage,
+but each stage's ``devices`` is not adjusted, so a stage can end up with e.g.
+tensor_parallel_size=4 while still holding a single-GPU deploy default. Without
+``--strategy-config`` the strategy-path device check never runs.
+"""
+
+import json
+import re
+import types
+from unittest import mock
+
+import pytest
+
+from vllm_omni.engine import stage_init_utils
+from vllm_omni.engine.stage_init_utils import (
+ _check_stage_device_layout,
+ build_vllm_config,
+ compute_replica_layout,
+ get_stage_devices_per_replica,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _stage(stage_id, devices, num_replicas=1, engine_args=None):
+ return types.SimpleNamespace(
+ stage_id=stage_id,
+ stage_type="llm",
+ engine_args=engine_args or {},
+ runtime=types.SimpleNamespace(devices=devices, num_replicas=num_replicas),
+ )
+
+
+def test_tp_broadcast_without_devices_fails_early():
+ """stage0 gets tensor_parallel_size=4 (broadcast) but only 1 device -> clear error."""
+ stage = _stage(0, devices="0")
+ engine_args = {
+ "tensor_parallel_size": 4,
+ "data_parallel_size": 1,
+ "pipeline_parallel_size": 1,
+ }
+ with pytest.raises(ValueError) as excinfo:
+ _check_stage_device_layout(stage, engine_args)
+ msg = str(excinfo.value)
+ # Message names the stage, the mismatch, and the actionable workaround.
+ assert "Stage 0" in msg
+ assert "1 device" in msg and "4" in msg
+ assert "--stage-overrides" in msg
+
+
+def test_workaround_example_sets_tp_on_every_stage():
+ """The JSON example in the error message must not reproduce the bug: every
+ stage entry it suggests has to set ``tensor_parallel_size`` (a single-GPU
+ stage that only overrides ``devices`` would still inherit the broadcast
+ tp and crash again — the original #5003 failure mode)."""
+ stage = _stage(0, devices="0")
+ with pytest.raises(ValueError) as excinfo:
+ _check_stage_device_layout(stage, {"tensor_parallel_size": 4})
+ msg = str(excinfo.value)
+
+ match = re.search(r"\{.*\}", msg, re.DOTALL)
+ assert match, f"no JSON example found in error message: {msg}"
+ example = json.loads(match.group(0))
+
+ assert example, "example must contain at least one stage"
+ for stage_id, overrides in example.items():
+ assert "tensor_parallel_size" in overrides, (
+ f"stage {stage_id} in the suggested workaround omits "
+ f"tensor_parallel_size, which reproduces #5003: {overrides}"
+ )
+
+
+def test_consistent_tp_and_devices_pass():
+ """tensor_parallel_size=4 with 4 assigned devices is valid."""
+ stage = _stage(0, devices="0,1,2,3")
+ _check_stage_device_layout(
+ stage,
+ {"tensor_parallel_size": 4, "data_parallel_size": 1, "pipeline_parallel_size": 1},
+ )
+
+
+def test_single_gpu_stage_passes():
+ """A TP=1 stage on a single device (talker/code2wav default) is valid."""
+ stage = _stage(1, devices="1")
+ _check_stage_device_layout(stage, {"tensor_parallel_size": 1})
+
+
+def test_missing_devices_is_skipped():
+ """No explicit devices -> vLLM assigns them; nothing to validate here."""
+ stage = _stage(0, devices=None)
+ _check_stage_device_layout(stage, {"tensor_parallel_size": 4})
+
+
+def test_replica_pool_layout_passes():
+ """Pool mode: num_replicas=2 x tp=2 => 4 devices is a valid pool shape."""
+ stage = _stage(0, devices="0,1,2,3", num_replicas=2)
+ _check_stage_device_layout(stage, {"tensor_parallel_size": 2})
+
+
+def test_multinode_dp_uses_local_width_for_device_validation():
+ """One local DP engine needs one device even when global DP spans four nodes."""
+ stage = _stage(0, devices="0")
+ _check_stage_device_layout(
+ stage,
+ {
+ "tensor_parallel_size": 1,
+ "data_parallel_size": 4,
+ "data_parallel_size_local": 1,
+ "pipeline_parallel_size": 1,
+ },
+ )
+
+
+def test_zero_local_dp_skips_local_device_validation():
+ """A head process with no local engines does not own a local device layout."""
+ stage = _stage(0, devices="0")
+ _check_stage_device_layout(
+ stage,
+ {
+ "tensor_parallel_size": 4,
+ "data_parallel_size": 4,
+ "data_parallel_size_local": 0,
+ "pipeline_parallel_size": 1,
+ },
+ )
+
+
+@pytest.mark.parametrize(
+ ("devices", "expected"),
+ [
+ ("0,1", ["0,1", "2,3"]),
+ ("0,1,2,3", ["0,1", "2,3"]),
+ ],
+)
+def test_replica_guard_and_splitter_share_local_world_size(devices, expected):
+ """PP and local DP contribute to the same per-replica size in both paths."""
+ engine_args = {
+ "tensor_parallel_size": 1,
+ "data_parallel_size": 4,
+ "data_parallel_size_local": 1,
+ "pipeline_parallel_size": 2,
+ }
+ stage = _stage(0, devices=devices, num_replicas=2, engine_args=engine_args)
+
+ _check_stage_device_layout(stage, engine_args)
+ assert get_stage_devices_per_replica(stage) == 2
+ _, replica_devices_map = compute_replica_layout([stage])
+ assert replica_devices_map == {0: expected}
+
+
+def test_non_tp_layout_error_uses_generic_guidance():
+ """A PP mismatch must not be reported as a top-level TP broadcast."""
+ stage = _stage(0, devices="0")
+ with pytest.raises(ValueError) as excinfo:
+ _check_stage_device_layout(
+ stage,
+ {
+ "tensor_parallel_size": 1,
+ "data_parallel_size": 1,
+ "pipeline_parallel_size": 2,
+ },
+ )
+
+ msg = str(excinfo.value)
+ assert "local world size" in msg
+ assert "top-level --tensor-parallel-size" not in msg
+
+
+def test_build_vllm_config_fails_before_engine_config_on_mismatch():
+ """Pin the production wiring: ``build_vllm_config`` must run the device-layout
+ guard *before* ``create_engine_config``. If the guard call is dropped or moved
+ after engine-config creation, #5003 regresses while the unit tests above stay
+ green — so assert here that a mismatched stage raises and that neither
+ ``create_engine_config`` nor ``Executor.get_class`` is reached."""
+ stage = _stage(0, devices="0")
+ with (
+ mock.patch.object(stage_init_utils.OmniEngineArgs, "create_engine_config") as create_engine_config,
+ mock.patch.object(stage_init_utils.Executor, "get_class") as get_class,
+ ):
+ with pytest.raises(ValueError) as excinfo:
+ build_vllm_config(
+ stage,
+ model="dummy-model",
+ engine_args_dict={
+ "tensor_parallel_size": 4,
+ "data_parallel_size": 1,
+ "pipeline_parallel_size": 1,
+ },
+ )
+ msg = str(excinfo.value)
+ assert "Stage 0" in msg
+ assert "--stage-overrides" in msg
+ # Guard fired early: worker/config construction was never reached.
+ create_engine_config.assert_not_called()
+ get_class.assert_not_called()
+
+
+def test_build_vllm_config_proceeds_on_consistent_layout():
+ """The guard must not false-positive: a consistent single-GPU stage flows past
+ it and reaches ``create_engine_config`` / ``Executor.get_class`` as usual."""
+ stage = _stage(1, devices="1")
+ fake_config = types.SimpleNamespace(
+ quant_config=None,
+ model_config=types.SimpleNamespace(hf_config=types.SimpleNamespace()),
+ )
+ sentinel_executor = object()
+ with (
+ mock.patch.object(
+ stage_init_utils.OmniEngineArgs, "create_engine_config", return_value=fake_config
+ ) as create_engine_config,
+ mock.patch.object(stage_init_utils.Executor, "get_class", return_value=sentinel_executor),
+ mock.patch.object(stage_init_utils.OmniINCConfig, "maybe_upgrade", side_effect=lambda quant: quant),
+ ):
+ vllm_config, executor_class = build_vllm_config(
+ stage,
+ model="dummy-model",
+ engine_args_dict={"tensor_parallel_size": 1},
+ )
+ create_engine_config.assert_called_once()
+ assert vllm_config is fake_config
+ assert executor_class is sentinel_executor
diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py
index 75660cbc7..0a54e2e80 100644
--- a/tests/engine/test_stage_engine_args.py
+++ b/tests/engine/test_stage_engine_args.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import copy
@@ -8,8 +11,23 @@ from pathlib import Path
import pytest
from pydantic.fields import FieldInfo
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
+from vllm.engine.arg_utils import EngineArgs
-from vllm_omni.config.omni_config import VllmOmniConfig
+from vllm_omni.config.omni_config import (
+ _LLM_STAGE_ENGINE_FIELDS,
+ OmniStageCacheConfig,
+ OmniStageLoadConfig,
+ OmniStageParallelConfig,
+ OmniStageSchedulerConfig,
+ VllmOmniARStageConfig,
+ VllmOmniConfig,
+)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config
from vllm_omni.config.stage_config import (
DeployConfig,
@@ -30,6 +48,7 @@ from vllm_omni.engine.stage_init_utils import (
build_engine_args_dict_from_omni_stage_config,
build_legacy_engine_args_dict,
)
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -53,6 +72,32 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str,
_LLM_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniEngineArgs))
_DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusionConfig))
+_TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"})
+_OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
+ {
+ "active_stream_window",
+ "codec_frame_rate_hz",
+ "custom_voice_dir",
+ "devices",
+ "disable_autocast",
+ "enable_multithread_weight_load",
+ "env",
+ "has_sampling_extra_args",
+ "log_level",
+ "log_stats",
+ "model_arch",
+ "model_subdir",
+ "num_gpus",
+ "num_replicas",
+ "num_weight_load_threads",
+ "omni_kv_config",
+ "parallel_config",
+ "silence_ban_frames",
+ "subtalker_sampling_params",
+ "task_type",
+ "tokenizer_subdir",
+ }
+)
@pytest.fixture(autouse=True)
@@ -80,6 +125,17 @@ def _stable_engine_arg_environment(monkeypatch, tmp_path):
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", _resolve_test_worker)
+ class _TestConnectorRunner(OmniConnectorModelRunnerMixin):
+ pass
+
+ class _TestWorker:
+ model_runner_cls = _TestConnectorRunner
+
+ for worker_type in ("ar", "generation"):
+ worker_module = types.ModuleType(f"test.{worker_type}")
+ worker_module.Worker = _TestWorker
+ monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module)
+
def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]:
parent_model = tmp_path / "parent-model"
@@ -139,6 +195,7 @@ def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, st
engine_extras={
"model": str(stage_model),
"attention_backend": "FLASHINFER",
+ "attention_config": {"use_trtllm_attention": False},
"hf_overrides": {"rope_scaling": {"factor": 2}},
"limit_mm_per_prompt": {"audio": 1},
"logits_processors": ["test.custom.LogitsProcessor"],
@@ -196,6 +253,100 @@ def _legacy_and_typed_stages(
return legacy_stages, omni_config
+def test_llm_stage_engine_field_schema_tracks_upstream_engine_args():
+ upstream_engine_fields = frozenset(field.name for field in fields(EngineArgs))
+
+ assert _LLM_STAGE_ENGINE_FIELDS - upstream_engine_fields == _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS
+
+
+def test_typed_llm_projection_discovers_explicit_upstream_config_fields():
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+ load_config=OmniStageLoadConfig(ignore_patterns=["*.bin"]),
+ cache_config=OmniStageCacheConfig(block_size=32, cache_dtype="fp8"),
+ scheduler_config=OmniStageSchedulerConfig(
+ prefill_schedule_interval=2,
+ policy="priority",
+ ),
+ parallel_config=OmniStageParallelConfig(
+ data_parallel_master_ip="10.0.0.1",
+ enable_dbo=True,
+ ),
+ )
+
+ engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+ assert engine_args["ignore_patterns"] == ["*.bin"]
+ assert engine_args["block_size"] == 32
+ assert engine_args["kv_cache_dtype"] == "fp8"
+ assert engine_args["prefill_schedule_interval"] == 2
+ assert engine_args["scheduling_policy"] == "priority"
+ assert engine_args["data_parallel_address"] == "10.0.0.1"
+ assert engine_args["enable_dbo"] is True
+
+
+def test_typed_llm_projection_does_not_emit_inherited_upstream_defaults():
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+ )
+
+ engine_args = stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+ inherited_defaults = {
+ "ignore_patterns",
+ "block_size",
+ "kv_cache_dtype",
+ "prefill_schedule_interval",
+ "scheduling_policy",
+ "data_parallel_address",
+ "data_parallel_rank",
+ "data_parallel_size_local",
+ "data_parallel_rpc_port",
+ "enable_dbo",
+ }
+ assert inherited_defaults.isdisjoint(engine_args)
+
+
+def test_typed_llm_projection_rejects_explicit_fields_owned_by_another_boundary():
+ stage_config = VllmOmniARStageConfig(
+ stage_pipeline_config=StagePipelineConfig(stage_id=0, model_stage="test"),
+ scheduler_config=OmniStageSchedulerConfig(scheduler_cls="test.CustomScheduler"),
+ )
+
+ with pytest.raises(ValueError, match=r"no EngineArgs projection: scheduler_cls"):
+ stage_init_utils._project_omni_stage_engine_args(stage_config)
+
+
+@pytest.mark.parametrize("stage_id", [0, 1], ids=["ar", "generation"])
+def test_typed_llm_engine_args_preserve_upstream_config_objects(tmp_path, stage_id):
+ pipeline, deploy, model = _engine_arg_inputs(tmp_path)
+ compilation_config = VllmCompilationConfig(backend="eager")
+ profiler_config = VllmProfilerConfig(profiler="cuda")
+ stage_prefix = f"stage_{stage_id}_"
+ omni_config = VllmOmniConfig.from_pipeline_config(
+ pipeline,
+ user_deploy_config=copy.deepcopy(deploy),
+ cli_overrides={
+ "model": model,
+ f"{stage_prefix}compilation_config": compilation_config,
+ f"{stage_prefix}profiler_config": profiler_config,
+ },
+ )
+
+ stage_config = omni_config.stage_by_id(stage_id)
+ typed_args = build_engine_args_dict_from_omni_stage_config(stage_config, model)
+
+ assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
+ assert isinstance(typed_args["compilation_config"], VllmCompilationConfig)
+ assert typed_args["compilation_config"] is not stage_config.compilation_config
+ assert typed_args["compilation_config"].backend == "eager"
+
+ assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
+ assert isinstance(typed_args["profiler_config"], VllmProfilerConfig)
+ assert typed_args["profiler_config"] is not stage_config.profiler_config
+ assert typed_args["profiler_config"].profiler == "cuda"
+
+
def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
pipeline, deploy, model = _engine_arg_inputs(tmp_path)
legacy_stages, omni_config = _legacy_and_typed_stages(pipeline, deploy, model)
@@ -218,9 +369,18 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
)
typed_args_by_stage[stage_id] = typed_args
- assert {name: typed_args[name] for name in legacy_args} == legacy_args
+ expected_args = {name: value for name, value in legacy_args.items() if name not in _TOPOLOGY_ONLY_ENGINE_ARGS}
+ assert {name: typed_args[name] for name in expected_args} == expected_args
+ assert _TOPOLOGY_ONLY_ENGINE_ARGS.isdisjoint(typed_args)
thinker_args = typed_args_by_stage[0]
+ inherited_vllm_fields = {
+ field.name
+ for config_cls in (VllmLoadConfig, VllmCacheConfig, VllmSchedulerConfig, VllmParallelConfig)
+ for field in fields(config_cls)
+ }
+ topology_projected_fields = {"scheduler_cls"}
+ assert (inherited_vllm_fields - _LLM_STAGE_ENGINE_FIELDS - topology_projected_fields).isdisjoint(thinker_args)
assert thinker_args["model"] == str(tmp_path / "stage-model" / "ar-model")
assert thinker_args["tokenizer"] == str(tmp_path / "stage-model" / "ar-tokenizer")
assert thinker_args["stage_id"] == 0
@@ -229,6 +389,7 @@ def test_typed_llm_engine_args_preserve_legacy_adapter_behavior(tmp_path):
assert thinker_args["omni_kv_config"] == {"need_send_cache": True}
assert thinker_args["worker_cls"] == "test.custom.Worker"
assert thinker_args["attention_backend"] == "FLASHINFER"
+ assert thinker_args["attention_config"] == {"use_trtllm_attention": False}
assert thinker_args["trust_remote_code"] is True
assert thinker_args["dtype"] == "float16"
assert thinker_args["distributed_executor_backend"] == "mp"
@@ -469,8 +630,24 @@ def test_typed_ming_image_engine_args_defer_diffusion_batch_default():
assert OmniDiffusionConfig(**typed_backend_args).max_num_seqs == 1
+def _fake_model_root(pipeline, tmp_path):
+ """A model directory carrying every subfolder the pipeline's stages declare.
+
+ Stage init fails closed when a declared ``model_subdir``/``tokenizer_subdir``
+ is not a real directory, because the joined path would otherwise reach
+ HuggingFace as a malformed repo id (issue #6638).
+ """
+ root = tmp_path / "model"
+ root.mkdir(exist_ok=True)
+ for stage in pipeline.stages:
+ for subdir in (stage.model_subdir, stage.tokenizer_subdir):
+ if subdir:
+ (root / subdir).mkdir(parents=True, exist_ok=True)
+ return str(root)
+
+
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
-def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
+def test_typed_engine_args_match_current_registry_backend_semantics(model_type, tmp_path):
pipeline = resolve_pipeline_config(model_type)
if pipeline is None:
pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve")
@@ -480,18 +657,19 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
if pipeline.default_deploy_config_name is not None
else DeployConfig()
)
+ model = _fake_model_root(pipeline, tmp_path)
legacy_stages, omni_config = _legacy_and_typed_stages(
pipeline,
deploy,
- model="/tmp",
+ model=model,
)
for legacy_stage in legacy_stages:
stage_id = legacy_stage.stage_id
- legacy_args = build_legacy_engine_args_dict(legacy_stage, model="/tmp")
+ legacy_args = build_legacy_engine_args_dict(legacy_stage, model=model)
typed_args = build_engine_args_dict_from_omni_stage_config(
omni_config.stage_by_id(stage_id),
- model="/tmp",
+ model=model,
)
if legacy_stage.stage_type == StageType.DIFFUSION:
backend_fields = _DIFFUSION_BACKEND_FIELDS
@@ -500,6 +678,7 @@ def test_typed_engine_args_match_current_registry_backend_semantics(model_type):
backend_fields = _LLM_BACKEND_FIELDS
backend_config_cls = OmniEngineArgs
+ backend_fields -= _TOPOLOGY_ONLY_ENGINE_ARGS
missing_fields = (legacy_args.keys() & backend_fields) - typed_args.keys()
assert not missing_fields, f"{model_type} stage {stage_id} lost backend fields: {sorted(missing_fields)}"
diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py
new file mode 100644
index 000000000..e8a228827
--- /dev/null
+++ b/tests/engine/test_stage_pool_collective_rpc.py
@@ -0,0 +1,124 @@
+"""Unit tests for StagePool.collective_rpc EngineCore control dispatch."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from vllm_omni.engine.stage_pool import StagePool
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_pool(*, stage_type: str = "llm", **client_methods: AsyncMock) -> tuple[StagePool, SimpleNamespace]:
+ client = SimpleNamespace(stage_type=stage_type, **client_methods)
+ if "collective_rpc_async" not in client_methods:
+ client.collective_rpc_async = AsyncMock(return_value={"via": "collective"})
+ pool = StagePool(0, [client]) # type: ignore[arg-type]
+ return pool, client
+
+
+@pytest.mark.cpu
+def test_collective_rpc_normalizes_none_args_on_control_helper():
+ async def run() -> None:
+ pause = AsyncMock(return_value="paused")
+ pool, client = _make_pool(pause_scheduler_async=pause)
+
+ result = await pool.collective_rpc(0, "pause_scheduler", args=None, kwargs={"mode": "abort"})
+
+ assert result == "paused"
+ pause.assert_awaited_once_with(mode="abort")
+ client.collective_rpc_async.assert_not_awaited()
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_unrelated_async_helper_uses_collective_path():
+ async def run() -> None:
+ other = AsyncMock(return_value="should-not-run")
+ pool, client = _make_pool(reset_prefix_cache_async=other)
+
+ result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1})
+
+ assert result == {"via": "collective"}
+ other.assert_not_awaited()
+ client.collective_rpc_async.assert_awaited_once_with(
+ method="reset_prefix_cache",
+ timeout=1.5,
+ args=("x",),
+ kwargs={"k": 1},
+ )
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_helper_honors_timeout():
+ async def run() -> None:
+ async def slow_sleep(*_args, **_kwargs):
+ await asyncio.sleep(1.0)
+ return "slept"
+
+ pool, _client = _make_pool(sleep_async=slow_sleep)
+ with pytest.raises(asyncio.TimeoutError):
+ await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,))
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_control_method_reraises_worker_error():
+ async def run() -> None:
+ pool, _client = _make_pool(wake_up_async=AsyncMock(side_effect=RuntimeError("worker died")))
+ with pytest.raises(RuntimeError, match="worker died"):
+ await pool.collective_rpc(0, "wake_up")
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_collective_rpc_non_control_method_still_returns_error_dict():
+ async def run() -> None:
+ pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed")))
+ result = await pool.collective_rpc(0, "reset_prefix_cache")
+ assert result["supported"] is False
+ assert "probe failed" in result["error"]
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
+ async def run() -> None:
+ class RecordingOutputProcessor:
+ def __init__(self) -> None:
+ self.collected = False
+ self.committed = False
+
+ def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True):
+ del internal
+ self.collected = True
+ assert commit_state is False
+ return list(request_ids), [SimpleNamespace(request_id=request_ids[0])]
+
+ def commit_aborted_request_state(self, request_ids, *, internal=False):
+ del request_ids, internal
+ self.committed = True
+
+ abort = AsyncMock(side_effect=RuntimeError("engine abort failed"))
+ output_processor = RecordingOutputProcessor()
+ client = SimpleNamespace(stage_type="llm", abort_requests_async=abort)
+ pool = StagePool(0, [client], output_processor=output_processor) # type: ignore[arg-type]
+ pool._request_bindings["req-1"] = 0
+
+ with pytest.raises(RuntimeError, match="engine abort failed"):
+ await pool.abort_requests(["req-1"])
+ assert output_processor.collected is True
+ assert output_processor.committed is False
+ abort.assert_awaited_once()
+
+ asyncio.run(run())
diff --git a/tests/engine/test_stage_request_token_budget.py b/tests/engine/test_stage_request_token_budget.py
new file mode 100644
index 000000000..045dfd31f
--- /dev/null
+++ b/tests/engine/test_stage_request_token_budget.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Token-budget clamping when the orchestrator builds a stage request."""
+
+from __future__ import annotations
+
+import pytest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StageModelConfig:
+ """Only the field ``build_engine_core_request_from_tokens`` reads."""
+
+ max_model_len = 4096
+
+
+class TestStageRequestTokenBudget:
+ def test_min_tokens_is_clamped_to_the_remaining_context(self):
+ # A min_tokens wider than the leftover window makes ``check_stop``
+ # return early forever: the request can neither stop nor be granted
+ # more tokens, so it stalls instead of being length-capped.
+ params = SamplingParams(max_tokens=4096, min_tokens=50)
+
+ request = build_engine_core_request_from_tokens(
+ request_id="req-clamped",
+ prompt={"prompt_token_ids": [0] * 4050},
+ params=params,
+ model_config=_StageModelConfig(),
+ )
+
+ assert request.sampling_params.min_tokens == 46
+ assert params.min_tokens == 50
+
+ def test_min_tokens_survives_when_the_prompt_leaves_room(self):
+ params = SamplingParams(max_tokens=4096, min_tokens=50)
+
+ request = build_engine_core_request_from_tokens(
+ request_id="req-untouched",
+ prompt={"prompt_token_ids": [0] * 68},
+ params=params,
+ model_config=_StageModelConfig(),
+ )
+
+ assert request.sampling_params.min_tokens == 50
+ assert request.sampling_params.max_tokens == 4096
+
+ def test_min_tokens_floors_at_zero_when_the_prompt_fills_the_context(self):
+ params = SamplingParams(max_tokens=16, min_tokens=8)
+
+ request = build_engine_core_request_from_tokens(
+ request_id="req-full",
+ prompt={"prompt_token_ids": [0] * 4096},
+ params=params,
+ model_config=_StageModelConfig(),
+ )
+
+ assert request.sampling_params.min_tokens == 0
+
+ def test_absent_max_tokens_still_defaults_to_the_remaining_context(self):
+ params = SamplingParams(min_tokens=0)
+ params.max_tokens = None
+
+ request = build_engine_core_request_from_tokens(
+ request_id="req-default",
+ prompt={"prompt_token_ids": [0] * 96},
+ params=params,
+ model_config=_StageModelConfig(),
+ )
+
+ assert request.sampling_params.max_tokens == 4000
diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py
index 3c3666026..9fd61095b 100644
--- a/tests/entrypoints/openai/test_duplex_protocol.py
+++ b/tests/entrypoints/openai/test_duplex_protocol.py
@@ -452,7 +452,7 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
assert caps["supports_core_resumable_request"] is True
assert caps["supports_stage_connector_handoff"] is True
assert caps["supports_audio_truncate"] is True
- assert caps["supports_barge_in"] is False
+ assert caps["supports_barge_in"] is True
assert caps["target_barge_in_latency_ms"] is None
assert caps["supports_multi_session"] is True
assert caps["supports_multi_session_same_replica"] is True
@@ -462,6 +462,15 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease():
assert caps["stage_handoff_transport"] == "scheduler_data_plane"
+def test_duplex_overlap_policy_auto_falls_back_to_listen_only():
+ from vllm_omni.experimental.fullduplex.openai.protocol import (
+ DuplexOverlapPolicy,
+ DuplexSessionConfig,
+ )
+
+ assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value
+
+
def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment():
caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict()
diff --git a/tests/entrypoints/openai/test_profiler_endpoints.py b/tests/entrypoints/openai/test_profiler_endpoints.py
new file mode 100644
index 000000000..02e48e19a
--- /dev/null
+++ b/tests/entrypoints/openai/test_profiler_endpoints.py
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.entrypoints.openai import api_server
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+ ("path", "engine_method"),
+ [
+ ("/start_profile", "start_profile"),
+ ("/stop_profile", "stop_profile"),
+ ],
+)
+def test_profiler_endpoint_forwards_selected_stages(path: str, engine_method: str) -> None:
+ engine_client = SimpleNamespace(
+ start_profile=AsyncMock(),
+ stop_profile=AsyncMock(),
+ )
+
+ app = FastAPI()
+ app.include_router(api_server.profiler_router)
+ app.state.engine_client = engine_client
+
+ response = TestClient(app).post(path, json={"stages": [1]})
+
+ assert response.status_code == 200
+ assert response.json() == {"status": "SUCCESS"}
+ getattr(engine_client, engine_method).assert_awaited_once_with(stages=[1])
diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py
new file mode 100644
index 000000000..2b9ed6ec0
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_api_server_guards.py
@@ -0,0 +1,882 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""API-server surface guards for the OpenAI Omni entrypoint.
+
+Purpose
+-------
+Freeze the *observable* server contract so helper moves (P0.2), route
+extraction (P0.3), and later endpoint-family PRs cannot silently change:
+
+* which routes exist,
+* how upstream routes are overridden at assembly time,
+* unavailable / missing-handler envelopes,
+* shared engine-error JSON fields,
+* app-state keys wired by ``omni_init_app_state``.
+
+CPU-only (fakes / TestClient). No checked-in OpenAPI golden JSON.
+
+Common wrong changes these tests are meant to catch
+----------------------------------------------------
+* Route accidentally removed, renamed, method-changed, or double-registered
+ after a move (census / assembled-app checks).
+* OpenAPI drops an Omni HTTP path while the router still has it (or the reverse).
+* Assembly forgets to strip overridden upstream chat / models / profiler
+ routes, or also deletes unrelated upstream routes.
+* Profiler / Omni router not mounted; storage not started; Omni engine
+ exception handlers not registered.
+* Timestamp middleware stops stamping HTTP, or starts mutating WebSocket scopes.
+* Missing-handler WS/HTTP responses change status, JSON shape, or wording
+ (clients treat these as capability probes).
+* Engine error payloads lose ``request_id`` / ``error_stage_id``.
+* Pure-diffusion vs multi-stage init drops an ``app.state`` key, or leaves a
+ handler as ``None`` while the route is still registered.
+
+Not covered here (add elsewhere if needed)
+------------------------------------------
+* Deep per-endpoint success behavior (family UTs).
+* Byte-identical helper bodies after a move (P0.2 AST integrity).
+* Live GPU / speech byte-identity checks.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.routing import APIRoute, APIWebSocketRoute
+from fastapi.testclient import TestClient
+from starlette.datastructures import State
+from starlette.requests import Request
+from starlette.websockets import WebSocketDisconnect
+from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+
+from vllm_omni.entrypoints.openai import api_server
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# FastAPI/Starlette auto-adds HEAD on every GET route. Including HEAD here
+# would invent ("HEAD", path) entries for every GET and fail the current census.
+# TODO: extend this set and the expected census if @router.head / @router.options appear.
+_HTTP_METHODS = {"GET", "POST", "DELETE", "PUT", "PATCH"}
+
+# Exact Omni router census on main. Update intentionally when routes change.
+_EXPECTED_ROUTER_ROUTES = {
+ ("POST", "/v1/chat/completions"),
+ ("POST", "/v1/chat/completions/batch"),
+ ("POST", "/v1/audio/speech"),
+ ("POST", "/v1/audio/speech/batch"),
+ ("POST", "/v1/audio/generate"),
+ ("GET", "/v1/audio/voices"),
+ ("POST", "/v1/audio/voices"),
+ ("DELETE", "/v1/audio/voices/{name}"),
+ ("WEBSOCKET", "/v1/audio/speech/stream"),
+ ("WEBSOCKET", "/v1/video/chat/stream"),
+ ("WEBSOCKET", "/v1/realtime/video"),
+ ("WEBSOCKET", "/v1/realtime"),
+ ("WEBSOCKET", "/v1/realtime/robot/openpi"),
+ ("WEBSOCKET", "/v1/duplex"),
+ ("GET", "/health"),
+ ("GET", "/v1/models"),
+ ("POST", "/v1/images/generations"),
+ ("POST", "/v1/images/edits"),
+ ("POST", "/v1/videos"),
+ ("POST", "/v1/videos/sync"),
+ ("GET", "/v1/videos"),
+ ("GET", "/v1/videos/{video_id}"),
+ ("DELETE", "/v1/videos/{video_id}"),
+ ("GET", "/v1/videos/{video_id}/content"),
+ ("POST", "/v1/omni/sleep"),
+ ("POST", "/v1/omni/wakeup"),
+}
+_EXPECTED_PROFILER_ROUTES = {
+ ("POST", "/start_profile"),
+ ("POST", "/stop_profile"),
+}
+
+# App-state keys that must remain available across refactor phases.
+# Presence is not enough: "not wired" is written as ``state.X = None``.
+_DIFFUSION_APP_STATE_KEYS = {
+ "engine_client",
+ "log_stats",
+ "args",
+ "sleeping_stages",
+ "stage_configs",
+ "vllm_config",
+ "diffusion_engine",
+ "openai_serving_models",
+ "serving_tokenization",
+ "openai_serving_chat",
+ "openai_serving_chat_batch",
+ "openai_serving_speech",
+ "openai_serving_audio_generate",
+ "openai_serving_video",
+ "openai_streaming_video_output",
+ "openai_serving_duplex",
+ "openai_streaming_speech",
+ "openai_streaming_video",
+ "openai_serving_realtime_robot",
+ "enable_server_load_tracking",
+ "server_load_metrics",
+}
+_DIFFUSION_MUST_BE_NONE = {
+ "vllm_config",
+ "serving_tokenization",
+ "openai_serving_duplex",
+ "openai_streaming_speech",
+ "openai_streaming_video",
+}
+_DIFFUSION_MUST_BE_WIRED = _DIFFUSION_APP_STATE_KEYS - _DIFFUSION_MUST_BE_NONE
+_MULTISTAGE_APP_STATE_KEYS = {
+ "engine_client",
+ "log_stats",
+ "args",
+ "sleeping_stages",
+ "stage_configs",
+ "vllm_config",
+ "openai_serving_models",
+ "serving_tokenization",
+ "online_renderer",
+ "openai_serving_chat",
+ "openai_serving_chat_batch",
+ "openai_serving_speech",
+ "openai_serving_audio_generate",
+ "openai_streaming_speech",
+ "openai_streaming_video",
+ "openai_serving_duplex",
+ "openai_serving_realtime",
+ "openai_serving_video",
+ "openai_serving_realtime_robot",
+ "enable_server_load_tracking",
+ "server_load_metrics",
+}
+_MULTISTAGE_MUST_BE_NONE = {
+ "openai_serving_duplex",
+ "openai_serving_realtime_robot",
+}
+_MULTISTAGE_MUST_BE_WIRED = _MULTISTAGE_APP_STATE_KEYS - _MULTISTAGE_MUST_BE_NONE
+
+
+def _route_entries(routes) -> list[tuple[str, str]]:
+ entries: list[tuple[str, str]] = []
+ for route in routes:
+ if isinstance(route, APIRoute):
+ for method in sorted((route.methods or set()) & _HTTP_METHODS):
+ entries.append((method, route.path))
+ elif isinstance(route, APIWebSocketRoute):
+ entries.append(("WEBSOCKET", route.path))
+ return entries
+
+
+def _route_manifest(routes) -> set[tuple[str, str]]:
+ return set(_route_entries(routes))
+
+
+def _assert_unique_route_entries(routes) -> None:
+ entries = _route_entries(routes)
+ duplicates = sorted({item for item in entries if entries.count(item) > 1})
+ assert not duplicates, f"duplicate route registrations: {duplicates}"
+
+
+def _matching_http_routes(routes, method: str, path: str) -> list[APIRoute]:
+ return [
+ route
+ for route in routes
+ if isinstance(route, APIRoute) and route.path == path and method in (route.methods or set())
+ ]
+
+
+def _single_http_route(routes, method: str, path: str) -> APIRoute:
+ matches = _matching_http_routes(routes, method, path)
+ assert len(matches) == 1, f"expected exactly one {method} {path}, got {len(matches)}"
+ return matches[0]
+
+
+def _assert_app_state_snapshot(
+ state,
+ *,
+ expected_keys: set[str],
+ must_be_wired: set[str],
+ must_be_none: set[str],
+) -> None:
+ assert must_be_wired.isdisjoint(must_be_none)
+ assert must_be_wired | must_be_none <= expected_keys
+ present = {key for key in expected_keys if hasattr(state, key)}
+ assert present == expected_keys
+ not_wired = sorted(key for key in must_be_wired if getattr(state, key) is None)
+ assert not not_wired, f"app.state keys registered but not wired: {not_wired}"
+ unexpectedly_set = sorted(key for key in must_be_none if getattr(state, key) is not None)
+ assert not unexpectedly_set, f"app.state keys expected to stay None: {unexpectedly_set}"
+
+
+def _minimal_args(**overrides) -> SimpleNamespace:
+ args = SimpleNamespace(
+ tool_parser_plugin=None,
+ reasoning_parser_plugin=None,
+ reasoning_parser=None,
+ structured_outputs_config=SimpleNamespace(reasoning_parser=None),
+ enable_ssl_refresh=False,
+ host="127.0.0.1",
+ port=8000,
+ uvicorn_log_level="info",
+ disable_uvicorn_access_log=True,
+ ssl_keyfile=None,
+ ssl_certfile=None,
+ ssl_ca_certs=None,
+ ssl_cert_reqs=0,
+ ssl_ciphers=None,
+ h11_max_incomplete_event_size=None,
+ h11_max_header_count=None,
+ log_config_file=None,
+ disable_access_log_for_endpoints=None,
+ disable_log_stats=True,
+ enable_log_requests=False,
+ max_log_len=None,
+ served_model_name=None,
+ model="demo-model",
+ enable_server_load_tracking=False,
+ log_error_stack=False,
+ chat_template=None,
+ chat_template_content_format="auto",
+ trust_request_chat_template=False,
+ tool_server=None,
+ lora_modules=None,
+ response_role="assistant",
+ default_chat_template_kwargs=None,
+ return_tokens_as_token_ids=False,
+ enable_auto_tool_choice=False,
+ exclude_tools_when_tool_choice_none=False,
+ tool_call_parser=None,
+ enable_prompt_tokens_details=False,
+ enable_force_include_usage=False,
+ enable_log_outputs=False,
+ enable_log_deltas=False,
+ tokens_only=False,
+ shutdown_timeout=0,
+ stage_configs_path=None,
+ deploy_config=None,
+ )
+ for key, value in overrides.items():
+ setattr(args, key, value)
+ return args
+
+
+def _request_for(app: FastAPI, *, method: str = "GET", path: str = "/") -> Request:
+ return Request(
+ {
+ "type": "http",
+ "method": method,
+ "path": path,
+ "headers": [],
+ "app": app,
+ "state": {},
+ }
+ )
+
+
+class _FakeSocket:
+ def __init__(self) -> None:
+ self.closed = False
+
+ def close(self) -> None:
+ self.closed = True
+
+
+class _FakeEngineClient:
+ def __init__(self, *, stage_configs=None, endpoint_restrictions=None, vllm_config=None) -> None:
+ self.stage_configs = stage_configs if stage_configs is not None else []
+ self.endpoint_restrictions = endpoint_restrictions if endpoint_restrictions is not None else {}
+ self.model_config = SimpleNamespace()
+ self.input_processor = object()
+ self.renderer = object()
+ self.errored = False
+ self.vllm_config = vllm_config
+
+ async def get_vllm_config(self):
+ return self.vllm_config
+
+ async def get_supported_tasks(self) -> tuple[str, ...]:
+ return ("generate",)
+
+ async def get_tokenizer(self):
+ return SimpleNamespace(chat_template="dummy")
+
+ async def check_health(self) -> None:
+ return None
+
+
+def _marker(name: str):
+ return SimpleNamespace(name=name)
+
+
+def test_router_route_manifest_is_explicit_and_complete() -> None:
+ """Lock the Omni + profiler router census.
+
+ Fails if a route is added/removed/renamed, a method changes, or a
+ WebSocket path disappears during helper/router moves. Also fails if the
+ same public route is registered more than once, which a set manifest alone
+ would collapse.
+ """
+ # TODO(P0.3): retarget this object-level census when routes move onto family
+ # routers. The assembled-app census is the client-facing lock that must survive.
+ _assert_unique_route_entries(api_server.router.routes)
+ _assert_unique_route_entries(api_server.profiler_router.routes)
+ assert _route_manifest(api_server.router.routes) == _EXPECTED_ROUTER_ROUTES
+ assert _route_manifest(api_server.profiler_router.routes) == _EXPECTED_PROFILER_ROUTES
+
+
+def test_router_openapi_paths_cover_http_manifest() -> None:
+ """Lock that OpenAPI advertises every HTTP census route/method.
+
+ Fails if a decorator still registers a route but OpenAPI loses the path
+ (or method), e.g. after moving route bodies / response models.
+ """
+ app = FastAPI()
+ app.include_router(api_server.router)
+ app.include_router(api_server.profiler_router)
+ paths = app.openapi().get("paths", {})
+
+ http_routes = {
+ (method, path) for method, path in _EXPECTED_ROUTER_ROUTES | _EXPECTED_PROFILER_ROUTES if method != "WEBSOCKET"
+ }
+ for method, path in sorted(http_routes):
+ assert path in paths, f"missing OpenAPI path {path}"
+ assert method.lower() in paths[path], f"missing OpenAPI method {method} {path}"
+
+
+@pytest.mark.asyncio
+async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_router(monkeypatch) -> None:
+ """Lock worker assembly: override, mount, storage, handlers, full census.
+
+ Fails if assembly stops replacing upstream chat/batch/models/profiler,
+ deletes unrelated upstream routes, forgets Omni/profiler mounts, skips
+ storage start, or drops Omni ``EngineDeadError`` / ``EngineGenerateError``
+ handlers.
+ """
+ captured: dict[str, object] = {}
+
+ def fake_build_openai_app(args, supported_tasks):
+ app = FastAPI()
+
+ @app.post("/v1/chat/completions")
+ async def upstream_chat():
+ return {"owner": "upstream"}
+
+ @app.post("/v1/chat/completions/batch")
+ async def upstream_batch():
+ return {"owner": "upstream"}
+
+ @app.get("/v1/models")
+ async def upstream_models():
+ return {"owner": "upstream"}
+
+ @app.post("/start_profile")
+ async def upstream_start_profile():
+ return {"owner": "upstream"}
+
+ @app.post("/stop_profile")
+ async def upstream_stop_profile():
+ return {"owner": "upstream"}
+
+ @app.get("/v1/upstream-only")
+ async def upstream_only():
+ return {"owner": "upstream"}
+
+ captured["supported_tasks"] = supported_tasks
+ return app
+
+ @asynccontextmanager
+ async def fake_build_async_omni(*_args, **_kwargs):
+ yield _FakeEngineClient(
+ stage_configs=[{"engine_args": {"profiler_config": {"profiler": "torch"}}}],
+ vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(_api_process_rank=0)),
+ )
+
+ async def fake_omni_init_app_state(engine_client, state, args):
+ state.engine_client = engine_client
+ state.initialized_by_omni = True
+ state.openai_serving_video = SimpleNamespace(
+ shutdown=lambda: captured.__setitem__("video_shutdown", True),
+ )
+ state.openai_serving_speech = SimpleNamespace(
+ shutdown=lambda: captured.__setitem__("speech_shutdown", True),
+ )
+
+ async def fake_storage_start():
+ captured["storage_started"] = True
+
+ async def fake_serve_http(app, **_kwargs):
+ captured["served_app"] = app
+ task = asyncio.create_task(asyncio.sleep(0))
+ await asyncio.sleep(0)
+ return task
+
+ monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+ monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+ monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+ monkeypatch.setattr(
+ api_server,
+ "shutdown_unsupported_routes",
+ lambda app, restrictions: captured.setdefault("restrictions", restrictions),
+ )
+ monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=fake_storage_start))
+ monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+ sock = _FakeSocket()
+ await api_server.omni_run_server_worker("127.0.0.1:8000", sock, _minimal_args())
+
+ served_wrapper = captured["served_app"]
+ # Timestamp middleware wraps the FastAPI app.
+ assert hasattr(served_wrapper, "_inner")
+ served_app = served_wrapper._inner
+ routes = served_app.routes
+ _assert_unique_route_entries(routes)
+ manifest = _route_manifest(routes)
+
+ assert captured["supported_tasks"] == ("generate",)
+ assert captured["storage_started"] is True
+ assert captured["restrictions"] == {}
+ assert captured["video_shutdown"] is True
+ assert captured["speech_shutdown"] is True
+ assert sock.closed is True
+ assert served_app.state.initialized_by_omni is True
+
+ # Upstream override: unrelated route kept; Omni owns overridden pairs once.
+ # Endpoint identity ensures the surviving route is the Omni handler, not
+ # the fake upstream handler with the same method/path.
+ assert ("GET", "/v1/upstream-only") in manifest
+ assert _single_http_route(routes, "POST", "/v1/chat/completions").endpoint is api_server.create_chat_completion
+ assert (
+ _single_http_route(routes, "POST", "/v1/chat/completions/batch").endpoint
+ is api_server.create_batch_chat_completion
+ )
+ assert _single_http_route(routes, "GET", "/v1/models").endpoint is api_server.show_available_models
+ assert _single_http_route(routes, "POST", "/start_profile").endpoint is api_server.start_profile
+ assert _single_http_route(routes, "POST", "/stop_profile").endpoint is api_server.stop_profile
+
+ # Assembled app includes Omni router HTTP/WS surface + profiler.
+ for item in _EXPECTED_ROUTER_ROUTES:
+ assert item in manifest, f"assembled app missing {item}"
+ assert ("POST", "/start_profile") in manifest
+ assert ("POST", "/stop_profile") in manifest
+
+ assert served_app.exception_handlers[EngineGenerateError]
+ assert served_app.exception_handlers[EngineDeadError]
+
+ # Light OpenAPI check on the assembled FastAPI app (not the ASGI wrapper).
+ openapi_paths = served_app.openapi().get("paths", {})
+ assert "/v1/chat/completions" in openapi_paths
+ assert "/v1/models" in openapi_paths
+ assert "/health" in openapi_paths
+
+
+@pytest.mark.asyncio
+async def test_timestamp_middleware_stamps_http_and_passes_websocket(monkeypatch) -> None:
+ """Lock outermost timestamp middleware behavior.
+
+ Fails if HTTP requests lose ``request_timestamp``, or WebSocket scopes
+ start getting stamped / mutated incorrectly.
+ """
+ captured: dict[str, object] = {}
+
+ def fake_build_openai_app(args, supported_tasks):
+ return FastAPI()
+
+ @asynccontextmanager
+ async def fake_build_async_omni(*_args, **_kwargs):
+ yield _FakeEngineClient()
+
+ async def fake_omni_init_app_state(engine_client, state, args):
+ state.engine_client = engine_client
+
+ async def fake_serve_http(app, **_kwargs):
+ captured["served_app"] = app
+ task = asyncio.create_task(asyncio.sleep(0))
+ await asyncio.sleep(0)
+ return task
+
+ monkeypatch.setattr(api_server, "build_openai_app", fake_build_openai_app)
+ monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni)
+ monkeypatch.setattr(api_server, "omni_init_app_state", fake_omni_init_app_state)
+ monkeypatch.setattr(api_server, "shutdown_unsupported_routes", lambda *_a, **_k: None)
+ monkeypatch.setattr(api_server, "STORAGE_MANAGER", SimpleNamespace(start=lambda: asyncio.sleep(0)))
+ monkeypatch.setattr(api_server, "serve_http", fake_serve_http)
+
+ await api_server.omni_run_server_worker("127.0.0.1:8000", _FakeSocket(), _minimal_args())
+ middleware = captured["served_app"]
+
+ http_scope = {"type": "http", "state": {}}
+ ws_scope = {"type": "websocket"}
+
+ async def _receive():
+ return {"type": "http.disconnect"}
+
+ async def _send(_message):
+ return None
+
+ called = {"http": False, "websocket": False}
+
+ async def inner(scope, receive, send):
+ called[scope["type"]] = True
+
+ middleware._inner = inner
+ await middleware(http_scope, _receive, _send)
+ await middleware(ws_scope, _receive, _send)
+
+ assert called["http"] is True
+ assert called["websocket"] is True
+ assert "request_timestamp" in http_scope["state"]
+ assert isinstance(http_scope["state"]["request_timestamp"], float)
+ assert "request_timestamp" not in ws_scope.get("state", {})
+
+
+@pytest.mark.parametrize(
+ ("path", "payload"),
+ [
+ ("/v1/audio/speech/stream", {"type": "error", "message": "Streaming speech is not available"}),
+ ("/v1/video/chat/stream", {"type": "error", "message": "Streaming video chat is not available"}),
+ ("/v1/realtime/video", {"type": "error", "message": "Streaming video generation is not available"}),
+ ("/v1/realtime", {"type": "error", "error": "Realtime API is not available", "code": "unsupported"}),
+ (
+ "/v1/realtime/robot/openpi",
+ {"type": "error", "error": "Robot policy not available", "code": "unsupported"},
+ ),
+ ("/v1/duplex", {"type": "error", "error": "Duplex API is not available", "code": "unsupported"}),
+ ],
+)
+def test_websocket_routes_emit_stable_unavailable_frames_and_close(path: str, payload: dict[str, str]) -> None:
+ """Lock missing-handler WebSocket capability probes.
+
+ Fails if unavailable JSON changes (keys/wording/code) or the socket stays
+ open after the error frame — common when WS wrappers move between modules.
+ """
+ app = FastAPI()
+ app.include_router(api_server.router)
+
+ with TestClient(app) as client:
+ with client.websocket_connect(path) as websocket:
+ assert websocket.receive_json() == payload
+ with pytest.raises(WebSocketDisconnect):
+ websocket.receive_text()
+
+
+def test_health_without_engine_returns_stable_unhealthy_response() -> None:
+ """Lock ``/health`` with no engine initialized.
+
+ Fails if status drifts from 503 or the unhealthy payload shape/reason
+ string changes during health-route extraction.
+ """
+ app = FastAPI()
+ response = asyncio.run(api_server.health(_request_for(app, path="/health")))
+
+ assert response.status_code == 503
+ assert json.loads(response.body) == {
+ "status": "unhealthy",
+ "reason": "No engine initialized",
+ }
+
+
+def test_models_without_handler_returns_empty_openai_list() -> None:
+ """Lock ``/v1/models`` when ``openai_serving_models`` is unset.
+
+ Fails if the empty listing becomes an error, or the OpenAI list envelope
+ changes (``object`` / ``data``).
+ """
+ app = FastAPI()
+ response = asyncio.run(api_server.show_available_models(_request_for(app, path="/v1/models")))
+
+ assert response.status_code == 200
+ assert json.loads(response.body) == {"object": "list", "data": []}
+
+
+def test_images_generation_without_engine_preserves_service_unavailable_error() -> None:
+ """Lock images generation when the multi-stage engine is missing.
+
+ Fails if the 503 detail string changes — clients/tests use it as a
+ capability signal when bootstrap failed.
+ """
+ app = FastAPI()
+ raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+ request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+
+ with pytest.raises(HTTPException) as exc_info:
+ asyncio.run(api_server.generate_images(request, raw_request))
+
+ assert exc_info.value.status_code == 503
+ assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model."
+
+
+def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None:
+ """Lock images generation when chat serving was not wired.
+
+ Fails if multi-stage images stop requiring ``openai_serving_chat``, or the
+ 503 detail changes after images/chat ownership splits.
+ """
+ app = FastAPI()
+ app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()])
+ app.state.stage_configs = app.state.engine_client.stage_configs
+ app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")])
+ app.state.openai_serving_chat = None
+ app.state.args = SimpleNamespace(max_generated_image_size=None)
+
+ raw_request = _request_for(app, method="POST", path="/v1/images/generations")
+ request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model")
+ monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion")
+
+ with pytest.raises(HTTPException) as exc_info:
+ asyncio.run(api_server.generate_images(request, raw_request))
+
+ assert exc_info.value.status_code == 503
+ assert exc_info.value.detail == "openai_serving_chat is not initialized for multi-stage image generation."
+
+
+def test_speech_without_handler_preserves_not_found_http_error() -> None:
+ """Lock speech HTTP when ``openai_serving_speech`` is unset.
+
+ Fails if missing speech becomes a different status/detail after TTS route
+ wrappers move out of ``api_server``.
+ """
+ app = FastAPI()
+ app.state.openai_serving_speech = None
+ raw_request = _request_for(app, method="POST", path="/v1/audio/speech")
+ request = api_server.OpenAICreateSpeechRequest(input="hello", model="demo-model", voice="alloy")
+
+ with pytest.raises(HTTPException) as exc_info:
+ asyncio.run(api_server.create_speech(request, raw_request))
+
+ assert exc_info.value.status_code == 404
+ assert exc_info.value.detail == "The model does not support Speech API"
+
+
+def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatch) -> None:
+ """Lock shared ``EngineGenerateError`` JSON fields used by Omni handlers.
+
+ Fails if ``request_id`` / ``error_stage_id`` are dropped when error helpers
+ move (P0.2) or exception-handler wiring changes.
+
+ Exercises the envelope through the registered handler, not the private
+ ``_create_engine_error_json_response`` helper, so relocating that helper
+ in P0.2 is not a false red.
+ """
+ app = FastAPI()
+ # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+ # leaves api_server (red here is a rename, not a surface regression).
+ api_server._register_omni_exception_handlers(app)
+ app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+ app.state.server = object()
+ app.state.args = SimpleNamespace(log_error_stack=False)
+
+ req = _request_for(app, method="POST", path="/v1/chat/completions")
+ req.state.request_metadata = SimpleNamespace(request_id="req-123")
+
+ monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+ exc = EngineGenerateError("boom")
+ exc.error_stage_id = 2 # type: ignore[attr-defined]
+ handler = app.exception_handlers[EngineGenerateError]
+ response = asyncio.run(handler(req, exc))
+ payload = json.loads(response.body)
+
+ assert "error" in payload
+ assert payload["error"]["request_id"] == "req-123"
+ assert payload["error"]["error_stage_id"] == 2
+ assert response.status_code >= 400
+
+
+def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None:
+ """Lock Omni ``EngineDeadError`` handler registration + JSON envelope.
+
+ Fails if registration is skipped or the handler stops returning OpenAI-style
+ JSON with ``request_id``.
+ """
+ app = FastAPI()
+ # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers``
+ # leaves api_server (red here is a rename, not a surface regression).
+ api_server._register_omni_exception_handlers(app)
+ app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False))
+ app.state.server = object()
+ app.state.args = SimpleNamespace(log_error_stack=False)
+
+ monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None)
+
+ handler = app.exception_handlers[EngineDeadError]
+ req = _request_for(app)
+ req.state.request_metadata = SimpleNamespace(request_id="dead-1")
+ response = asyncio.run(handler(req, EngineDeadError()))
+ payload = json.loads(response.body)
+ assert payload["error"]["request_id"] == "dead-1"
+ assert response.status_code >= 400
+
+
+@pytest.mark.asyncio
+async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None:
+ """Lock pure-diffusion ``app.state`` keys after init, including live vs None.
+
+ Fails if bootstrap drops keys route owners still read (e.g. video/speech/
+ models), wires a handler that must stay None, or leaves a required
+ handler as None.
+ """
+ stage = SimpleNamespace(engine_args={})
+ engine = _FakeEngineClient(stage_configs=[stage])
+
+ def _for_diffusion_factory(label: str):
+ @classmethod
+ def _factory(cls, *args, **kwargs):
+ return _marker(label)
+
+ return _factory
+
+ monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
+ monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
+ monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
+ monkeypatch.setattr(
+ api_server.OmniOpenAIServingAudioGenerate,
+ "for_diffusion",
+ _for_diffusion_factory("audio_generate"),
+ )
+ monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
+ monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
+ monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _for_diffusion_factory("speech"))
+ monkeypatch.setattr(
+ api_server.ServingRealtimeRobotOpenPI,
+ "create_policy_server",
+ classmethod(lambda cls, *a, **k: _marker("openpi")),
+ )
+
+ state = State()
+ await api_server.omni_init_app_state(
+ engine,
+ state,
+ _minimal_args(),
+ )
+
+ _assert_app_state_snapshot(
+ state,
+ expected_keys=_DIFFUSION_APP_STATE_KEYS,
+ must_be_wired=_DIFFUSION_MUST_BE_WIRED,
+ must_be_none=_DIFFUSION_MUST_BE_NONE,
+ )
+ assert state.diffusion_engine is engine
+
+
+@pytest.mark.asyncio
+async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None:
+ stage = SimpleNamespace(engine_args={})
+ engine = _FakeEngineClient(stage_configs=[stage])
+ speech_kwargs = {}
+
+ def _for_diffusion_factory(label: str):
+ @classmethod
+ def _factory(cls, *args, **kwargs):
+ return _marker(label)
+
+ return _factory
+
+ @classmethod
+ def _speech_factory(cls, *args, **kwargs):
+ speech_kwargs.update(kwargs)
+ return _marker("speech")
+
+ monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion")
+ monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat"))
+ monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch"))
+ monkeypatch.setattr(
+ api_server.OmniOpenAIServingAudioGenerate,
+ "for_diffusion",
+ _for_diffusion_factory("audio_generate"),
+ )
+ monkeypatch.setattr(api_server.OmniOpenAIServingVideo, "for_diffusion", _for_diffusion_factory("video"))
+ monkeypatch.setattr(api_server.OmniStreamingVideoOutputHandler, "__init__", lambda self, *a, **k: None)
+ monkeypatch.setattr(api_server.OmniOpenAIServingSpeech, "for_diffusion", _speech_factory)
+ monkeypatch.setattr(
+ api_server.ServingRealtimeRobotOpenPI,
+ "create_policy_server",
+ classmethod(lambda cls, *a, **k: _marker("openpi")),
+ )
+
+ state = State()
+ await api_server.omni_init_app_state(
+ engine,
+ state,
+ _minimal_args(
+ allowed_local_media_path="/allowed/media",
+ allowed_media_domains=["media.example.com"],
+ ),
+ )
+
+ assert speech_kwargs["allowed_local_media_path"] == "/allowed/media"
+ assert speech_kwargs["allowed_media_domains"] == ["media.example.com"]
+
+
+@pytest.mark.asyncio
+async def test_multistage_app_state_key_snapshot(monkeypatch) -> None:
+ """Lock multi-stage ``app.state`` keys after init, including live vs None.
+
+ Fails if chat/speech/video/realtime/tokenization keys disappear or are
+ left as None after init refactors — classic “route still registered,
+ handler never wired” failure mode.
+ """
+ engine = _FakeEngineClient(
+ stage_configs=[object(), object()],
+ vllm_config=SimpleNamespace(
+ lora_config=None,
+ model_config=SimpleNamespace(),
+ parallel_config=SimpleNamespace(_api_process_rank=0),
+ ),
+ )
+
+ class _FakeModels:
+ def __init__(self, *args, **kwargs):
+ self.base_model_paths = kwargs.get("base_model_paths") or []
+
+ async def init_static_loras(self):
+ return None
+
+ class _FakeCtor:
+ def __init__(self, *args, **kwargs):
+ pass
+
+ def warmup(self):
+ return None
+
+ class _FakeSpeech(_FakeCtor):
+ async def warmup(self):
+ return None
+
+ monkeypatch.setattr(api_server, "load_chat_template", lambda *_a, **_k: None)
+ monkeypatch.setattr(api_server, "process_lora_modules", lambda modules, _defaults: modules or [])
+ monkeypatch.setattr(api_server, "OpenAIServingModels", _FakeModels)
+ monkeypatch.setattr(api_server, "OnlineRenderer", _FakeCtor)
+ monkeypatch.setattr(api_server, "OpenAIServingResponses", _FakeCtor)
+ monkeypatch.setattr(api_server, "OmniOpenAIServingChat", _FakeCtor)
+ monkeypatch.setattr(api_server, "OmniOpenAIServingChatBatch", _FakeCtor)
+ monkeypatch.setattr(api_server, "OpenAIServingCompletion", _FakeCtor)
+ monkeypatch.setattr(api_server, "ServingPooling", _FakeCtor)
+ monkeypatch.setattr(api_server, "OpenAIServingEmbedding", _FakeCtor)
+ monkeypatch.setattr(api_server, "ServingClassification", _FakeCtor)
+ monkeypatch.setattr(api_server, "ServingScores", _FakeCtor)
+ monkeypatch.setattr(api_server, "ServingTokenization", _FakeCtor)
+ monkeypatch.setattr(api_server, "OpenAIServingTranscription", _FakeCtor)
+ monkeypatch.setattr(api_server, "OpenAIServingTranslation", _FakeCtor)
+ monkeypatch.setattr(api_server, "AnthropicServingMessages", _FakeCtor)
+ monkeypatch.setattr(api_server, "ServingTokens", _FakeCtor)
+ monkeypatch.setattr(api_server, "OmniOpenAIServingSpeech", _FakeSpeech)
+ monkeypatch.setattr(api_server, "OmniOpenAIServingAudioGenerate", _FakeCtor)
+ monkeypatch.setattr(api_server, "OmniStreamingSpeechHandler", _FakeCtor)
+ monkeypatch.setattr(api_server, "create_streaming_video_handler", lambda **_k: _marker("streaming_video"))
+ monkeypatch.setattr(api_server, "OpenAIServingRealtime", _FakeCtor)
+ monkeypatch.setattr(api_server, "OmniOpenAIServingVideo", _FakeCtor)
+ monkeypatch.setattr(api_server, "should_enable_duplex_endpoint", lambda *_a, **_k: False)
+
+ state = State()
+ await api_server.omni_init_app_state(engine, state, _minimal_args())
+
+ _assert_app_state_snapshot(
+ state,
+ expected_keys=_MULTISTAGE_APP_STATE_KEYS,
+ must_be_wired=_MULTISTAGE_MUST_BE_WIRED,
+ must_be_none=_MULTISTAGE_MUST_BE_NONE,
+ )
diff --git a/tests/entrypoints/openai_api/test_audex_serving_guards.py b/tests/entrypoints/openai_api/test_audex_serving_guards.py
index 68661865f..27ed11356 100644
--- a/tests/entrypoints/openai_api/test_audex_serving_guards.py
+++ b/tests/entrypoints/openai_api/test_audex_serving_guards.py
@@ -26,8 +26,13 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
+ from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter
+ from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext
+
serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
serving._tts_model_type = model_type
+ adapter_cls = resolve_adapter(model_type)
+ serving._adapter = adapter_cls(SpeechServingContext(server=serving)) if adapter_cls is not None else None
serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
serving._discard_ref_audio_artifact_warmup = lambda request_id: None
return serving
@@ -154,7 +159,7 @@ class TestAudexCfgValidation:
assert _adapter().validate(_speech_request(cfg_scale=None)) is None
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
params = SimpleNamespace(extra_args={"cfg_scale": None})
- out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
assert out.extra_args == {}
@pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
@@ -177,21 +182,21 @@ class TestAudexCfgInjection:
def test_no_extra_args_is_untouched(self):
serving = self._serving_with_fake_tokenizer()
params = SimpleNamespace(extra_args=None)
- out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
assert out.extra_args is None
assert params.extra_args is None
def test_scale_one_leaves_params_byte_identical(self):
serving = self._serving_with_fake_tokenizer()
params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
- out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
assert out.extra_args == {}
def test_guided_request_gains_pair_contract(self):
serving = self._serving_with_fake_tokenizer()
params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
cond = self._cond_prompt("Some transcription text here.")
- out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
extra = out.extra_args
assert extra["cfg_scale"] == 1.5
@@ -207,14 +212,14 @@ class TestAudexCfgInjection:
serving = self._serving_with_fake_tokenizer()
params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
with pytest.raises(ValueError, match="adapter-built text prompt"):
- serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
+ serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
def test_injection_never_mutates_input_params(self):
"""The input may be the engine's SHARED default params (review P2):
pair state must land only on the returned clone."""
serving = self._serving_with_fake_tokenizer()
shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
- out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
assert out is not shared
assert shared.extra_args == {"cfg_scale": 1.5}
assert shared.temperature == 0.1
@@ -252,11 +257,11 @@ def test_guided_injection_sets_guided_temperature():
serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
cond = TestAudexCfgInjection()._cond_prompt()
params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
- out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
assert out.temperature == 0.05
unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
- out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
+ out = serving._adapter._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
assert out.temperature == 0.1
@@ -340,7 +345,7 @@ class TestAudexTTAInjection:
def test_default_injection_applies_rvq_and_cfg3(self):
serving = self._serving()
params = SimpleNamespace(extra_args=None)
- out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
+ out = serving._adapter._inject_audex_tta_args("req-9", self._prompt(), params)
extra = out.extra_args
rvq = extra["tta_rvq"]
@@ -355,15 +360,15 @@ class TestAudexTTAInjection:
def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
serving = self._serving()
params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
- out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
+ out = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), params)
assert "tta_rvq" in out.extra_args
assert "cfg_scale" not in out.extra_args
assert "cfg_role" not in out.extra_args
def test_rvq_contract_cached_per_process(self):
serving = self._serving()
- out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
- out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
+ out_a = serving._adapter._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
+ out_b = serving._adapter._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
def test_injection_never_mutates_input_params(self):
@@ -372,8 +377,8 @@ class TestAudexTTAInjection:
per-request pair state must land only on the returned clone."""
serving = self._serving()
shared = SimpleNamespace(extra_args=None)
- out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
- out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
+ out_1 = serving._adapter._inject_audex_tta_args("req-1", self._prompt(), shared)
+ out_2 = serving._adapter._inject_audex_tta_args("req-2", self._prompt(), shared)
assert shared.extra_args is None
assert out_1 is not shared and out_2 is not shared
assert out_1.extra_args["cfg_pair_id"] == "req-1"
diff --git a/tests/entrypoints/openai_api/test_decode_audio_url.py b/tests/entrypoints/openai_api/test_decode_audio_url.py
index 2a5c5be47..5dff3b052 100644
--- a/tests/entrypoints/openai_api/test_decode_audio_url.py
+++ b/tests/entrypoints/openai_api/test_decode_audio_url.py
@@ -5,6 +5,7 @@
import asyncio
import base64
import os
+import tempfile
import httpx
import pytest
@@ -101,7 +102,7 @@ class TestDecodeAudioUrlSuffixSanitization:
try:
# ".." contains non-alnum dots, so falls back to .wav
assert path.endswith(".wav")
- assert os.path.dirname(path) == "/tmp"
+ assert os.path.dirname(path) == tempfile.gettempdir()
finally:
_cleanup(path)
diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py
index a37d9d2b1..c75f1e599 100644
--- a/tests/entrypoints/openai_api/test_duplex_handler.py
+++ b/tests/entrypoints/openai_api/test_duplex_handler.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import asyncio
@@ -30,6 +33,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.runtime import (
from vllm_omni.experimental.fullduplex.minicpmo45.session import (
MiniCPMO45ServingSessionState,
)
+from vllm_omni.experimental.fullduplex.openai import vad as realtime_vad
from vllm_omni.experimental.fullduplex.openai.protocol import (
DuplexCapabilities,
DuplexOverlapPolicy,
@@ -46,6 +50,9 @@ from vllm_omni.experimental.fullduplex.openai.serving import (
)
from vllm_omni.experimental.fullduplex.openai.websocket import DuplexWebSocketActor
from vllm_omni.experimental.fullduplex.output import attach_duplex_output_decision
+from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import (
+ PersonaPlexServingRuntimeAdapter,
+)
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -105,16 +112,16 @@ class FakeEngineClient:
self.close_result = close_result
self.opened: list[str] = []
self.opened_fences: list[DuplexFence | None] = []
- self.appended: list[tuple[str, str, object, bool]] = []
+ self.appended: list[tuple[str, str, Any, bool]] = []
self.append_operation_ids: list[str | None] = []
self.appended_fences: list[DuplexFence | None] = []
- self.opened_configs: list[dict[str, object]] = []
- self.opened_runtime_configs: list[dict[str, object]] = []
+ self.opened_configs: list[dict[str, Any]] = []
+ self.opened_runtime_configs: list[dict[str, Any]] = []
self.signals: list[tuple[str, str]] = []
self.signal_fences: list[DuplexFence | None] = []
self.signal_next_fences: list[DuplexFence | None] = []
- self.signal_session_configs: list[dict[str, object] | None] = []
- self.signal_runtime_configs: list[dict[str, object] | None] = []
+ self.signal_session_configs: list[dict[str, Any] | None] = []
+ self.signal_runtime_configs: list[dict[str, Any] | None] = []
self.closed: list[tuple[str, str]] = []
self.touched: list[tuple[str, DuplexLeaseActivity]] = []
self.touch_fences: list[DuplexFence] = []
@@ -136,7 +143,7 @@ class FakeEngineClient:
runtime_config: dict[str, object] | None = None,
timeout: float | None = None,
fence: DuplexFence | None = None,
- ) -> None:
+ ) -> Any:
del timeout
if self.fail_open:
raise RuntimeError("open failed")
@@ -157,7 +164,7 @@ class FakeEngineClient:
timeout: float | None = None,
collect_outputs: bool = True,
fence: DuplexFence | None = None,
- ) -> None:
+ ) -> Any:
del timeout, collect_outputs
self.appended.append((session_id, mode, payload, final))
self.append_operation_ids.append(operation_id)
@@ -191,7 +198,7 @@ class FakeEngineClient:
next_fence: DuplexFence | None = None,
session_config: dict[str, object] | None = None,
runtime_config: dict[str, object] | None = None,
- ) -> None:
+ ) -> Any:
del timeout
if self.fail_signal or event in self.fail_signal_events:
raise RuntimeError("signal failed")
@@ -209,7 +216,7 @@ class FakeEngineClient:
reason: str = "client_close",
timeout: float | None = None,
fence: DuplexFence | None = None,
- ) -> None:
+ ) -> Any:
del timeout
if self.fail_close:
raise RuntimeError("close failed")
@@ -344,7 +351,7 @@ def _project_data_plane(
result: object,
*,
session: DuplexSession | None = None,
-) -> list[dict[str, object]]:
+) -> list[dict[str, Any]]:
return list(data_plane.project(result, context=_data_plane_context(session)))
@@ -545,6 +552,54 @@ async def test_native_realtime_protocol_conversation_item_create_commits_user_te
}
+@pytest.mark.asyncio
+async def test_function_output_is_recorded_only_after_runtime_delivery():
+ ws = TimedWebSocket()
+ protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
+ protocol.bind_sender(ws.send_json)
+ protocol.encode_outbound_event(
+ {
+ "type": "conversation.item.created",
+ "item": {
+ "id": "item-call",
+ "type": "function_call",
+ "call_id": "call-1",
+ "name": "weather",
+ "arguments": "{}",
+ "status": "completed",
+ },
+ }
+ )
+ event = {
+ "type": "conversation.item.create",
+ "item": {
+ "id": "item-output",
+ "type": "function_call_output",
+ "call_id": "call-1",
+ "output": '{"temperature":20}',
+ },
+ }
+
+ first = await protocol._to_duplex_event(event)
+ retry = await protocol._to_duplex_event(event)
+
+ assert first is not None
+ assert retry is not None
+ assert "item-output" not in protocol._conversation_items
+
+ delivered_item = first["payload"]["item"]
+ protocol.encode_outbound_event(
+ {
+ "type": "conversation.item.created",
+ "item": delivered_item,
+ }
+ )
+
+ assert protocol._conversation_items["item-output"] == delivered_item
+ assert await protocol._to_duplex_event(event) is None
+ assert ws.sent[-1]["error"]["code"] == "invalid_function_call_output"
+
+
@pytest.mark.asyncio
async def test_native_realtime_protocol_audio_commit_requires_non_empty_buffer():
ws = TimedWebSocket()
@@ -652,7 +707,7 @@ async def test_native_realtime_protocol_preserves_input_turn_policy_hints():
@pytest.mark.asyncio
-async def test_native_realtime_protocol_rejects_unimplemented_server_vad():
+async def test_native_realtime_protocol_rejects_server_vad_without_interrupt():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
@@ -675,7 +730,11 @@ async def test_native_realtime_protocol_rejects_unimplemented_server_vad():
error = ws.sent[-1]
assert error["type"] == "error"
assert error["error"]["code"] == "unsupported_turn_detection"
- assert "server_vad" in error["error"]["message"]
+ assert "interrupt_response=false" in error["error"]["message"]
+ threshold_error = protocol._validate_realtime_turn_detection(
+ {"turn_detection": {"type": "server_vad", "threshold": 0.15}}
+ )
+ assert threshold_error == "turn_detection.threshold must be greater than 0.15 and at most 1"
@pytest.mark.asyncio
@@ -699,7 +758,7 @@ async def test_native_realtime_protocol_accepts_disabled_turn_detection():
@pytest.mark.asyncio
-async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_is_disabled():
+async def test_native_realtime_protocol_prefers_nested_vad_over_top_level():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
protocol.bind_sender(ws.send_json)
@@ -719,8 +778,55 @@ async def test_native_realtime_protocol_rejects_nested_vad_even_when_top_level_i
}
)
- assert translated is None
- assert ws.sent[-1]["error"]["code"] == "unsupported_turn_detection"
+ assert translated["session"]["extra_body"]["realtime_session_payload"]["overlap_policy"] == "barge_in_on_speech"
+
+
+@pytest.mark.asyncio
+async def test_realtime_vad_edge_order_and_tiny_chunk_resampling(monkeypatch):
+ ws = TimedWebSocket()
+ protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
+ protocol.bind_sender(ws.send_json)
+ protocol._server_vad = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig())
+ result = realtime_vad.StreamingVADResult(True, True, True, True)
+ monkeypatch.setattr(protocol._server_vad, "process_base64", lambda *args, **kwargs: result)
+ protocol._input_speech_started = True
+ protocol._hold_realtime_output_until_session_created = False
+ await protocol._to_duplex_event(
+ {"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"}
+ )
+ assert [event["type"] for event in ws.sent] == [
+ "input_audio_buffer.speech_stopped",
+ "input_audio_buffer.speech_started",
+ ]
+ ws.sent.clear()
+ protocol._input_speech_started = False
+ result = realtime_vad.StreamingVADResult(True, False, True, True)
+ await protocol._to_duplex_event({"type": "input_audio_buffer.append", "audio": "AAAA", "format": "pcm_f32le"})
+ assert [event["type"] for event in ws.sent] == [
+ "input_audio_buffer.speech_started",
+ "input_audio_buffer.speech_stopped",
+ ]
+ detector = realtime_vad.SileroStreamingVAD(realtime_vad.SileroVADConfig(), frame_scorer=lambda _: 0.0)
+ for _ in range(1024):
+ detector.process_base64(_pcm_f32_b64(1), fmt="pcm_f32le", sample_rate_hz=48_000)
+ assert detector._processed_samples + detector._pending.size == 341
+
+
+@pytest.mark.asyncio
+async def test_realtime_server_vad_runtime_failures_are_terminal(monkeypatch):
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
+ ws = TimedWebSocket()
+ ws.put(_server_vad_update("generic-vad", model="test-model", session_id="generic-vad"))
+ await handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({}))
+ assert ws.sent[-1]["error"]["code"] == "server_vad_requires_native_duplex"
+ monkeypatch.setattr(realtime_vad.SileroStreamingVAD, "process_base64", lambda *args, **kwargs: 1 / 0)
+ ws = TimedWebSocket()
+ create = _native_realtime_session_update("vad-inference-failure")
+ create["session"]["turn_detection"] = {"type": "server_vad"}
+ ws.put(create)
+ ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(1), "format": "pcm_f32le"})
+ await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=NativeRealtimeSessionProtocol({})), timeout=2)
+ assert any(e.get("error", {}).get("code") == "realtime_input_failed" for e in ws.sent)
def test_native_duplex_handler_has_no_fixed_session_admission_cap():
@@ -988,6 +1094,178 @@ async def test_minicpmo_native_session_update_requires_ref_audio_before_enabling
assert ("sid-runtime-update-ref-required", "session.update") not in engine.signals
+def test_personaplex_candidate_update_does_not_leak_minicpmo_ref_audio_check():
+ handler = OmniDuplexSessionHandler(
+ chat_service=SimpleNamespace(engine_client=SimpleNamespace()),
+ serving_runtime_adapter=PersonaPlexServingRuntimeAdapter(lambda *_: None),
+ )
+ session = DuplexSession(
+ session_id="sid-personaplex-voice-change",
+ config=DuplexSessionConfig(modalities=["audio"], voice="NATF2.pt"),
+ capabilities=PersonaPlexServingRuntimeAdapter.capabilities(max_sessions=1),
+ )
+ candidate_config = DuplexSessionConfig(modalities=["audio"], voice="OtherVoice.pt")
+
+ error = handler._runtime_session_candidate_update_error(session, candidate_config)
+
+ assert error is None
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_native_session_update_rejects_instructions_after_native_context_locked():
+ engine = FakeEngineClient()
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(engine),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ ws = TimedWebSocket()
+ ws.put(_native_session_create("sid-instructions-locked", modalities=["text"]))
+ ws.put(
+ {
+ "type": "input_audio_buffer.append",
+ "audio": _pcm_f32_b64(16000, value=0.05),
+ "format": "pcm_f32le",
+ "sample_rate_hz": 16000,
+ "duration_ms": 1000,
+ "is_speech": True,
+ }
+ )
+ ws.put(
+ {
+ "type": "turn.signal",
+ "event": "session.update",
+ "payload": {"instructions": "You are now a pirate."},
+ }
+ )
+ ws.put({"type": "session.close"})
+
+ await handler.handle_session(ws)
+
+ error = next(message for message in ws.sent if message.get("type") == "error")
+ assert error["code"] == "instructions_update_unsupported"
+ assert "session.updated" not in ws.sent_types()
+ assert ("sid-instructions-locked", "session.update") not in engine.signals
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_native_session_update_allows_unchanged_instructions_after_lock():
+ engine = FakeEngineClient()
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(engine),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ ws = TimedWebSocket()
+ ws.put(_native_session_create("sid-instructions-unchanged", modalities=["text"]))
+ ws.put(
+ {
+ "type": "input_audio_buffer.append",
+ "audio": _pcm_f32_b64(16000, value=0.05),
+ "format": "pcm_f32le",
+ "sample_rate_hz": 16000,
+ "duration_ms": 1000,
+ "is_speech": True,
+ }
+ )
+ ws.put(
+ {
+ "type": "turn.signal",
+ "event": "session.update",
+ "payload": {"instructions": "You are a concise assistant.", "temperature": 0.3},
+ }
+ )
+ ws.put({"type": "session.close"})
+
+ await handler.handle_session(ws)
+
+ assert "session.updated" in ws.sent_types()
+ assert not any(
+ message.get("type") == "error" and message.get("code") == "instructions_update_unsupported"
+ for message in ws.sent
+ )
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_native_session_update_rejects_instructions_before_first_real_append():
+ engine = FakeEngineClient()
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(engine),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ ws = TimedWebSocket()
+ ws.put(_native_session_create("sid-buffered-before-lock", modalities=["text"]))
+ ws.put(
+ {
+ "type": "input_audio_buffer.append",
+ "audio": _pcm_f32_b64(3200, value=0.05),
+ "format": "pcm_f32le",
+ "sample_rate_hz": 16000,
+ "duration_ms": 200,
+ "is_speech": True,
+ }
+ )
+ ws.put(
+ {
+ "type": "turn.signal",
+ "event": "session.update",
+ "payload": {
+ "instructions": (
+ "You are now a swashbuckling pirate captain who speaks entirely in "
+ "nautical slang, references buried treasure constantly, and never "
+ "breaks character no matter what the user asks."
+ )
+ },
+ }
+ )
+ ws.put({"type": "input_audio_buffer.commit", "final": True, "response_create": True})
+ ws.put({"type": "session.close"})
+
+ await handler.handle_session(ws)
+
+ error = next(message for message in ws.sent if message.get("type") == "error")
+ assert error["code"] == "instructions_update_unsupported"
+ assert "session.updated" not in ws.sent_types()
+ assert engine.appended, "the commit must still flush the buffered chunk despite the rejected update"
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip():
+ engine = FakeEngineClient()
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(engine),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ ws = TimedWebSocket()
+ ws.put(_native_session_create("sid-native-duplex-flag-flip", modalities=["text"]))
+ ws.put(
+ {
+ "type": "turn.signal",
+ "event": "session.update",
+ "payload": {"extra_body": {"minicpmo45_native_duplex": False}},
+ }
+ )
+ ws.put(
+ {
+ "type": "turn.signal",
+ "event": "session.update",
+ "payload": {"instructions": "You are now a pirate."},
+ }
+ )
+ ws.put({"type": "session.close"})
+
+ await handler.handle_session(ws)
+
+ errors = [message for message in ws.sent if message.get("type") == "error"]
+ assert [error["code"] for error in errors] == [
+ "native_duplex_mode_update_unsupported",
+ "instructions_update_unsupported",
+ ]
+ assert "session.updated" not in ws.sent_types()
+
+
def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz():
ws = TimedWebSocket()
protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
@@ -1325,6 +1603,14 @@ def _native_realtime_session_update(
}
+def _server_vad_update(event_id: str, **session: object) -> dict[str, object]:
+ return {
+ "type": "session.update",
+ "event_id": event_id,
+ "session": {"turn_detection": {"type": "server_vad"}, **session},
+ }
+
+
def test_native_stage0_request_id_matches_engine_after_same_session_id_reopen():
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
first = handler._registry.create(session_id="sid-reused")
@@ -1525,6 +1811,27 @@ def test_auto_response_playback_overlap_keeps_model_owned_listen_speak_decision(
assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False
+def test_legacy_auto_overlap_policy_falls_back_to_listen_only():
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
+ session = DuplexSession(
+ session_id="sid-auto-policy-fallback",
+ config=DuplexSessionConfig(
+ overlap_policy=DuplexSessionConfig._normalize_overlap_policy("auto"),
+ ),
+ )
+ session.capabilities = DuplexCapabilities.minicpmo45_native()
+
+ decision = handler._overlap_decision(
+ session,
+ {"duration_ms": session.config.overlap_barge_in_ms, "is_speech": True},
+ _native_audio_payload(is_speech=True),
+ )
+
+ assert session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
+ assert decision["action"] == "listen"
+ assert decision["reason"] == "policy_listen_only"
+
+
def test_auto_response_overlap_silence_advances_model_unit_and_preserves_realtime_input():
handler, session = _auto_response_context(
"sid-auto-silent-wire-buffer",
@@ -2922,7 +3229,7 @@ async def test_minicpmo_pre_response_continuation_drops_after_model_turn_ends():
@pytest.mark.asyncio
-async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
+async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundary():
request_id = "duplex-sid-long-response-e0-stage0"
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
@@ -2938,7 +3245,7 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
_install_direct_silence_scheduler(handler, session)
ws = TimedWebSocket()
- for _ in range(9):
+ for _ in range(handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS):
await handler._maybe_continue_native_response(
ws.send_json,
session=session,
@@ -2946,9 +3253,160 @@ async def test_minicpmo_auto_response_continuation_has_no_semantic_unit_cap():
)
await asyncio.sleep(0.02)
- assert len(engine.appended) == 9
+ response_id = session.active_response_id
+ assert response_id is not None
+ assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
assert all(payload["duplex_turn_id"] == 0 for _, _, payload, _ in engine.appended)
- assert all("force_speak" not in payload for _, _, payload, _ in engine.appended)
+ assert all(not {"force_speak", "force_listen"} & payload.keys() for _, _, payload, _ in engine.appended[:-1])
+ assert engine.appended[-1][2]["force_listen"] is True
+ assert handler._native_response_continuations_remaining(session, response_id) is False
+
+ await handler._maybe_continue_native_response(
+ ws.send_json,
+ session=session,
+ expected_epoch=session.epoch,
+ )
+
+ assert len(engine.appended) == handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
+ assert ws.sent_types()[-2:] == ["response.listen", "response.done"]
+ assert session.active_response_id is None
+ assert session.active_request_id == request_id
+ assert session.turn_id == 1
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_auto_response_boundary_listen_closes_response():
+ request_id = "duplex-sid-boundary-listen-e0-stage0"
+ engine = FakeEngineClient()
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
+ session = DuplexSession(
+ session_id="sid-boundary-listen",
+ config=DuplexSessionConfig(extra_body={"auto_response": True}),
+ )
+ session.capabilities = DuplexCapabilities.minicpmo45_native()
+ response_id = session.begin_response(turn_id=0)
+ session.bind_request(request_id)
+ native = handler._minicpmo_session_state(session)
+ native.continuation_owner_id = f"response:{response_id}"
+ native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
+ _install_direct_silence_scheduler(handler, session)
+ ws = TimedWebSocket()
+
+ close_reason, emitted = await handler._send_one_native_duplex_event(
+ ws.send_json,
+ {
+ "supported": True,
+ "stage_role": "llm",
+ "is_listen": True,
+ "model_listen": True,
+ "data_plane_request_id": request_id,
+ "end_of_turn": False,
+ "uses_model_runner_scheduler": True,
+ "runner_kv_backed": True,
+ "runtime_impl": "scheduler_data_plane",
+ },
+ session=session,
+ expected_epoch=session.epoch,
+ )
+
+ assert close_reason is None
+ assert emitted is True
+ assert engine.appended == []
+ assert ws.sent_types() == ["response.listen", "response.done"]
+ assert session.active_response_id is None
+ assert session.active_request_id == request_id
+ assert session.turn_id == 1
+ assert native.continuation_owner_id is None
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_auto_response_boundary_allows_next_model_turn_response():
+ request_id = "duplex-sid-boundary-next-turn-e0-stage0"
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
+ session = DuplexSession(
+ session_id="sid-boundary-next-turn",
+ config=DuplexSessionConfig(extra_body={"auto_response": True}),
+ )
+ session.capabilities = DuplexCapabilities.minicpmo45_native()
+ session.begin_response(turn_id=0)
+ session.bind_request(request_id)
+ native = handler._minicpmo_session_state(session)
+ native.continuation_owner_id = f"response:{session.active_response_id}"
+ native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
+ ws = TimedWebSocket()
+
+ await handler._maybe_continue_native_response(
+ ws.send_json,
+ session=session,
+ expected_epoch=session.epoch,
+ )
+
+ assert session.active_response_id is None
+ assert session.turn_id == 1
+
+ close_reason, emitted = await handler._send_one_native_duplex_event(
+ ws.send_json,
+ {
+ "supported": True,
+ "stage_role": "tts",
+ "is_listen": False,
+ "data_plane_request_id": request_id,
+ "text": "next response",
+ "audio_data": "next-audio",
+ "audio_format": "pcm16",
+ "audio_duration_ms": 100,
+ "end_of_turn": False,
+ "model_turn_id": 1,
+ },
+ session=session,
+ expected_epoch=session.epoch,
+ )
+
+ assert close_reason is None
+ assert emitted is True
+ assert ws.sent_types() == [
+ "response.listen",
+ "response.done",
+ "response.created",
+ "response.speak",
+ "response.output_audio.delta",
+ ]
+ assert session.active_response_turn_id == 1
+
+
+@pytest.mark.asyncio
+async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response():
+ request_id = "duplex-sid-boundary-race-e0-stage0"
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient()))
+ session = DuplexSession(
+ session_id="sid-boundary-race",
+ config=DuplexSessionConfig(extra_body={"auto_response": True}),
+ )
+ session.capabilities = DuplexCapabilities.minicpmo45_native()
+ old_response_id = session.begin_response(turn_id=0)
+ session.bind_request(request_id)
+ native = handler._minicpmo_session_state(session)
+ native.continuation_owner_id = f"response:{old_response_id}"
+ native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS
+ sent: list[dict[str, Any]] = []
+ new_response_ids: list[str] = []
+
+ async def _barge_in_on_listen(payload: dict[str, Any]) -> None:
+ sent.append(payload)
+ if payload["type"] == "response.listen":
+ session.barge_in()
+ session.bind_request("duplex-sid-boundary-race-e1-stage0")
+ new_response_ids.append(session.begin_response(turn_id=1))
+
+ await handler._maybe_continue_native_response(
+ _barge_in_on_listen,
+ session=session,
+ expected_epoch=session.epoch,
+ )
+
+ assert [payload["type"] for payload in sent] == ["response.listen"]
+ assert len(new_response_ids) == 1
+ assert session.active_response_id == new_response_ids[0]
@pytest.mark.asyncio
@@ -3641,6 +4099,55 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event():
]
+@pytest.mark.asyncio
+async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot():
+ class SnapshotChatService(FakeChatService):
+ async def create_chat_completion(self, request, raw_request=None):
+ del raw_request
+ self.seen_request_ids.append(request.request_id)
+ snapshots = [
+ {"num_tokens_out": 1, "vllm_tpot_ms": 0.0, "vllm_itls_ms": []},
+ {"num_tokens_out": 2, "vllm_tpot_ms": 10.0, "vllm_itls_ms": [10.0]},
+ {"num_tokens_out": 3, "vllm_tpot_ms": 11.0, "vllm_itls_ms": [10.0, 12.0]},
+ ]
+
+ async def generate():
+ for index, snapshot in enumerate(snapshots):
+ payload = {
+ "modality": "audio",
+ "choices": [{"delta": {"content": f"audio-{index}"}}],
+ "metrics": {"stage_metrics": {"0": snapshot}},
+ }
+ yield f"data: {json.dumps(payload)}\n\n"
+
+ return generate()
+
+ handler = OmniDuplexSessionHandler(
+ chat_service=SnapshotChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-chat-snapshots",
+ config=DuplexSessionConfig(instructions="speak"),
+ )
+ sent: list[dict[str, Any]] = []
+
+ async def send_json(data: dict[str, Any]) -> None:
+ sent.append(data)
+
+ await handler._run_response(session, send_json)
+
+ audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"]
+ done = next(event for event in sent if event.get("type") == "response.done")
+ assert all("vllm_omni" not in event for event in audio_deltas)
+ assert done["vllm_omni"]["stage_metrics"]["0"] == {
+ "num_tokens_out": 3,
+ "vllm_tpot_ms": 11.0,
+ "vllm_itls_ms": [10.0, 12.0],
+ }
+
+
@pytest.mark.asyncio
async def test_duplex_chat_rejects_unknown_output_modality():
handler = OmniDuplexSessionHandler(
@@ -3840,6 +4347,385 @@ async def test_playback_ack_response_id_selects_matching_pending_history_item():
assert second_item_id not in session.pending_history_item_ids
+@pytest.mark.asyncio
+async def test_playback_ack_recovers_missing_response_history_registration():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-playback-missing-registration",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ response_id = session.begin_response()
+ session.append_assistant_text("response text")
+ session.mark_audio_sent(1000, text_chars=len("response text"))
+ session.end_response(commit_text=True, preserve_request=True)
+ assert not session.history_item_ids
+ assert f"item_{response_id}" in session.pending_history_item_ids
+
+ ws = TimedWebSocket()
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": f"item_{response_id}",
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged")
+ assert ack["history_committed"] is True
+ assert f"item_{response_id}" in session.history_item_ids
+ assert session.history[-1] == {"role": "assistant", "content": "response text"}
+
+
+@pytest.mark.asyncio
+async def test_playback_ack_rejects_response_after_followup_user_commit():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-playback-after-followup-user",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ session.append_text("user A")
+ session.commit_user_input()
+ response_id = session.begin_response()
+ session.append_assistant_text("assistant A")
+ session.mark_audio_sent(1000, text_chars=len("assistant A"))
+ session.end_response(commit_text=False, preserve_request=True)
+ session.append_text("user B")
+ session.commit_user_input()
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": f"item_{response_id}",
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ assert ws.sent_types() == ["error"]
+ assert ws.sent[0]["code"] == "playback_ack_too_late"
+ assert session.history == (
+ {"role": "user", "content": "user A"},
+ {"role": "user", "content": "user B"},
+ )
+
+
+@pytest.mark.asyncio
+async def test_precommit_zero_audio_ack_reserves_history_for_response_that_finishes_after_commit():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-zero-audio-checkpoint",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ session.mark_user_input_activity()
+ response_id = session.begin_response()
+ item_id = f"item_{response_id}"
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 0,
+ "committed_ms": 0,
+ },
+ ws.send_json,
+ )
+ session.commit_native_audio_input(transcript="later user input")
+ session.append_assistant_text("hello world")
+ session.mark_audio_sent(1000, text_chars=len("hello world"))
+ committed_message = session.end_response(commit_text=True, preserve_request=True)
+ session.register_history_item(item_id, committed_message)
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ assert not any(message.get("type") == "error" for message in ws.sent)
+ assert session.history == (
+ {"role": "assistant", "content": "hello world"},
+ {
+ "role": "user",
+ "content": [
+ {
+ "type": "audio_url",
+ "audio_url": {"url": "native-duplex:input-audio"},
+ "transcript": "later user input",
+ }
+ ],
+ "transcript": "later user input",
+ },
+ )
+
+
+@pytest.mark.asyncio
+async def test_precommit_partial_ack_extends_history_after_response_finishes():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-partial-checkpoint",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ session.mark_user_input_activity()
+ response_id = session.begin_response()
+ item_id = f"item_{response_id}"
+ session.append_assistant_text("hello world")
+ session.mark_audio_sent(1000, text_chars=len("hello world"))
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 500,
+ "committed_ms": 500,
+ },
+ ws.send_json,
+ )
+ session.commit_native_audio_input(transcript="later user input")
+ committed_message = session.end_response(commit_text=True, preserve_request=True)
+ session.register_history_item(item_id, committed_message)
+
+ assert session.history[0] == {"role": "assistant", "content": "hello"}
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ assert not any(message.get("type") == "error" for message in ws.sent)
+ assert session.history[0] == {"role": "assistant", "content": "hello world"}
+ assert session.history[1]["role"] == "user"
+
+
+@pytest.mark.asyncio
+async def test_explicit_truncate_seals_snapshot_against_later_playback_ack():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-hard-playback-truncate",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ response_id = session.begin_response()
+ item_id = f"item_{response_id}"
+ session.append_assistant_text("hello world")
+ session.mark_audio_sent(1000, text_chars=len("hello world"))
+ session.end_response(commit_text=False, preserve_request=True)
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 500,
+ "committed_ms": 500,
+ "truncate": True,
+ },
+ ws.send_json,
+ )
+ assert session.history == ({"role": "assistant", "content": "hello"},)
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": response_id,
+ "item_id": item_id,
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ assert not any(message.get("type") == "error" for message in ws.sent)
+ assert session.history == ({"role": "assistant", "content": "hello"},)
+
+
+@pytest.mark.asyncio
+async def test_cancel_does_not_append_old_assistant_after_followup_user():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(
+ session_id="sid-cancel-after-followup-user",
+ config=DuplexSessionConfig(playback_commit_policy="ack_only"),
+ )
+ session.append_text("user A")
+ session.commit_user_input()
+ session.begin_response()
+ session.append_assistant_text("assistant A")
+ session.mark_audio_sent(1000, text_chars=len("assistant A"))
+ session.acknowledge_playback(500, 500)
+ session.append_text("user B")
+ session.commit_user_input()
+ ws = TimedWebSocket()
+
+ cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="new_response")
+
+ assert cancelled is True
+ assert session.history == (
+ {"role": "user", "content": "user A"},
+ {"role": "user", "content": "user B"},
+ )
+
+
+@pytest.mark.asyncio
+async def test_playback_ack_rejects_item_from_another_response():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(session_id="sid-playback-mismatch", config=DuplexSessionConfig())
+ first_response_id = session.begin_response()
+ session.append_assistant_text("first response")
+ session.mark_audio_sent(1000)
+ session.end_response(commit_text=False, preserve_request=True)
+ session.register_history_item(f"item_{first_response_id}", None)
+ second_response_id = session.begin_response()
+ session.mark_audio_sent(1000)
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": second_response_id,
+ "item_id": f"item_{first_response_id}",
+ "played_ms": 500,
+ "committed_ms": 500,
+ },
+ ws.send_json,
+ )
+
+ assert ws.sent_types() == ["error"]
+ assert ws.sent[0]["code"] == "playback_item_mismatch"
+ assert session.active_response_id == second_response_id
+ assert session.playback.as_dict() == {
+ "generated_ms": 1000,
+ "sent_ms": 1000,
+ "played_ms": 0,
+ "committed_ms": 0,
+ }
+
+
+@pytest.mark.asyncio
+async def test_playback_ack_rejects_user_history_item():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(session_id="sid-playback-user-item", config=DuplexSessionConfig())
+ user_message = {"role": "user", "content": "user input"}
+ session.append_history_message(user_message)
+ session.register_history_item("item_user", user_message)
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": "user",
+ "item_id": "item_user",
+ "played_ms": 500,
+ "committed_ms": 500,
+ },
+ ws.send_json,
+ )
+
+ assert ws.sent_types() == ["error"]
+ assert ws.sent[0]["code"] == "playback_item_not_found"
+ assert session.history == ({"role": "user", "content": "user input"},)
+
+
+@pytest.mark.asyncio
+async def test_late_playback_ack_recovers_response_local_history_after_next_response_begins():
+ handler = OmniDuplexSessionHandler(
+ chat_service=FakeChatService(FakeEngineClient()),
+ config_timeout_s=0.1,
+ idle_timeout_s=1,
+ )
+ session = DuplexSession(session_id="sid-playback-response-local", config=DuplexSessionConfig())
+ first_response_id = session.begin_response()
+ session.append_assistant_text("first response")
+ session.mark_audio_sent(1000, text_chars=len("first response"))
+ session.end_response(commit_text=False, preserve_request=True)
+ second_response_id = session.begin_response()
+ session.mark_audio_sent(2000, text_chars=len("second response"))
+ ws = TimedWebSocket()
+
+ await handler._handle_playback_ack(
+ session,
+ {
+ "type": "playback.ack",
+ "response_id": first_response_id,
+ "item_id": f"item_{first_response_id}",
+ "played_ms": 1000,
+ "committed_ms": 1000,
+ },
+ ws.send_json,
+ )
+
+ ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged")
+ assert ack["history_committed"] is True
+ assert session.history[-1] == {"role": "assistant", "content": "first response"}
+ assert session.active_response_id == second_response_id
+ assert session.playback.as_dict() == {
+ "generated_ms": 2000,
+ "sent_ms": 2000,
+ "played_ms": 0,
+ "committed_ms": 0,
+ }
+
+
@pytest.mark.asyncio
async def test_late_playback_ack_does_not_advance_new_response_cursor():
handler = OmniDuplexSessionHandler(
@@ -4194,6 +5080,7 @@ async def test_realtime_resume_preserves_append_tail_order_across_connections():
await asyncio.sleep(0.01)
else:
pytest.fail("resumable session did not open")
+ assert created is not None
session = handler._registry.get("sid-resume-append-order")
assert session is not None
native_state = handler._minicpmo_session_state(session)
@@ -4323,6 +5210,7 @@ async def test_realtime_takeover_does_not_clear_new_attachment_pending_turns():
await asyncio.sleep(0.001)
else:
pytest.fail("first attachment did not create the session")
+ assert created is not None
second = TimedWebSocket(receive_timeout_s=1.0)
second.put(
@@ -4386,6 +5274,7 @@ async def test_realtime_takeover_drops_input_queued_by_replaced_attachment():
await asyncio.sleep(0.001)
else:
pytest.fail("first attachment did not create the session")
+ assert created is not None
second = TimedWebSocket(receive_timeout_s=1.0)
second.put(
@@ -4988,7 +5877,7 @@ async def test_audio_clear_preserves_later_wire_order_append():
@pytest.mark.asyncio
-async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
+async def test_native_append_is_cancelled_by_later_wire_order_input_cancel(mocker):
class SlowAppendEngine(FakeEngineClient):
def __init__(self) -> None:
super().__init__()
@@ -5002,11 +5891,12 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
*,
mode: str,
payload: object,
+ operation_id: str | None = None,
final: bool = False,
timeout: float | None = None,
collect_outputs: bool = True,
fence: DuplexFence | None = None,
- ):
+ ) -> Any:
self.append_started.set()
try:
await self.release_append.wait()
@@ -5017,12 +5907,14 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
session_id,
mode=mode,
payload=payload,
+ operation_id=operation_id,
final=final,
timeout=timeout,
collect_outputs=collect_outputs,
fence=fence,
)
+ rollback = mocker.spy(MiniCPMO45PcmAppendBuffer, "_rollback_reservation")
engine = SlowAppendEngine()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
@@ -5053,6 +5945,7 @@ async def test_native_append_is_cancelled_by_later_wire_order_input_cancel():
await asyncio.wait_for(handler_task, timeout=1)
assert engine.append_cancelled.is_set()
+ rollback.assert_called_once()
assert engine.signal_fences[-1] == DuplexFence("sid-cancel-slow-append")
assert engine.signal_next_fences[-1] == DuplexFence("sid-cancel-slow-append", epoch=1)
@@ -5128,7 +6021,9 @@ async def test_duplex_handler_session_update_control_failure_rolls_back_config()
error = next(m for m in ws.sent if m.get("type") == "error")
assert error["code"] == "runtime_signal_failed"
assert error["runtime_control"]["error_count"] == 1
- assert engine.signal_session_configs[-1]["instructions"] == "must roll back"
+ signaled_config = engine.signal_session_configs[-1]
+ assert signaled_config is not None
+ assert signaled_config["instructions"] == "must roll back"
@pytest.mark.asyncio
@@ -5146,33 +6041,59 @@ async def test_native_session_update_commits_config_only_after_runtime_ack():
return await super().signal_duplex_turn_async(session_id, **kwargs)
engine = BlockingUpdateEngine()
- handler = OmniDuplexSessionHandler(
- chat_service=FakeChatService(engine),
- config_timeout_s=0.1,
- idle_timeout_s=1,
- )
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
ws = TimedWebSocket()
- ws.put(_session_create("sid-update-two-phase"))
- ws.put(
- {
- "type": "turn.signal",
- "event": "session.update",
- "payload": {"instructions": "candidate instructions"},
- }
- )
+ protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
+ ws.put(_native_realtime_session_update("sid-update-two-phase"))
+ ws.put(_server_vad_update("update-two-phase"))
- handler_task = asyncio.create_task(handler.handle_session(ws))
+ handler_task = asyncio.create_task(handler.handle_session(ws, realtime_protocol=protocol))
await asyncio.wait_for(engine.update_started.wait(), timeout=1)
live_session = handler._registry.get("sid-update-two-phase")
assert live_session is not None
- assert live_session.config.instructions != "candidate instructions"
+ assert live_session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value
+ assert protocol._pending_turn_detection_update is not None
+ assert protocol._server_vad is None
+
+ protocol.encode_outbound_event({"type": "error", "code": "unrelated", "error": "unrelated"})
+ assert protocol._pending_turn_detection_update is not None
engine.release_update.set()
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
- updated = next(message for message in ws.sent if message.get("type") == "session.updated")
- assert updated["session"]["instructions"] == "candidate instructions"
+ updated = next(message for message in reversed(ws.sent) if message.get("type") == "session.updated")
+ assert updated["session"]["overlap_policy"] == DuplexOverlapPolicy.BARGE_IN_ON_SPEECH.value
+ assert isinstance(protocol._server_vad, realtime_vad.SileroStreamingVAD)
+
+
+@pytest.mark.parametrize(
+ ("append_fails", "expected_code"),
+ [(True, "session_update_aborted"), (False, "instructions_update_unsupported")],
+)
+@pytest.mark.asyncio
+async def test_realtime_vad_update_rejection_does_not_stall(append_fails, expected_code):
+ append_result = {"operation": "append", "ok": False, "error_count": 1} if append_fails else None
+ handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient(append_result=append_result)))
+ ws = TimedWebSocket()
+ protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type]
+ create = _native_realtime_session_update("sid-rejected-vad-update")
+ create["session"]["extra_body"]["auto_response"] = append_fails
+ ws.put(create)
+ if append_fails:
+ ws.put({"type": "input_audio_buffer.append", "audio": _pcm_f32_b64(16_000), "format": "pcm_f32le"})
+ update = _server_vad_update("rejected-vad-update")
+ if not append_fails:
+ update["session"]["instructions"] = "You are now a pirate."
+ ws.put(update)
+ ws.put({"type": "session.close"})
+ await asyncio.wait_for(handler.handle_session(ws, realtime_protocol=protocol), timeout=2)
+
+ errors = [message["error"] for message in ws.sent if isinstance(message.get("error"), dict)]
+ error = next(error for error in errors if error.get("code") == expected_code)
+ assert error["event_id"] == "rejected-vad-update"
+ assert protocol._pending_turn_detection_update is None
+ assert protocol._server_vad is None
@pytest.mark.asyncio
@@ -5296,7 +6217,7 @@ async def test_failed_native_append_prevents_queued_append_from_running():
@pytest.mark.asyncio
-async def test_failed_final_append_keeps_committed_audio_for_response_retry():
+async def test_failed_final_append_discards_retained_audio_without_retry():
failed_result = {
"operation": "append",
"session_id": "sid-final-append-retry",
@@ -5305,15 +6226,7 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
"error_count": 1,
"stage_results": [],
}
- successful_result = {
- "operation": "append",
- "session_id": "sid-final-append-retry",
- "ok": True,
- "unsupported_count": 0,
- "error_count": 0,
- "stage_results": [],
- }
- engine = FakeEngineClient(append_results=[failed_result, successful_result])
+ engine = FakeEngineClient(append_result=failed_result)
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
@@ -5338,15 +6251,18 @@ async def test_failed_final_append_keeps_committed_audio_for_response_retry():
handler_task = asyncio.create_task(handler.handle_session(ws))
for _ in range(100):
- if len(engine.appended) >= 2:
+ state = handler._serving_runtime_adapter.session_states.get("sid-final-append-retry")
+ if len(engine.appended) >= 1 and state is not None and state.committed_audio_payload is None:
break
await asyncio.sleep(0.01)
+ state = handler._serving_runtime_adapter.session_states["sid-final-append-retry"]
+ assert state.committed_audio_payload is None
+ assert state.committed_audio_reserved_bytes == 0
+ assert handler._registry.get("sid-final-append-retry").pending_input_bytes == 0
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
- assert len(engine.appended) == 2
- assert engine.appended[0][2]["audio"] == engine.appended[1][2]["audio"]
- assert engine.append_operation_ids[0] == engine.append_operation_ids[1]
+ assert len(engine.appended) == 1
@pytest.mark.asyncio
@@ -5429,7 +6345,7 @@ async def test_minicpmo_native_duplex_open_unsupported_fails_session_create():
@pytest.mark.asyncio
async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
class LegacyEngineClient(FakeEngineClient):
- async def open_duplex_session_async(
+ async def open_duplex_session_async( # type: ignore[override]
self,
session_id: str,
*,
@@ -5437,7 +6353,7 @@ async def test_minicpmo_native_duplex_rejects_engine_without_fence_contract():
capabilities: dict[str, object] | None = None,
session_config: dict[str, object] | None = None,
timeout: float | None = None,
- ) -> None:
+ ) -> Any:
return await super().open_duplex_session_async(
session_id,
session_mode=session_mode,
@@ -5534,6 +6450,9 @@ async def test_minicpmo_native_duplex_separates_public_and_runtime_config(monkey
assert opened_runtime_config["duplex_stage_max_tokens"] == {"0": 20, "1": 8192}
assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_k"] == 20
assert opened_runtime_config["duplex_stage_sampling_params"]["0"]["top_p"] == 0.8
+ # The Talker YAML carries upstream's min_new_token=50 for chat TTS; a duplex
+ # chunk is 26 codec samples, so check_stop would never release the request.
+ assert opened_runtime_config["duplex_stage_sampling_params"]["1"] == {"min_tokens": 0}
assert base64.b64decode(opened_runtime_config["ref_audio_data"]) == struct.pack("<1600f", *([0.25] * 1600))
@@ -5674,7 +6593,7 @@ async def test_minicpmo_native_duplex_rejects_ref_audio_path():
assert engine.opened == []
-def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen():
+def test_minicpmo_native_duplex_explicit_barge_in_request_interrupts():
engine = FakeEngineClient()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
@@ -5701,12 +6620,10 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen(
)
assert decision == {
- "action": "listen",
- "reason": "barge_in_unsupported",
+ "action": "barge_in",
+ "reason": "client_force_barge_in",
"duration_ms": 1000,
- "overlap_speech_ms": 1000,
"buffer_audio": True,
- "defer_runtime_append": True,
}
@@ -5718,7 +6635,7 @@ def test_minicpmo_native_duplex_explicit_barge_in_request_is_deferred_as_listen(
{"type": "turn.signal", "event": "barge_in"},
],
)
-async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control(
+async def test_minicpmo_native_duplex_accepts_explicit_barge_in_control(
barge_in_event: dict[str, object],
):
engine = FakeEngineClient()
@@ -5734,9 +6651,8 @@ async def test_minicpmo_native_duplex_rejects_unadvertised_barge_in_control(
await handler.handle_session(ws)
- error = next(message for message in ws.sent if message.get("code") == "barge_in_unsupported")
- assert error["session_id"] == "sid-native-barge-control-disabled"
- assert ("sid-native-barge-control-disabled", "barge_in") not in engine.signals
+ assert not any(message.get("code") == "barge_in_unsupported" for message in ws.sent)
+ assert ("sid-native-barge-control-disabled", "barge_in") in engine.signals
@pytest.mark.asyncio
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 2136d84f8..b2eb2edfc 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -117,10 +117,23 @@ def test_encode_image_base64():
class MockGenerationResult:
"""Mock result object compatible with current diffusion output shape."""
- def __init__(self, images):
+ def __init__(self, images, stage_durations=None, peak_memory_mb=0.0):
self.images = images
- self.stage_durations = {}
- self.peak_memory_mb = 0.0
+ self.stage_durations = {} if stage_durations is None else stage_durations
+ self.peak_memory_mb = peak_memory_mb
+
+
+def _stub_engine_generate_with_metrics(engine, *, stage_durations, peak_memory_mb):
+ """Stub engine.generate to return profiler metrics."""
+
+ async def generate(*args, **kwargs):
+ yield MockGenerationResult(
+ [Image.new("RGB", (16, 16), color="blue")],
+ stage_durations=stage_durations,
+ peak_memory_mb=peak_memory_mb,
+ )
+
+ engine.generate = generate
class MockStageResult:
@@ -972,6 +985,53 @@ def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_clie
assert "multi-stage" in response.json()["detail"]
+def test_generate_images_returns_metrics_single_stage(test_client):
+ """Single-stage /v1/images/generations copies getattr metrics onto the response."""
+ stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0}
+ peak_memory_mb = 1024.0
+ _stub_engine_generate_with_metrics(
+ test_client.app.state.engine_client,
+ stage_durations=stage_durations,
+ peak_memory_mb=peak_memory_mb,
+ )
+
+ response = test_client.post(
+ "/v1/images/generations",
+ json={"prompt": "a cat", "n": 1, "size": "256x256"},
+ )
+ assert response.status_code == 200
+ assert response.json()["metrics"] == {
+ "stage_durations": stage_durations,
+ "peak_memory_mb": peak_memory_mb,
+ }
+
+
+def test_generate_images_returns_metrics_multistage(async_omni_test_client):
+ """Multi-stage /v1/images/generations unpacks generate_diffusion_images metrics."""
+ stage_durations = {
+ "queue_wait_ms": 1.0,
+ "preprocess_ms": 2.0,
+ "stage_0_gen_ms": 3.0,
+ "stage_1_gen_ms": 4.0,
+ }
+ peak_memory_mb = 2048.0
+ _stub_engine_generate_with_metrics(
+ async_omni_test_client.app.state.engine_client,
+ stage_durations=stage_durations,
+ peak_memory_mb=peak_memory_mb,
+ )
+
+ response = async_omni_test_client.post(
+ "/v1/images/generations",
+ json={"prompt": "a cat", "n": 1, "size": "256x256"},
+ )
+ assert response.status_code == 200
+ assert response.json()["metrics"] == {
+ "stage_durations": stage_durations,
+ "peak_memory_mb": peak_memory_mb,
+ }
+
+
def test_generate_images_max_size_rejected(async_omni_test_client):
"""Test that a size exceeding max_generated_image_size returns 400."""
response = async_omni_test_client.post(
@@ -1987,6 +2047,56 @@ def test_image_edit_with_seed_zero(async_omni_test_client):
)
+def test_image_edits_returns_metrics_single_stage(test_client):
+ """Single-stage /v1/images/edits copies getattr metrics onto the response."""
+ stage_durations = {"queue_wait_ms": 1.0, "stage_0_gen_ms": 2.0}
+ peak_memory_mb = 1024.0
+ _stub_engine_generate_with_metrics(
+ test_client.app.state.engine_client,
+ stage_durations=stage_durations,
+ peak_memory_mb=peak_memory_mb,
+ )
+
+ response = test_client.post(
+ "/v1/images/edits",
+ files=[("image", make_test_image_bytes((16, 16)))],
+ data={"prompt": "edit me"},
+ )
+ assert response.status_code == 200
+ assert response.json()["metrics"] == {
+ "stage_durations": stage_durations,
+ "peak_memory_mb": peak_memory_mb,
+ }
+
+
+def test_image_edits_returns_metrics_multistage(async_omni_test_client):
+ """Multi-stage /v1/images/edits unpacks generate_diffusion_images metrics."""
+ stage_durations = {
+ "queue_wait_ms": 1.0,
+ "preprocess_ms": 2.0,
+ "ar2diffusion_ms": 3.0,
+ "stage_0_gen_ms": 4.0,
+ "stage_1_gen_ms": 5.0,
+ }
+ peak_memory_mb = 2048.0
+ _stub_engine_generate_with_metrics(
+ async_omni_test_client.app.state.engine_client,
+ stage_durations=stage_durations,
+ peak_memory_mb=peak_memory_mb,
+ )
+
+ response = async_omni_test_client.post(
+ "/v1/images/edits",
+ files=[("image", make_test_image_bytes((16, 16)))],
+ data={"prompt": "edit me"},
+ )
+ assert response.status_code == 200
+ assert response.json()["metrics"] == {
+ "stage_durations": stage_durations,
+ "peak_memory_mb": peak_memory_mb,
+ }
+
+
def test_image_edit_with_seed_zero_single_stage(test_client):
"""Test that seed=0 is correctly handled in image editing (single stage).
diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
index 325daabdd..7664b2402 100644
--- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
+++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out).
@@ -38,6 +38,10 @@ MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
REALTIME_SYNTH_PHRASE_TEXT = (
"Translate into Chinese: Beijing is the Capital of China. It is the center of culture and politics"
)
+ISSUE_6474_SYNTH_PHRASE_TEXT = (
+ "Can you tell me the current temperature and weather conditions in New York City? "
+ "What about Los Angeles? Please compare them in detail using at least eight complete sentences."
+)
# Simulate realtime upload pacing (``openai_realtime_client.py --send-delay-ms``).
SEND_DELAY_MS = 200
@@ -100,6 +104,7 @@ async def _run_realtime_audio_roundtrip(
*,
chunk_ms: int = 100,
send_delay_ms: int = 0,
+ completion_timeout_s: float = 600,
) -> dict:
uri = f"ws://{host}:{port}/v1/realtime"
incremental: list[bytes] = []
@@ -131,7 +136,7 @@ async def _run_realtime_audio_roundtrip(
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
while True:
- message = await asyncio.wait_for(ws.recv(), timeout=600)
+ message = await asyncio.wait_for(ws.recv(), timeout=completion_timeout_s)
if isinstance(message, bytes):
continue
@@ -178,12 +183,16 @@ async def _run_realtime_audio_roundtrip(
}
-def _synthetic_pcm16_input() -> bytes:
+def _synthetic_pcm16_input(
+ *,
+ phrase_text: str = REALTIME_SYNTH_PHRASE_TEXT,
+ duration_s: int = 10,
+) -> bytes:
syn = generate_synthetic_audio(
- 10,
+ duration_s,
1,
sample_rate=16000,
- phrase_text=REALTIME_SYNTH_PHRASE_TEXT,
+ phrase_text=phrase_text,
)
wav_bytes = base64.b64decode(syn["base64"])
return _pcm16_mono_16k_from_wav_bytes(wav_bytes)
@@ -260,6 +269,33 @@ class TestQwen3OmniRealtimeWebSocket:
_assert_realtime_smoke(result)
_assert_realtime_accuracy(result)
+ @pytest.mark.advanced_model
+ @pytest.mark.omni
+ @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+ @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True)
+ def test_long_audio_response_completes_async_chunk(self, omni_server) -> None:
+ """Regression for #6474: long async-chunk audio must emit response.audio.done."""
+ issue_6474_pcm16 = _synthetic_pcm16_input(
+ phrase_text=ISSUE_6474_SYNTH_PHRASE_TEXT,
+ )
+ issue_6474_result = asyncio.run(
+ _run_realtime_audio_roundtrip(
+ omni_server.host,
+ omni_server.port,
+ omni_server.model,
+ issue_6474_pcm16,
+ chunk_ms=100,
+ send_delay_ms=SEND_DELAY_MS,
+ completion_timeout_s=180,
+ )
+ )
+
+ _assert_realtime_smoke(issue_6474_result)
+ output_duration_s = len(issue_6474_result["output_pcm"]) / (2 * issue_6474_result["output_sample_rate"])
+ assert output_duration_s > 5, (
+ f"Expected an issue-like audio response longer than 5 seconds, got {output_duration_s:.2f}s"
+ )
+
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py
index 0647c40a3..2fdf9f646 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py
@@ -2,9 +2,18 @@
"""Unit tests for OmniChatCompletionResponse/StreamResponse metrics field."""
from types import SimpleNamespace
+from unittest.mock import MagicMock
import pytest
+from tests.helpers.serving_chat import (
+ build_serving_chat,
+ collect_stream,
+ make_request,
+ make_text_omni_output,
+ parse_sse_chunks,
+)
+
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -54,6 +63,65 @@ def test_omni_chat_completion_stream_response_metrics():
assert response.metrics == {"stage_latency": 0.5}
+@pytest.mark.asyncio
+async def test_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens():
+ serving_chat = build_serving_chat()
+ serving_chat.enable_prompt_tokens_details = True
+ request = make_request(modalities=["text"], include_usage=True)
+
+ async def result_generator():
+ yield make_text_omni_output(
+ token_ids=[10],
+ finish_reason="stop",
+ num_cached_tokens=0,
+ )
+
+ raw_lines = await collect_stream(
+ serving_chat.chat_completion_stream_generator(
+ request=request,
+ result_generator=result_generator(),
+ request_id="test-req",
+ model_name="test-model",
+ conversation=[],
+ tokenizer=MagicMock(),
+ request_metadata=MagicMock(),
+ mm_token_counts={"image": 7},
+ )
+ )
+
+ chunks = parse_sse_chunks(raw_lines)
+ usage = next(chunk["usage"] for chunk in chunks if chunk.get("usage"))
+
+ assert usage["prompt_tokens_details"] == {
+ "cached_tokens": 0,
+ "multimodal_tokens": {"image": 7},
+ }
+
+
+def test_non_stream_prompt_token_details_include_zero_cached_and_multimodal_tokens():
+ serving_chat = build_serving_chat()
+ serving_chat.enable_prompt_tokens_details = True
+ request = make_request(modalities=["text"], stream=False)
+ omni_output = make_text_omni_output(
+ token_ids=[10],
+ finish_reason="stop",
+ num_cached_tokens=0,
+ )
+
+ _, usage, _, _, _ = serving_chat._create_text_choice(
+ request=request,
+ omni_outputs=omni_output,
+ tokenizer=MagicMock(),
+ conversation=[],
+ role="assistant",
+ mm_token_counts={"audio": 5},
+ )
+
+ assert usage.prompt_tokens_details is not None
+ assert usage.prompt_tokens_details.cached_tokens == 0
+ assert usage.prompt_tokens_details.multimodal_tokens == {"audio": 5}
+
+
def test_create_image_choice_exposes_diffusion_metrics():
"""Ensure image chat content exposes profiler metrics for clients."""
from PIL import Image
diff --git a/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
new file mode 100644
index 000000000..59d070a39
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_serving_chat_minicpmo45_reference_audio.py
@@ -0,0 +1,207 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MiniCPM-o 4.5 chat reference audio transport."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import cast
+
+import numpy as np
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.inputs import TokensPrompt
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
+from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import (
+ llm2tts,
+ tts2code2wav_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _EngineClientStub:
+ stage_configs: list[dict[str, object]]
+
+
+@dataclass
+class _ModelConfigStub:
+ allowed_local_media_path: str = ""
+ allowed_media_domains: list[str] | None = None
+
+
+@dataclass
+class _ThinkerCompletionStub:
+ token_ids: list[int]
+ text: str
+ hidden_states: torch.Tensor
+ multimodal_output: dict[str, object] = field(default_factory=dict)
+
+
+@dataclass
+class _ThinkerRequestOutputStub:
+ request_id: str
+ prompt_token_ids: list[int]
+ outputs: list[_ThinkerCompletionStub]
+ multimodal_output: dict[str, object] | None = None
+
+
+@dataclass
+class _Code2WavRequestStub:
+ external_req_id: str
+ request_id: str
+ model_intermediate_buffer: dict[str, object]
+
+
+@dataclass
+class _TransferManagerStub:
+ request_payload: dict[str, object] = field(default_factory=dict)
+
+
+@pytest.fixture
+def serving_chat() -> OmniOpenAIServingChat:
+ return object.__new__(OmniOpenAIServingChat)
+
+
+def _stage_config(model_arch: str) -> dict[str, object]:
+ return {"engine_args": {"model_arch": model_arch}}
+
+
+def _chat_request(reference_audio_url: str) -> ChatCompletionRequest:
+ return ChatCompletionRequest(
+ model="test-model",
+ messages=[{"role": "user", "content": "target text"}],
+ ref_audio=reference_audio_url,
+ )
+
+
+@pytest.mark.asyncio
+async def test_chat_reference_audio_reaches_code2wav_transport(
+ serving_chat: OmniOpenAIServingChat,
+ mocker: MockerFixture,
+) -> None:
+ reference_audio_url = "data:audio/wav;base64,AAAA"
+ reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+ serving_chat.engine_client = _EngineClientStub(
+ [_stage_config("MiniCPMO45OmniForConditionalGeneration")],
+ )
+ serving_chat.model_config = _ModelConfigStub()
+ engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+ media_connector_class = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+ )
+ media_connector = media_connector_class.return_value
+ media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+ await serving_chat._attach_minicpmo45_reference_audio(
+ engine_prompt,
+ _chat_request(reference_audio_url),
+ )
+
+ media_connector_class.assert_called_once_with(
+ media_io_kwargs=None,
+ allowed_local_media_path="",
+ allowed_media_domains=None,
+ )
+ media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+ preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY]
+ assert isinstance(preserved_reference, tuple)
+ preserved_waveform, preserved_sample_rate = preserved_reference
+ assert isinstance(preserved_waveform, np.ndarray)
+ assert isinstance(preserved_sample_rate, int)
+ np.testing.assert_array_equal(preserved_waveform, reference_waveform)
+ assert preserved_sample_rate == 16000
+
+ thinker_output = _ThinkerRequestOutputStub(
+ request_id="request-0",
+ prompt_token_ids=[10],
+ outputs=[
+ _ThinkerCompletionStub(
+ token_ids=[20],
+ text="target text",
+ hidden_states=torch.zeros((2, 4)),
+ )
+ ],
+ )
+ talker_prompt = llm2tts([thinker_output], prompt=engine_prompt)[0]
+ talker_request = _Code2WavRequestStub(
+ external_req_id="request-0",
+ request_id="request-0",
+ model_intermediate_buffer=talker_prompt["model_intermediate_buffer"],
+ )
+ transfer_manager = _TransferManagerStub()
+
+ code2wav_payload = tts2code2wav_async_chunk(
+ transfer_manager,
+ {"codes": {"audio": torch.arange(7).reshape(-1, 1)}},
+ talker_request,
+ True,
+ )
+
+ assert code2wav_payload is not None
+ torch.testing.assert_close(code2wav_payload.codes.ref, torch.from_numpy(reference_waveform))
+ assert code2wav_payload.meta.ref_audio_sr == 16000
+
+
+@pytest.mark.parametrize("request_shape", ["extra_body", "model_extra"])
+@pytest.mark.asyncio
+async def test_reference_audio_accepts_request_extra_mappings(
+ serving_chat: OmniOpenAIServingChat,
+ mocker: MockerFixture,
+ request_shape: str,
+) -> None:
+ reference_audio_url = "data:audio/wav;base64,AAAA"
+ reference_waveform = np.array([0.25, -0.5], dtype=np.float32)
+ serving_chat.engine_client = _EngineClientStub(
+ [_stage_config("MiniCPMO45OmniForConditionalGeneration")],
+ )
+ serving_chat.model_config = _ModelConfigStub()
+ engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+ if request_shape == "extra_body":
+ request = ChatCompletionRequest(
+ model="test-model",
+ messages=[{"role": "user", "content": "target text"}],
+ extra_body={"ref_audio": reference_audio_url},
+ )
+ else:
+ request = SimpleNamespace(model_extra={"ref_audio": reference_audio_url})
+ assert not hasattr(request, "ref_audio")
+ media_connector = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+ ).return_value
+ media_connector.fetch_audio_async = mocker.AsyncMock(return_value=(reference_waveform, 16000))
+
+ await serving_chat._attach_minicpmo45_reference_audio(engine_prompt, request)
+
+ media_connector.fetch_audio_async.assert_awaited_once_with(reference_audio_url)
+ preserved_reference = cast(dict[str, object], engine_prompt)[MINICPMO45_REFERENCE_AUDIO_KEY]
+ assert isinstance(preserved_reference, tuple)
+ np.testing.assert_array_equal(preserved_reference[0], reference_waveform)
+ assert preserved_reference[1] == 16000
+
+
+@pytest.mark.asyncio
+async def test_reference_audio_is_ignored_for_other_models(
+ serving_chat: OmniOpenAIServingChat,
+ mocker: MockerFixture,
+) -> None:
+ reference_audio_url = "data:audio/wav;base64,AAAA"
+ serving_chat.engine_client = _EngineClientStub([_stage_config("SomeOtherOmniForConditionalGeneration")])
+ engine_prompt: TokensPrompt = {"type": "tokens", "prompt_token_ids": [1]}
+ media_connector_class = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_chat.MediaConnector",
+ )
+
+ await serving_chat._attach_minicpmo45_reference_audio(
+ engine_prompt,
+ _chat_request(reference_audio_url),
+ )
+
+ media_connector_class.assert_not_called()
+ assert MINICPMO45_REFERENCE_AUDIO_KEY not in cast(dict[str, object], engine_prompt)
diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py
index e7f014a3b..f81e61cf1 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py
@@ -14,6 +14,8 @@ from pytest_mock import MockerFixture
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.sampling_params import SamplingParams
+from tests.helpers.serving_chat import build_serving_chat
+
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -77,18 +79,8 @@ def mock_engine_client(
return engine_client
-@pytest.fixture
-def serving_chat(mock_engine_client):
- """Create OmniOpenAIServingChat instance with mocked dependencies."""
- from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
- # Create instance without calling __init__
- instance = object.__new__(OmniOpenAIServingChat)
- instance.engine_client = mock_engine_client
- return instance
-
-
-def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_chat):
+def test_serving_boundary_normalizes_declared_root_and_nested_extras(mock_engine_client):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_extra_body_params = frozenset({"cfg_text_scale", "negative_prompt"})
request = ChatCompletionRequest(
model="test",
@@ -118,7 +110,8 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
assert diffusion_request_args["modalities"] == ["image"]
-def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
+def test_unknown_root_extra_does_not_claim_canonical_extra(mock_engine_client):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_extra_body_params = frozenset()
request = ChatCompletionRequest(
model="test",
@@ -135,7 +128,8 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
assert "pipeline_option" not in diffusion_request_args
-def test_serving_boundary_rejects_invalid_quality(serving_chat):
+def test_serving_boundary_rejects_invalid_quality(mock_engine_client):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_extra_body_params = frozenset()
request = ChatCompletionRequest(
model="test",
@@ -148,7 +142,8 @@ def test_serving_boundary_rejects_invalid_quality(serving_chat):
serving_chat._normalize_diffusion_request_args(request)
-def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
+def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(mock_engine_client):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_extra_body_params = frozenset()
request = ChatCompletionRequest(
model="test",
@@ -170,7 +165,8 @@ def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
],
ids=["registered-model-extra", "unregistered-common-alias"],
)
-def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, request_kwargs):
+def test_cfg_scale_owner_conflicts_are_rejected(mock_engine_client, registered, request_kwargs):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_extra_body_params = frozenset(registered)
request = ChatCompletionRequest(model="test", messages=[], **request_kwargs)
@@ -180,10 +176,11 @@ def test_cfg_scale_owner_conflicts_are_rejected(serving_chat, registered, reques
@pytest.mark.parametrize("diffusion_mode", [True, False], ids=["pure", "mixed"])
def test_duplicate_extras_return_the_same_bad_request_before_dispatch(
- serving_chat,
+ mock_engine_client,
mocker: MockerFixture,
diffusion_mode: bool,
):
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
serving_chat._diffusion_mode = diffusion_mode
serving_chat._diffusion_extra_body_params = frozenset({"sample_solver"})
serving_chat.engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
@@ -211,9 +208,10 @@ def test_duplicate_extras_return_the_same_bad_request_before_dispatch(
pure_dispatch.assert_not_awaited()
-def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat, mocker: MockerFixture):
+def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(mock_engine_client, mocker: MockerFixture):
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
captured: dict[str, object] = {}
async def generate(**kwargs):
@@ -261,19 +259,20 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
assert response.error.message == "No output generated from AsyncOmni"
-def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat, mocker: MockerFixture):
+def test_mixed_consumer_keeps_root_common_args_with_nested_extras(mock_engine_client, mocker: MockerFixture):
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
- serving_chat.engine_client.stage_configs = [
+ mock_engine_client.stage_configs = [
SimpleNamespace(stage_type="llm", is_comprehension=True),
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
]
- serving_chat.engine_client.default_sampling_params_list = [
+ mock_engine_client.default_sampling_params_list = [
SamplingParams(),
OmniDiffusionSamplingParams(),
]
- serving_chat.engine_client.output_modalities = ["image"]
- serving_chat.engine_client.errored = False
+ mock_engine_client.output_modalities = ["image"]
+ mock_engine_client.errored = False
+ mock_engine_client.renderer = SimpleNamespace(get_tokenizer=lambda: object())
captured: dict[str, object] = {}
@@ -285,21 +284,15 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat,
captured.update(kwargs)
return results()
- serving_chat.engine_client.generate = generate
- serving_chat.__dict__.update(
- _diffusion_mode=False,
- _diffusion_extra_body_params=frozenset(),
+ mock_engine_client.generate = generate
+ serving_chat = build_serving_chat(
+ engine_client=mock_engine_client,
models=SimpleNamespace(model_name=lambda _: "test"),
- renderer=SimpleNamespace(get_tokenizer=lambda: object()),
online_renderer=SimpleNamespace(validate_chat_template=lambda **_: None),
- parser_cls=None,
- use_harmony=False,
- enable_auto_tools=False,
- exclude_tools_when_tool_choice_none=False,
trust_request_chat_template=True,
- chat_template=None,
- chat_template_content_format="auto",
)
+ serving_chat._diffusion_mode = False
+ serving_chat._diffusion_extra_body_params = frozenset()
mocker.patch.multiple(
serving_chat,
_check_model=mocker.AsyncMock(return_value=None),
@@ -386,8 +379,9 @@ def test_openai_sampling_fields_contains_expected_fields():
# =============================================================================
-def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_request):
+def test_preserves_yaml_defaults_when_no_request_params(mock_engine_client, mock_request):
"""Test that YAML defaults are preserved when request has no params."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
result = serving_chat._build_sampling_params_list_from_request(mock_request)
assert len(result) == 2
@@ -400,8 +394,9 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque
assert comprehension_params.repetition_penalty == 1.05 # YAML custom param preserved
-def test_request_temperature_overrides_yaml_default(serving_chat, mock_request):
+def test_request_temperature_overrides_yaml_default(mock_engine_client, mock_request):
"""Test that request temperature overrides YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.temperature = 0.8
mock_request.model_fields_set = {"temperature"}
@@ -413,8 +408,9 @@ def test_request_temperature_overrides_yaml_default(serving_chat, mock_request):
assert comprehension_params.top_k == 1 # YAML custom param preserved
-def test_request_top_p_overrides_yaml_default(serving_chat, mock_request):
+def test_request_top_p_overrides_yaml_default(mock_engine_client, mock_request):
"""Test that request top_p overrides YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.top_p = 0.95
mock_request.model_fields_set = {"top_p"}
@@ -425,8 +421,9 @@ def test_request_top_p_overrides_yaml_default(serving_chat, mock_request):
assert comprehension_params.temperature == 0.4 # Preserved from YAML
-def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request):
+def test_request_max_tokens_overrides_yaml_default(mock_engine_client, mock_request):
"""Test that request max_tokens overrides YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.max_tokens = 100
mock_request.model_fields_set = {"max_tokens"}
@@ -435,15 +432,17 @@ def test_request_max_tokens_overrides_yaml_default(serving_chat, mock_request):
assert result[0].max_tokens == 100
-def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_request):
+def test_max_tokens_uses_yaml_default_when_not_specified(mock_engine_client, mock_request):
"""Test that max_tokens falls back to YAML default when not in request."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
result = serving_chat._build_sampling_params_list_from_request(mock_request)
assert result[0].max_tokens == 4353
-def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
+def test_request_seed_overrides_yaml_default(mock_engine_client, mock_request):
"""Test that request seed overrides YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.seed = 123
mock_request.model_fields_set = {"seed"}
@@ -454,8 +453,9 @@ def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
assert comprehension_params.temperature == 0.4 # Preserved from YAML
-def test_request_frequency_penalty_overrides(serving_chat, mock_request):
+def test_request_frequency_penalty_overrides(mock_engine_client, mock_request):
"""Test that request frequency_penalty is applied."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.frequency_penalty = 0.5
mock_request.model_fields_set = {"frequency_penalty"}
@@ -464,8 +464,9 @@ def test_request_frequency_penalty_overrides(serving_chat, mock_request):
assert result[0].frequency_penalty == 0.5
-def test_request_presence_penalty_overrides(serving_chat, mock_request):
+def test_request_presence_penalty_overrides(mock_engine_client, mock_request):
"""Test that request presence_penalty is applied."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.presence_penalty = 0.3
mock_request.model_fields_set = {"presence_penalty"}
@@ -474,8 +475,9 @@ def test_request_presence_penalty_overrides(serving_chat, mock_request):
assert result[0].presence_penalty == 0.3
-def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request):
+def test_non_comprehension_stages_use_cloned_defaults(mock_engine_client, mock_request):
"""Test that non-comprehension stages always use cloned YAML defaults."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.max_tokens = 50
mock_request.temperature = 0.1
@@ -488,8 +490,9 @@ def test_non_comprehension_stages_use_cloned_defaults(serving_chat, mock_request
assert other_params.seed == 42 # YAML default
-def test_multiple_params_override_together(serving_chat, mock_request):
+def test_multiple_params_override_together(mock_engine_client, mock_request):
"""Test that multiple request params can override together."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.max_tokens = 200
mock_request.temperature = 0.7
mock_request.top_p = 0.85
@@ -514,15 +517,17 @@ def test_multiple_params_override_together(serving_chat, mock_request):
# =============================================================================
-def test_apply_request_overrides_clones_params(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_clones_params(mock_engine_client, mock_request, default_comprehension_params):
"""Test that _apply_request_overrides returns a cloned object."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request)
assert result is not default_comprehension_params # Different object
-def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_preserves_defaults(mock_engine_client, mock_request, default_comprehension_params):
"""Test that _apply_request_overrides preserves defaults when request has None."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
result = serving_chat._apply_request_overrides(default_comprehension_params, mock_request)
assert result.temperature == 0.4
@@ -531,8 +536,9 @@ def test_apply_request_overrides_preserves_defaults(serving_chat, mock_request,
assert result.top_k == 1 # YAML custom param
-def test_apply_request_overrides_applies_values(serving_chat, mock_request, default_comprehension_params):
+def test_apply_request_overrides_applies_values(mock_engine_client, mock_request, default_comprehension_params):
"""Test that _apply_request_overrides applies non-None request values."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.temperature = 0.8
mock_request.seed = 123
mock_request.model_fields_set = {"temperature", "seed"}
@@ -550,8 +556,9 @@ def test_apply_request_overrides_applies_values(serving_chat, mock_request, defa
# =============================================================================
-def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker):
+def test_apply_overrides_empty_stop_list_preserves_default(mock_engine_client, mocker):
"""Test that request.stop=[] does NOT override YAML default stop words."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"])
request = mocker.MagicMock()
request.temperature = None
@@ -573,8 +580,9 @@ def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker)
assert result.stop == ["<|im_end|>"] # YAML default preserved
-def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mocker):
+def test_apply_overrides_nonempty_stop_list_overrides_default(mock_engine_client, mocker):
"""Test that request.stop=["\\n"] overrides YAML default stop words."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"])
request = mocker.MagicMock()
request.temperature = None
@@ -596,8 +604,9 @@ def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mock
assert result.stop == ["\n"] # Overridden by request
-def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mocker):
+def test_apply_overrides_empty_stop_token_ids_preserves_default(mock_engine_client, mocker):
"""Test that request.stop_token_ids=[] does NOT override YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3])
request = mocker.MagicMock()
request.temperature = None
@@ -617,8 +626,9 @@ def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mo
assert result.stop_token_ids == [2, 3] # YAML default preserved
-def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, mocker):
+def test_apply_overrides_nonempty_stop_token_ids_overrides_default(mock_engine_client, mocker):
"""Test that request.stop_token_ids=[100] overrides YAML default."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3])
request = mocker.MagicMock()
request.temperature = None
@@ -640,8 +650,9 @@ def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat,
assert result.stop_token_ids == [100] # Overridden by request
-def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker):
+def test_apply_overrides_mixed_empty_and_nonempty_lists(mock_engine_client, mocker):
"""Test mixing empty and non-empty list fields with scalar fields."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(
temperature=0.4,
stop=["<|end|>"],
@@ -669,8 +680,9 @@ def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker):
assert result.stop_token_ids == [100, 200] # Non-empty list DID override
-def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocker):
+def test_apply_overrides_none_scalar_still_preserves_default(mock_engine_client, mocker):
"""Regression: ensure None scalar values still don't override defaults."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(temperature=0.5, max_tokens=100, seed=42)
request = mocker.MagicMock()
request.temperature = None
@@ -694,8 +706,9 @@ def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocke
assert result.seed == 42
-def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocker):
+def test_apply_overrides_both_lists_empty_preserves_defaults(mock_engine_client, mocker):
"""Test that both stop=[] and stop_token_ids=[] preserve YAML defaults."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
default_params = SamplingParams(
temperature=0.5,
stop=["<|end|>", "\\n"],
@@ -722,9 +735,10 @@ def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocke
assert result.stop_token_ids == [2, 32000]
-def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock_request):
+def test_build_sampling_params_list_empty_stop_preserves_yaml(mock_engine_client, mock_request):
"""Test that empty stop list in request preserves YAML defaults via
_build_sampling_params_list_from_request."""
+ serving_chat = build_serving_chat(engine_client=mock_engine_client)
mock_request.stop = []
mock_request.stop_token_ids = []
@@ -736,21 +750,18 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock
assert comprehension_params.stop_token_ids == []
-def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults():
+def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(mocker: MockerFixture):
"""AURA callers may pass 3 semantic model params for a 4-stage engine pipeline."""
- from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
- instance = object.__new__(OmniOpenAIServingChat)
default_params = [
SamplingParams(max_tokens=10),
SamplingParams(max_tokens=20),
SamplingParams(max_tokens=30),
SamplingParams(max_tokens=40),
]
- instance.engine_client = SimpleNamespace(
- stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)],
- default_sampling_params_list=default_params,
- )
+ engine_client = mocker.MagicMock()
+ engine_client.stage_configs = [SimpleNamespace(stage_type="llm") for _ in range(4)]
+ engine_client.default_sampling_params_list = default_params
+ instance = build_serving_chat(engine_client=engine_client)
result = instance._to_sampling_params_list(
[
@@ -772,44 +783,35 @@ def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults():
def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client):
"""Test finding comprehension stage when it's at index 0."""
- from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
- instance = object.__new__(OmniOpenAIServingChat)
- instance.engine_client = mock_engine_client
+ instance = build_serving_chat(engine_client=mock_engine_client)
assert instance._get_comprehension_stage_index() == 0
def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture):
"""Test finding comprehension stage when it's at index 1."""
- from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
- instance = object.__new__(OmniOpenAIServingChat)
-
other = mocker.MagicMock()
other.is_comprehension = False
comprehension = mocker.MagicMock()
comprehension.is_comprehension = True
- instance.engine_client = mocker.MagicMock()
- instance.engine_client.stage_configs = [other, comprehension]
+ engine_client = mocker.MagicMock()
+ engine_client.stage_configs = [other, comprehension]
+ instance = build_serving_chat(engine_client=engine_client)
assert instance._get_comprehension_stage_index() == 1
def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture):
"""Test that ValueError is raised when no comprehension stage exists."""
- from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
- instance = object.__new__(OmniOpenAIServingChat)
-
stage1 = mocker.MagicMock()
stage1.is_comprehension = False
stage2 = mocker.MagicMock()
stage2.is_comprehension = False
- instance.engine_client = mocker.MagicMock()
- instance.engine_client.stage_configs = [stage1, stage2]
+ engine_client = mocker.MagicMock()
+ engine_client.stage_configs = [stage1, stage2]
+ instance = build_serving_chat(engine_client=engine_client)
with pytest.raises(ValueError, match="No comprehension stage"):
instance._get_comprehension_stage_index()
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index 94d84d041..19d383504 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/tests/entrypoints/openai_api/test_serving_speech.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
# tests/entrypoints/openai/test_serving_speech.py
import asyncio
import base64
@@ -8,6 +11,7 @@ import logging
import os
import struct
import wave
+from dataclasses import FrozenInstanceError, replace
from inspect import Signature, signature
from pathlib import Path
from types import SimpleNamespace
@@ -33,15 +37,24 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
OpenAICreateAudioGenerateRequest,
OpenAICreateSpeechRequest,
SpeechBatchItem,
+ SpeechInputTokenDetails,
+ SpeechTokenUsage,
StreamingSpeechSessionConfig,
)
from vllm_omni.entrypoints.openai.serving_speech import (
- _TTS_LANGUAGES,
OmniOpenAIServingSpeech,
_create_wav_header,
)
-from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+ DEFAULT_TTS_LANGUAGES,
+ PreparedRequest,
+ SpeechServingContext,
+ TTSGenerationError,
+)
+from vllm_omni.entrypoints.openai.tts_adapters.capabilities import load_supported_speakers
from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
+from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
FISH_TEXT_ONLY_SYSTEM_PROMPT,
build_fish_voice_clone_prompt_ids,
@@ -52,6 +65,7 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
KEY_SPEAKER_WAVEFORM,
KEY_SPEAKER_WAVEFORM_LENGTHS,
SPEAKER_EMBEDDING_DIM,
+ TEXT_EOS_TOKEN_ID,
)
from vllm_omni.outputs import OmniRequestOutput
@@ -154,6 +168,7 @@ class TestAudioMixin:
# Helper to create mock model output for endpoint tests
def create_mock_audio_output_for_test(
request_id: str = "speech-mock-123",
+ metrics: dict | None = None,
) -> OmniRequestOutput:
class MockCompletionOutput:
def __init__(self, index: int = 0):
@@ -180,11 +195,16 @@ def create_mock_audio_output_for_test(
audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
- return OmniRequestOutput.from_stage_output(
+ if metrics is None:
+ metrics = {"stage_metrics": {"0": {"num_tokens_in": 0, "num_tokens_out": 5}}}
+
+ output = OmniRequestOutput.from_stage_output(
mock_request_output,
stage_id=0,
final_output_type="audio",
)
+ output.metrics = metrics
+ return output
def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
@@ -209,7 +229,8 @@ def _wav_data_url(samples: np.ndarray, sample_rate: int) -> str:
@pytest.fixture
-def test_app(mocker: MockerFixture):
+def test_app(mocker: MockerFixture, tmp_path, monkeypatch):
+ monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path))
# Mock the engine client
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
@@ -232,6 +253,12 @@ def test_app(mocker: MockerFixture):
models=mock_models,
request_logger=mock_request_logger,
)
+ # Voice endpoint tests exercise generic embedding parsing/persistence with
+ # multiple valid dimensions. Supply the adapter contract without imposing
+ # one real model's fixed embedding dimension on this shared fixture.
+ speech_server._adapter = mocker.MagicMock()
+ speech_server._adapter.validate_tts_embedding_dim.return_value = None
+ speech_server._tts_tokenizer = lambda text, padding=False: {"input_ids": text.split()}
# Skip TTS validation in tests (mock doesn't set up supported_speakers)
speech_server._validate_tts_request = mocker.MagicMock(return_value=None)
@@ -253,31 +280,9 @@ def test_app(mocker: MockerFixture):
speech_server.create_speech = awaitable_patched_create_speech
app = FastAPI()
+ app.state.openai_serving_speech = speech_server
app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
-
- # Add list_voices endpoint
- async def list_voices():
- speakers = sorted(speech_server.supported_speakers) if speech_server.supported_speakers else []
- uploaded_voices = []
- if hasattr(speech_server, "uploaded_speakers"):
- for voice_name, info in speech_server.uploaded_speakers.items():
- voice_entry = {
- "name": info.get("name", voice_name),
- "consent": info.get("consent", ""),
- "created_at": info.get("created_at", 0),
- "file_size": info.get("file_size", 0),
- "mime_type": info.get("mime_type", ""),
- "embedding_source": info.get("embedding_source", "audio"),
- "embedding_dim": info.get("embedding_dim"),
- }
- if info.get("ref_text"):
- voice_entry["ref_text"] = info["ref_text"]
- if info.get("speaker_description"):
- voice_entry["speaker_description"] = info["speaker_description"]
- uploaded_voices.append(voice_entry)
- return {"voices": speakers, "uploaded_voices": uploaded_voices}
-
- app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"])
+ app.add_api_route("/v1/audio/voices", api_server_module.list_voices, methods=["GET"])
app.add_api_route("/v1/audio/speech/batch", speech_server.create_speech_batch, methods=["POST"])
# Add upload_voice endpoint
@@ -369,6 +374,37 @@ class TestSpeechAPI:
assert response.headers["content-type"] == "audio/wav"
assert len(response.content) > 0
+ def test_create_speech_includes_token_usage_and_detail_headers(self, client):
+ payload = {
+ "input": "Hello world",
+ "model": "tts-model",
+ "voice": "alloy",
+ "response_format": "wav",
+ }
+ response = client.post("/v1/audio/speech", json=payload)
+ assert response.status_code == 200
+ assert response.headers["x-vllm-omni-input-tokens"] == "2"
+ assert response.headers["x-vllm-omni-output-tokens"] == "5"
+ assert response.headers["x-vllm-omni-total-tokens"] == "7"
+ assert response.headers["x-vllm-omni-input-text-tokens"] == "2"
+ assert response.headers["x-vllm-omni-input-audio-tokens"] == "0"
+
+ def test_build_speech_usage_headers_uses_usage_and_detail_field_names(self):
+ usage = SpeechTokenUsage(
+ input_tokens=7,
+ output_tokens=11,
+ total_tokens=18,
+ input_token_details=SpeechInputTokenDetails(text_tokens=3, audio_tokens=4),
+ )
+ headers = OmniOpenAIServingSpeech._build_speech_usage_headers(usage)
+ assert headers == {
+ "X-VLLM-OMNI-INPUT-TOKENS": "7",
+ "X-VLLM-OMNI-OUTPUT-TOKENS": "11",
+ "X-VLLM-OMNI-TOTAL-TOKENS": "18",
+ "X-VLLM-OMNI-INPUT-TEXT-TOKENS": "3",
+ "X-VLLM-OMNI-INPUT-AUDIO-TOKENS": "4",
+ }
+
def test_create_speech_mp3_format(self, client):
payload = {
"input": "Hello world",
@@ -418,9 +454,31 @@ class TestSpeechAPI:
assert audio_obj.speed == 2.5
def test_list_voices_endpoint(self, client):
+ handler = client.app.state.openai_serving_speech
+ handler._adapter.capabilities.supported_speakers = frozenset({"alice"})
+ handler._adapter.capabilities.precomputed_speakers = {"bob": {"name": "Bob"}}
+ handler.uploaded_speakers = {
+ "carol": {
+ "name": "Carol",
+ "created_at": 1,
+ }
+ }
+
response = client.get("/v1/audio/voices")
+
assert response.status_code == 200
- assert "voices" in response.json()
+ assert response.json()["voices"] == ["alice", "bob", "carol", "default"]
+ assert response.json()["uploaded_voices"] == [
+ {
+ "name": "Carol",
+ "consent": "",
+ "created_at": 1,
+ "file_size": 0,
+ "mime_type": "",
+ "embedding_source": "audio",
+ "embedding_dim": None,
+ }
+ ]
def test_upload_voice_success(self, client, tmp_path):
"""Test successful voice upload without ref_text."""
@@ -678,6 +736,20 @@ class TestSpeechAPI:
assert response.status_code == 400
assert "finite" in response.json()["detail"]
+ @pytest.mark.asyncio
+ async def test_diffusion_create_speech_with_unknown_voice(self, mocker: MockerFixture):
+ engine_client = mocker.MagicMock()
+ server = OmniOpenAIServingSpeech.for_diffusion(
+ diffusion_engine=engine_client,
+ model_name="test-model",
+ )
+ assert server._adapter is None
+
+ response = await server.create_speech(OpenAICreateSpeechRequest(input="test-input", voice="test-voice"))
+
+ assert response.status_code == 400
+ assert b"Invalid voice" in response.body
+
@pytest.mark.asyncio
async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
"""Test public diffusion speech success and extra_params propagation."""
@@ -726,7 +798,8 @@ class TestTTSMethods:
"""Unit tests for TTS validation and parameter building."""
@pytest.fixture
- def speech_server(self, mocker: MockerFixture):
+ def speech_server(self, mocker, tmp_path, monkeypatch):
+ monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path))
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_configs = []
@@ -949,12 +1022,18 @@ class TestTTSMethods:
def test_speaker_embedding_valid_base_task(self, speech_server):
"""speaker_embedding with Base task, x_vector_only_mode, and no ref_audio is accepted."""
+ speech_server.engine_client.model_config = SimpleNamespace(
+ hf_config=SimpleNamespace(talker_config=SimpleNamespace(hidden_size=1024))
+ )
emb = [0.1] * 1024
req = OpenAICreateSpeechRequest(input="Hello", task_type="Base", speaker_embedding=emb, x_vector_only_mode=True)
assert speech_server._validate_tts_request(req) is None
def test_speaker_embedding_auto_sets_x_vector_only_mode(self, speech_server):
"""speaker_embedding auto-implies x_vector_only_mode, so validation passes."""
+ speech_server.engine_client.model_config = SimpleNamespace(
+ hf_config=SimpleNamespace(talker_config=SimpleNamespace(hidden_size=1024))
+ )
emb = [0.1] * 1024
req = OpenAICreateSpeechRequest(input="Hello", task_type="Base", speaker_embedding=emb)
result = speech_server._validate_tts_request(req)
@@ -1026,6 +1105,7 @@ class TestTTSMethods:
"""Embedding uploads must match the loaded Qwen3-TTS model before being stored."""
speech_server._tts_model_type = "qwen3_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(
talker_config=SimpleNamespace(hidden_size=2048),
@@ -1044,6 +1124,7 @@ class TestTTSMethods:
def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
"""Ming embedding uploads must match the model dimension before replacement."""
speech_server._tts_model_type = "ming_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
invalid_dim = SPEAKER_EMBEDDING_DIM - 1
existing = {
"name": "bad_emb_voice",
@@ -1071,9 +1152,13 @@ class TestTTSMethods:
def test_upload_ming_audio_voice_defers_speaker_extraction(self, speech_server, mocker: MockerFixture, tmp_path):
"""Ming audio uploads stay model-agnostic and defer extraction to use."""
speech_server._tts_model_type = "ming_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.uploaded_speakers_dir = tmp_path
speech_server.uploaded_speakers = {}
- speech_server.supported_speakers = set()
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ supported_speakers=frozenset(),
+ )
speech_server._speaker_cache.clear()
left = np.full(32000, 0.5, dtype=np.float32)
@@ -1161,23 +1246,421 @@ class TestTTSMethods:
assert second[1] == 24000
assert first[0] is second[0]
assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4)
+ cache_key = first[2]
assert speech_server._get_resolved_ref_audio_artifact_key(
- ref_audio
+ cache_key
) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
+ @pytest.mark.asyncio
+ async def test_diffusion_ref_audio_uses_media_access_config(
+ self,
+ mocker: MockerFixture,
+ monkeypatch,
+ tmp_path,
+ ):
+ monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path / "speakers"))
+ connector = mocker.MagicMock()
+ connector.fetch_audio_async = mocker.AsyncMock(return_value=(np.zeros(24000, dtype=np.float32), 24000))
+ connector_cls = mocker.patch.object(
+ serving_speech_module,
+ "MediaConnector",
+ return_value=connector,
+ )
+ media_dir = tmp_path / "media"
+ media_dir.mkdir()
+ ref_audio = media_dir / "reference.wav"
+ ref_audio.write_bytes(b"v1")
+
+ server = OmniOpenAIServingSpeech.for_diffusion(
+ diffusion_engine=mocker.MagicMock(),
+ model_name="test-model",
+ allowed_local_media_path=str(media_dir),
+ allowed_media_domains=["media.example.com"],
+ )
+ try:
+ first = await server._resolve_ref_audio(ref_audio.as_uri())
+ ref_audio.write_bytes(b"v2-with-different-size")
+ second = await server._resolve_ref_audio(ref_audio.as_uri())
+ finally:
+ server.shutdown()
+
+ connector_cls.assert_called_once_with(
+ allowed_local_media_path=str(media_dir),
+ allowed_media_domains=["media.example.com"],
+ )
+ assert first[2] != second[2]
+ assert connector.fetch_audio_async.await_count == 2
+
+ # ── ref-audio cache key tests (static helper) ──
+
+ def test_local_file_cache_key_invalidation(self):
+ """Modified file must produce a different cache key (mtime change)."""
+ import pathlib
+ import tempfile
+ import time as _time
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ test_file = pathlib.Path(tmp_dir) / "test_audio.wav"
+ test_file.write_text("dummy audio data 1")
+ file_uri = f"file://{test_file}"
+ key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ file_uri,
+ allowed_local_media_path=tmp_dir,
+ )
+ _time.sleep(0.01)
+ test_file.write_text("dummy audio data 2")
+ key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ file_uri,
+ allowed_local_media_path=tmp_dir,
+ )
+ assert key1 != key2, "Cache key should change when file is modified!"
+
+ def test_remote_url_cache_key(self):
+ """Remote URLs must produce a stable, repeatable cache key."""
+ url = "https://example.com/audio.wav"
+ key1 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+ key2 = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+ assert key1 == key2, "Cache key for URLs should remain constant!"
+
+ def test_missing_file_cache_key(self, caplog):
+ """A missing file falls back to the URI-string hash and emits a log
+ so the stale-cache risk is visible."""
+ file_uri = "file:///path/to/nonexistent/file.wav"
+ target_logger = logging.getLogger("vllm_omni.entrypoints.openai.serving_speech")
+ target_logger.addHandler(caplog.handler)
+ prev_level = target_logger.level
+ target_logger.setLevel(logging.DEBUG)
+ try:
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ file_uri,
+ allowed_local_media_path="/path/to/nonexistent",
+ )
+ finally:
+ target_logger.removeHandler(caplog.handler)
+ target_logger.setLevel(prev_level)
+ expected_key = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+ assert key == expected_key, "Missing files should fallback to the string hash."
+ assert any("stale cache" in r.getMessage() for r in caplog.records), (
+ "A log message about stale cache must be emitted when os.stat fails"
+ )
+
+ def test_percent_encoded_file_uri(self):
+ """Percent-encoded file:// URIs must be decoded so os.stat hits the real file."""
+ import pathlib
+ import tempfile
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ test_file = pathlib.Path(tmp_dir) / "my audio.wav"
+ test_file.write_text("dummy audio data")
+ encoded_uri = f"file://{str(test_file).replace(' ', '%20')}"
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ encoded_uri,
+ allowed_local_media_path=tmp_dir,
+ )
+ string_only_key = hashlib.sha1(encoded_uri.encode("utf-8")).hexdigest()
+ assert key != string_only_key, (
+ "Percent-encoded URI should be decoded and stat'd, not fall back to string-only cache key!"
+ )
+
+ def test_no_allowlist_skips_stat_for_local_paths(self):
+ """When allowed_local_media_path is None, local paths must NOT be stat'd
+ — the key falls back to the string-only hash."""
+ import pathlib
+ import tempfile
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ test_file = pathlib.Path(tmp_dir) / "test.wav"
+ test_file.write_text("data")
+ file_uri = f"file://{test_file}"
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ file_uri,
+ allowed_local_media_path=None,
+ )
+ # With no allowlist the stat is skipped, so the key is just sha1(uri)
+ expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+ assert key == expected, "No allowlist should skip stat and use string-only key"
+
+ def test_empty_string_allowlist_skips_stat_for_local_paths(self):
+ """vLLM's default allowed_local_media_path is '' not None; that must
+ also skip stat rather than treating cwd as the allowlist."""
+ import pathlib
+ import tempfile
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ test_file = pathlib.Path(tmp_dir) / "test.wav"
+ test_file.write_text("data")
+ file_uri = f"file://{test_file}"
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ file_uri,
+ allowed_local_media_path="",
+ )
+ expected = hashlib.sha1(file_uri.encode("utf-8")).hexdigest()
+ assert key == expected, "Empty-string allowlist should skip stat and use string-only key"
+
+ def test_uppercase_file_uri_uses_local_metadata(self):
+ """FILE:// must be treated as a local file, not a string-only locator."""
+ import pathlib
+ import tempfile
+
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ test_file = pathlib.Path(tmp_dir) / "test.wav"
+ test_file.write_text("data")
+ upper_uri = f"FILE://{test_file}"
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(
+ upper_uri,
+ allowed_local_media_path=tmp_dir,
+ )
+ string_only = hashlib.sha1(upper_uri.encode("utf-8")).hexdigest()
+ assert key != string_only
+
+ def test_uppercase_http_uri_is_remote(self, speech_server):
+ url = "HTTP://example.com/audio.wav"
+ key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(url)
+ assert key == hashlib.sha1(url.encode("utf-8")).hexdigest()
+ assert speech_server._validate_ref_audio_format(url) is None
+ assert speech_server._validate_ref_audio_format("FILE:///tmp/a.wav") is None
+ assert speech_server._validate_ref_audio_format("/tmp/a.wav") is not None
+
+ @pytest.mark.asyncio
+ async def test_resolve_ref_audio_retries_when_file_changes_during_fetch(self, speech_server, mocker):
+ import pathlib
+ import tempfile
+
+ wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32)
+ with tempfile.TemporaryDirectory() as tmp_dir:
+ path = pathlib.Path(tmp_dir) / "spk.wav"
+ path.write_bytes(b"v1")
+ uri = f"file://{path}"
+ speech_server.model_config.allowed_local_media_path = tmp_dir
+ speech_server.model_config.allowed_media_domains = None
+ speech_server._diffusion_mode = False
+
+ fetches = {"n": 0}
+
+ async def fake_fetch(_ref):
+ fetches["n"] += 1
+ if fetches["n"] == 1:
+ path.write_bytes(b"v2-longer-content")
+ return wav, 24000
+
+ mock_connector = mocker.MagicMock()
+ mock_connector.fetch_audio_async = fake_fetch
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_speech.MediaConnector",
+ return_value=mock_connector,
+ )
+
+ wav_list, sr, cache_key = await speech_server._resolve_ref_audio(uri)
+ expected_key = OmniOpenAIServingSpeech._get_ref_audio_cache_key(uri, allowed_local_media_path=tmp_dir)
+ assert sr == 24000
+ assert len(wav_list) == 48000
+ assert fetches["n"] == 2
+ assert cache_key == expected_key
+ assert cache_key in speech_server._ref_audio_resolve_cache
+
+ def test_conditioning_cache_salt_changes_with_ref_audio_cache_key(self):
+ """Prefix-cache salt must fold the content-aware resolve key so a
+ same-path file rewrite cannot reuse KV from the previous reference."""
+ from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt
+
+ req = OpenAICreateSpeechRequest(input="hello", ref_audio="file:///data/spk.wav")
+ salt_a = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_aaa"})
+ salt_b = _conditioning_cache_salt(req, {"ref_audio_cache_key": "key_bbb"})
+ salt_locator_only = _conditioning_cache_salt(req, {})
+ assert salt_a != salt_b
+ assert salt_a != salt_locator_only
+
+ @pytest.mark.asyncio
+ async def test_higgs_v3_cache_salt_changes_with_ref_audio_cache_key(self, speech_server, mocker):
+ """Higgs v3 voice clone uses placeholder prompt ids + prefix caching.
+ The salt must move when the resolve key moves, or a same-path rewrite
+ reuses KV from the previous reference."""
+ from unittest.mock import AsyncMock, MagicMock
+
+ adapter = MagicMock()
+ adapter.build_prompt.return_value = [1, 1, 1, 1]
+ speech_server._resolve_higgs_audio_v3_adapter = AsyncMock(return_value=adapter)
+ codes = torch.arange(8, dtype=torch.long).view(8, 1)
+ speech_server._resolve_higgs_audio_v3_ref_codes = AsyncMock(return_value=(codes, False, False))
+ speech_server._get_resolved_ref_audio_artifact_key = MagicMock(return_value="art")
+
+ req = OpenAICreateSpeechRequest(
+ input="hello",
+ ref_audio="file:///data/spk.wav",
+ ref_text="transcript",
+ )
+ speech_server._resolve_ref_audio = AsyncMock(return_value=([0.1] * 48000, 24000, "key_aaa"))
+ prompt_a = await speech_server._build_higgs_audio_v3_params(req)
+ speech_server._resolve_ref_audio = AsyncMock(return_value=([0.9] * 48000, 24000, "key_bbb"))
+ prompt_b = await speech_server._build_higgs_audio_v3_params(req)
+
+ assert prompt_a["prompt_token_ids"] == prompt_b["prompt_token_ids"]
+ assert prompt_a["cache_salt"] != prompt_b["cache_salt"]
+ assert prompt_a["additional_information"]["ref_audio_cache_key"] == "key_aaa"
+ assert prompt_b["additional_information"]["ref_audio_cache_key"] == "key_bbb"
+
+ # ── encode_reference_codes: speaker cache invalidation ──
+
+ @pytest.mark.asyncio
+ async def test_encode_reference_codes_reencodes_on_key_change(self):
+ """Changing the resolve_cache_key must trigger a fresh encode, not serve
+ the stale cached tensor. This is the user-visible fix for delay-family
+ voice clones that were reproducing the old speaker."""
+ from unittest.mock import MagicMock
+
+ from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes
+ from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+ speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024)
+
+ call_count = 0
+ codes_v1 = torch.arange(10, dtype=torch.int64)
+ codes_v2 = torch.arange(10, 20, dtype=torch.int64)
+
+ # Stub processor that tracks call count and returns different codes
+ processor = MagicMock()
+
+ def fake_encode(wavs, sampling_rate, n_vq):
+ nonlocal call_count
+ call_count += 1
+ return [codes_v1 if call_count == 1 else codes_v2]
+
+ processor.encode_audios_from_wav = fake_encode
+
+ # resolve_ref_audio returns different cache keys on successive calls
+ resolve_keys = iter(["key_aaa", "key_bbb"])
+
+ async def mock_resolve(ref_str):
+ return ([0.1, 0.2, 0.3], 24000, next(resolve_keys))
+
+ # First call: cold miss, encodes and stores
+ result1 = await encode_reference_codes(
+ "data:audio/wav;base64,fake",
+ processor=processor,
+ resolve_ref_audio=mock_resolve,
+ speaker_cache=speaker_cache,
+ variant="tts",
+ n_vq=32,
+ sr_target=24000,
+ )
+ assert call_count == 1
+ assert torch.equal(result1, codes_v1)
+
+ # Second call: different resolve key → must re-encode, not serve cached
+ result2 = await encode_reference_codes(
+ "data:audio/wav;base64,fake",
+ processor=processor,
+ resolve_ref_audio=mock_resolve,
+ speaker_cache=speaker_cache,
+ variant="tts",
+ n_vq=32,
+ sr_target=24000,
+ )
+ assert call_count == 2, "Different resolve key should trigger re-encode"
+ assert torch.equal(result2, codes_v2)
+
+ @pytest.mark.asyncio
+ async def test_encode_reference_codes_named_voice_skips_resolve_on_hit(self):
+ """Named-voice cache hit must NOT call resolve_ref_audio."""
+ from unittest.mock import MagicMock
+
+ from vllm_omni.model_executor.models.moss_tts.reference_encoder import encode_reference_codes
+ from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+ speaker_cache = SpeakerEmbeddingCache(max_bytes=64 * 1024 * 1024)
+
+ processor = MagicMock()
+ processor.encode_audios_from_wav = lambda wavs, sampling_rate, n_vq: [torch.zeros(5)]
+
+ resolve_called = 0
+
+ async def mock_resolve(ref_str):
+ nonlocal resolve_called
+ resolve_called += 1
+ return ([0.1], 24000, "some_key")
+
+ # First call: cold miss, resolves + encodes
+ await encode_reference_codes(
+ "data:audio/wav;base64,fake",
+ processor=processor,
+ resolve_ref_audio=mock_resolve,
+ speaker_cache=speaker_cache,
+ variant="tts",
+ n_vq=32,
+ sr_target=24000,
+ voice_name="alice",
+ voice_created_at=100,
+ )
+ assert resolve_called == 1
+
+ # Second call: warm hit, must NOT resolve
+ await encode_reference_codes(
+ "data:audio/wav;base64,fake",
+ processor=processor,
+ resolve_ref_audio=mock_resolve,
+ speaker_cache=speaker_cache,
+ variant="tts",
+ n_vq=32,
+ sr_target=24000,
+ voice_name="alice",
+ voice_created_at=100,
+ )
+ assert resolve_called == 1, "Named-voice cache hit should skip resolve_ref_audio"
+
+ @pytest.mark.asyncio
+ async def test_ttsd_second_reference_does_not_use_named_voice(self, speech_server, mocker):
+ """Uploaded voice names speaker 1 only; speaker 2 must stay anonymous."""
+ speech_server._moss_variant = "ttsd"
+ speech_server.uploaded_speakers = {"alice": {}}
+ mocker.patch.object(speech_server, "_voice_created_at", return_value=42)
+
+ proc = mocker.MagicMock()
+ proc.model_config.n_vq = 8
+ proc.model_config.sampling_rate = 24000
+ proc.build_user_message.return_value = "msg"
+ proc.return_value = {"input_ids": [torch.zeros((4, 9), dtype=torch.int64)]}
+ mocker.patch.object(speech_server, "_get_moss_processor", return_value=proc)
+
+ seen: list[tuple[str, str | None]] = []
+
+ async def fake_encode(ref_str, **kwargs):
+ seen.append((ref_str, kwargs.get("voice_name")))
+ return torch.zeros(3, dtype=torch.int64)
+
+ mocker.patch(
+ "vllm_omni.model_executor.models.moss_tts.reference_encoder.encode_reference_codes",
+ fake_encode,
+ )
+
+ req = OpenAICreateSpeechRequest(
+ input="hello",
+ voice="alice",
+ ref_audio="data:audio/wav;base64,aaa",
+ ref_audio_2="data:audio/wav;base64,bbb",
+ )
+ await speech_server._build_moss_tts_params(req, has_inline_ref_audio=False)
+ assert seen == [
+ ("data:audio/wav;base64,aaa", "alice"),
+ ("data:audio/wav;base64,bbb", None),
+ ]
+
def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server):
"""Precomputed Qwen3 voices are reusable by name without per-request ref_audio."""
speech_server._tts_model_type = "qwen3_tts"
- speech_server.precomputed_speakers = {
- "alice": {
- "name": "Alice",
- "model_type": "qwen3_tts",
- "mode": "icl",
- "ref_text": "reference transcript",
- "ref_code_length": 3,
- }
- }
- speech_server.supported_speakers = {"alice"}
+ speech_server._adapter = speech_server._get_tts_adapter()
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={
+ "alice": {
+ "name": "Alice",
+ "model_type": "qwen3_tts",
+ "mode": "icl",
+ "ref_text": "reference transcript",
+ "ref_code_length": 3,
+ }
+ },
+ supported_speakers=frozenset({"alice"}),
+ )
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is None
@@ -1198,6 +1681,7 @@ class TestTTSMethods:
uploaded_path = tmp_path / "alice.safetensors"
save_file({"speaker_embedding": torch.tensor([0.1] * 4)}, str(uploaded_path))
speech_server._tts_model_type = "qwen3_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.uploaded_speakers_dir = tmp_path
speech_server.uploaded_speakers = {
"alice": {
@@ -1207,15 +1691,18 @@ class TestTTSMethods:
"embedding_source": "direct",
}
}
- speech_server.precomputed_speakers = {
- "alice": {
- "name": "Alice",
- "model_type": "qwen3_tts",
- "mode": "icl",
- "ref_text": "precomputed transcript",
- "ref_code_length": 3,
- }
- }
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={
+ "alice": {
+ "name": "Alice",
+ "model_type": "qwen3_tts",
+ "mode": "icl",
+ "ref_text": "precomputed transcript",
+ "ref_code_length": 3,
+ }
+ },
+ )
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is None
@@ -1242,13 +1729,16 @@ class TestTTSMethods:
talker_config=SimpleNamespace(hidden_size=4),
)
)
+ speech_server._adapter = speech_server._get_tts_adapter()
- profiles = speech_server._load_precomputed_speakers()
+ profiles = speech_server._adapter._load_precomputed_speakers()
assert profiles == {}
- speech_server.precomputed_speakers = profiles
- speech_server.supported_speakers = set(profiles)
- assert "alice" not in speech_server.supported_speakers
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers=profiles,
+ )
+ assert "alice" not in speech_server._adapter.capabilities.precomputed_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is not None
@@ -1276,13 +1766,16 @@ class TestTTSMethods:
talker_config=SimpleNamespace(hidden_size=4),
)
)
+ speech_server._adapter = speech_server._get_tts_adapter()
- profiles = speech_server._load_precomputed_speakers()
+ profiles = speech_server._adapter._load_precomputed_speakers()
assert profiles == {}
- speech_server.precomputed_speakers = profiles
- speech_server.supported_speakers = set(profiles)
- assert "alice" not in speech_server.supported_speakers
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers=profiles,
+ )
+ assert "alice" not in speech_server._adapter.capabilities.precomputed_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
assert speech_server._validate_tts_request(req) is not None
@@ -1297,13 +1790,16 @@ class TestTTSMethods:
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(custom_voice_dir=str(tmp_path))
)
+ speech_server._adapter = speech_server._get_tts_adapter()
- profiles = speech_server._load_precomputed_speakers()
+ profiles = speech_server._adapter._load_precomputed_speakers()
assert profiles == {}
- speech_server.precomputed_speakers = profiles
- speech_server.supported_speakers = set(profiles)
- assert "bob" not in speech_server.supported_speakers
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers=profiles,
+ )
+ assert "bob" not in speech_server._adapter.capabilities.precomputed_speakers
req = OpenAICreateSpeechRequest(input="Hello", voice="Bob")
assert speech_server._validate_tts_request(req) is not None
assert speech_server._validate_tts_request(OpenAICreateSpeechRequest(input="Hello")) is None
@@ -1311,9 +1807,12 @@ class TestTTSMethods:
def test_prepare_voxcpm2_rejects_supported_speaker_without_voice_profile(self, speech_server, mocker):
"""VoxCPM2 named voices must be uploaded or precomputed on the real request path."""
speech_server._tts_model_type = "voxcpm2"
- speech_server.supported_speakers = {"bob", "default"}
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.uploaded_speakers = {}
- speech_server.precomputed_speakers = {}
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={},
+ )
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value="generator")
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
@@ -1329,9 +1828,12 @@ class TestTTSMethods:
def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker):
"""VoxCPM2 default voice preserves the built-in zero-shot request path."""
speech_server._tts_model_type = "voxcpm2"
- speech_server.supported_speakers = {"default"}
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.uploaded_speakers = {}
- speech_server.precomputed_speakers = {}
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={},
+ )
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
@@ -1346,23 +1848,26 @@ class TestTTSMethods:
def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker):
"""VoxCPM2 precomputed voices must carry voice metadata to the model cache lookup."""
speech_server._tts_model_type = "voxcpm2"
- speech_server.supported_speakers = {"alice"}
+ speech_server._adapter = speech_server._get_tts_adapter()
speech_server.uploaded_speakers = {}
- speech_server.precomputed_speakers = {
- "alice": {
- "name": "Alice",
- "model_type": "voxcpm2",
- "mode": "reference",
- "ref_audio_feat_len": 2,
- }
- }
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={
+ "alice": {
+ "name": "Alice",
+ "model_type": "voxcpm2",
+ "mode": "reference",
+ "ref_audio_feat_len": 2,
+ }
+ },
+ )
speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
return_value={
"prompt_token_ids": [1],
"additional_information": {
- "voice_profile": speech_server.precomputed_speakers["alice"],
+ "voice_profile": speech_server._adapter.capabilities.precomputed_speakers["alice"],
},
}
)
@@ -1441,6 +1946,10 @@ class TestTTSMethods:
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
+ mock_stage = mocker.MagicMock()
+ mock_stage.engine_args.model_stage = "qwen3_tts"
+ mock_engine_client.stage_configs = [mock_stage]
+
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
@@ -1448,7 +1957,10 @@ class TestTTSMethods:
)
# Verify speakers are normalized to lowercase
- assert server.supported_speakers == {"ryan", "vivian", "aiden"}
+ assert server._adapter.capabilities.supported_speakers == {"ryan", "vivian", "aiden"}
+ assert isinstance(server._adapter.capabilities.supported_speakers, frozenset)
+ with pytest.raises(FrozenInstanceError):
+ setattr(server._adapter.capabilities, "codec_frame_rate", 25.0)
def test_load_supported_speakers_skips_non_tts_omni_model(
self,
@@ -1461,7 +1973,7 @@ class TestTTSMethods:
)
assert speech_server._is_tts is False
- assert speech_server._load_supported_speakers() == set()
+ assert load_supported_speakers(speech_server.engine_client) == set()
warning.assert_not_called()
def test_load_supported_languages_from_config(self, speech_server):
@@ -1474,7 +1986,8 @@ class TestTTSMethods:
)
)
)
- assert speech_server._load_supported_languages() == {
+ adapter = speech_server._get_tts_adapter()
+ assert adapter._load_supported_languages() == {
"Chinese",
"English",
"Beijing_Dialect",
@@ -1487,11 +2000,17 @@ class TestTTSMethods:
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}})
)
- assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"}
+ adapter = speech_server._get_tts_adapter()
+ assert adapter._load_supported_languages() == {"Chinese", "English", "Auto"}
def test_validate_language_custom_dialect_accepted(self, speech_server):
"""A language present in the model config passes validation, case-insensitively."""
- speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"}
+ speech_server._tts_model_type = "qwen3_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ supported_languages=frozenset({"Chinese", "English", "Beijing_Dialect", "Auto"}),
+ )
for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"):
req = OpenAICreateSpeechRequest(input="Hello", language=language)
result = speech_server._validate_tts_request(req)
@@ -1501,7 +2020,12 @@ class TestTTSMethods:
def test_validate_language_unknown_rejected(self, speech_server):
"""A language not in the configured set is rejected."""
- speech_server.supported_languages = {"Chinese", "English", "Auto"}
+ speech_server._tts_model_type = "qwen3_tts"
+ speech_server._adapter = speech_server._get_tts_adapter()
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ supported_languages=frozenset({"Chinese", "English", "Auto"}),
+ )
for language in ("Klingon", "klingon"):
req = OpenAICreateSpeechRequest(input="Hello", language=language)
assert "Invalid language" in speech_server._validate_tts_request(req)
@@ -1512,18 +2036,21 @@ class TestTTSMethods:
speech_server.engine_client.model_config = SimpleNamespace(
hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={}))
)
- assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+ adapter = speech_server._get_tts_adapter()
+ assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES
def test_load_supported_languages_default_on_config_error(self, speech_server):
"""If the model config cannot be read, fall back to the default list."""
speech_server._tts_model_type = "qwen3_tts"
speech_server.engine_client = SimpleNamespace() # no model_config -> AttributeError
- assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+ speech_server._adapter = None
+ adapter = speech_server._get_tts_adapter()
+ assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES
def test_load_supported_languages_default_for_non_qwen(self, speech_server):
"""Non-qwen3_tts model types get the default language set."""
- speech_server._tts_model_type = None
- assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+ adapter = MingTTSAdapter(SpeechServingContext(server=speech_server, engine_client=speech_server.engine_client))
+ assert adapter._load_supported_languages() == DEFAULT_TTS_LANGUAGES
def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture):
"""Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only)."""
@@ -1536,7 +2063,6 @@ class TestTTSMethods:
"created_at": 1711234567,
}
}
- speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv"
@@ -1560,7 +2086,6 @@ class TestTTSMethods:
"created_at": 1711234567,
}
}
- speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv"
@@ -1577,7 +2102,6 @@ class TestTTSMethods:
"""Test _build_tts_params does not auto-set ref_audio for non-uploaded voices."""
# No uploaded speakers
speech_server.uploaded_speakers = {}
- speech_server.supported_speakers = {"ryan", "vivian"}
req = OpenAICreateSpeechRequest(input="Hello", voice="ryan", task_type="Base")
@@ -1597,7 +2121,6 @@ class TestTTSMethods:
"mime_type": "audio/wav",
}
}
- speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
req = OpenAICreateSpeechRequest(
input="Hello", voice="custom_voice", task_type="Base", ref_audio="data:audio/wav;base64,ZXhwbGljaXQ="
@@ -1703,7 +2226,6 @@ class TestTTSMethods:
"cache_file": "/tmp/voice_samples/emb_voice.safetensors",
}
}
- speech_server.supported_speakers = {"ryan", "vivian", "emb_voice"}
fake_embedding = [0.1] * 1024
mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
@@ -1777,7 +2299,7 @@ class TestTTSMethods:
mocker.patch.object(
speech_server,
"_resolve_ref_audio",
- new=mocker.AsyncMock(return_value=ref_audio_data),
+ new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")),
)
mock_prompt = mocker.patch.object(
speech_server,
@@ -1894,7 +2416,7 @@ class TestTTSMethods:
mocker.patch.object(
speech_server,
"_resolve_ref_audio",
- new=mocker.AsyncMock(return_value=ref_audio_data),
+ new=mocker.AsyncMock(return_value=(*ref_audio_data, "fake_cache_key")),
)
mock_prompt = mocker.patch.object(
speech_server,
@@ -2162,6 +2684,28 @@ class TestTTSMethods:
assert request.ref_audio is not None
assert call.kwargs["has_inline_ref_audio"] is False
+ @pytest.mark.asyncio
+ async def test_warmup_with_no_adapter(self, speech_server):
+ assert speech_server._adapter is None
+ await speech_server.warmup()
+
+ @pytest.mark.asyncio
+ async def test_voxcpm2_warmup_successfully(self, speech_server, mocker):
+ speech_server._tts_model_type = "voxcpm2"
+ speech_server._adapter = speech_server._get_tts_adapter()
+ speech_server.uploaded_speakers = {}
+ speech_server._adapter.capabilities = replace(
+ speech_server._adapter.capabilities,
+ precomputed_speakers={},
+ )
+ speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+ speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
+ speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+ return_value={"prompt_token_ids": [1], "additional_information": {}}
+ )
+
+ await speech_server.warmup()
+
class TestFileValidationFunctions:
"""Unit tests for file validation helper functions."""
@@ -2390,6 +2934,40 @@ class TestStreamingResponse:
assert "text/event-stream" not in response.headers["content-type"]
assert len(response.content) > 0
+ def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monkeypatch):
+ """The raw route must not bypass the shared Qwen3-TTS EOS guard."""
+ speech_server = streaming_app.state.speech_server
+ finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]}
+ captured: dict = {}
+
+ async def prepare(_request, request_id=None):
+ return request_id, object(), finalized_tts_params
+
+ async def generate_chunks(
+ _generator,
+ _request_id,
+ _response_format="pcm",
+ raw_request=None,
+ request_start_s=None,
+ include_sample_rate=False,
+ usage_acc=None,
+ tts_params=None,
+ collect=None,
+ ):
+ captured["tts_params"] = tts_params
+ yield b"\x00\x00"
+
+ monkeypatch.setattr(speech_server, "_prepare_speech_generation", prepare)
+ monkeypatch.setattr(speech_server, "_generate_audio_chunks", generate_chunks)
+
+ response = TestClient(streaming_app).post(
+ "/v1/audio/speech",
+ json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+ )
+
+ assert response.status_code == 200
+ assert captured["tts_params"] is finalized_tts_params
+
def test_sse_streaming(self, streaming_app):
"""stream_format=sse without stream=True returns audio deltas as SSE."""
client = TestClient(streaming_app)
@@ -2932,6 +3510,39 @@ def test_api_server_upload_voice_exception_returns_500(mocker: MockerFixture):
)
+def test_api_server_upload_voice_with_no_adapter(
+ mocker: MockerFixture,
+ monkeypatch: pytest.MonkeyPatch,
+ tmp_path: Path,
+):
+ monkeypatch.setenv("SPEAKER_SAMPLES_DIR", str(tmp_path))
+ _patch_api_server_base(mocker)
+ engine_client = mocker.MagicMock()
+ engine_client.errored = False
+ engine_client.stage_configs = []
+ models = mocker.MagicMock()
+ models.is_base_model.return_value = True
+ handler = OmniOpenAIServingSpeech(
+ engine_client=engine_client,
+ models=models,
+ request_logger=mocker.MagicMock(),
+ )
+ assert handler._adapter is None
+ raw_request = _make_api_server_request(handler)
+
+ response = asyncio.run(
+ api_server_module.upload_voice(
+ raw_request,
+ audio_sample=None,
+ speaker_embedding=json.dumps([0.1] * 1024),
+ consent="cons_test",
+ name="probe",
+ )
+ )
+
+ assert response.status_code == 200
+
+
def test_api_server_delete_voice_without_speech_handler_returns_404(mocker: MockerFixture):
_patch_api_server_base(mocker)
raw_request = _make_api_server_request(None, method="DELETE", path="/v1/audio/voices/probe")
@@ -3261,8 +3872,14 @@ class _FakeFishTokenizer:
@pytest.fixture
def fish_speech_server(mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
@@ -3532,8 +4149,14 @@ class TestWAVStreaming:
@pytest.fixture
def cosyvoice3_server(mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
@@ -3607,18 +4230,31 @@ class TestCosyVoice3Serving:
error = cosyvoice3_server._validate_tts_request(request)
assert error is None
- def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server):
- request = OpenAICreateSpeechRequest(
- input="Hello",
- ref_audio="data:audio/wav;base64,abc",
- ref_text="hello",
- max_new_tokens=0,
- )
- error = cosyvoice3_server._validate_tts_request(request)
- assert error is not None
- assert "max_new_tokens" in error
+ def test_validate_cosyvoice3_max_new_tokens_range(self):
+ """Ensure max_new_tokens below the minimum is rejected during request validation."""
+ with pytest.raises(ValidationError, match="max_new_tokens"):
+ OpenAICreateSpeechRequest(
+ input="Hello",
+ ref_audio="data:audio/wav;base64,abc",
+ ref_text="hello",
+ max_new_tokens=0,
+ )
- def test_prepare_speech_generation_cosyvoice3(self, cosyvoice3_server, mocker: MockerFixture):
+ @pytest.mark.parametrize(
+ ("max_new_tokens", "expected_min_tokens", "expected_max_tokens"),
+ [
+ (None, 10, 2000),
+ (5, 5, 5),
+ ],
+ )
+ def test_prepare_speech_generation_cosyvoice3(
+ self,
+ cosyvoice3_server,
+ mocker: MockerFixture,
+ max_new_tokens: int | None,
+ expected_min_tokens: int,
+ expected_max_tokens: int,
+ ):
cosyvoice3_server._build_cosyvoice3_prompt = mocker.AsyncMock(
return_value={
"prompt": "Hello",
@@ -3626,15 +4262,23 @@ class TestCosyVoice3Serving:
"mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
}
)
- cosyvoice3_server._apply_cosyvoice3_dynamic_tokens = mocker.MagicMock(side_effect=lambda spl, req: spl)
+ cosyvoice3_server.model_config.hf_config = SimpleNamespace(
+ min_token_text_ratio=1, max_token_text_ratio=200, allowed_special=True
+ )
+ cosyvoice3_server._cosyvoice3_tokenizer = mocker.MagicMock()
+ mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=[None, 10])
request = OpenAICreateSpeechRequest(
input="Hello",
ref_audio="data:audio/wav;base64,abc",
ref_text="Reference text",
+ max_new_tokens=max_new_tokens,
)
request_id, generator, tts_params = asyncio.run(cosyvoice3_server._prepare_speech_generation(request))
+ sampling_params = cosyvoice3_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+ assert sampling_params.min_tokens == expected_min_tokens
+ assert sampling_params.max_tokens == expected_max_tokens
assert request_id.startswith("speech-")
assert generator == "generator"
assert tts_params == {}
@@ -3646,8 +4290,14 @@ class TestCosyVoice3Serving:
@pytest.fixture
def glm_tts_server(mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
@@ -3669,7 +4319,6 @@ def glm_tts_server(mocker: MockerFixture):
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
-
return OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
@@ -3718,11 +4367,107 @@ class TestGLMTTSServing:
assert text_token_len == 3
load_tokenizer.assert_called_once()
+ def test_prepare_speech_generation_glm_tts(self, glm_tts_server, mocker: MockerFixture):
+ glm_tts_server._build_glm_tts_prompt = mocker.AsyncMock(
+ return_value={
+ "prompt": "Hello",
+ "multi_modal_data": {"audio": (np.zeros(24000), 24000)},
+ "mm_processor_kwargs": {"prompt_text": "ref text", "sample_rate": 24000},
+ }
+ )
+ request = OpenAICreateSpeechRequest(
+ input="Hello",
+ ref_audio="data:audio/wav;base64,abc",
+ ref_text="Reference text",
+ )
+ glm_tts_server.model_config.hf_config = SimpleNamespace(min_token_text_ratio=1, max_token_text_ratio=200)
+ glm_tts_server._estimate_glm_tts_text_token_len = mocker.MagicMock(return_value=10)
+ request_id, generator, tts_params = asyncio.run(glm_tts_server._prepare_speech_generation(request))
+ sampling_params = glm_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+
+ assert sampling_params.min_tokens == 10
+ assert sampling_params.max_tokens == 2000
+ assert request_id.startswith("speech-")
+ assert generator == "generator"
+ assert tts_params == {}
+ glm_tts_server._build_glm_tts_prompt.assert_awaited_once()
+
+
+@pytest.fixture
+def ming_tts_server(mocker: MockerFixture):
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
+
+ mock_engine_client = mocker.MagicMock()
+ mock_engine_client.errored = False
+ mock_engine_client.model_config = mocker.MagicMock(
+ model="inclusionAI/Ming-omni-tts-0.5B",
+ hf_config=SimpleNamespace(min_token_text_ratio=2, max_token_text_ratio=20),
+ )
+ mock_engine_client.default_sampling_params_list = [
+ SimpleNamespace(max_tokens=2048, min_tokens=None, extra_args=None)
+ ]
+ mock_engine_client.tts_batch_max_items = 32
+ mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+ mock_engine_client.stage_configs = [
+ SimpleNamespace(
+ engine_args=SimpleNamespace(
+ model_stage="llm", model_arch="MingTTSForConditionalGeneration", worker_type="ar"
+ ),
+ tts_args={},
+ )
+ ]
+
+ mock_models = mocker.MagicMock()
+ mock_models.is_base_model.return_value = True
+ return OmniOpenAIServingSpeech(
+ engine_client=mock_engine_client,
+ models=mock_models,
+ request_logger=mocker.MagicMock(),
+ )
+
+
+class TestMingTTSServing:
+ def test_prepare_speech_generation_ming(self, ming_tts_server, mocker: MockerFixture):
+ ming_tts_server._adapter.build = mocker.AsyncMock(
+ return_value=PreparedRequest(
+ prompt={"prompt": "hello"},
+ tts_params={},
+ model_type="ming_tts",
+ )
+ )
+
+ request = OpenAICreateSpeechRequest(
+ input="Hello",
+ ref_audio="data:audio/wav;base64,abc",
+ ref_text="Reference text",
+ max_new_tokens=7,
+ )
+ asyncio.run(ming_tts_server._prepare_speech_generation(request))
+
+ assert ming_tts_server._tts_model_type == "ming_tts"
+ ming_tts_server._adapter.build.assert_awaited_once()
+ sampling_params = ming_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+ assert sampling_params.stop_token_ids == [TEXT_EOS_TOKEN_ID]
+ assert sampling_params.max_tokens == 8
+
@pytest.fixture
def ming_flash_omni_tts_server(mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
@@ -3785,8 +4530,14 @@ class TestTTSAsyncOffloading:
@pytest.fixture
def voxtral_server(self, mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.voxtral.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.model_config = mocker.MagicMock(model="mistralai/Voxtral")
@@ -3811,8 +4562,14 @@ class TestTTSAsyncOffloading:
@pytest.fixture
def qwen3_tts_server(self, mocker: MockerFixture):
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.model_config = mocker.MagicMock(model="Qwen/Qwen3-TTS", hf_config=mocker.MagicMock())
@@ -3836,6 +4593,32 @@ class TestTTSAsyncOffloading:
yield server
server.shutdown()
+ def test_voxtral_loads_supported_speakers(self, mocker: MockerFixture):
+ engine_client = mocker.MagicMock()
+ engine_client.model_config.hf_config.audio_config = {
+ "speaker_id": {
+ "Alice": 0,
+ "Bob": 1,
+ }
+ }
+
+ server = mocker.MagicMock()
+ adapter = VoxtralTTSAdapter(
+ SpeechServingContext(
+ server=server,
+ engine_client=engine_client,
+ )
+ )
+
+ adapter.load_capabilities()
+
+ assert adapter.capabilities.supported_speakers == frozenset(
+ {
+ "alice",
+ "bob",
+ }
+ )
+
def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mocker: MockerFixture):
"""Voxtral path in _prepare_speech_generation should call the async wrapper."""
voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock(
@@ -3901,6 +4684,9 @@ class TestTTSAsyncOffloading:
model_type=adapter_model_type,
)
+ def apply_sampling_overrides(self, sampling_params_list, request, prompt, request_id=None):
+ return sampling_params_list
+
voxtral_server._tts_model_type = legacy_tts_model_type
mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")
@@ -3909,9 +4695,7 @@ class TestTTSAsyncOffloading:
assert adapter_model_type != legacy_tts_model_type
assert any(
- call.args
- and call.args[0] == "TTS speech request %s: text=%r, model=%s"
- and call.args[3] == adapter_model_type
+ call.args and call.args[0] == "TTS speech request %s: model=%s" and call.args[2] == adapter_model_type
for call in log_info.call_args_list
)
@@ -4034,7 +4818,7 @@ class TestTTSAsyncOffloading:
):
"""Base explicit true should reach the model prompt additional_information."""
qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
- qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000))
+ qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000, "fake_cache_key"))
qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None)
qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
@@ -4066,7 +4850,10 @@ class TestTTSAsyncOffloading:
artifact_key,
)
qwen3_tts_server._ref_audio_model_artifact_ready.add((artifact_key, False))
- qwen3_tts_server._codec_frame_rate = 25.0
+ qwen3_tts_server._adapter.capabilities = replace(
+ qwen3_tts_server._adapter.capabilities,
+ codec_frame_rate=25.0,
+ )
qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))}
qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace(
codec_language_id={},
@@ -4101,6 +4888,57 @@ class TestTTSAsyncOffloading:
assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
+ @pytest.mark.asyncio
+ async def test_qwen3_tts_nonstream_retries_codec_limit_once(self, qwen3_tts_server, mocker: MockerFixture) -> None:
+ qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+ qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(
+ side_effect=[
+ Qwen3TTSCodecLimitError("codec limit"),
+ (b"RIFF" + b"\x00" * 32, "audio/wav"),
+ ]
+ )
+ request = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
+
+ response = await qwen3_tts_server.create_speech(request)
+
+ assert response.status_code == 200
+ assert qwen3_tts_server._generate_audio_bytes.await_count == 2
+ first_call, retry_call = qwen3_tts_server._generate_audio_bytes.await_args_list
+ assert first_call.args[0] is request
+ assert first_call.kwargs["request_id"] != retry_call.kwargs["request_id"]
+ assert retry_call.args[0].seed is not None
+ assert request.seed is None
+
+ @pytest.mark.asyncio
+ async def test_nonstream_does_not_retry_nonretryable_generation_error(
+ self, qwen3_tts_server, mocker: MockerFixture
+ ) -> None:
+ qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+ qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=TTSGenerationError("invalid generation"))
+
+ response = await qwen3_tts_server.create_speech(OpenAICreateSpeechRequest(input="Hello", task_type="Base"))
+
+ assert response.status_code == 500
+ qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("request_kwargs", [{"seed": 100423}, {"max_new_tokens": 192}])
+ async def test_qwen3_tts_nonstream_does_not_retry_explicit_sampling_controls(
+ self,
+ qwen3_tts_server,
+ mocker: MockerFixture,
+ request_kwargs: dict[str, int],
+ ) -> None:
+ qwen3_tts_server._check_model = mocker.AsyncMock(return_value=None)
+ qwen3_tts_server._generate_audio_bytes = mocker.AsyncMock(side_effect=Qwen3TTSCodecLimitError("codec limit"))
+
+ response = await qwen3_tts_server.create_speech(
+ OpenAICreateSpeechRequest(input="Hello", task_type="Base", **request_kwargs)
+ )
+
+ assert response.status_code == 500
+ qwen3_tts_server._generate_audio_bytes.assert_awaited_once()
+
@pytest.mark.asyncio
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
async def failing_generator():
@@ -4115,6 +4953,80 @@ class TestTTSAsyncOffloading:
assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready
+ @pytest.mark.asyncio
+ async def test_generate_audio_chunks_rejects_qwen3_codec_limit_before_success(self, qwen3_tts_server):
+ async def length_limited_generator():
+ yield SimpleNamespace(
+ multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+ metrics={
+ "stage_metrics": {
+ "0": {
+ "num_tokens_out": 191,
+ "finish_reason": "length",
+ }
+ }
+ },
+ )
+
+ chunks = qwen3_tts_server._generate_audio_chunks(
+ length_limited_generator(),
+ "req-codec-limit",
+ tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+ )
+ with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+ async for _ in chunks:
+ pass
+
+ @pytest.mark.asyncio
+ async def test_sse_codec_limit_marks_partial_audio_for_discard(self, qwen3_tts_server):
+ async def length_limited_generator():
+ yield SimpleNamespace(
+ multimodal_output={"audio": torch.zeros(16, dtype=torch.float32), "sr": 24000},
+ metrics={
+ "stage_metrics": {
+ "0": {
+ "num_tokens_out": 191,
+ "finish_reason": "length",
+ }
+ }
+ },
+ )
+
+ events = [
+ event
+ async for event in qwen3_tts_server._generate_audio_sse_events(
+ length_limited_generator(),
+ "req-sse-codec-limit",
+ tts_params={"task_type": ["Base"], "_qwen3_tts_effective_max_tokens": [192]},
+ )
+ ]
+
+ assert any("event: speech.audio.delta" in event for event in events)
+ assert not any("event: speech.audio.done" in event for event in events)
+ error_event = next(event for event in events if "event: speech.audio.error" in event)
+ payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+ assert payload["error"]["partial_audio"] is True
+ assert payload["error"]["action"] == "discard"
+
+ @pytest.mark.asyncio
+ async def test_sse_error_before_audio_is_not_marked_partial(self, qwen3_tts_server):
+ async def failing_generator():
+ raise RuntimeError("boom")
+ yield # pragma: no cover
+
+ events = [
+ event
+ async for event in qwen3_tts_server._generate_audio_sse_events(
+ failing_generator(),
+ "req-sse-no-audio",
+ )
+ ]
+
+ error_event = next(event for event in events if "event: speech.audio.error" in event)
+ payload = json.loads(next(line for line in error_event.splitlines() if line.startswith("data: "))[6:])
+ assert "partial_audio" not in payload["error"]
+ assert "action" not in payload["error"]
+
@pytest.mark.asyncio
async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
async def pcm_generator():
@@ -4171,8 +5083,14 @@ class TestTTSAsyncOffloading:
def test_shutdown_is_idempotent(self, mocker: MockerFixture):
"""Calling shutdown() twice should not raise."""
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
- mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_supported_speakers",
+ return_value=set(),
+ )
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.tts_adapters.base.load_codec_frame_rate",
+ return_value=None,
+ )
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_configs = []
diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py
index 37e3509f5..7c553af8f 100644
--- a/tests/entrypoints/openai_api/test_serving_speech_stream.py
+++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
import base64
from types import SimpleNamespace
@@ -44,7 +47,9 @@ def _build_test_app(
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
speech_service.forced_aligner_enabled = False
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+ async def mock_generate_pcm_chunks(
+ _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+ ):
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
yield (chunk, 24000) if include_sample_rate else chunk
@@ -244,6 +249,7 @@ class TestStreamingSpeechWebSocket:
def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
captured_requests = []
+ captured_tts_params = []
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
@@ -253,11 +259,12 @@ class TestStreamingSpeechWebSocket:
async def mock_prepare_speech_generation(request):
captured_requests.append(request)
- return "req-stream", object(), {}
+ return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
speech_service._prepare_speech_generation = mock_prepare_speech_generation
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
+ captured_tts_params.append(tts_params)
for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
yield (chunk, 24000) if include_sample_rate else chunk
@@ -302,6 +309,7 @@ class TestStreamingSpeechWebSocket:
assert captured_requests[0].stream is True
assert captured_requests[0].response_format == "pcm"
assert captured_requests[0].initial_codec_chunk_frames == 12
+ assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
assert speech_service._generate_audio_bytes.await_count == 0
def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
@@ -344,7 +352,9 @@ class TestStreamingSpeechWebSocket:
# The forced-aligner stage rides the same generator: its pooling output
# is surfaced via the ``collect`` channel once the audio has streamed.
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+ async def mock_generate_pcm_chunks(
+ _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+ ):
for chunk in (first_chunk, second_chunk):
yield (chunk, 1000) if include_sample_rate else chunk
if collect is not None:
@@ -427,7 +437,9 @@ class TestStreamingSpeechWebSocket:
speech_service.forced_aligner_enabled = True
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
+ async def mock_generate_pcm_chunks(
+ _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None
+ ):
chunk = b"\x01" * 1000
yield (chunk, 1000) if include_sample_rate else chunk
if collect is not None:
@@ -608,7 +620,7 @@ class TestStreamingSpeechWebSocket:
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
yield b"\x01\x02"
raise RuntimeError("stream boom")
@@ -633,6 +645,8 @@ class TestStreamingSpeechWebSocket:
assert ws.receive_json() == {
"type": "error",
"message": "Generation failed for utterance 0, sentence 0: stream boom",
+ "partial_audio": True,
+ "action": "discard",
}
assert ws.receive_json() == {
"type": "audio.done",
@@ -702,7 +716,7 @@ class TestStreamingSpeechWebSocket:
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
- async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None):
yield b"\x01\x02"
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py
index 27f1bed1e..e57e3e279 100644
--- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py
+++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py
@@ -47,7 +47,13 @@ def _build_test_app(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
- engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+ engine_client.stage_configs = [
+ SimpleNamespace(
+ stage_type="diffusion",
+ final_output=True,
+ final_output_type="video",
+ )
+ ]
engine_client.generate = mock_generate
class FakeStreamingVideoEncoder:
@@ -72,10 +78,6 @@ def _build_test_app(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
- mocker.patch(
- "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
- return_value="diffusion",
- )
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
@@ -398,7 +400,13 @@ def _build_handler_for_async_tests(
engine_client.abort = mocker.AsyncMock()
engine_client.submit_interaction_async = mocker.AsyncMock()
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()]
- engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+ engine_client.stage_configs = [
+ SimpleNamespace(
+ stage_type="diffusion",
+ final_output=True,
+ final_output_type="video",
+ )
+ ]
engine_client.generate = mock_generate
chunks = streaming_chunks or [(b"mp4-chunk-0", True)]
@@ -421,10 +429,6 @@ def _build_handler_for_async_tests(
"vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder",
encoder_factory,
)
- mocker.patch(
- "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type",
- return_value="diffusion",
- )
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list",
return_value=[OmniDiffusionSamplingParams()],
diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py
index 2fba10516..56b74c05b 100644
--- a/tests/entrypoints/openai_api/test_speech_usage.py
+++ b/tests/entrypoints/openai_api/test_speech_usage.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for speech token-usage accounting (issue #4646).
These exercise the pure usage logic without a model/server: input-token
@@ -205,6 +206,14 @@ def test_output_tokens_from_stage_metrics():
assert acc.total() == 77
+def test_output_tokens_preserves_stage0_finish_reason():
+ acc = SpeechOutputTokenCounter()
+ res = _final_res(stage0_out=128)
+ res.metrics["stage_metrics"]["0"]["finish_reason"] = "length"
+ acc.observe(res)
+ assert acc.stage0_finish_reason == "length"
+
+
def test_output_tokens_takes_max_across_stages():
acc = SpeechOutputTokenCounter()
acc.observe(_final_res(stage0_out=42, stage1_out=0))
diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py
index 9ab744c83..03faacfcc 100644
--- a/tests/entrypoints/openai_api/test_tts_adapter.py
+++ b/tests/entrypoints/openai_api/test_tts_adapter.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the TTS serving adapter registry (RFC #4327).
Pure-Python registry/resolution logic; no model or GPU resources are loaded.
@@ -6,23 +7,36 @@ Pure-Python registry/resolution logic; no model or GPU resources are loaded.
import asyncio
from types import SimpleNamespace
+from typing import Any
import pytest
+from vllm.sampling_params import SamplingParams
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
from vllm_omni.entrypoints.openai.tts_adapters import (
TTS_ADAPTER_REGISTRY,
ARTTSAdapter,
DiffusionTTSAdapter,
+ SpeechServingContext,
all_tts_model_types,
detect_tts_model_type,
resolve_adapter,
)
+from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter
from vllm_omni.entrypoints.openai.tts_adapters.indextts2 import (
IndexTTS2Adapter,
IndexTTS25Adapter,
indextts2_conditioning_cache_salt,
)
-from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import (
+ MossTTSAdapter,
+ MossTTSNanoAdapter,
+)
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import (
+ QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY,
+ Qwen3TTSAdapter,
+ Qwen3TTSCodecLimitError,
+)
from vllm_omni.model_executor.models.indextts2 import prompt_utils
from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
INDEXTTS25_TOKENIZER_FILE,
@@ -94,11 +108,114 @@ def test_all_adapters_are_ar_or_diffusion():
assert cls.backend in ("ar", "diffusion")
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+def test_moss_tts_applies_request_max_new_tokens(adapter_cls):
+ adapter = adapter_cls(SimpleNamespace(server=object()))
+ stage_defaults = [SimpleNamespace(max_tokens=4096)]
+
+ overridden = adapter.apply_sampling_overrides(
+ stage_defaults,
+ SimpleNamespace(max_new_tokens=512),
+ )
+
+ assert overridden[0].max_tokens == 512
+ assert stage_defaults[0].max_tokens == 4096
+
+
+def _build_moss_tts_request(adapter_cls, mocker, *, request_seed):
+ server = mocker.Mock()
+ server._build_moss_tts_params = mocker.AsyncMock(return_value={})
+ adapter = adapter_cls(SpeechServingContext(server=server))
+ request = OpenAICreateSpeechRequest(input="hello", seed=request_seed)
+
+ return asyncio.run(
+ adapter.build(
+ request,
+ [SamplingParams(seed=42)],
+ has_inline_ref_audio=False,
+ )
+ )
+
+
+# Full-family coverage pins the adapter contract; only Nano consumes this seed end to end today.
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+@pytest.mark.parametrize("request_seed", [0, 1234])
+def test_moss_tts_request_seed_overrides_stage_default(adapter_cls, request_seed, mocker):
+ prepared = _build_moss_tts_request(
+ adapter_cls,
+ mocker,
+ request_seed=request_seed,
+ )
+
+ assert prepared.tts_params["seed"] == [request_seed]
+
+
+@pytest.mark.parametrize("adapter_cls", [MossTTSAdapter, MossTTSNanoAdapter])
+def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker):
+ prepared = _build_moss_tts_request(
+ adapter_cls,
+ mocker,
+ request_seed=None,
+ )
+
+ assert prepared.tts_params["seed"] == [42]
+
+
def test_qwen3_tts_metadata():
assert Qwen3TTSAdapter.backend == "ar"
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
+@pytest.mark.parametrize(
+ ("task_type", "text_tokens", "request_cap", "expected_cap"),
+ [
+ ("Base", 0, None, 4096),
+ ("Base", 10, None, 192),
+ ("Base", 23, None, 276),
+ ("Base", 23, 128, 128),
+ ("Base", 23, 512, 512),
+ ("Base", 400, None, 4096),
+ ("CustomVoice", 10, None, 4096),
+ ("CustomVoice", 10, 256, 256),
+ ],
+)
+def test_qwen3_tts_applies_text_scaled_codec_safety_limit(task_type, text_tokens, request_cap, expected_cap, mocker):
+ server = mocker.Mock()
+ server._count_usage_text_tokens.return_value = text_tokens
+ adapter = Qwen3TTSAdapter(SpeechServingContext(server=server))
+ stage_defaults = [SamplingParams(max_tokens=4096, min_tokens=2)]
+ request = OpenAICreateSpeechRequest(
+ input="test text",
+ task_type=task_type,
+ max_new_tokens=request_cap,
+ )
+ prompt: dict[str, Any] = {"additional_information": {}}
+
+ overridden = adapter.apply_sampling_overrides(stage_defaults, request, prompt)
+
+ assert overridden[0].max_tokens == expected_cap
+ assert prompt["additional_information"][QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY] == [expected_cap]
+ assert stage_defaults[0].max_tokens == 4096
+
+
+def test_qwen3_tts_rejects_only_length_finished_base_audio(mocker):
+ adapter = Qwen3TTSAdapter(SpeechServingContext(server=mocker.Mock()))
+ params = {
+ "task_type": ["Base"],
+ QWEN3_TTS_EFFECTIVE_MAX_TOKENS_KEY: [192],
+ }
+
+ # EOS at the exact budget is valid; the token count alone is not a
+ # sufficient failure signal.
+ adapter.validate_generation(params, stage0_finish_reason="stop", output_tokens=192)
+ adapter.validate_generation(params, stage0_finish_reason=None, output_tokens=192)
+
+ # Some frontends expose 191 decoded frames for max_new_tokens=192. The
+ # engine terminal reason still makes this an unambiguous limit failure.
+ with pytest.raises(Qwen3TTSCodecLimitError, match="191/192"):
+ adapter.validate_generation(params, stage0_finish_reason="length", output_tokens=191)
+
+
def test_indextts_adapters_are_versioned():
assert resolve_adapter("indextts2") is IndexTTS2Adapter
assert resolve_adapter("indextts2_5") is IndexTTS25Adapter
@@ -162,6 +279,13 @@ def test_indextts25_speed_does_not_change_conditioning_cache_salt():
) == indextts2_conditioning_cache_salt(slow_request, slow_params)
+def test_indextts2_conditioning_cache_salt_changes_with_ref_audio_cache_key():
+ request = SimpleNamespace(input="hello", ref_audio="file:///data/spk.wav")
+ salt_a = indextts2_conditioning_cache_salt(request, {"ref_audio_cache_key": ["key_aaa"]})
+ salt_b = indextts2_conditioning_cache_salt(request, {"ref_audio_cache_key": ["key_bbb"]})
+ assert salt_a != salt_b
+
+
@pytest.mark.parametrize(
("hf_config", "expected_tokenizer_file"),
[
@@ -216,5 +340,58 @@ def test_diffusion_adapter_extra_body_params_fallback():
assert _DiffAdapter.extra_body_params() == frozenset()
+def _higgs_v2_adapter() -> HiggsAudioV2Adapter:
+ server = SimpleNamespace(
+ _apply_uploaded_speaker=lambda request: None,
+ uploaded_speakers={},
+ )
+ return HiggsAudioV2Adapter(SimpleNamespace(server=server))
+
+
+def _higgs_v2_request(**overrides: Any) -> SimpleNamespace:
+ fields: dict[str, Any] = {
+ "input": "Hello world.",
+ "ref_audio": None,
+ "ref_text": None,
+ "voice": None,
+ "x_vector_only_mode": None,
+ "speaker_embedding": None,
+ "instructions": None,
+ "task_type": None,
+ "language": None,
+ "speed": None,
+ "max_new_tokens": None,
+ }
+ fields.update(overrides)
+ return SimpleNamespace(**fields)
+
+
+@pytest.mark.parametrize(
+ "overrides, err_substr",
+ [
+ pytest.param({"ref_audio": "data:audio/wav;base64,AA=="}, "ref_text", id="ref_audio_without_ref_text"),
+ pytest.param({"ref_text": "some transcript"}, "ref_audio", id="ref_text_without_ref_audio"),
+ pytest.param({"task_type": "Base"}, "task_type", id="task_type"),
+ pytest.param({"language": "Chinese"}, "language", id="language_override"),
+ pytest.param({"input": "[SPEAKER0] hi"}, "multi-speaker", id="multi_speaker_tag"),
+ pytest.param({"input": " "}, "empty", id="input_whitespace_only"),
+ ],
+)
+def test_higgs_audio_v2_validate_rejects_out_of_scope_fields(overrides: dict[str, object], err_substr: str) -> None:
+ """Adapter-only policy checks formerly covered by invalid_param e2e on a live V2 server."""
+ adapter = _higgs_v2_adapter()
+ err = adapter.validate(_higgs_v2_request(**overrides))
+ assert err is not None
+ assert err_substr.lower() in err.lower()
+
+
+def test_higgs_audio_v2_validate_accepts_plain_text_and_paired_clone() -> None:
+ adapter = _higgs_v2_adapter()
+ assert adapter.validate(_higgs_v2_request()) is None
+ assert (
+ adapter.validate(_higgs_v2_request(ref_audio="data:audio/wav;base64,AA==", ref_text="some transcript")) is None
+ )
+
+
if __name__ == "__main__":
pytest.main([__file__, "-v"])
diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py
index c63e5c8e4..6b8b97849 100644
--- a/tests/entrypoints/openai_api/test_tts_detection.py
+++ b/tests/entrypoints/openai_api/test_tts_detection.py
@@ -14,7 +14,6 @@ import itertools
import pytest
from vllm_omni.entrypoints.openai.tts_adapters import (
- LEGACY_TTS_DETECTORS,
TTS_ADAPTER_REGISTRY,
all_tts_stage_keys,
detect_tts_model_type,
@@ -154,6 +153,7 @@ _PIPELINE_STAGES = [
"indextts2_talker",
"latent_generator",
"llm",
+ "minimax_music3_ar",
"ming_tts",
"moss_tts",
"moss_tts_codec",
@@ -282,27 +282,17 @@ def test_entry_stage_archs_is_ming_only():
assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
-def test_every_detected_type_is_adapter_backed_or_declared_legacy():
- """Detection may only name a model that has an adapter or an explicit
- :data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string."""
- legacy_names = {d.name for d in LEGACY_TTS_DETECTORS}
+def test_every_detected_type_is_adapter_backed():
+ """Detection may only name a model that has a registered adapter."""
for stage, arch in itertools.product(_STAGES, _ARCHS): # full cross product on purpose
detected = detect_tts_model_type(stage, arch)
if detected is None:
continue
- assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, (
+ assert detected in TTS_ADAPTER_REGISTRY, (
f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}"
)
-def test_legacy_detectors_have_no_adapter():
- """A legacy entry that gained an adapter must be deleted from the list."""
- for detector in LEGACY_TTS_DETECTORS:
- assert resolve_adapter(detector.name) is None, (
- f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS"
- )
-
-
def test_no_ambiguous_detectors_at_equal_priority():
"""Two detectors that can match the same deployment must be ordered.
diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py
index 08d223c68..9f643d8ca 100644
--- a/tests/entrypoints/openai_api/test_video_api_utils.py
+++ b/tests/entrypoints/openai_api/test_video_api_utils.py
@@ -1,17 +1,127 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for OpenAI-compatible video API encoding helpers."""
+import base64
+import threading
+from io import BytesIO
+from typing import Any
+
+import av
+import httpx
import numpy as np
import pytest
import torch
+from PIL import Image
+from vllm import envs
from vllm_omni.diffusion.postprocess import rife_interpolator
+from vllm_omni.diffusion.utils import media_utils
from vllm_omni.entrypoints.openai import video_api_utils
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def _png_bytes() -> bytes:
+ buffer = BytesIO()
+ Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG")
+ return buffer.getvalue()
+
+
+def _install_http_transport(monkeypatch, handler, client_kwargs):
+ async_client = httpx.AsyncClient
+
+ def _client_factory(*args, **kwargs):
+ client_kwargs.append(kwargs.copy())
+ return async_client(*args, transport=httpx.MockTransport(handler), **kwargs)
+
+ monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_follows_redirects_when_allowed(monkeypatch):
+ requested_paths = []
+ client_kwargs: list[dict[str, object]] = []
+
+ def _handler(request):
+ requested_paths.append(request.url.path)
+ if request.url.path == "/redirect.png":
+ return httpx.Response(302, headers={"location": "/image.png"})
+ return httpx.Response(200, content=_png_bytes(), headers={"content-type": "image/png"})
+
+ monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+ _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+ image = await video_api_utils.decode_image_url("https://example.com/redirect.png")
+
+ assert image.size == (2, 1)
+ assert image.mode == "RGB"
+ assert requested_paths == ["/redirect.png", "/image.png"]
+ assert client_kwargs == [{"timeout": 60, "follow_redirects": True}]
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_rejects_redirects_when_disabled(monkeypatch):
+ requested_paths = []
+ client_kwargs: list[dict[str, object]] = []
+
+ def _handler(request):
+ requested_paths.append(request.url.path)
+ return httpx.Response(302, headers={"location": "/image.png"})
+
+ monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", False)
+ _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+ with pytest.raises(InvalidInputReferenceError, match="redirect.*VLLM_MEDIA_URL_ALLOW_REDIRECTS"):
+ await video_api_utils.decode_image_url("https://example.com/redirect.png")
+
+ assert requested_paths == ["/redirect.png"]
+ assert client_kwargs == [{"timeout": 60, "follow_redirects": False}]
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_reports_http_status(monkeypatch):
+ client_kwargs: list[dict[str, object]] = []
+
+ def _handler(request):
+ return httpx.Response(404)
+
+ monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+ _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+ with pytest.raises(InvalidInputReferenceError, match="server returned HTTP 404"):
+ await video_api_utils.decode_image_url("https://example.com/missing.png")
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_reports_connection_failure(monkeypatch):
+ client_kwargs: list[dict[str, object]] = []
+
+ def _handler(request):
+ raise httpx.ConnectError("connection refused", request=request)
+
+ monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+ _install_http_transport(monkeypatch, _handler, client_kwargs)
+
+ with pytest.raises(InvalidInputReferenceError, match="failed to download image"):
+ await video_api_utils.decode_image_url("https://example.com/unreachable.png")
+
+
+@pytest.mark.asyncio
+async def test_decode_image_url_keeps_data_urls_local(monkeypatch):
+ def _unexpected_http_client(*args, **kwargs):
+ pytest.fail("data URLs must not create an HTTP client")
+
+ monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _unexpected_http_client)
+ encoded_image = base64.b64encode(_png_bytes()).decode()
+
+ image = await video_api_utils.decode_image_url(f"data:image/png;base64,{encoded_image}")
+
+ assert image.size == (2, 1)
+ assert image.mode == "RGB"
+
+
def _install_fake_video_mux(monkeypatch, mux_calls):
def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None):
mux_calls.append(
@@ -32,7 +142,7 @@ def _install_fake_video_mux(monkeypatch, mux_calls):
def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
- mux_calls = []
+ mux_calls: list[dict[str, Any]] = []
_install_fake_video_mux(monkeypatch, mux_calls)
frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)]
@@ -113,6 +223,686 @@ def test_channel_first_video_tensor_is_converted_to_channel_last():
np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy())
+def test_contiguous_uint8_video_reaches_mux_without_conversion(monkeypatch):
+ video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
+
+ def fail_prepare(*args, **kwargs):
+ raise AssertionError("ready-to-encode uint8 video must not be normalized")
+
+ monkeypatch.setattr(video_api_utils, "_prepare_video_frames", fail_prepare)
+
+ frames = video_api_utils._coerce_video_to_uint8_frames(video)
+
+ assert frames is video
+
+
+def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch):
+ calls = []
+
+ def fake_planar_mux(frames, width, height, audio_waveform, **kwargs):
+ calls.append((list(frames), width, height, audio_waveform, kwargs))
+ return b"planar-video"
+
+ def fail_compat_mux(*args, **kwargs):
+ raise AssertionError("CFHW-backed FHWC frames should use the planar path")
+
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux)
+ monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fail_compat_mux)
+
+ video = torch.linspace(-1, 1, 3 * 2 * 4 * 6, dtype=torch.float32).reshape(3, 2, 4, 6)
+ encoded = video_api_utils._encode_video_bytes(
+ video,
+ fps=12,
+ )
+
+ assert encoded == b"planar-video"
+ assert len(calls) == 1
+ frames, width, height, audio, options = calls[0]
+ assert (len(frames), width, height) == (2, 6, 4)
+ assert all(frame.format.name == "gbrp" for frame in frames)
+ assert audio is None
+ assert options["fps"] == 12.0
+
+
+def test_planar_converter_serial_and_parallel_outputs_match():
+ frames = [np.full((2, 3, 3), fill_value=value / 10, dtype=np.float32) for value in (1, 4, 7)]
+
+ serial_converter = video_api_utils._PlanarFrameConverter(max_workers=1)
+ parallel_converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+ try:
+ serial = [
+ frame.to_ndarray(format="rgb24") for frame in serial_converter.iter_frames(frames, np.dtype(np.float32))
+ ]
+ parallel = [
+ frame.to_ndarray(format="rgb24") for frame in parallel_converter.iter_frames(frames, np.dtype(np.float32))
+ ]
+ finally:
+ serial_converter.shutdown()
+ parallel_converter.shutdown()
+
+ np.testing.assert_array_equal(np.stack(serial), np.stack(parallel))
+
+
+def test_planar_converter_preserves_fifo_order_and_bounds_pending_futures(monkeypatch):
+ converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+ started = threading.Event()
+ release = threading.Event()
+ four_submitted = threading.Event()
+ outstanding = 0
+ peak_outstanding = 0
+ lock = threading.Lock()
+ real_submit = converter._executor.submit
+
+ def fake_build(frame, _common_dtype):
+ started.set()
+ release.wait(timeout=2)
+ return int(frame[0, 0, 0])
+
+ def submit(*args, **kwargs):
+ nonlocal outstanding, peak_outstanding
+ with lock:
+ outstanding += 1
+ peak_outstanding = max(peak_outstanding, outstanding)
+ if peak_outstanding >= 4:
+ four_submitted.set()
+ future = real_submit(*args, **kwargs)
+
+ def done(_future):
+ nonlocal outstanding
+ with lock:
+ outstanding -= 1
+
+ future.add_done_callback(done)
+ return future
+
+ monkeypatch.setattr(converter, "_build_frame", fake_build)
+ monkeypatch.setattr(converter._executor, "submit", submit)
+ frames = [np.full((1, 1, 3), value, dtype=np.uint8) for value in range(8)]
+ result: list[int] = []
+ iterator = converter.iter_frames(frames, np.dtype(np.uint8))
+ consumer = threading.Thread(target=lambda: result.extend(iterator))
+ consumer.start()
+ try:
+ assert started.wait(timeout=2)
+ assert four_submitted.wait(timeout=2)
+ release.set()
+ consumer.join(timeout=2)
+ assert not consumer.is_alive()
+ finally:
+ release.set()
+ iterator.close()
+ converter.shutdown()
+
+ assert result == list(range(8))
+ assert peak_outstanding <= 4
+
+
+def test_planar_converter_reuses_executor_and_shutdowns_after_cancel():
+ converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+ executor = converter._executor
+ frames = [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(4)]
+ try:
+ list(converter.iter_frames(frames, np.dtype(np.uint8)))
+ iterator = converter.iter_frames(frames, np.dtype(np.uint8))
+ next(iterator)
+ iterator.close()
+ assert converter._executor is executor
+ finally:
+ converter.shutdown()
+
+ assert executor._shutdown
+
+
+@pytest.mark.parametrize(
+ ("video", "expected_path", "expected_reason", "expected_mux"),
+ [
+ (
+ np.transpose(np.zeros((3, 2, 4, 4), dtype=np.float32), (1, 2, 3, 0)),
+ "direct_planar",
+ None,
+ "direct",
+ ),
+ (
+ np.zeros((2, 4, 6, 3), dtype=np.float32),
+ "legacy_fallback",
+ "non_contiguous_rgb_planes",
+ "legacy",
+ ),
+ ],
+ ids=["automatic-direct", "interleaved-legacy-fallback"],
+)
+def test_audio_sample_rate_defaults_and_path_log(monkeypatch, video, expected_path, expected_reason, expected_mux):
+ mux_calls = []
+ path_logs = []
+ default_rate_logs = []
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info",
+ lambda message, *args, **kwargs: path_logs.append(message % args),
+ )
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info_once",
+ lambda message, *args, **kwargs: default_rate_logs.append(message % args),
+ )
+
+ def fake_planar_mux(frames, width, height, audio_waveform, **kwargs):
+ mux_calls.append({"kind": "direct", "audio_sample_rate": kwargs["audio_sample_rate"]})
+ return b"planar-video"
+
+ def fake_compat_mux(frames, audio, **kwargs):
+ mux_calls.append({"kind": "legacy", "audio_sample_rate": kwargs["audio_sample_rate"]})
+ return b"legacy-video"
+
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux)
+ monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+
+ encoded = video_api_utils._encode_video_bytes(
+ video,
+ fps=12,
+ audio=np.zeros((1, 8), dtype=np.float32),
+ )
+
+ assert encoded == (b"planar-video" if expected_mux == "direct" else b"legacy-video")
+ assert mux_calls == [{"kind": expected_mux, "audio_sample_rate": 24000}]
+ assert len(default_rate_logs) == 1
+ assert "default sample rate of 24000 Hz" in default_rate_logs[0]
+ assert len(path_logs) == 1
+ assert f"selected_path={expected_path}" in path_logs[0]
+ if expected_reason is not None:
+ assert f"reason={expected_reason}" in path_logs[0]
+ assert "frames=2" in path_logs[0]
+ expected_shape = "(4, 4, 3)" if expected_mux == "direct" else "(4, 6, 3)"
+ assert f"frame_shape={expected_shape}" in path_logs[0]
+ assert "dtype=float32" in path_logs[0]
+ assert "fps=12" in path_logs[0]
+ assert "audio_present=True" in path_logs[0]
+ assert "effective_audio_sample_rate=24000" in path_logs[0]
+
+
+def test_explicit_audio_sample_rate_bypasses_default_log(monkeypatch):
+ log_messages = []
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info_once",
+ lambda message, *args, **kwargs: log_messages.append(message % args),
+ )
+
+ audio = [[0.0, 1.0], [1.0, 0.0]]
+ assert video_api_utils._resolve_audio_sample_rate(audio, 32000) == 32000
+ assert log_messages == []
+
+
+def test_video_only_does_not_log_default_audio_sample_rate(monkeypatch):
+ log_messages = []
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info_once",
+ lambda message, *args, **kwargs: log_messages.append(message % args),
+ )
+ monkeypatch.setattr(
+ media_utils,
+ "mux_video_audio_bytes",
+ lambda frames, audio, **kwargs: b"legacy-video",
+ )
+
+ video_api_utils._encode_video_bytes(np.zeros((1, 2, 2, 3), dtype=np.uint8), fps=12)
+
+ assert log_messages == []
+
+
+def test_planar_mux_accepts_none_rate_and_defaults_audio_to_44100():
+ def make_frame():
+ return next(
+ video_api_utils._iter_planar_video_frames(
+ [np.zeros((2, 2, 3), dtype=np.uint8)],
+ np.dtype(np.uint8),
+ )
+ )
+
+ video_only = media_utils.mux_av_video_audio_bytes(
+ [make_frame()],
+ width=2,
+ height=2,
+ audio_sample_rate=None,
+ )
+ with av.open(BytesIO(video_only), mode="r", format="mp4") as container:
+ assert not container.streams.audio
+
+ with_audio = media_utils.mux_av_video_audio_bytes(
+ [make_frame()],
+ width=2,
+ height=2,
+ audio_waveform=np.zeros((1, 1024), dtype=np.float32),
+ audio_sample_rate=None,
+ )
+ with av.open(BytesIO(with_audio), mode="r", format="mp4") as container:
+ assert container.streams.audio[0].rate == 44100
+
+
+def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
+ calls = []
+ path_logs = []
+ coerce_calls = []
+
+ real_coerce_video_to_frames = video_api_utils._coerce_video_to_frames
+
+ def count_coerce_video_to_frames(video):
+ coerce_calls.append(video)
+ return real_coerce_video_to_frames(video)
+
+ def fail_planar_mux(*args, **kwargs):
+ raise AssertionError("interleaved FHWC frames must use the legacy path")
+
+ def fake_compat_mux(frames, audio, **kwargs):
+ calls.append((frames, audio, kwargs))
+ return b"legacy-video"
+
+ monkeypatch.setattr(
+ video_api_utils.logger, "info", lambda message, *args, **kwargs: path_logs.append(message % args)
+ )
+ monkeypatch.setattr(video_api_utils, "_coerce_video_to_frames", count_coerce_video_to_frames)
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_planar_mux)
+ monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+
+ video = np.zeros((2, 4, 6, 3), dtype=np.float32)
+ encoded = video_api_utils._encode_video_bytes(video, fps=12)
+
+ assert encoded == b"legacy-video"
+ assert calls[0][0].shape == (2, 4, 6, 3)
+ assert calls[0][0].dtype == np.uint8
+ assert len(path_logs) == 1
+ assert "selected_path=legacy_fallback" in path_logs[0]
+ assert "reason=non_contiguous_rgb_planes" in path_logs[0]
+ assert "audio_present=False" in path_logs[0]
+ assert "effective_audio_sample_rate=None" in path_logs[0]
+ assert "effective_frame_conversion_workers=0" in path_logs[0]
+ assert len(coerce_calls) == 1
+
+
+def test_uint8_frames_reach_the_muxer_without_a_float_round_trip():
+ """uint8 is the muxer's own dtype, so normalisation must leave it alone.
+
+ MiniMax-H3 quantises on the accelerator to keep the response small across
+ its two process hops. Re-normalising to float here would undo that at the
+ last step and cost a full-size conversion in each direction.
+ """
+ video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
+
+ frames, frame_shape, common_dtype = video_api_utils._prepare_video_frames(video)
+
+ assert common_dtype == np.dtype(np.uint8)
+ assert frame_shape == (4, 5, 3)
+ np.testing.assert_array_equal(np.stack(frames), video)
+
+
+def test_non_uint8_integer_frames_are_still_normalised():
+ """Only uint8 is the muxer's dtype; other integer payloads still scale."""
+ video = np.full((2, 4, 5, 3), 255, dtype=np.int32)
+
+ frames, _, common_dtype = video_api_utils._prepare_video_frames(video)
+
+ assert np.issubdtype(common_dtype, np.floating)
+ np.testing.assert_allclose(frames[0], 1.0)
+
+
+@pytest.mark.parametrize(
+ "video",
+ [
+ np.array(
+ [[[[0, 1, 127], [128, 254, 255]], [[255, 127, 1], [2, 3, 4]]]],
+ dtype=np.uint8,
+ ),
+ np.array(
+ [[[[-1.0, -0.5, 0.0], [0.5, 1.0, 2.0]], [[1.5, -2.0, 0.25], [0.75, 0.125, 0.875]]]],
+ dtype=np.float32,
+ ),
+ np.array(
+ [[[[-512, -1, 0], [128, 255, 512]], [[1024, -256, 64], [127, 300, -300]]]],
+ dtype=np.int16,
+ ),
+ ],
+ ids=["uint8", "float-bounded-and-out-of-range", "signed-int-out-of-range"],
+)
+def test_prepared_automatic_fallback_preserves_legacy_quantization(monkeypatch, video):
+ mux_inputs = []
+
+ def fail_planar_mux(*args, **kwargs):
+ raise AssertionError("interleaved frames must use the legacy mux")
+
+ def fake_compat_mux(frames, audio, **kwargs):
+ mux_inputs.append(frames.copy())
+ return b"legacy-video"
+
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_planar_mux)
+ monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+
+ prepared_frames, frame_shape, _ = video_api_utils._prepare_video_frames(video)
+ stacked = np.stack([frame[..., :3] for frame in prepared_frames])
+ # uint8 frames are the muxer's own dtype and reach it unchanged; everything
+ # else is still quantised out of the normalised [0, 1] range.
+ reference = stacked if stacked.dtype == np.uint8 else np.rint(np.clip(stacked, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+ assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video"
+ assert len(mux_inputs) == 1
+ assert mux_inputs[0].dtype == np.uint8
+ assert mux_inputs[0].shape == (len(prepared_frames), *frame_shape[:-1], 3)
+ np.testing.assert_array_equal(mux_inputs[0], reference)
+
+
+@pytest.mark.parametrize(
+ ("worker_count", "expected_path", "expected_reason"),
+ [
+ (None, "legacy_fallback", "non_contiguous_rgb_planes"),
+ (1, "legacy_fallback", "non_contiguous_rgb_planes"),
+ (8, "direct_planar", None),
+ ],
+)
+def test_interleaved_video_routing(monkeypatch, worker_count, expected_path, expected_reason):
+ routes = []
+
+ def fake_planar_mux(*args, **kwargs):
+ return b"planar-video"
+
+ def fake_compat_mux(frames, audio, **kwargs):
+ return b"legacy-video"
+
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fake_planar_mux)
+ monkeypatch.setattr(media_utils, "mux_video_audio_bytes", fake_compat_mux)
+ monkeypatch.setattr(video_api_utils, "_log_video_encoding_path", lambda **kwargs: routes.append(kwargs))
+
+ converter = None if worker_count is None else video_api_utils._PlanarFrameConverter(max_workers=worker_count)
+ try:
+ encoded = video_api_utils._encode_video_bytes(
+ np.zeros((2, 4, 6, 3), dtype=np.float32),
+ fps=12,
+ frame_converter=converter,
+ )
+ finally:
+ if converter is not None:
+ converter.shutdown()
+
+ assert encoded == (b"planar-video" if expected_path == "direct_planar" else b"legacy-video")
+ assert len(routes) == 1
+ assert routes[0]["selected_path"] == expected_path
+ assert routes[0].get("reason") == expected_reason
+
+
+@pytest.mark.parametrize(
+ ("video", "reason"),
+ [
+ (np.zeros((2, 4, 6), dtype=np.float32), "unsupported_shape"),
+ (np.empty((2, 0, 6, 3), dtype=np.float32), "unsupported_shape"),
+ (np.empty((2, 4, 0, 3), dtype=np.float32), "unsupported_shape"),
+ (np.zeros((2, 2, 2, 3), dtype=np.complex64), "unsupported_dtype"),
+ (np.zeros((2, 2, 2, 3), dtype=object), "unsupported_dtype"),
+ (np.zeros((2, 2, 2, 3), dtype="U1"), "unsupported_dtype"),
+ (np.zeros((2, 2, 2, 3), dtype="datetime64[D]"), "unsupported_dtype"),
+ ],
+)
+def test_unsupported_direct_capability_uses_legacy_with_reason(monkeypatch, video, reason):
+ events = []
+
+ def fake_legacy(*args, **kwargs):
+ events.append("legacy")
+ return b"legacy-video"
+
+ monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fake_legacy)
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info",
+ lambda message, *args, **kwargs: events.append(message % args),
+ )
+
+ assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video"
+ assert events[-1] == "legacy"
+ assert f"selected_path=legacy_fallback reason={reason}" in events[0]
+ assert "frames=" in events[0]
+ assert "fps=12" in events[0]
+ assert "audio_present=False" in events[0]
+ assert "effective_audio_sample_rate=None" in events[0]
+
+
+def test_non_contiguous_fallback_is_logged_before_legacy_mux(monkeypatch):
+ events = []
+ video = np.zeros((2, 4, 6, 3), dtype=np.float32)
+
+ def fake_legacy(*args, **kwargs):
+ events.append("legacy")
+ return b"legacy-video"
+
+ monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fake_legacy)
+ monkeypatch.setattr(
+ video_api_utils.logger,
+ "info",
+ lambda message, *args, **kwargs: events.append(message % args),
+ )
+
+ assert video_api_utils._encode_video_bytes(video, fps=12) == b"legacy-video"
+ assert events[0].startswith("Video response encoding route selected: selected_path=legacy_fallback")
+ assert events[1] == "legacy"
+
+
+def test_direct_mux_error_is_propagated_without_legacy_retry(monkeypatch):
+ legacy_calls = []
+
+ def fail_direct(*args, **kwargs):
+ raise RuntimeError("direct mux failed")
+
+ def fail_legacy(*args, **kwargs):
+ legacy_calls.append((args, kwargs))
+ raise AssertionError("direct failures must not retry through legacy")
+
+ monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", fail_direct)
+ monkeypatch.setattr(video_api_utils, "_encode_prepared_video_bytes_legacy", fail_legacy)
+ planar = np.transpose(np.zeros((3, 2, 4, 4), dtype=np.float32), (1, 2, 3, 0))
+
+ with pytest.raises(RuntimeError, match="direct mux failed"):
+ video_api_utils._encode_video_bytes(planar, fps=12)
+
+ assert legacy_calls == []
+
+
+@pytest.mark.parametrize("dtype", [np.float16, np.float32])
+@pytest.mark.parametrize("channels", [3, 4])
+def test_planar_frames_decode_with_bounded_rgb_quantization(dtype, channels):
+ channel_values = np.array(
+ [
+ [[-0.1, 0.0, 0.1, 0.5, 0.9, 1.0, 1.1]],
+ [[1.0, 0.9, 0.5, 0.1, 0.0, -0.1, 1.1]],
+ [[0.25, 0.75, 0.501, 0.499, 0.125, 0.875, 0.5]],
+ [[0.2, 0.4, 0.6, 0.8, 1.0, 0.0, 0.5]],
+ ],
+ dtype=dtype,
+ )[:channels]
+ fhwc = np.transpose(channel_values[:, None, :, :], (1, 2, 3, 0))
+ frames = list(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(dtype)))
+
+ decoded = frames[0].to_ndarray(format="rgb24")
+ expected = np.rint(np.clip(fhwc[0, ..., :3], 0.0, 1.0) * 255.0).astype(np.uint8)
+ np.testing.assert_array_equal(decoded, expected)
+
+
+def test_planar_frame_writes_gbr_planes_and_clears_padding():
+ width = 17
+ planar = np.stack(
+ [
+ np.full((2, width), 11, dtype=np.uint8),
+ np.full((2, width), 22, dtype=np.uint8),
+ np.full((2, width), 33, dtype=np.uint8),
+ ]
+ )
+ fhwc = np.transpose(planar[:, None, :, :], (1, 2, 3, 0))
+
+ frame = next(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(np.uint8)))
+
+ for plane, expected in zip(frame.planes, (22, 33, 11)):
+ plane_data = np.frombuffer(plane, dtype=np.uint8).reshape(plane.height, plane.line_size)
+ np.testing.assert_array_equal(plane_data[:2, :width], expected)
+ np.testing.assert_array_equal(plane_data[:2, width:], 0)
+ np.testing.assert_array_equal(plane_data[2:], 0)
+
+
+def test_planar_bool_frames_match_bounded_compatible_output():
+ planar = np.array(
+ [
+ [[[False, True], [True, False]]],
+ [[[True, False], [True, False]]],
+ [[[False, False], [True, True]]],
+ ],
+ dtype=np.bool_,
+ )
+ direct_video = np.transpose(planar, (1, 2, 3, 0))
+ compatible_video = np.ascontiguousarray(direct_video)
+
+ prepared_frames, _, common_dtype = video_api_utils._prepare_video_frames(direct_video)
+ direct_frames = list(video_api_utils._iter_planar_video_frames(prepared_frames, common_dtype))
+ decoded_direct = np.stack([frame.to_ndarray(format="rgb24") for frame in direct_frames])
+ bounded_compatible = video_api_utils._coerce_video_to_uint8_frames(compatible_video)
+
+ np.testing.assert_array_equal(decoded_direct, bounded_compatible)
+
+
+@pytest.mark.parametrize(
+ ("plane_height", "line_size"),
+ [(1, 2), (2, 1)],
+)
+def test_planar_frame_rejects_undersized_plane(monkeypatch, plane_height, line_size):
+ class FakePlane:
+ height: int
+ line_size: int
+
+ plane = FakePlane()
+ plane.height = plane_height
+ plane.line_size = line_size
+
+ class FakeFrame:
+ planes = [plane]
+
+ monkeypatch.setattr(av, "VideoFrame", lambda *args, **kwargs: FakeFrame())
+ planar = np.zeros((3, 1, 2, 2), dtype=np.uint8)
+ fhwc = np.transpose(planar, (1, 2, 3, 0))
+
+ with pytest.raises(ValueError, match="smaller than"):
+ next(video_api_utils._iter_planar_video_frames(list(fhwc), np.dtype(np.uint8)))
+
+
+def test_unequal_frame_shapes_fail_before_output_allocation(monkeypatch):
+ def fail_empty(*args, **kwargs):
+ raise AssertionError("output allocation must happen after shape validation")
+
+ monkeypatch.setattr(video_api_utils.np, "empty", fail_empty)
+
+ with pytest.raises(ValueError, match="same shape"):
+ video_api_utils._encode_video_bytes(
+ [np.zeros((2, 3, 3), dtype=np.float32), np.zeros((3, 3, 3), dtype=np.float32)],
+ fps=12,
+ )
+
+
+@pytest.mark.parametrize("with_audio", [False, True])
+def test_direct_and_compatible_paths_produce_equivalent_mp4(with_audio):
+ rng = np.random.default_rng(7)
+ planar = rng.random((3, 3, 32, 32), dtype=np.float32)
+ direct_video = np.transpose(planar, (1, 2, 3, 0))
+ compatible_video = np.ascontiguousarray(direct_video)
+ audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None
+
+ direct_bytes = video_api_utils._encode_video_bytes(
+ direct_video,
+ fps=12,
+ audio=audio,
+ audio_sample_rate=24000,
+ video_codec_options={"preset": "ultrafast", "threads": "1"},
+ )
+ compatible_bytes = video_api_utils._encode_video_bytes(
+ compatible_video,
+ fps=12,
+ audio=audio,
+ audio_sample_rate=24000,
+ video_codec_options={"preset": "ultrafast", "threads": "1"},
+ )
+
+ def decode(mp4_bytes):
+ with av.open(BytesIO(mp4_bytes), mode="r", format="mp4") as container:
+ stream_types = [stream.type for stream in container.streams]
+ video_stream = container.streams.video[0]
+ frames = np.stack([frame.to_ndarray(format="rgb24") for frame in container.decode(video_stream)])
+ return stream_types, frames
+
+ direct_streams, direct_frames = decode(direct_bytes)
+ compatible_streams, compatible_frames = decode(compatible_bytes)
+ assert direct_bytes == compatible_bytes
+ assert direct_streams == compatible_streams
+ assert ("audio" in direct_streams) is with_audio
+ np.testing.assert_array_equal(direct_frames, compatible_frames)
+
+
+@pytest.mark.parametrize("with_audio", [False, True])
+def test_parallel_interleaved_and_legacy_paths_produce_equivalent_mp4(with_audio):
+ rng = np.random.default_rng(7)
+ video = rng.random((3, 32, 32, 3), dtype=np.float32)
+ audio = np.zeros((2, 2400), dtype=np.float32) if with_audio else None
+ kwargs = {
+ "fps": 12,
+ "audio": audio,
+ "audio_sample_rate": 24000,
+ "video_codec_options": {"preset": "ultrafast", "threads": "1"},
+ }
+
+ converter = video_api_utils._PlanarFrameConverter(max_workers=8)
+ try:
+ parallel_bytes = video_api_utils._encode_video_bytes(
+ video,
+ **kwargs,
+ frame_converter=converter,
+ )
+ finally:
+ converter.shutdown()
+ legacy_bytes = video_api_utils._encode_video_bytes_legacy(video, **kwargs)
+
+ assert parallel_bytes == legacy_bytes
+
+
+def test_mux_closes_container_and_preserves_generator_error(monkeypatch):
+ class GeneratorError(RuntimeError):
+ pass
+
+ class FakeStream:
+ options = None
+
+ def encode(self, frame=None):
+ return []
+
+ class FakeContainer:
+ def __init__(self):
+ self.closed = False
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ self.closed = True
+ return False
+
+ def add_stream(self, codec, rate):
+ return FakeStream()
+
+ def mux(self, packet):
+ pass
+
+ container = FakeContainer()
+ monkeypatch.setattr(media_utils.av, "open", lambda *args, **kwargs: container)
+
+ def failing_frames():
+ yield object()
+ raise GeneratorError("frame generation failed")
+
+ with pytest.raises(GeneratorError, match="frame generation failed"):
+ media_utils.mux_av_video_audio_bytes(failing_frames(), width=2, height=2)
+ assert container.closed
+
+
def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
muxers = []
diff --git a/tests/entrypoints/openai_api/test_video_pipeline_capability.py b/tests/entrypoints/openai_api/test_video_pipeline_capability.py
new file mode 100644
index 000000000..cfa11f96d
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_video_pipeline_capability.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.io_support import get_diffusion_output_type
+from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_video_pipeline_requires_declared_final_video_stage():
+ assert is_video_generation_pipeline(
+ [
+ SimpleNamespace(
+ stage_type="llm",
+ final_output=False,
+ final_output_type=None,
+ ),
+ SimpleNamespace(
+ stage_type="diffusion",
+ final_output=True,
+ final_output_type="video",
+ ),
+ ]
+ )
+
+
+@pytest.mark.parametrize(
+ "stage_configs",
+ [
+ [SimpleNamespace(stage_type="diffusion")],
+ [
+ SimpleNamespace(
+ stage_type="diffusion",
+ final_output=True,
+ final_output_type="image",
+ )
+ ],
+ [
+ {
+ "stage_type": "diffusion",
+ "final_output": False,
+ "final_output_type": "video",
+ }
+ ],
+ ],
+)
+def test_video_pipeline_rejects_non_video_final_outputs(stage_configs):
+ assert not is_video_generation_pipeline(stage_configs)
+
+
+@pytest.mark.parametrize(
+ "model_class_name",
+ [
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledOneStagePipeline",
+ "WanDMDPipeline",
+ "LingBotWorldCausalDMDPipeline",
+ "LongCatVideoAvatarPipeline",
+ ],
+)
+def test_registered_video_aliases_declare_video_output(model_class_name):
+ assert get_diffusion_output_type(model_class_name) == "video"
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index bcda12e2d..6d8977c55 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -16,15 +16,17 @@ from pathlib import Path
from types import SimpleNamespace
import av
+import httpx
import numpy as np
import pytest
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from PIL import Image
from pytest_mock import MockerFixture
+from vllm import envs
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.openai import api_server
+from vllm_omni.entrypoints.openai import api_server, video_api_utils
from vllm_omni.entrypoints.openai.api_server import router
from vllm_omni.entrypoints.openai.protocol.videos import (
VideoGenerationRequest,
@@ -32,7 +34,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
VideoParams,
VideoResponse,
)
-from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
+from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
from vllm_omni.entrypoints.openai.storage import LocalStorageManager
from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
from vllm_omni.errors import GuardrailViolationError
@@ -62,7 +64,13 @@ class MockVideoResult:
class FakeAsyncOmni:
def __init__(self):
- self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
+ self.stage_configs = [
+ SimpleNamespace(
+ stage_type="diffusion",
+ final_output=True,
+ final_output_type="video",
+ )
+ ]
self.default_sampling_params_list = [OmniDiffusionSamplingParams()]
self.model_class_name = "WanPipeline"
self.captured_prompt = None
@@ -87,6 +95,58 @@ class FakeAsyncOmni:
yield MockVideoResult(videos)
+def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture):
+ engine = FakeAsyncOmni()
+ handler = OmniOpenAIServingVideo.for_diffusion(
+ engine,
+ model_name="test-model",
+ )
+ assert handler._video_frame_converter.max_workers == 8
+ raw_encoder = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+ return_value=b"encoded-video",
+ )
+ base64_encoder = mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+ return_value="encoded-video",
+ )
+
+ async def _generate_both_response_types():
+ request = VideoGenerationRequest(prompt="test prompt")
+ await handler.generate_video_bytes(request, "raw-request")
+ await handler.generate_videos(request, "base64-request")
+
+ asyncio.run(_generate_both_response_types())
+
+ assert raw_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
+ assert base64_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
+ handler.shutdown()
+
+
+def test_resolve_diffusion_od_config_falls_back_to_attribute():
+ od_config = SimpleNamespace(model_class_name="WanPipeline")
+ handler = OmniOpenAIServingVideo.for_diffusion(
+ SimpleNamespace(od_config=od_config),
+ model_name="test-model",
+ )
+
+ assert handler._resolve_diffusion_od_config() is od_config
+
+
+def test_resolve_diffusion_od_config_prefers_getter_over_attribute():
+ attribute_config = SimpleNamespace(model_class_name="WanPipeline")
+ getter_config = SimpleNamespace(model_class_name="MiniMaxH3Pipeline")
+ handler = OmniOpenAIServingVideo.for_diffusion(
+ SimpleNamespace(
+ od_config=attribute_config,
+ get_diffusion_od_config=lambda: getter_config,
+ ),
+ model_name="test-model",
+ )
+
+ assert handler._resolve_diffusion_od_config() is getter_config
+
+
class BlockingVideoHandler:
def __init__(self):
self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
@@ -461,6 +521,51 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
assert input_image.size == (96, 64)
+def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch):
+ engine = FakeAsyncOmni()
+ engine.get_diffusion_od_config = lambda: SimpleNamespace(
+ model="org/model",
+ model_class_name="ExamplePipeline",
+ revision="pinned-revision",
+ )
+ captured = {}
+
+ def fake_policy(model_class_name, *, model, revision=None):
+ captured.update(
+ model_class_name=model_class_name,
+ model=model,
+ revision=revision,
+ )
+ return True
+
+ monkeypatch.setattr(
+ "vllm_omni.entrypoints.openai.serving_video.should_preserve_reference_image_size",
+ fake_policy,
+ )
+ handler = OmniOpenAIServingVideo.for_diffusion(
+ diffusion_engine=engine,
+ model_name="fallback/model",
+ )
+ image = Image.new("RGB", (48, 32))
+
+ asyncio.run(
+ handler._run_and_extract(
+ VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64),
+ "pipeline-owned-resize",
+ reference_image=ReferenceImage(image),
+ )
+ )
+
+ input_image = engine.captured_prompt["multi_modal_data"]["image"]
+ assert isinstance(input_image, Image.Image)
+ assert input_image.size == (48, 32)
+ assert captured == {
+ "model_class_name": "ExamplePipeline",
+ "model": "org/model",
+ "revision": "pinned-revision",
+ }
+
+
def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture):
image_bytes = _make_test_image_bytes((48, 48))
mocker.patch(
@@ -554,6 +659,45 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
assert input_image.size == (40, 24)
+def test_i2v_video_generation_follows_allowed_image_redirect(test_client, mocker: MockerFixture, monkeypatch):
+ requested_paths = []
+ async_client = httpx.AsyncClient
+
+ def _handler(request):
+ requested_paths.append(request.url.path)
+ if request.url.path == "/redirect.png":
+ return httpx.Response(302, headers={"location": "/image.png"})
+ return httpx.Response(200, content=_make_test_image_bytes((40, 24)))
+
+ def _client_factory(*args, **kwargs):
+ return async_client(*args, transport=httpx.MockTransport(_handler), **kwargs)
+
+ monkeypatch.setattr(envs, "VLLM_MEDIA_URL_ALLOW_REDIRECTS", True)
+ monkeypatch.setattr(video_api_utils.httpx, "AsyncClient", _client_factory)
+ mocker.patch(
+ "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+ return_value=b"fake-video",
+ )
+
+ response = test_client.post(
+ "/v1/videos",
+ data={
+ "prompt": "A fox running through snow.",
+ "image_reference": json.dumps({"image_url": "https://example.com/redirect.png"}),
+ },
+ )
+
+ assert response.status_code == 200
+ video_id = response.json()["id"]
+ _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+ assert requested_paths == ["/redirect.png", "/image.png"]
+
+ engine = test_client.app.state.openai_serving_video._engine_client
+ input_image = engine.captured_prompt["multi_modal_data"]["image"]
+ assert isinstance(input_image, Image.Image)
+ assert input_image.size == (40, 24)
+
+
def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
@@ -692,6 +836,18 @@ def test_multi_video_generation_preserves_uploaded_files_until_generation(
assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0
+def test_mixed_reference_capability_uses_model_metadata_when_config_defaults_false(test_client):
+ handler = test_client.app.state.openai_serving_video
+ handler._engine_client.model_class_name = None
+ handler._engine_client.stage_configs = [
+ SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3TextEncoder"}),
+ SimpleNamespace(engine_args={"model_class_name": "MiniMaxH3Pipeline"}),
+ ]
+ handler._stage_configs = handler._engine_client.stage_configs
+
+ assert handler.supports_mixed_reference_inputs
+
+
def test_decode_video_bytes_can_keep_first_frames():
from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
@@ -1071,8 +1227,15 @@ def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker:
def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
audio_sample_rates = []
- def _fake_encode(video, fps, audio=None, audio_sample_rate=None, video_codec_options=None):
- del video, fps, audio, video_codec_options
+ def _fake_encode(
+ video,
+ fps,
+ audio=None,
+ audio_sample_rate=None,
+ video_codec_options=None,
+ frame_converter=None,
+ ):
+ del video, fps, audio, video_codec_options, frame_converter
audio_sample_rates.append(audio_sample_rate)
return b"fake-video"
@@ -2082,6 +2245,31 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
assert captured.quality == "high"
+def test_sync_sana_wm_extra_params_payload_passes_to_engine_prompt(test_client, mocker: MockerFixture):
+ _mock_encode_video_bytes(mocker)
+ response = test_client.post(
+ "/v1/videos/sync",
+ data={
+ "prompt": "drive forward",
+ "extra_params": json.dumps(
+ {
+ "sana_wm": {"action": "d-4", "rotation_speed_deg": 1.5},
+ "sana_wm_native_max_tokens": 30000,
+ }
+ ),
+ },
+ )
+
+ assert response.status_code == 200
+ engine = test_client.app.state.openai_serving_video._engine_client
+ captured_params = engine.captured_sampling_params_list[0]
+ # The camera block reaches the model through extra_args; the Sana preprocess
+ # hook is what lifts it onto the prompt.
+ assert captured_params.extra_args["sana_wm"]["action"] == "d-4"
+ assert captured_params.extra_args["sana_wm"]["rotation_speed_deg"] == 1.5
+ assert captured_params.extra_args["sana_wm_native_max_tokens"] == 30000
+
+
def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
"""Frame interpolation parameters should be forwarded on the sync path."""
encode_mock = _mock_encode_video_bytes(mocker)
diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py
index 6120d5ea4..2549889b3 100644
--- a/tests/entrypoints/test_async_omni.py
+++ b/tests/entrypoints/test_async_omni.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
import re
from types import SimpleNamespace
@@ -8,6 +11,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
from tests.helpers.mark import hardware_test
from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.outputs import OmniRequestOutput
@@ -77,7 +81,7 @@ def test_generate_submits_randomized_id_to_engine():
"""Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""
async def run():
- submitted_ids = []
+ submitted_ids: list[str] = []
omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))
req_id = "my-req-1"
@@ -106,8 +110,8 @@ def test_generate_forwards_lora_request_to_engine():
"""
async def run():
- submitted_ids = []
- submitted_loras = []
+ submitted_ids: list[str] = []
+ submitted_loras: list[LoRARequest | None] = []
omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras))
lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
@@ -143,7 +147,7 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
user provided request ID will be aborted."""
async def run():
- aborted_batches = []
+ aborted_batches: list[list[str]] = []
omni = get_async_omni_instance(
fake_abort_request=get_fake_abort(aborted_batches),
)
@@ -176,7 +180,6 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
assert len(aborted_ids) == expected_cancel_count
assert len(set(aborted_ids)) == expected_cancel_count
- assert len(omni.request_states) == len(req_ids) - expected_cancel_count
for t in tasks:
t.cancel()
@@ -185,11 +188,152 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi
asyncio.run(run())
+@pytest.mark.cpu
+def test_abort_keeps_request_states_until_generate_cleanup():
+ """Frontend abort must not drop request_states; generate() owns cleanup."""
+
+ async def run():
+ release = asyncio.Event()
+ seen_during_wait: list[int] = []
+
+ async def slow_abort_async(request_ids):
+ del request_ids
+ seen_during_wait.append(len(omni.request_states))
+ await release.wait()
+
+ omni = get_async_omni_instance(fake_abort_request=slow_abort_async)
+ omni.request_states["req-1-aaaa"] = SimpleNamespace(
+ request_id="req-1-aaaa",
+ external_request_id="req-1",
+ input_stream_task=None,
+ )
+
+ task = asyncio.create_task(omni.abort("req-1"))
+ await asyncio.sleep(0)
+ assert not task.done()
+ assert "req-1-aaaa" in omni.request_states
+ assert seen_during_wait == [1]
+
+ release.set()
+ await task
+ assert "req-1-aaaa" in omni.request_states
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_propagates_engine_errors_without_popping_state():
+ async def run():
+ async def failing_abort_async(request_ids):
+ del request_ids
+ raise RuntimeError("orchestrator abort failed")
+
+ omni = get_async_omni_instance(fake_abort_request=failing_abort_async)
+ omni.request_states["req-1-bbbb"] = SimpleNamespace(
+ request_id="req-1-bbbb",
+ external_request_id="req-1",
+ input_stream_task=None,
+ )
+
+ with pytest.raises(RuntimeError, match="orchestrator abort failed"):
+ await omni.abort("req-1")
+ assert "req-1-bbbb" in omni.request_states
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_enqueues_prefix_tokens_from_engine():
+ """AR abort must deliver the generated prefix on the request queue (CPU)."""
+
+ async def run():
+ from vllm.outputs import CompletionOutput
+
+ from vllm_omni.engine.messages import OutputMessage
+
+ queue: asyncio.Queue = asyncio.Queue()
+
+ async def abort_with_prefix(request_ids):
+ rid = request_ids[0]
+ engine_output = OmniRequestOutput(
+ request_id=rid,
+ finished=True,
+ stage_id=0,
+ final_output_type="text",
+ outputs=[
+ CompletionOutput(
+ index=0,
+ text="partial",
+ token_ids=[7, 8, 9],
+ cumulative_logprob=None,
+ logprobs=None,
+ finish_reason="abort",
+ stop_reason=None,
+ )
+ ],
+ )
+ return [
+ OutputMessage(
+ request_id=rid,
+ stage_id=0,
+ replica_id=None,
+ engine_outputs=engine_output,
+ metrics=None,
+ finished=True,
+ stage_submit_ts=None,
+ )
+ ]
+
+ omni = get_async_omni_instance(fake_abort_request=abort_with_prefix)
+ omni.request_states["req-1-cccc"] = SimpleNamespace(
+ request_id="req-1-cccc",
+ external_request_id="req-1",
+ input_stream_task=None,
+ queue=queue,
+ )
+ await omni._abort(["req-1-cccc"])
+ msg = queue.get_nowait()
+ assert msg.finished is True
+ assert list(msg.engine_outputs.outputs[0].token_ids) == [7, 8, 9]
+ assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+ assert "req-1-cccc" in omni.request_states
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_abort_enqueues_synthetic_finished_when_engine_returns_empty():
+ """Empty abort_async must still unblock generate() with finish_reason=abort."""
+
+ async def run():
+ queue: asyncio.Queue = asyncio.Queue()
+
+ async def empty_abort_async(request_ids):
+ del request_ids
+ return []
+
+ omni = get_async_omni_instance(fake_abort_request=empty_abort_async)
+ omni.request_states["req-1-dddd"] = SimpleNamespace(
+ request_id="req-1-dddd",
+ external_request_id="req-1",
+ input_stream_task=None,
+ queue=queue,
+ )
+ await omni._abort(["req-1-dddd"])
+ msg = queue.get_nowait()
+ assert msg.finished is True
+ assert msg.request_id == "req-1-dddd"
+ assert msg.engine_outputs.outputs[0].finish_reason == "abort"
+ assert "req-1-dddd" in omni.request_states
+
+ asyncio.run(run())
+
+
@pytest.mark.cpu
def test_generate_accepts_request_after_repeated_cancellations():
async def run_test():
- submitted_request_ids = []
- aborted_request_batches = []
+ submitted_request_ids: list[str] = []
+ aborted_request_batches: list[list[str]] = []
async def collect_outputs(request_id):
outputs = []
@@ -234,6 +378,43 @@ def test_generate_accepts_request_after_repeated_cancellations():
asyncio.run(run_test())
+@pytest.mark.cpu
+def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker):
+ async def run_test():
+ submitted_request_ids: list[str] = []
+ engine = object.__new__(AsyncOmniEngine)
+ engine._shutdown_called = True
+ engine.request_queue = mocker.Mock()
+ engine.num_stages = 1
+ engine.add_request_async = get_fake_add_request(submitted_request_ids)
+
+ omni = get_async_omni_instance()
+ omni.engine = engine
+
+ async def collect_outputs():
+ async for _ in omni.generate(
+ prompt={"prompt": "prompt"},
+ request_id="cancel-shutdown",
+ sampling_params_list=[SimpleNamespace()],
+ output_modalities=["image"],
+ ):
+ pass
+
+ task = asyncio.create_task(collect_outputs())
+ await asyncio.sleep(0)
+ task.cancel()
+
+ with pytest.raises(asyncio.CancelledError):
+ await task
+
+ assert len(submitted_request_ids) == 1
+ assert submitted_request_ids[0].startswith("cancel-shutdown-")
+ assert omni.request_states == {}
+ engine.request_queue.sync_q.put.assert_not_called()
+
+ asyncio.run(run_test())
+
+
@pytest.mark.cpu
def test_generate_yields_streaming_diffusion_chunks_before_final():
"""AsyncOmni.generate yields every intermediate diffusion chunk before the final one."""
@@ -364,7 +545,7 @@ async def test_omni_generate_request_id():
Text modality only; asserts caller-visible ids are preserved across the
multi-stage orchestrator path on H100.
"""
- engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
+ engine = AsyncOmni(model=OMNI_MODEL, deploy_config=OMNI_STAGE_CONFIG)
try:
for req_id in _OMNI_REQ_IDS:
async for output in engine.generate(
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 2fa8c92bd..71a6e00a5 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/tests/entrypoints/test_async_omni_diffusion_config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -33,6 +33,77 @@ def test_default_stage_config_includes_cache_backend():
assert engine_args["model_stage"] == "diffusion"
+def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ {
+ "ulysses_degree": 4,
+ "ulysses_a2a_permute": True,
+ }
+ )[0]
+
+ parallel_config = stage_cfg["engine_args"]["parallel_config"]
+ assert parallel_config.ulysses_degree == 4
+ assert parallel_config.ulysses_a2a_permute is True
+
+
+def test_default_stage_config_preserves_model_extras():
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
+
+ assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+
+
+def test_default_stage_config_preserves_and_overrides_promoted_extras():
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ {
+ "extras": {
+ "auxiliary_text_encoder": "/models/extras-llama",
+ "default_llama_model_id": "extras/default-llama",
+ },
+ "auxiliary_text_encoder": None,
+ }
+ )[0]
+
+ extras = stage_cfg["engine_args"]["extras"]
+ assert extras["auxiliary_text_encoder"] == "/models/extras-llama"
+ assert extras["default_llama_model_id"] == "extras/default-llama"
+
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ {
+ "extras": {
+ "auxiliary_text_encoder": "/models/extras-llama",
+ "default_llama_model_id": "extras/default-llama",
+ },
+ "auxiliary_text_encoder": "/models/top-level-llama",
+ "default_llama_model_id": "top-level/default-llama",
+ }
+ )[0]
+
+ extras = stage_cfg["engine_args"]["extras"]
+ assert extras["auxiliary_text_encoder"] == "/models/top-level-llama"
+ assert extras["default_llama_model_id"] == "top-level/default-llama"
+
+
+def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker):
+ """Local/unregistered Diffusers checkpoints still honor stage-0 extras."""
+
+ def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs):
+ return None, default_stage_cfg_factory(), None
+
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+ side_effect=resolve_with_default,
+ )
+ engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+ _, stage_configs = engine._resolve_stage_configs(
+ "/models/LTX-2.5-Diffusers",
+ {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'},
+ trust_remote_code=False,
+ )
+
+ assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+
+
def test_default_cache_config_used_when_missing():
"""Ensure default cache_config is synthesized when only backend is given."""
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -58,6 +129,23 @@ def test_default_stage_devices_from_sequence_parallel():
assert stage_cfg["runtime"]["devices"] == "0,1,2,3"
+def test_default_stage_devices_and_dp_from_num_gpus():
+ """Resolve omitted DP before deriving devices from an explicit GPU count."""
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+ {
+ "num_gpus": 8,
+ "tensor_parallel_size": 2,
+ "ulysses_degree": 2,
+ }
+ )[0]
+
+ parallel_config = stage_cfg["engine_args"]["parallel_config"]
+ assert parallel_config.data_parallel_size == 2
+ assert parallel_config.world_size == 8
+ assert stage_cfg["engine_args"]["num_gpus"] == 8
+ assert stage_cfg["runtime"]["devices"] == "0,1,2,3,4,5,6,7"
+
+
def test_default_stage_config_uses_parallel_size_kwargs():
"""Ensure default diffusion parallel_config uses CLI/API parallel sizes."""
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
@@ -76,8 +164,8 @@ def test_default_stage_config_uses_parallel_size_kwargs():
assert parallel_config.enable_expert_parallel is True
-def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
- """Ensure nullified diffusion parallel-size kwargs fall back to defaults."""
+def test_default_stage_config_preserves_omitted_dp_for_runtime_inference():
+ """Keep omitted DP unresolved until the runtime WORLD size is known."""
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
{
"pipeline_parallel_size": None,
@@ -92,7 +180,7 @@ def test_default_stage_config_defaults_nullified_parallel_size_kwargs():
parallel_config = stage_cfg["engine_args"]["parallel_config"]
assert parallel_config.pipeline_parallel_size == 1
- assert parallel_config.data_parallel_size == 1
+ assert parallel_config.data_parallel_size is None
assert parallel_config.tensor_parallel_size == 1
assert parallel_config.enable_expert_parallel is False
assert stage_cfg["engine_args"]["enforce_eager"] is False
@@ -208,7 +296,7 @@ def test_default_stage_config_whitelist_none_fallback():
engine_args = stage_cfg["engine_args"]
assert engine_args["trust_remote_code"] is False
- assert engine_args["distributed_executor_backend"] == "mp"
+ assert engine_args["distributed_executor_backend"] is None
assert engine_args["dtype"] == "auto"
assert engine_args["enforce_eager"] is False
@@ -240,6 +328,29 @@ def test_serve_cli_accepts_ulysses_mode():
assert parallel_config.ulysses_mode == "advanced_uaa"
+def test_serve_cli_accepts_text_encoder_tp_size():
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="command")
+ OmniServeCommand().subparser_init(subparsers)
+
+ args = parser.parse_args(
+ [
+ "serve",
+ "MiniMaxAI/MiniMax-H3",
+ "--omni",
+ "--text-encoder-tp-size",
+ "4",
+ ]
+ )
+
+ explicit_kwargs = args.get_explicit_kwargs_dict()
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ parallel_config = stage_cfg["engine_args"]["parallel_config"]
+
+ assert args.text_encoder_tp_size == 4
+ assert parallel_config.text_encoder_tp_size == 4
+
+
def test_serve_cli_forwards_model_defined_task_type_to_diffusion_stage():
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
@@ -345,6 +456,39 @@ def test_serve_cli_forwards_distributed_offload_residency():
assert engine_args["dlo_resident_layers"] == 20
+def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo():
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="command")
+ OmniServeCommand().subparser_init(subparsers)
+
+ args = parser.parse_args(
+ [
+ "serve",
+ "MiniMaxAI/MiniMax-H3",
+ "--omni",
+ "--enable-distributed-layerwise-offload",
+ "--dlo-no-use-allgather",
+ "--host-weight-runtime-mode",
+ "preferred",
+ "--host-weight-runtime-root",
+ "/var/cache/vllm-omni/hwr",
+ "--dlo-host-registration-limit-gib",
+ "80",
+ ]
+ )
+
+ explicit_kwargs = args.get_explicit_kwargs_dict()
+ stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+ engine_args = stage_cfg["engine_args"]
+
+ assert explicit_kwargs["host_weight_runtime_mode"] == "preferred"
+ assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+ assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80
+ assert engine_args["host_weight_runtime_mode"] == "preferred"
+ assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr"
+ assert engine_args["dlo_host_registration_limit_gib"] == 80
+
+
def test_serve_cli_accepts_diffusion_compile_controls():
"""Ensure both compile controls reach the diffusion stage."""
parser = TrackingArgumentParser()
@@ -383,7 +527,9 @@ def test_serve_cli_accepts_diffusion_attention_backend():
"Qwen/Qwen-Image",
"--omni",
"--diffusion-attention-backend",
- "FLASH_ATTN",
+ "FASTVIDEO_VSA",
+ "--fastvideo-vsa-topk",
+ "96",
]
)
@@ -391,10 +537,12 @@ def test_serve_cli_accepts_diffusion_attention_backend():
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
diffusion_attention_config = stage_cfg["engine_args"]["diffusion_attention_config"]
- assert args.diffusion_attention_backend == "FLASH_ATTN"
+ assert args.diffusion_attention_backend == "FASTVIDEO_VSA"
+ assert args.fastvideo_vsa_topk == 96
assert isinstance(diffusion_attention_config, AttentionConfig)
assert diffusion_attention_config.default is not None
- assert diffusion_attention_config.default.backend == "FLASH_ATTN"
+ assert diffusion_attention_config.default.backend == "FASTVIDEO_VSA"
+ assert diffusion_attention_config.default.backend_kwargs() == {"topk": 96}
def test_serve_cli_accepts_request_batch_max_wait_ms():
@@ -462,6 +610,66 @@ def test_serve_cli_accepts_additional_config():
assert engine_args["additional_config"] == {"torchair_graph_config": {"enabled": True}}
+def test_default_stage_resolves_video_output_from_checkpoint(mocker):
+ captured = {}
+
+ def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
+ del args, kwargs
+ captured.update(default_stage_cfg_factory()[0])
+ stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
+ return ("", [stage], None)
+
+ resolver = mocker.patch(
+ "vllm_omni.engine.async_omni_engine.resolve_model_class_name",
+ return_value="MiniMaxH3Pipeline",
+ )
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+ side_effect=resolve_with_default,
+ )
+ engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+ engine._strip_single_engine_args = lambda kwargs: kwargs
+
+ engine._resolve_stage_configs(
+ "/models/MiniMax-H3/FL2VA",
+ {},
+ trust_remote_code=False,
+ )
+
+ resolver.assert_called_once_with("/models/MiniMax-H3/FL2VA", "default")
+ assert captured["final_output_type"] == "video"
+
+
+def test_default_diffusers_stage_preserves_video_model_identity(mocker):
+ captured = {}
+
+ def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
+ del args, kwargs
+ captured.update(default_stage_cfg_factory()[0])
+ stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
+ return ("", [stage], None)
+
+ mocker.patch(
+ "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+ return_value={"_class_name": "WanImageToVideoPipeline"},
+ )
+ mocker.patch(
+ "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+ side_effect=resolve_with_default,
+ )
+ engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+ engine._strip_single_engine_args = lambda kwargs: kwargs
+
+ engine._resolve_stage_configs(
+ "/models/Wan2.2-I2V",
+ {"diffusion_load_format": "diffusers"},
+ trust_remote_code=False,
+ )
+
+ assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline"
+ assert captured["final_output_type"] == "video"
+
+
def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
"""Ensure YAML/deploy stage resolution forwards top-level additional_config."""
fake_diffusion_stage = SimpleNamespace(
@@ -478,12 +686,11 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
- engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
- "stage_configs_path": "dummy.yaml",
+ "deploy_config": "dummy.yaml",
"additional_config": {"torchair_graph_config": {"enabled": True}},
},
trust_remote_code=False,
@@ -493,6 +700,24 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
+@pytest.mark.parametrize(
+ ("legacy_arg", "value"),
+ [
+ ("stage_configs_path", "legacy.yaml"),
+ ("stage_configs", [{"stage_id": 0}]),
+ ],
+)
+def test_resolve_stage_configs_rejects_legacy_config_arguments(legacy_arg, value):
+ engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+ with pytest.raises(ValueError, match=rf"`{legacy_arg}`.*`deploy_config`"):
+ engine._resolve_stage_configs(
+ "dummy-model",
+ {legacy_arg: value},
+ trust_remote_code=False,
+ )
+
+
def test_default_stage_config_includes_quantization_config():
"""Ensure structured quantization_config survives default diffusion-stage creation."""
quantization_config = {
@@ -516,12 +741,11 @@ def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(
)
engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
- engine._strip_single_engine_args = lambda kwargs: kwargs
_, stage_configs = engine._resolve_stage_configs(
"dummy-model",
{
- "stage_configs_path": "dummy.yaml",
+ "deploy_config": "dummy.yaml",
"quantization_config": {"method": "bitsandbytes"},
},
trust_remote_code=False,
diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py
new file mode 100644
index 000000000..9600b990d
--- /dev/null
+++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py
@@ -0,0 +1,441 @@
+"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, call
+
+import pytest
+
+from vllm_omni.diffusion.data import CuMemTag, OmniACK
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model]
+
+
+def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
+ omni = object.__new__(AsyncOmni)
+ omni._pause_cond = asyncio.Condition()
+ omni._paused = False
+ omni._admitting = 0
+ omni._hold_admission_until_resume = False
+ omni._sleeping_tags = set()
+ omni._stage_sleeping_tags = {}
+ omni._level2_sleeping = False
+ omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
+ omni._final_output_handler = lambda: None
+ omni.reset_mm_cache = AsyncMock()
+
+ stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
+ omni.engine = SimpleNamespace(
+ stage_clients=stage_clients,
+ stage_vllm_configs=[None] * len(stage_types),
+ collective_rpc_async=AsyncMock(return_value=[True]),
+ )
+ omni.collective_rpc = AsyncMock(return_value=[True])
+ return omni
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type():
+ omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
+ ar_ids, diff_ids = omni._split_stage_ids_by_type()
+ assert ar_ids == [0, 2]
+ assert diff_ids == [1]
+
+
+@pytest.mark.cpu
+def test_split_stage_ids_by_type_rejects_out_of_range():
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"):
+ omni._split_stage_ids_by_type([0, 2])
+ with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"):
+ omni._split_stage_ids_by_type([-1])
+
+
+@pytest.mark.cpu
+def test_pause_generation_routes_ar_via_collective_rpc():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ omni.reset_prefix_cache = AsyncMock(return_value=True)
+ omni.reset_mm_cache = AsyncMock()
+ omni.reset_encoder_cache = AsyncMock()
+
+ await omni.pause_generation(mode="abort", clear_cache=True)
+
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+ omni.collective_rpc.assert_awaited_once_with(
+ method="pause_scheduler",
+ args=(),
+ kwargs={"mode": "abort", "clear_cache": True},
+ stage_ids=[0],
+ )
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_generation_still_rpcs_when_already_paused():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ omni.reset_prefix_cache = AsyncMock(return_value=True)
+ omni.reset_mm_cache = AsyncMock()
+ omni.reset_encoder_cache = AsyncMock()
+ omni._paused = True
+
+ await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0])
+
+ omni.collective_rpc.assert_awaited_once_with(
+ method="pause_scheduler",
+ args=(),
+ kwargs={"mode": "abort", "clear_cache": True},
+ stage_ids=[0],
+ )
+ omni.reset_prefix_cache.assert_awaited_once()
+ omni.reset_mm_cache.assert_awaited_once()
+ omni.reset_encoder_cache.assert_awaited_once()
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_resume_generation_resumes_ar_then_clears_frontend_pause():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ omni._paused = True
+
+ await omni.resume_generation()
+
+ omni.collective_rpc.assert_awaited_once_with(
+ method="resume_scheduler",
+ args=(),
+ kwargs=None,
+ stage_ids=[0],
+ )
+ assert omni._paused is False
+ assert omni._hold_admission_until_resume is False
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+ omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+ acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+
+ omni.collective_rpc.assert_awaited_once_with(
+ method="sleep",
+ args=(1, "abort"),
+ kwargs=None,
+ stage_ids=[0],
+ )
+ omni._sleep_diffusion.assert_awaited_once_with([1], 1)
+ assert {ack.stage_id for ack in acks} == {0, 1}
+ assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0)
+ assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+ assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+ # Sleep gates frontend admission for the trainer resume contract.
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_blocks_admission_before_engine_core_rpc():
+ """Problem 3: _paused must be set before awaiting sleep RPC."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+ paused_at_rpc: list[bool] = []
+
+ async def rpc_side_effect(**kwargs):
+ if kwargs.get("method") == "sleep":
+ paused_at_rpc.append(omni._paused)
+ return [True]
+
+ omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+ assert omni._paused is False
+ await omni.sleep(level=1, mode="abort")
+
+ assert paused_at_rpc == [True]
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+ # Sleep must not also call pause_scheduler (EngineCore.sleep pauses).
+ assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list)
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_waits_for_in_flight_generate_admission():
+ """Sleep must not offload while generate() is still in add_request."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+ omni._admitting = 1
+ rpc_started = asyncio.Event()
+
+ async def rpc_side_effect(**kwargs):
+ if kwargs.get("method") == "sleep":
+ rpc_started.set()
+ return [True]
+
+ omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+ sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+ await asyncio.sleep(0.05)
+ assert not sleep_task.done()
+ assert not rpc_started.is_set()
+ await omni._release_generate_admission()
+ await asyncio.wait_for(sleep_task, timeout=1)
+ assert rpc_started.is_set()
+ assert omni._admitting == 0
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "diffusion"])
+ omni._paused = True
+ omni._hold_admission_until_resume = True
+ omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
+ diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
+ omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack])
+
+ acks = await omni.wake_up(stage_ids=[0, 1])
+
+ omni.collective_rpc.assert_awaited_once()
+ wake_kwargs = omni.collective_rpc.await_args.kwargs
+ assert wake_kwargs["method"] == "wake_up"
+ assert wake_kwargs["stage_ids"] == [0]
+ assert set(wake_kwargs["kwargs"]["tags"]) == {
+ CuMemTag.WEIGHTS.value,
+ CuMemTag.KV_CACHE.value,
+ }
+ omni._wake_diffusion.assert_awaited_once()
+ assert {ack.stage_id for ack in acks} == {0, 1}
+ assert not omni._sleeping_tags
+ # Mixed/AR wake restores memory but does not resume frontend admission.
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_does_not_resume_frontend_admission():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+ await omni.sleep(level=1, mode="abort")
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+
+ await omni.wake_up()
+
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+ assert not omni._sleeping_tags
+ # Explicit resume is required after AR sleep/wake.
+ await omni.resume_generation()
+ assert omni._paused is False
+ assert omni._hold_admission_until_resume is False
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags():
+ """sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+
+ await omni.sleep(level=1, mode="abort")
+ assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+ assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
+
+ await omni.wake_up()
+ assert not omni._sleeping_tags
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
+ async def run() -> None:
+ omni = _make_omni(stage_types=["diffusion"])
+ omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+
+ await omni.sleep(level=1)
+
+ # AR EngineCore sleep path must not run for diffusion-only.
+ assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in (
+ omni.collective_rpc.await_args_list
+ )
+ omni._sleep_diffusion.assert_awaited_once_with([0], 1)
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is False
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_streaming_wait_for_first_chunk_does_not_block_sleep():
+ """Waiting for the next client chunk must not hold an admission slot."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+ rpc_started = asyncio.Event()
+ first_chunk = asyncio.Event()
+
+ async def rpc_side_effect(**kwargs):
+ if kwargs.get("method") == "sleep":
+ rpc_started.set()
+ return [True]
+
+ omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+ async def wait_then_submit() -> None:
+ await first_chunk.wait()
+ await omni._submit_with_admission(asyncio.sleep(0))
+
+ wait_task = asyncio.create_task(wait_then_submit())
+ sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+ await asyncio.sleep(0.05)
+ assert not wait_task.done()
+ assert omni._admitting == 0
+ await asyncio.wait_for(sleep_task, timeout=1)
+ assert rpc_started.is_set()
+
+ wait_task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await wait_task
+ assert omni._admitting == 0
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_streaming_in_flight_add_blocks_sleep():
+ """Sleep must wait while a streaming ADD/update holds an admission slot."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm"])
+ add_started = asyncio.Event()
+ add_release = asyncio.Event()
+ rpc_started = asyncio.Event()
+
+ async def rpc_side_effect(**kwargs):
+ if kwargs.get("method") == "sleep":
+ rpc_started.set()
+ return [True]
+
+ omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
+
+ async def slow_add() -> None:
+ add_started.set()
+ await add_release.wait()
+
+ add_task = asyncio.create_task(omni._submit_with_admission(slow_add()))
+ await add_started.wait()
+ sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
+ await asyncio.sleep(0.05)
+ assert not sleep_task.done()
+ assert not rpc_started.is_set()
+ assert omni._admitting == 1
+
+ add_release.set()
+ await asyncio.wait_for(add_task, timeout=1)
+ await asyncio.wait_for(sleep_task, timeout=1)
+ assert rpc_started.is_set()
+ assert omni._admitting == 0
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_wake_up_restores_admission_for_diffusion_only():
+ """Pure diffusion must keep sleep → wake → generate (no resume)."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["diffusion"])
+ omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+ omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+ await omni.sleep(level=1)
+ assert omni._paused is True
+
+ async def wait_for_admission() -> None:
+ async with omni._pause_cond:
+ await omni._pause_cond.wait_for(lambda: not omni._paused)
+
+ waiter = asyncio.create_task(wait_for_admission())
+ await omni.wake_up()
+ await asyncio.wait_for(waiter, timeout=1.0)
+ assert omni._paused is False
+ assert omni._hold_admission_until_resume is False
+ assert not omni._sleeping_tags
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
+ """Explicit pause_generation still requires resume after diffusion wake."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["diffusion"])
+ omni.reset_prefix_cache = AsyncMock(return_value=True)
+ omni.reset_mm_cache = AsyncMock()
+ omni.reset_encoder_cache = AsyncMock()
+ omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
+ omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
+
+ await omni.pause_generation()
+ await omni.sleep(level=1)
+ await omni.wake_up()
+
+ assert omni._paused is True
+ assert omni._hold_admission_until_resume is True
+ await omni.resume_generation()
+ assert omni._paused is False
+
+ asyncio.run(run())
+
+
+@pytest.mark.cpu
+def test_partial_wake_does_not_skip_remaining_sleeping_stage():
+ """sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1])."""
+
+ async def run() -> None:
+ omni = _make_omni(stage_types=["llm", "llm"])
+ await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
+ assert omni._stage_sleeping_tags.keys() == {0, 1}
+
+ await omni.wake_up(stage_ids=[0])
+ assert 0 not in omni._stage_sleeping_tags
+ assert 1 in omni._stage_sleeping_tags
+ assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
+
+ omni.collective_rpc.reset_mock()
+ acks = await omni.wake_up(stage_ids=[1])
+ assert acks
+ omni.collective_rpc.assert_awaited_once()
+ wake_kwargs = omni.collective_rpc.await_args.kwargs
+ assert wake_kwargs["method"] == "wake_up"
+ assert wake_kwargs["stage_ids"] == [1]
+ assert not omni._sleeping_tags
+ assert not omni._stage_sleeping_tags
+
+ asyncio.run(run())
diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py
index 3f793e89e..827ab67b6 100644
--- a/tests/entrypoints/test_omni_base_profiler.py
+++ b/tests/entrypoints/test_omni_base_profiler.py
@@ -46,6 +46,16 @@ class TestOmniBaseProfiler:
stage_ids=None,
)
+ def test_removed_diffusion_batch_size_fails_loudly(self, mock_engine, mocker: MockerFixture):
+ """Legacy diffusion_batch_size callers receive an actionable error."""
+ mocker.patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine)
+ mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x)
+ mocker.patch("vllm_omni.entrypoints.omni_base.weakref.finalize")
+ from vllm_omni.entrypoints.omni_base import OmniBase
+
+ with pytest.raises(TypeError, match="diffusion_batch_size.*max_num_seqs"):
+ OmniBase(model="test-model", diffusion_batch_size=8)
+
def test_start_profile_with_prefix(self, omni_base_instance, mock_engine):
"""Test that start_profile passes profile_prefix to collective_rpc."""
omni_base_instance.start_profile(profile_prefix="test_trace")
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index 161b3fe2a..4f9fc592b 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -12,12 +12,13 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
+from vllm_omni.config.stage_config import StageConfig
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
from vllm_omni.engine.messages import ErrorMessage, OutputMessage
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.client_request_state import ClientRequestState
from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
+from vllm_omni.entrypoints.omni_base import OmniBase, OmniEngineDeadError
from vllm_omni.errors import (
OmniClientError,
client_error_from_metadata,
@@ -188,14 +189,39 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine)
def _make_base():
- from vllm_omni.entrypoints.omni_base import OmniBase
-
obj = object.__new__(OmniBase)
obj.engine = MagicMock()
obj.request_states = {}
return obj
+def test_resolve_sampling_params_list_preserves_stage_constraints():
+ base = _make_base()
+ base.engine.num_stages = 1
+ base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])]
+ base.engine.stage_configs = [
+ StageConfig(
+ stage_id=0,
+ model_stage="dummy-model",
+ sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+ ).to_omegaconf()
+ ]
+ base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
+ assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}]
+ caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7])
+
+ resolved = base.resolve_sampling_params_list(caller_params)
+
+ assert resolved[0] is not caller_params
+ assert resolved[0].seed == 1234
+ assert resolved[0].max_tokens == 7
+ assert resolved[0].detokenize is False
+ assert resolved[0].stop_token_ids == [42]
+ assert 42 in resolved[0]._all_stop_token_ids
+ assert caller_params.detokenize is True
+ assert caller_params.stop_token_ids == [7]
+
+
def _stage_spec(
stage_id: int,
*,
@@ -585,7 +611,13 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
# external ID to avoid collisions, so this also tests mapping
external_req_id = "req-1"
req_id = "req-1-12345678"
+ recorded_failures = []
try:
+ monkeypatch.setattr(
+ app,
+ "_record_request_failure_once",
+ lambda request_id, reason: recorded_failures.append((request_id, reason)),
+ )
app.request_states[req_id] = ClientRequestState(
request_id=req_id,
external_request_id=external_req_id,
@@ -596,6 +628,7 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat
assert engine.aborted == [[req_id]]
assert external_req_id not in app.request_states
+ assert recorded_failures == [(req_id, "client_abort")]
@pytest.mark.asyncio
@@ -788,7 +821,13 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
_patch_engine(monkeypatch, engine)
app = Omni("dummy-model")
+ recorded_failures = []
try:
+ monkeypatch.setattr(
+ app,
+ "_record_request_failure_once",
+ lambda request_id, reason: recorded_failures.append((request_id, reason)),
+ )
app.request_states["req-1"] = object()
app.abort("req-1")
finally:
@@ -796,6 +835,7 @@ def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
assert engine.aborted == [["req-1"]]
assert "req-1" not in app.request_states
+ assert recorded_failures == [("req-1", "client_abort")]
def test_omni_forces_final_only_on_llm_stages(monkeypatch: pytest.MonkeyPatch):
diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py
index 0377e1a9a..9567cbc96 100644
--- a/tests/entrypoints/test_omni_sleep_mode.py
+++ b/tests/entrypoints/test_omni_sleep_mode.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Omni sleep mode: entrypoint-level VRAM/ACK tests plus multi-stage e2e.
@@ -109,11 +109,15 @@ def get_dynamic_devices(stage_idx: int, num_stages: int, tp_size: int) -> str:
async def _ensure_awake(engine: AsyncOmni, stage_ids: list[int]) -> None:
- """Best-effort full wake so the next shared-engine test starts active."""
+ """Best-effort full wake + resume so the next shared-engine test starts active."""
try:
await engine.wake_up(stage_ids=stage_ids)
except Exception as e:
logger.warning("ensure_awake failed (stage_ids=%s): %s", stage_ids, e)
+ try:
+ await engine.resume_generation(stage_ids=stage_ids)
+ except Exception as e:
+ logger.warning("ensure_resume failed (stage_ids=%s): %s", stage_ids, e)
def _build_llm_stages() -> tuple[list[dict], list[dict]]:
@@ -148,13 +152,13 @@ def _build_llm_stages() -> tuple[list[dict], list[dict]]:
@pytest.fixture(scope="module", autouse=True)
def _module_device_cleanup():
"""One device cleanup pass around the module (engines are class-scoped)."""
- from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+ from tests.helpers.clean import cleanup_test_environment
print("\n=== PRE-MODULE DEVICE CLEANUP (sleep_mode) ===")
- run_pre_test_cleanup()
+ cleanup_test_environment()
yield
print("\n=== POST-MODULE DEVICE CLEANUP (sleep_mode) ===")
- run_post_test_cleanup()
+ cleanup_test_environment()
@pytest_asyncio.fixture(scope="class", loop_scope="class")
@@ -240,6 +244,9 @@ class TestOmniLlmSleepMode:
try:
await llm_engine.sleep(stage_ids=[0], level=1)
await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+ # wake_up restores weights but keeps the AR admission hold.
+ # Resume so generate() can reach the sleeping-tags rejection.
+ await llm_engine.resume_generation(stage_ids=[0])
with pytest.raises(RuntimeError, match="partially or fully asleep"):
async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
pass
@@ -364,6 +371,8 @@ class TestOmniDiffusionSleepMode:
# --- wake + VRAM restore to initial (±3 GiB) ---
logger.info("Waking up (Reloading Weights)...")
await diffusion_engine.wake_up(stage_ids=[0])
+ # AR stages keep the frontend admission gate held after wake.
+ await diffusion_engine.resume_generation(stage_ids=[0])
await asyncio.sleep(2.0)
gc.collect()
get_vram_info(device_id)
@@ -490,6 +499,7 @@ async def test_multistage_llm_diffusion_sleep_wake(tp_size: int):
assert len(acks) == 2
await engine.wake_up(stage_ids=[0, 1])
+ await engine.resume_generation(stage_ids=[0, 1])
async for _ in engine.generate("verify", sampling_params_list=[SamplingParams(), sp]):
pass
finally:
@@ -522,6 +532,7 @@ async def test_pure_diffusion_sleep_wake():
try:
await engine.sleep(level=1)
await engine.wake_up()
+ await engine.resume_generation()
async for _ in engine.generate(
"test",
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
diff --git a/tests/entrypoints/test_omni_snapshot_download.py b/tests/entrypoints/test_omni_snapshot_download.py
index d2c76a249..eae774178 100644
--- a/tests/entrypoints/test_omni_snapshot_download.py
+++ b/tests/entrypoints/test_omni_snapshot_download.py
@@ -1,17 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""``omni_snapshot_download`` must honor vLLM's ``VLLM_USE_MODELSCOPE`` semantics.
-vLLM treats the flag as enabled only for the literal string ``"true"``
-(case-insensitive). Reading ``os.environ`` directly made every non-empty value
-truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0`` still took the
-ModelScope path.
+vLLM treats the flag as enabled for the literal strings ``"1"`` or ``"true"``
+(case-insensitive; see vllm.envs). Reading ``os.environ`` directly made every
+non-empty value truthy, so an explicit opt-out such as ``VLLM_USE_MODELSCOPE=0``
+still took the ModelScope path.
"""
import sys
import types
+from pathlib import Path
import pytest
+from pytest_mock import MockerFixture
from vllm import envs
from vllm_omni.entrypoints import omni_base
@@ -21,6 +23,64 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
MODEL_ID = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+@pytest.mark.parametrize(
+ "model_uri",
+ [
+ "s3://bucket/model",
+ "gs://bucket/model",
+ "az://bucket/model",
+ ],
+)
+def test_omni_snapshot_download_preserves_object_storage_uri(
+ model_uri: str,
+ mocker: MockerFixture,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
+ hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
+
+ assert omni_base.omni_snapshot_download(model_uri) == model_uri
+ hf_download.assert_not_called()
+
+
+def test_omni_snapshot_download_preserves_existing_local_path(
+ tmp_path: Path,
+ mocker: MockerFixture,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
+ absolute_model_path = tmp_path / "absolute-model"
+ absolute_model_path.mkdir()
+ relative_model_path = "relative-model"
+ (tmp_path / relative_model_path).mkdir()
+ monkeypatch.chdir(tmp_path)
+ hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
+
+ assert omni_base.omni_snapshot_download(str(absolute_model_path)) == str(absolute_model_path)
+ assert omni_base.omni_snapshot_download(relative_model_path) == relative_model_path
+ hf_download.assert_not_called()
+
+
+def test_omni_snapshot_download_uses_hf_for_relative_repo_id(
+ mocker: MockerFixture,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.delenv("VLLM_USE_MODELSCOPE", raising=False)
+ model_id = "org/model"
+ # Stub the modular-index probe so the test never reaches huggingface.co.
+ hf_file_probe = mocker.patch.object(omni_base, "file_or_path_exists", return_value=False)
+ hf_download = mocker.patch.object(omni_base, "download_weights_from_hf_specific")
+
+ assert omni_base.omni_snapshot_download(model_id) == model_id
+ hf_file_probe.assert_called_once_with(model_id, "modular_model_index.json", revision=None)
+ hf_download.assert_called_once_with(
+ model_name_or_path=model_id,
+ cache_dir=None,
+ allow_patterns=["*"],
+ require_all=True,
+ )
+
+
@pytest.fixture
def download_backend(monkeypatch: pytest.MonkeyPatch):
"""Run ``omni_snapshot_download`` and report which backend it selected.
@@ -67,7 +127,7 @@ def download_backend(monkeypatch: pytest.MonkeyPatch):
return run
-@pytest.mark.parametrize("value", ["0", "1", "False", "false", "no", "off"])
+@pytest.mark.parametrize("value", ["0", "False", "false", "no", "off"])
def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend, value):
monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
@@ -75,7 +135,7 @@ def test_non_true_values_do_not_enable_modelscope(monkeypatch, download_backend,
assert download_backend() == "huggingface"
-@pytest.mark.parametrize("value", ["true", "True", "TRUE"])
+@pytest.mark.parametrize("value", ["1", "true", "True", "TRUE"])
def test_true_values_enable_modelscope(monkeypatch, download_backend, value):
monkeypatch.setenv("VLLM_USE_MODELSCOPE", value)
diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py
index cc459c269..660203710 100644
--- a/tests/entrypoints/test_resolve_dreamzero_config.py
+++ b/tests/entrypoints/test_resolve_dreamzero_config.py
@@ -49,7 +49,9 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
monkeypatch.setattr(
"vllm.transformers_utils.config.get_hf_file_to_dict",
- lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
+ lambda path, _model, **_kwargs: (
+ None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
+ ),
)
od_config = OmniDiffusionConfig(
diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py
index 04a29edd1..e759e0f15 100644
--- a/tests/entrypoints/test_serve.py
+++ b/tests/entrypoints/test_serve.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Unit tests for the Omni serve CLI helpers."""
from __future__ import annotations
@@ -76,6 +79,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
assert args.cfg_parallel_size == 4
+def test_serve_parser_accepts_ulysses_a2a_permute() -> None:
+ parser = TrackingArgumentParser()
+ subparsers = parser.add_subparsers(dest="subcommand")
+ OmniServeCommand().subparser_init(subparsers)
+
+ args = parser.parse_args(["serve", "fake-model", "--omni", "--ulysses-a2a-permute"])
+
+ assert args.ulysses_a2a_permute is True
+ assert args.get_explicit_kwargs_dict()["ulysses_a2a_permute"] is True
+
+
def _make_headless_args(**kwargs) -> TrackingNamespace:
defaults = {
"model": "fake-model",
@@ -193,7 +207,7 @@ def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture)
with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
run_headless(args)
- assert captured["args"][1] is None
+ assert len(captured["args"]) == 2
assert captured["deploy_config_path"] == "/tmp/deploy.yaml"
assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py
index e2890797f..73702df3a 100644
--- a/tests/entrypoints/test_stage_utils.py
+++ b/tests/entrypoints/test_stage_utils.py
@@ -216,6 +216,20 @@ def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytes
assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4"
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_resolve_stage_physical_devices_captured_unset_ignores_live_env(monkeypatch: pytest.MonkeyPatch):
+ """Parallel init captures an unset CUDA_VISIBLE_DEVICES as None.
+
+ A sibling thread may then narrow the live env for spawn. The captured
+ None must not fall back to that live value, or logical device 0 would
+ map onto the other stage's GPU.
+ """
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
+ assert resolve_stage_physical_devices(0, "0", visible_baseline=None) == "0"
+ assert resolve_stage_physical_devices(1, "1", visible_baseline=None) == "1"
+
+
@pytest.mark.core_model
@pytest.mark.cpu
def test_map_device_list_raises_on_non_numeric():
diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py
index b02936838..97ed2370b 100644
--- a/tests/entrypoints/test_stream_finish_reason.py
+++ b/tests/entrypoints/test_stream_finish_reason.py
@@ -16,19 +16,20 @@ Key invariants tested:
- voice/speaker parameter compatibility in chat completions
"""
-import json
import time
from unittest.mock import MagicMock
import pytest
-from vllm.entrypoints.openai.chat_completion.protocol import (
- ChatCompletionRequest,
- ChatCompletionResponseStreamChoice,
-)
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
from vllm.outputs import CompletionOutput, RequestOutput
+from tests.helpers.serving_chat import (
+ build_serving_chat,
+ collect_stream,
+ make_request,
+ make_text_omni_output,
+ parse_sse_chunks,
+)
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
from vllm_omni.outputs import OmniRequestOutput
@@ -40,43 +41,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ---------------------------------------------------------------------------
-def _make_text_omni_output(
- request_id: str = "test-req",
- text: str = "hello",
- token_ids: list[int] | None = None,
- finish_reason: str | None = None,
- index: int = 0,
- num_prompt_tokens: int = 3,
-) -> OmniRequestOutput:
- """Build an OmniRequestOutput wrapping a text RequestOutput."""
- if token_ids is None:
- token_ids = [10, 11, 12]
- res = RequestOutput(
- request_id=request_id,
- prompt="test",
- prompt_token_ids=list(range(num_prompt_tokens)),
- prompt_logprobs=None,
- outputs=[
- CompletionOutput(
- index=index,
- text=text,
- token_ids=token_ids,
- cumulative_logprob=0.0,
- logprobs=None,
- finish_reason=finish_reason,
- stop_reason=None,
- )
- ],
- finished=finish_reason is not None,
- )
- return OmniRequestOutput.from_stage_output(
- res,
- request_id=request_id,
- final_output_type="text",
- finished=finish_reason is not None,
- )
-
-
def _make_audio_omni_output(
request_id: str = "test-req",
index: int = 0,
@@ -118,82 +82,6 @@ def _make_audio_omni_output(
)
-def _mock_audio_choices(index: int = 0, role: str = "assistant"):
- return [
- ChatCompletionResponseStreamChoice(
- index=index,
- delta=DeltaMessage(role=role, content="dGVzdA=="),
- logprobs=None,
- finish_reason="stop",
- )
- ]
-
-
-def _build_serving_chat():
- """Create a minimal OmniOpenAIServingChat for testing."""
- mock_engine = MagicMock()
- mock_engine.errored = False
-
- models = OpenAIServingModels(
- engine_client=mock_engine,
- base_model_paths=[],
- )
- mock_render = MagicMock()
-
- instance = OmniOpenAIServingChat(
- engine_client=mock_engine,
- models=models,
- response_role="assistant",
- online_renderer=mock_render,
- request_logger=None,
- chat_template=None,
- chat_template_content_format="auto",
- )
- instance._create_audio_choice = MagicMock(
- side_effect=lambda omni_res, role, request, stream=False: _mock_audio_choices(
- index=omni_res.outputs[0].index,
- role=role,
- )
- )
- return instance
-
-
-def _make_request(modalities: list[str], n: int = 1) -> ChatCompletionRequest:
- req = ChatCompletionRequest(
- model="test-model",
- messages=[{"role": "user", "content": "hello"}],
- n=n,
- stream=True,
- )
- req.modalities = modalities # type: ignore[attr-defined]
- return req
-
-
-def _parse_sse_chunks(lines: list[str]) -> list[dict]:
- """Parse SSE lines into JSON dicts."""
- prefix = "data: "
- chunks = []
- for line in lines:
- line = line.strip()
- if not line.startswith(prefix):
- continue
- payload = line[len(prefix) :].strip()
- if payload == "[DONE]":
- continue
- try:
- chunks.append(json.loads(payload))
- except json.JSONDecodeError:
- pass
- return chunks
-
-
-async def _collect_stream(gen):
- result = []
- async for item in gen:
- result.append(item)
- return result
-
-
# ---------------------------------------------------------------------------
# Tests: finish_reason correctness
# ---------------------------------------------------------------------------
@@ -202,14 +90,14 @@ async def _collect_stream(gen):
@pytest.mark.asyncio
async def test_single_modality_text_only_one_stop():
"""Text-only streaming: exactly one chunk has finish_reason='stop'."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text"])
async def result_generator():
- yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
- yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
+ yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
+ yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -221,7 +109,7 @@ async def test_single_modality_text_only_one_stop():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [c["choices"][0]["finish_reason"] for c in chunks if c.get("choices")]
assert finish_reasons[-1] == "stop"
@@ -233,15 +121,15 @@ async def test_single_modality_text_only_one_stop():
@pytest.mark.asyncio
async def test_multi_modal_text_audio_only_last_stop():
"""text+audio: text finish sends finish_reason=null, audio sends stop."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"])
async def result_generator():
- yield _make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
- yield _make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
+ yield make_text_omni_output(text="he", token_ids=[10, 11], finish_reason=None)
+ yield make_text_omni_output(text="llo", token_ids=[12], finish_reason="stop")
yield _make_audio_omni_output()
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -253,7 +141,7 @@ async def test_multi_modal_text_audio_only_last_stop():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
assert finish_reasons.count("stop") == 1
@@ -269,18 +157,18 @@ async def test_multi_modal_text_audio_only_last_stop():
@pytest.mark.asyncio
async def test_multi_modal_n2_independent_per_choice():
"""n=2 with text+audio: each choice gets exactly one stop, at the end."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"], n=2)
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"], n=2)
async def result_generator():
- yield _make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0)
- yield _make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1)
- yield _make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0)
- yield _make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1)
+ yield make_text_omni_output(text="A", token_ids=[10], finish_reason=None, index=0)
+ yield make_text_omni_output(text="B", token_ids=[20], finish_reason=None, index=1)
+ yield make_text_omni_output(text="", token_ids=[11], finish_reason="stop", index=0)
+ yield make_text_omni_output(text="", token_ids=[21], finish_reason="stop", index=1)
yield _make_audio_omni_output(index=0)
yield _make_audio_omni_output(index=1)
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -292,7 +180,7 @@ async def test_multi_modal_n2_independent_per_choice():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
per_choice: dict[int, list] = {}
for c in chunks:
for ch in c.get("choices", []):
@@ -306,13 +194,13 @@ async def test_multi_modal_n2_independent_per_choice():
@pytest.mark.asyncio
async def test_single_modality_audio_only_one_stop():
"""Audio-only streaming: the audio chunk carries finish_reason='stop'."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["audio"])
async def result_generator():
yield _make_audio_omni_output()
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -324,7 +212,7 @@ async def test_single_modality_audio_only_one_stop():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
assert finish_reasons.count("stop") == 1
@@ -349,8 +237,8 @@ async def test_streaming_audio_metrics_resolve_replica_id(
import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod
from vllm_omni.entrypoints.client_request_state import ClientRequestState
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["audio"])
req_state = ClientRequestState(
request_id="internal-req",
external_request_id="test-req",
@@ -377,7 +265,7 @@ async def test_streaming_audio_metrics_resolve_replica_id(
async def result_generator():
yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400)
- await _collect_stream(
+ await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -408,15 +296,15 @@ async def test_streaming_audio_metrics_resolve_replica_id(
async def test_declared_modality_not_produced_emits_fallback_stop():
"""If request.modalities declares ["text","audio"] but engine only produces
text, a fallback stop chunk is emitted at stream end."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"])
async def result_generator():
# Engine only produces text, no audio output at all
- yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
- yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+ yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+ yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -428,7 +316,7 @@ async def test_declared_modality_not_produced_emits_fallback_stop():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
# Text finish is suppressed (audio not seen yet), but fallback stop
@@ -441,15 +329,15 @@ async def test_declared_modality_not_produced_emits_fallback_stop():
async def test_declared_modality_not_produced_text_finish_suppressed():
"""When text finishes but audio (declared in modalities) never appears,
the text finish chunk has finish_reason=null (suppressed)."""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"])
async def result_generator():
- yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
- yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+ yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+ yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
# No audio output — stream ends
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -461,7 +349,7 @@ async def test_declared_modality_not_produced_text_finish_suppressed():
)
)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
# Find the text finish chunk (content "!")
for c in chunks:
@@ -485,8 +373,8 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
yields (field, value) tuples, so the loop raised
``AttributeError: 'tuple' object has no attribute 'finish_reason'``.
"""
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"])
empty_audio = _make_audio_omni_output()
empty_audio.outputs[0].multimodal_output = {"audio": []}
@@ -497,11 +385,11 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
serving_chat._create_audio_choice = create_audio_choice
async def result_generator():
- yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
- yield _make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
+ yield make_text_omni_output(text="hi", token_ids=[10], finish_reason=None)
+ yield make_text_omni_output(text="!", token_ids=[11], finish_reason="stop")
yield empty_audio
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -514,7 +402,7 @@ async def test_audio_chunk_without_waveform_keeps_stream_alive():
)
assert not any("Error in chat completion stream generator" in line for line in raw_lines)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}"
assert finish_reasons[-1] == "stop"
@@ -525,18 +413,18 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices():
"""Any ErrorResponse from the audio path is skipped, not iterated."""
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
- serving_chat = _build_serving_chat()
- request = _make_request(modalities=["text", "audio"])
+ serving_chat = build_serving_chat()
+ request = make_request(modalities=["text", "audio"])
serving_chat._create_audio_choice = MagicMock(
side_effect=lambda omni_res, role, request, stream=False: serving_chat._create_error_response("boom")
)
async def result_generator():
- yield _make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop")
+ yield make_text_omni_output(text="hi", token_ids=[10], finish_reason="stop")
yield _make_audio_omni_output()
- raw_lines = await _collect_stream(
+ raw_lines = await collect_stream(
serving_chat.chat_completion_stream_generator(
request=request,
result_generator=result_generator(),
@@ -550,7 +438,7 @@ async def test_audio_choice_error_response_is_not_iterated_as_choices():
assert isinstance(serving_chat._create_error_response("boom"), ErrorResponse)
assert not any("AttributeError" in line for line in raw_lines)
- chunks = _parse_sse_chunks(raw_lines)
+ chunks = parse_sse_chunks(raw_lines)
finish_reasons = [ch["finish_reason"] for c in chunks for ch in c.get("choices", [])]
assert finish_reasons.count("stop") == 1, f"Expected 1 stop, got {finish_reasons}"
diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py
index c6d8e0354..1fc4779a7 100644
--- a/tests/entrypoints/test_utils.py
+++ b/tests/entrypoints/test_utils.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Unit tests for vllm_omni.entrypoints.utils module."""
import os
@@ -22,11 +25,12 @@ from vllm_omni.engine.stage_init_utils import (
from vllm_omni.entrypoints.utils import (
_convert_dataclasses_to_dict,
_filter_dict_like_object,
+ _try_resolve_omni_model_type,
coerce_param_message_types,
filter_dataclass_kwargs,
filter_stages,
load_and_resolve_stage_configs,
- load_stage_configs_from_yaml,
+ load_stage_configs_from_model,
resolve_model_config_path,
)
@@ -371,17 +375,60 @@ class TestResolveModelConfigPath:
assert sender["extra"]["connector_get_max_wait_first_chunk"] == 3000
assert sender["extra"]["connector_get_max_wait"] == 300
+ def test_object_storage_uri_resolves_via_materialized_configs(
+ self,
+ mocker: MockerFixture,
+ tmp_path,
+ ):
+ """An object-storage URI must not reach HF helpers; resolution reads the
+ locally materialized config files instead."""
+ hf_config = SimpleNamespace(
+ model_type="qwen3_omni_moe",
+ architectures=["Qwen3OmniMoeForConditionalGeneration"],
+ )
+ uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
+ materialize = mocker.patch(
+ "vllm_omni.entrypoints.utils._materialize_object_storage_configs",
+ return_value=str(tmp_path),
+ )
+ get_config = mocker.patch(
+ "vllm_omni.entrypoints.utils.get_config",
+ return_value=hf_config,
+ )
+
+ result = resolve_model_config_path(uri)
+
+ materialize.assert_called_once_with(uri)
+ get_config.assert_called_once()
+ assert get_config.call_args[0][0] == str(tmp_path)
+ assert result is not None
+ assert Path(result).as_posix().endswith("qwen3_omni_moe.yaml")
+
+
+class TestTryResolveOmniModelType:
+ """Name matching must only scan the last path component."""
+
+ def test_uri_bucket_name_cannot_hijack_pipeline(self):
+ # "qwen3-tts-models" normalizes to "qwen3ttsmodels", which contains the
+ # registered key "qwen3tts" — only the basename may participate, so a
+ # misleadingly named bucket selects nothing.
+ assert _try_resolve_omni_model_type("s3://qwen3-tts-models/plain-checkpoint") is None
+
+ def test_basename_still_matches_pipeline(self):
+ assert _try_resolve_omni_model_type("gs://any-bucket/Fun-CosyVoice3-0.5B-2512") == "cosyvoice3"
+
class TestLoadAndResolveStageConfigs:
def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
- """Ensure that dtype survives default stage creation."""
- kwargs = {"dtype": torch.float32}
+ """Ensure that kwargs survive default stage creation."""
+ engine_backend = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+ kwargs = {"dtype": torch.float32, "engine_backend": engine_backend}
+ kwargs["revision"] = "pinned-revision"
mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="black-forest-labs/FLUX.2-klein-4B",
- stage_configs_path=None,
kwargs=kwargs,
trust_remote_code=False,
default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
@@ -389,10 +436,10 @@ class TestLoadAndResolveStageConfigs:
assert config_path is None
assert len(stage_configs) == 1
assert "dtype" in stage_configs[0]["engine_args"]
+ assert stage_configs[0]["engine_args"]["engine_backend"] == engine_backend
+ assert stage_configs[0]["engine_args"]["revision"] == "pinned-revision"
- def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
- self, tmp_path, mocker: MockerFixture
- ):
+ def test_deploy_config_preserves_cli_overrides_and_replicas(self, tmp_path, mocker: MockerFixture):
deploy_path = tmp_path / "qwen3_multi.yaml"
deploy_path.write_text(
'stages:\n - stage_id: 0\n devices: "0"\n - stage_id: 1\n devices: "1,2,3"\n num_replicas: 3\n',
@@ -413,18 +460,23 @@ class TestLoadAndResolveStageConfigs:
"vllm_omni.entrypoints.utils.load_stage_configs_from_model",
return_value=(returned_stage_configs, None),
)
+ cli_overrides = {
+ "interleave_mm_strings": True,
+ "media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}},
+ "gpu_memory_utilization": 0.55,
+ }
config_path, stage_configs, _ = load_and_resolve_stage_configs(
model="dummy-model",
- stage_configs_path=str(deploy_path),
- kwargs={},
+ kwargs=cli_overrides,
trust_remote_code=True,
+ deploy_config_path=str(deploy_path),
)
load_stage_configs.assert_called_once_with(
"dummy-model",
trust_remote_code=True,
- base_engine_args={},
+ base_engine_args=cli_overrides,
deploy_config_path=str(deploy_path),
stage_overrides=None,
strategy_config_path=None,
@@ -473,71 +525,20 @@ class TestLoadAndResolveStageConfigs:
assert filtered[0].final_output_type is None
-class TestLoadStageConfigsFromYaml:
- """Regression tests for stage-config loading and merging."""
-
- def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
- yaml_config = create_config(
- {
- "async_chunk": True,
- "stage_args": [
- {
- "stage_id": 0,
- "runtime": {"device": 0},
- "engine_args": {
- "parallel_config": {"tensor_parallel_size": 4},
- },
- }
- ],
- }
- )
+class TestLoadStageConfigsFromModel:
+ def test_unresolved_model_does_not_fall_back_to_yaml(self, mocker: MockerFixture, caplog):
mocker.patch(
- "vllm_omni.entrypoints.utils.load_yaml_config",
- return_value=yaml_config,
+ "vllm_omni.entrypoints.utils.StageConfigFactory.create_legacy_stage_configs_from_model",
+ return_value=(None, None),
)
+ resolve_path = mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path")
- stages = load_stage_configs_from_yaml(
- "fake.yaml",
- base_engine_args={
- "parallel_config": {
- "tensor_parallel_size": 1,
- "pipeline_parallel_size": 2,
- },
- "model": "base-model",
- },
- )
+ result = load_stage_configs_from_model("unregistered-model", trust_remote_code=False)
- merged_engine_args = stages[0]["engine_args"]
- assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
- assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
- assert merged_engine_args["model"] == "base-model"
- assert merged_engine_args["async_chunk"] is True
-
- def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
- yaml_config = create_config(
- {
- "stage_args": [
- {
- "stage_id": 0,
- "engine_args": {
- "nested": {"override": 2},
- },
- }
- ],
- }
- )
- mocker.patch(
- "vllm_omni.entrypoints.utils.load_yaml_config",
- return_value=yaml_config,
- )
-
- stages = load_stage_configs_from_yaml(
- "fake.yaml",
- base_engine_args={"nested": {"base": 1}},
- )
-
- assert stages[0]["engine_args"]["nested"]["base"] == 1
- assert stages[0]["engine_args"]["nested"]["override"] == 2
+ assert result == ([], None)
+ resolve_path.assert_not_called()
+ assert "No registered PipelineConfig resolved" in caplog.text
+ assert "deploy_config" in caplog.text
class TestCumulativeStreamingCoercion:
diff --git a/tests/examples/offline_inference/test_image_task_prompts.py b/tests/examples/offline_inference/test_image_task_prompts.py
index 57e8f1560..054756e33 100644
--- a/tests/examples/offline_inference/test_image_task_prompts.py
+++ b/tests/examples/offline_inference/test_image_task_prompts.py
@@ -1,17 +1,29 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
+import functools
+import importlib.util
+
import pytest
from PIL import Image
-from examples.offline_inference.image_to_video.image_to_video import build_image_to_video_prompt
-from examples.offline_inference.text_to_image.text_to_image import build_text_to_image_prompt
+from tests.examples.helpers import EXAMPLES
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@functools.cache
+def _load_example_module(relative_path: str, module_name: str):
+ path = EXAMPLES / relative_path
+ spec = importlib.util.spec_from_file_location(module_name, path)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
@@ -20,7 +32,11 @@ def test_text_to_image_builds_canonical_prompt(
negative_prompt: str | None,
expected_negative_prompt: str | None,
) -> None:
- result = build_text_to_image_prompt("a red fox", negative_prompt)
+ mod = _load_example_module(
+ "offline_inference/text_to_image/text_to_image.py",
+ "text_to_image_example",
+ )
+ result = mod.build_text_to_image_prompt("a red fox", negative_prompt)
assert result["prompt"] == "a red fox"
assert result["modalities"] == ["image"]
@@ -31,9 +47,13 @@ def test_text_to_image_builds_canonical_prompt(
def test_image_to_video_builds_canonical_prompt() -> None:
+ mod = _load_example_module(
+ "offline_inference/image_to_video/image_to_video.py",
+ "image_to_video_example",
+ )
image = Image.new("RGB", (32, 16), "red")
- result = build_image_to_video_prompt(
+ result = mod.build_image_to_video_prompt(
prompt="the fox turns toward the camera",
negative_prompt="flicker",
media_inputs={"image": image},
@@ -45,3 +65,38 @@ def test_image_to_video_builds_canonical_prompt() -> None:
"multi_modal_data": {"image": image},
"negative_prompt": "flicker",
}
+
+
+@pytest.mark.parametrize(
+ ("relative_path", "module_name"),
+ [
+ ("offline_inference/text_to_video/text_to_video.py", "text_to_video_output_example"),
+ ("offline_inference/image_to_video/image_to_video.py", "image_to_video_output_example"),
+ ],
+)
+@pytest.mark.parametrize("output_type", ["video", "image"])
+def test_video_examples_accept_semantic_and_legacy_output_types(
+ relative_path: str,
+ module_name: str,
+ output_type: str,
+) -> None:
+ mod = _load_example_module(relative_path, module_name)
+
+ mod._validate_video_output_type(output_type)
+
+
+@pytest.mark.parametrize(
+ ("relative_path", "module_name"),
+ [
+ ("offline_inference/text_to_video/text_to_video.py", "text_to_video_invalid_output_example"),
+ ("offline_inference/image_to_video/image_to_video.py", "image_to_video_invalid_output_example"),
+ ],
+)
+def test_video_examples_reject_unrelated_output_types(
+ relative_path: str,
+ module_name: str,
+) -> None:
+ mod = _load_example_module(relative_path, module_name)
+
+ with pytest.raises(ValueError, match="Unexpected output type 'audio'"):
+ mod._validate_video_output_type("audio")
diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py
index 9b5597c3e..5ee2a2fae 100644
--- a/tests/examples/offline_inference/test_text_to_image.py
+++ b/tests/examples/offline_inference/test_text_to_image.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Offline inference tests: text-to-image.
See examples/offline_inference/text_to_image/README.md
@@ -26,9 +29,8 @@ EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_t2i"
def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]:
if h2_title == "Web UI Demo":
return True, f"README section '{h2_title}' is intentionally excluded for examples tests"
- # Krea-2 Turbo/Raw README snippets currently fail in CI (more_cli_examples_006/007).
- if "krea/Krea-2-Turbo" in code or "krea/Krea-2-Raw" in code:
- return True, "Krea-2 README snippets are temporarily skipped (CI failures)"
+ if "krea/Krea-2-Raw" in code:
+ return True, "krea/Krea-2-Raw is gated on Hugging Face Hub; CI has no access"
return False, ""
diff --git a/tests/examples/offline_inference/test_text_to_video_prompts.py b/tests/examples/offline_inference/test_text_to_video_prompts.py
index ff7244d65..c588b7a38 100644
--- a/tests/examples/offline_inference/test_text_to_video_prompts.py
+++ b/tests/examples/offline_inference/test_text_to_video_prompts.py
@@ -3,19 +3,27 @@
from __future__ import annotations
+import functools
+import importlib.util
+
import pytest
import torch
-from examples.offline_inference.text_to_video.text_to_video import (
- _MODEL_PRESETS,
- _detect_preset,
- _normalize_float_tensor,
- build_text_to_video_prompt,
-)
+from tests.examples.helpers import EXAMPLES
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@functools.cache
+def _load_text_to_video():
+ path = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.py"
+ spec = importlib.util.spec_from_file_location("text_to_video_example", path)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
@pytest.mark.parametrize(
("negative_prompt", "expected_negative_prompt"),
[(None, None), ("", ""), ("blurry", "blurry")],
@@ -24,7 +32,8 @@ def test_text_to_video_builds_canonical_prompt(
negative_prompt: str | None,
expected_negative_prompt: str | None,
) -> None:
- result = build_text_to_video_prompt("a robot waves", negative_prompt)
+ mod = _load_text_to_video()
+ result = mod.build_text_to_video_prompt("a robot waves", negative_prompt)
assert result["prompt"] == "a robot waves"
assert result["modalities"] == ["video"]
@@ -37,6 +46,9 @@ def test_text_to_video_builds_canonical_prompt(
@pytest.mark.parametrize(
("model", "model_class_name", "preset_name"),
[
+ ("Efficient-Large-Model/SANA-Video_2B_480p_diffusers", None, "sana_480p"),
+ ("Efficient-Large-Model/SANA-Video_2B_720p_diffusers", None, "sana_720p"),
+ ("/models/custom-checkpoint", "SanaVideoPipeline", "sana_480p"),
("robbyant/lingbot-video-dense-1.3b", None, "lingbot"),
("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"),
("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"),
@@ -50,11 +62,13 @@ def test_detect_preset_is_scoped_to_model_family(
model_class_name: str | None,
preset_name: str,
) -> None:
- assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name]
+ mod = _load_text_to_video()
+ assert mod._detect_preset(model, model_class_name) is mod._MODEL_PRESETS[preset_name]
def test_lingbot_preset_matches_lingbot_defaults() -> None:
- assert _MODEL_PRESETS["lingbot"] == {
+ mod = _load_text_to_video()
+ assert mod._MODEL_PRESETS["lingbot"] == {
"model_class_name": "LingBotVideoPipeline",
"height": 192,
"width": 320,
@@ -68,24 +82,27 @@ def test_lingbot_preset_matches_lingbot_defaults() -> None:
def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None:
+ mod = _load_text_to_video()
frames = torch.tensor([0.0, 0.25, 1.0, 1.5])
- assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
+ assert torch.equal(mod._normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None:
+ mod = _load_text_to_video()
frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5])
assert torch.equal(
- _normalize_float_tensor(frames, "negative_one_to_one"),
+ mod._normalize_float_tensor(frames, "negative_one_to_one"),
torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
)
def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None:
+ mod = _load_text_to_video()
frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])]
- normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
+ normalized = [mod._normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0]))
assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0]))
diff --git a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py
index 93fc4db38..dfb7b32b1 100644
--- a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py
+++ b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py
@@ -1,3 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
import base64
import importlib.util
import sys
@@ -169,6 +173,87 @@ def test_open_streaming_response_requires_post_commit_drain():
)
+def test_final_input_commit_checkpoints_playback_first():
+ demo = _load_demo_module()
+ client = demo.RealtimeDuplexClient("ws://unused")
+ client.events.add(
+ {
+ "type": "response.created",
+ "response": {"id": "resp-zero-audio"},
+ }
+ )
+ calls = []
+
+ async def send(event):
+ calls.append(event)
+
+ client.send = send
+ commit_sent_at_s = asyncio.run(demo._commit_input_after_playback_checkpoint(client))
+
+ assert [event["type"] for event in calls] == ["playback.ack", "input_audio_buffer.commit"]
+ assert calls[0]["response_id"] == "resp-zero-audio"
+ assert calls[0]["played_ms"] == 0
+ assert isinstance(commit_sent_at_s, float)
+
+
+def _asymmetric_image():
+ from PIL import Image
+
+ image = Image.new("RGB", (4, 2), (0, 0, 0))
+ image.putpixel((0, 0), (255, 0, 0))
+ return image
+
+
+@pytest.mark.parametrize(
+ "rotation,expected_corner",
+ [
+ (0, (0, 0)),
+ (90, (0, 3)),
+ (180, (3, 1)),
+ (-90, (1, 0)),
+ (270, (1, 0)),
+ (-270, (0, 3)),
+ ],
+)
+def test_upright_applies_the_display_matrix_quarter_turn(rotation, expected_corner):
+ demo = _load_demo_module()
+
+ rotated = demo._upright(_asymmetric_image(), rotation)
+
+ assert rotated.getpixel(expected_corner) == (255, 0, 0)
+
+
+def test_decoded_frames_apply_the_container_display_matrix(monkeypatch, tmp_path):
+ """A -90 phone clip must reach the model upright, not on its side."""
+ demo = _load_demo_module()
+ source = _asymmetric_image()
+
+ class _Frame:
+ rotation = -90
+
+ def to_image(self):
+ return source
+
+ class _Container:
+ streams = type("_Streams", (), {"video": [type("_Stream", (), {})()]})()
+
+ def decode(self, stream):
+ yield _Frame()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ return False
+
+ monkeypatch.setitem(sys.modules, "av", type("_Av", (), {"open": staticmethod(lambda path: _Container())}))
+
+ (image,) = demo._decode_video_frames_rgb(tmp_path / "clip.mp4", [0])
+
+ assert image.size == (2, 4)
+ assert image.getpixel((1, 0)) == (255, 0, 0)
+
+
def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, capsys):
demo = _load_demo_module()
writer = demo._StreamingOutputWriter(tmp_path)
diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py
index 9929beab8..9fa313409 100644
--- a/tests/helpers/__init__.py
+++ b/tests/helpers/__init__.py
@@ -1,6 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Shared, importable test helper utilities.
-Submodules (``assertions``, ``env``, ``media``, ``runtime``, …) are imported
+Submodules (``assertions``, ``clean``, ``client``, ``media``, ``runtime``, …) are imported
explicitly by callers. Unit tests for these helpers live under ``tests/``.
Avoid star-importing everything here: that ran before refactor only inside the
old monolithic ``conftest``; a greedy ``__init__`` changes import order and can
@@ -11,11 +14,16 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
import pytest
-def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
+def skip_if_gated_repo_inaccessible(
+ repo_id: str,
+ *,
+ revision: str | None = None,
+ filename: str = "config.json",
+) -> None:
"""Skip the test if a gated HuggingFace repo is not accessible.
- Tries to download the model's config.json via ``hf_hub_download``,
- which performs an actual file-access check (unlike ``HfApi().model_info()``
+ Tries to download ``filename`` via ``hf_hub_download``, which performs an
+ actual file-access check (unlike ``HfApi().model_info()``
that only checks metadata). If the token has metadata access but not
file-download access, ``hf_hub_download`` will raise ``GatedRepoError``
and we skip cleanly.
@@ -26,7 +34,7 @@ def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
except Exception:
return
try:
- hf_hub_download(repo_id=repo_id, filename="config.json")
+ hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
except GatedRepoError as exc:
pytest.skip(
f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py
index 3b6fbe137..755cf0a4e 100644
--- a/tests/helpers/assertions.py
+++ b/tests/helpers/assertions.py
@@ -1,16 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Assertion and response validation helpers for tests."""
+import base64
import io
import json
+import math
import tempfile
import threading
import wave
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
if TYPE_CHECKING:
- from tests.helpers.runtime import DiffusionResponse
+ from tests.helpers.client import DiffusionResponse
import av
import numpy as np
@@ -20,12 +28,112 @@ from PIL import Image
from tests.helpers.media import (
convert_audio_bytes_to_text,
cosine_similarity_text,
- decode_b64_image,
preprocess_text,
)
+_ResponseT = TypeVar("_ResponseT")
+
_GENDER_PIPELINE = None
_GENDER_PIPELINE_LOCK = threading.Lock()
+
+# Assertions that sample a random variable rather than testing whether serving works:
+# whether TTS audio is intelligible enough to transcribe back, and whether a preset voice's
+# timbre estimates as the expected gender. Both are decided by the model and the estimator,
+# so a caller measuring a success rate has to tell them apart from a serving failure. The
+# messages are constants because a caller matching the wording by hand would silently stop
+# matching if it were reworded, reclassifying every quality failure as a hard failure.
+AUDIO_MISMATCH_MESSAGE = "The audio content is not same as the text"
+GENDER_MISMATCH_MESSAGE = "estimated gender is"
+_QUALITY_FAILURE_MESSAGES = (AUDIO_MISMATCH_MESSAGE, GENDER_MISMATCH_MESSAGE)
+
+
+@dataclass(frozen=True)
+class SuccessRateGate:
+ """Gate a sampled quality check on how many requests pass, not on each one passing.
+
+ Built by the send helpers from their own arguments, so a test states a policy and never
+ touches this class. See ``collect_at_success_rate``.
+
+ Attributes:
+ min_successes: Fail below this many successes out of ``request_num``. Predeclare it
+ from a measured baseline; tuning it until CI passes defeats the gate.
+ concurrency: In-flight requests, in batches until the sample count is reached.
+ Defaults to sending all of them at once.
+ """
+
+ min_successes: int
+ concurrency: int | None = None
+
+
+def is_quality_failure(exc: BaseException) -> bool:
+ """True when ``exc`` is one of the sampled quality assertions above.
+
+ A caller absorbing a few of these into a success-rate budget must not absorb an HTTP
+ failure, missing audio or a text-keyword miss as well, or a broken server passes.
+ """
+ if not isinstance(exc, AssertionError):
+ return False
+ return any(message in str(exc) for message in _QUALITY_FAILURE_MESSAGES)
+
+
+def collect_at_success_rate(
+ sample: Callable[[], _ResponseT],
+ gate: SuccessRateGate,
+ *,
+ request_num: int,
+ report: Callable[[str], None] = print,
+) -> list[_ResponseT]:
+ """Send ``request_num`` samples and gate on how many passed, returning the successes.
+
+ ``sample`` sends one request and asserts it. Only the sampled quality assertions are
+ counted; anything else propagates immediately. Does not retry and does not exit early,
+ since both would bias the rate being measured.
+ """
+ if gate.min_successes > request_num:
+ raise ValueError(f"min_successes={gate.min_successes} exceeds request_num={request_num}, gate can never pass")
+ concurrency = min(gate.concurrency or request_num, request_num)
+ responses: list[_ResponseT] = []
+ failures: list[str] = []
+ lock = threading.Lock()
+
+ def _one_sample(_: int) -> None:
+ try:
+ response = sample()
+ except AssertionError as exc:
+ if not is_quality_failure(exc):
+ raise
+ with lock:
+ failures.append(str(exc))
+ return
+ with lock:
+ responses.append(response)
+
+ for start in range(0, request_num, concurrency):
+ size = min(concurrency, request_num - start)
+ with ThreadPoolExecutor(max_workers=size) as executor:
+ for _ in executor.map(_one_sample, range(size)):
+ pass
+
+ successes = request_num - len(failures)
+ report(f"success rate: {successes}/{request_num} succeeded")
+ assert successes >= gate.min_successes, (
+ f"success rate {successes}/{request_num} is below the "
+ f"{gate.min_successes}/{request_num} gate "
+ f"(95% CI {wilson_interval(successes, request_num)}). Failures:\n"
+ + "\n".join(f" - {failure}" for failure in failures)
+ )
+ return responses
+
+
+def wilson_interval(successes: int, total: int, z: float = 1.96) -> str:
+ """Format a 95% Wilson score interval, so a rate failure shows how marginal it is."""
+ p = successes / total
+ denom = 1 + z * z / total
+ centre = (p + z * z / (2 * total)) / denom
+ half = z / denom * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total))
+ return f"{max(0.0, centre - half):.1%}-{min(1.0, centre + half):.1%}"
+
+
# Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
# short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
# leading SFX token), which flakes the deterministic similarity gate. Short
@@ -66,7 +174,7 @@ def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
def assert_image_diffusion_response(
response: "DiffusionResponse",
request_config: dict[str, Any],
- run_level: str = None,
+ run_level: str | None = None,
) -> None:
"""
Validate image diffusion response.
@@ -138,14 +246,15 @@ def assert_images_generations_response(
assert isinstance(item, dict), "Image generation data entries must be objects"
b64_json = item.get("b64_json")
assert isinstance(b64_json, str) and b64_json, "Image generation response is missing b64_json"
- image = decode_b64_image(b64_json)
+ image = Image.open(io.BytesIO(base64.b64decode(b64_json)))
+ image.load()
assert_image_valid(image, width=width, height=height)
def assert_video_diffusion_response(
response: "DiffusionResponse",
request_config: dict[str, Any],
- run_level: str = None,
+ run_level: str | None = None,
) -> None:
"""
Validate video diffusion response.
@@ -214,7 +323,7 @@ def assert_video_first_frame_matches(
def assert_audio_diffusion_response(
response: "DiffusionResponse",
request_config: dict[str, Any],
- run_level: str = None,
+ run_level: str | None = None,
) -> None:
"""
Validate audio diffusion response.
@@ -474,7 +583,7 @@ def _assert_preset_voice_gender_from_audio(
print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
if estimated_gender != "unknown":
assert estimated_gender == expected_gender, (
- f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
+ f"{voice_name!r} is expected {expected_gender}, but {GENDER_MISMATCH_MESSAGE} {estimated_gender!r}"
)
@@ -529,7 +638,7 @@ def _response_has_audio_output(response: Any) -> bool:
return bool(audio_data)
-def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool:
+def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str | None) -> bool:
if run_level not in {"advanced_model", "full_model"}:
return False
modalities = request_config.get("modalities", ["text", "audio"])
@@ -551,18 +660,18 @@ def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_l
return "text" in modalities
-def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool:
+def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str | None) -> bool:
if run_level not in {"advanced_model", "full_model"}:
return False
if request_config.get("response_format") == "pcm":
return False
- return bool(request_config.get("input"))
+ return bool(request_config.get("transcript_expected_text", request_config.get("input")))
def _resolve_audio_transcript(
response: Any,
request_config: dict[str, Any],
- run_level: str,
+ run_level: str | None,
*,
speech_api: bool,
) -> str | None:
@@ -655,7 +764,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
assert shorter_clean and (shorter_clean in longer_clean), (
- f"The audio content is not same as the text "
+ f"{AUDIO_MISMATCH_MESSAGE} "
f"(short-text containment check failed: "
f"text={text_output!r}, transcript={transcript!r})"
)
@@ -666,7 +775,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
text_output.lower(),
)
print(f"similarity is: {similarity}")
- assert similarity > similarity_threshold, "The audio content is not same as the text"
+ assert similarity > similarity_threshold, AUDIO_MISMATCH_MESSAGE
if audio_ref_text:
assert transcript is not None, "No audio transcript for reference-text validation"
audio_similarity = cosine_similarity_text(
@@ -737,28 +846,13 @@ def _assert_transcript_matches(
)
-def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
- """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.
+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`.
- When ``request_config`` carries ``status_code`` and/or ``err_message``, the
- request is expected to be rejected: assert it failed and that the HTTP status
- / error text match. Otherwise the normal success-path checks run.
+ Success-path checks only. Negative / contract cases belong on
+ :meth:`~tests.helpers.client.OnlineOmniClient.send_audio_speech_http_request`
+ with :func:`assert_http_error`.
"""
- expected_status = request_config.get("status_code")
- expected_err = request_config.get("err_message")
- if expected_status is not None or expected_err is not None:
- assert not response.success, "Expected an error response, but the request succeeded."
- if expected_status is not None:
- allowed = expected_status if isinstance(expected_status, (list, tuple)) else (expected_status,)
- assert response.status_code in allowed, f"Expected HTTP status in {allowed}, got {response.status_code}"
- if expected_err is not None:
- alternatives = expected_err if isinstance(expected_err, (list, tuple)) else (expected_err,)
- error_text = response.error_message or ""
- assert any(alt in error_text for alt in alternatives), (
- f"Expected one of {alternatives} in error text, got: {error_text!r}"
- )
- return
-
assert response.success, "The request failed."
# Optional floor on decoded audio size (models with very short clips may use a lower value).
@@ -786,7 +880,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True)
if transcript is not None:
- expected_text = request_config.get("input")
+ expected_text = request_config.get("transcript_expected_text", request_config.get("input"))
if expected_text:
print(f"audio content is: {transcript}")
print(f"input text is: {expected_text}")
@@ -805,7 +899,9 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
)
-def assert_diffusion_response(response: "DiffusionResponse", request_config: dict[str, Any], run_level: str = None):
+def assert_diffusion_response(
+ response: "DiffusionResponse", request_config: dict[str, Any], run_level: str | None = None
+):
assert response.success, "The request failed."
has_any_content = any(content is not None for content in (response.images, response.videos, response.audios))
assert has_any_content, "Response contains no images, videos, or audios"
@@ -871,9 +967,9 @@ def assert_http_error(
err_message: str | tuple[str, ...] | list[str] | None = None,
websocket_json_message: bool = False,
) -> dict[str, Any] | None:
- """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object.
+ """Validate a raw-HTTP :class:`~tests.helpers.client.HttpResponse`-like object.
- Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when
+ Used by :class:`~tests.helpers.client.OnlineOmniClient` ``send_*_http_request`` helpers when
``request_config`` contains optional ``err_code`` and/or ``err_message``.
When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame).
diff --git a/tests/helpers/clean.py b/tests/helpers/clean.py
new file mode 100644
index 000000000..17fef825d
--- /dev/null
+++ b/tests/helpers/clean.py
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Test cleanup helpers (device memory, distributed state, and memory monitoring).
+
+``vllm_omni.platforms`` is imported only inside functions that need it so importing this module
+at pytest plugin load does not run before session autouse fixtures.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+import os
+import subprocess
+import time
+
+from vllm_omni.platforms import current_omni_platform
+
+logger = logging.getLogger(__name__)
+
+
+def get_physical_device_indices(devices):
+ """Map logical device indices to physical IDs via the platform visibility env.
+
+ Uses ``current_omni_platform.device_control_env_var`` (e.g. ``CUDA_VISIBLE_DEVICES``,
+ ``ASCEND_RT_VISIBLE_DEVICES``) rather than hard-coding CUDA.
+ """
+ visible_devices = os.environ.get(current_omni_platform.device_control_env_var)
+ if visible_devices is None:
+ return devices
+ visible_indices = [int(x) for x in visible_devices.split(",") if x.strip() != ""]
+ index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
+ return [index_mapping[i] for i in devices if i in index_mapping]
+
+
+def pick_least_used_device_indices(num_devices: int) -> list[int]:
+ """Return physical device indices for the least-used accelerators.
+
+ Args:
+ num_devices: Number of device indices to return.
+
+ Returns:
+ Physical device indices sorted from least to most used memory.
+ """
+ if num_devices <= 0:
+ raise ValueError(f"num_devices must be positive, got {num_devices}.")
+
+ logical_count = current_omni_platform.get_device_count()
+ if logical_count < num_devices:
+ raise RuntimeError(f"Need {num_devices} devices, but only {logical_count} are available.")
+
+ device_usage: list[tuple[int, int]] = []
+ for device_index in range(logical_count):
+ with current_omni_platform.device(device_index):
+ free_bytes, total_bytes = current_omni_platform.get_device_memory()
+ device_usage.append((total_bytes - free_bytes, device_index))
+ device_usage.sort()
+ logical_indices = [device_index for _, device_index in device_usage[:num_devices]]
+ return get_physical_device_indices(logical_indices)
+
+
+def wait_for_gpu_memory_to_clear(
+ *,
+ devices: list[int],
+ threshold_bytes: int | None = None,
+ threshold_ratio: float | None = None,
+ timeout_s: float = 120,
+) -> None:
+ assert threshold_bytes is not None or threshold_ratio is not None
+ devices = get_physical_device_indices(devices)
+ start_time = time.time()
+
+ device_list = ", ".join(str(d) for d in devices)
+ if threshold_bytes is not None:
+ condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB"
+
+ def is_free(used, total):
+ return used <= threshold_bytes / 2**30
+ else:
+ ratio = threshold_ratio
+ assert ratio is not None
+ condition_str = f"Memory usage ratio ≤ {ratio * 100:.1f}%"
+
+ def is_free(used, total):
+ return used / total <= ratio
+
+ print(f"[Device Memory Monitor] Waiting for device(s) {device_list} to free memory, Condition: {condition_str}")
+
+ def get_mem_gib(device: int) -> tuple[float, float]:
+ with current_omni_platform.device(device):
+ free_bytes, total_bytes = current_omni_platform.mem_get_info()
+ return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30
+
+ while True:
+ output_raw = {d: get_mem_gib(d) for d in devices}
+ output = {
+ d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)"
+ for d, (used, total) in output_raw.items()
+ }
+
+ print("[Device Memory Status] Current usage:")
+ for device_id, mem_info in output.items():
+ print(f" Device {device_id}: {mem_info}")
+
+ dur_s = time.time() - start_time
+ if all(is_free(used, total) for used, total in output_raw.values()):
+ print(f"[Device Memory Freed] Device(s) {device_list} meet memory condition")
+ print(f" Condition: {condition_str}")
+ print(f" Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)")
+ break
+
+ if dur_s >= timeout_s:
+ raise ValueError(
+ f"[Device Memory Timeout] Device(s) {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
+ f"Condition: {condition_str}\n"
+ f"Current status:\n" + "\n".join(f" Device {d}: {output[d]}" for d in devices)
+ )
+
+ gc.collect()
+ current_omni_platform.empty_cache()
+ time.sleep(5)
+
+
+def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None:
+ print("\n" + "=" * 80)
+ print(label)
+ print("=" * 80)
+ try:
+ result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+ if result.returncode == 0 and result.stdout.strip():
+ lines = result.stdout.strip().split("\n")
+ for line in lines[:head_lines]:
+ print(line)
+ if len(lines) > head_lines:
+ print(f"... (showing first {head_lines} of {len(lines)} lines)")
+ else:
+ print(f"{cmd[0]} command failed or produced no output")
+ except (subprocess.TimeoutExpired, FileNotFoundError):
+ print(f"{cmd[0]} not available or timed out")
+ except Exception as e:
+ print(f"Error running {' '.join(cmd)}: {e}")
+
+
+def _print_device_processes() -> None:
+ """Print device information via platform SMI tools."""
+ if current_omni_platform.is_cuda():
+ _run_smi("NVIDIA GPU Information (nvidia-smi)", ["nvidia-smi"], 20)
+ _run_smi("Detailed GPU Processes (nvidia-smi pmon)", ["nvidia-smi", "pmon", "-c", "1"], 100, timeout=3)
+ elif current_omni_platform.is_npu():
+ _run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40)
+ elif current_omni_platform.is_rocm():
+ # amd-smi can enter uninterruptible sleep in the amdgpu CPER ioctl
+ # (amdgpu_cper_ring_write). In that state subprocess.run's timeout
+ # cannot reap it, causing this diagnostic hook to hang the entire test.
+ # rocm-smi provides the information needed here without that ioctl.
+ _run_smi(
+ "AMD GPU Information (rocm-smi)",
+ ["rocm-smi", "--showuse", "--showmeminfo", "vram"],
+ 60,
+ )
+ _run_smi("Detailed AMD GPU Processes (rocm-smi)", ["rocm-smi", "--showpids"], 100, timeout=3)
+ elif current_omni_platform.is_xpu():
+ _run_smi("Intel XPU Information (xpu-smi discovery)", ["xpu-smi", "discovery"], 40)
+ elif current_omni_platform.is_musa():
+ _run_smi("Moore Threads GPU Information (mthreads-gmi)", ["mthreads-gmi"], 30)
+ else:
+ print("\n" + "=" * 80)
+ print("WARNING: No supported device platform detected")
+ print("=" * 80)
+
+
+def _cleanup_stale_device_locks() -> None:
+ """Remove stale device-initialization lock files whose recorded PID is dead.
+
+ Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a
+ crashed / killed test run and block subsequent orchestrator startups.
+ """
+ import glob as _glob
+
+ for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"):
+ try:
+ with open(lock_file) as fh:
+ content = fh.read().strip()
+ if not content:
+ continue
+ pid = int(content)
+ except (OSError, ValueError):
+ continue
+
+ try:
+ os.kill(pid, 0)
+ except ProcessLookupError:
+ print(f" Removing stale device lock {lock_file} (PID {pid} is dead)")
+ try:
+ os.unlink(lock_file)
+ except OSError:
+ pass
+ except PermissionError:
+ pass
+
+
+def cleanup_test_environment(*, shutdown_ray: bool = False) -> None:
+ """Tear down distributed state and reset device memory for tests."""
+ from vllm import envs
+ from vllm.distributed.parallel_state import (
+ destroy_distributed_environment,
+ destroy_model_parallel,
+ )
+
+ envs.disable_envs_cache()
+
+ if current_omni_platform.is_rocm():
+ from vllm._aiter_ops import rocm_aiter_ops
+
+ rocm_aiter_ops.refresh_env_variables()
+
+ gc.unfreeze()
+ destroy_model_parallel()
+ destroy_distributed_environment()
+ if shutdown_ray:
+ import ray
+
+ ray.shutdown()
+
+ print("Pre-test device status:")
+ _cleanup_stale_device_locks()
+
+ num_devices = current_omni_platform.device_count()
+ if num_devices > 0:
+ try:
+ wait_for_gpu_memory_to_clear(
+ devices=list(range(num_devices)),
+ threshold_ratio=0.05,
+ timeout_s=60,
+ )
+ except Exception as e:
+ print(f"Device cleanup note: {e}")
+
+ gc.collect()
+ if not current_omni_platform.is_cpu():
+ current_omni_platform.empty_cache()
+ try:
+ import torch
+
+ torch._C._host_emptyCache()
+ except AttributeError:
+ logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
+
+ if current_omni_platform.is_available():
+ print("Post-test device status:")
+ _print_device_processes()
+
+
+__all__ = [
+ "cleanup_test_environment",
+ "get_physical_device_indices",
+ "pick_least_used_device_indices",
+ "wait_for_gpu_memory_to_clear",
+]
diff --git a/tests/helpers/client.py b/tests/helpers/client.py
new file mode 100644
index 000000000..9ef4a7a0c
--- /dev/null
+++ b/tests/helpers/client.py
@@ -0,0 +1,2148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Online/offline request clients and response types for tests."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import concurrent.futures
+import copy
+import io
+import json
+import time
+from dataclasses import asdict, dataclass
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import quote
+
+import numpy as np
+import requests
+import soundfile as sf
+import torch
+from openai import OpenAI, omit
+from PIL import Image
+from vllm.logger import init_logger
+
+from tests.helpers.assertions import (
+ SuccessRateGate,
+ assert_audio_speech_response,
+ assert_diffusion_response,
+ assert_http_error,
+ assert_images_generations_response,
+ assert_omni_response,
+ collect_at_success_rate,
+)
+from tests.helpers.media import (
+ _merge_base64_audio_to_segment,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+ from tests.helpers.runtime import OmniRunner
+
+logger = init_logger(__name__)
+
+
+def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None:
+ try:
+ data = r.json()
+ if isinstance(data, (dict, list)):
+ return data
+ except Exception:
+ pass
+ return None
+
+
+def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
+ """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
+ eb = cfg.get("extra_body")
+ if not eb:
+ return None
+ h, w = eb.get("height"), eb.get("width")
+ if (isinstance(h, (list, tuple)) or isinstance(w, (list, tuple))) and not (
+ isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))
+ ):
+ raise ValueError("extra_body height and width must both be lists or both be scalars")
+ if not (isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))):
+ return None
+ if len(h) != len(w):
+ raise ValueError(f"height and width lists must have equal length; got {len(h)=} {len(w)=}")
+ n = len(h)
+ n_out = eb.get("num_outputs_per_prompt")
+ if n_out is not None:
+ n_out = int(n_out)
+ if n_out != n:
+ raise ValueError(
+ "When height/width are lists, num_outputs_per_prompt must equal their length; "
+ f"got num_outputs_per_prompt={n_out}, len(lists)={n}"
+ )
+ splits: list[dict[str, Any]] = []
+ for i in range(n):
+ sub = copy.deepcopy(cfg)
+ sub_eb = dict(sub.get("extra_body") or {})
+ sub_eb["height"] = int(h[i])
+ sub_eb["width"] = int(w[i])
+ sub_eb["num_outputs_per_prompt"] = 1
+ sub["extra_body"] = sub_eb
+ splits.append(sub)
+ return splits
+
+
+@dataclass
+class OmniResponse:
+ """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
+
+ text_content: str | None = None
+ audio_data: list[str] | None = None
+ audio_content: str | None = None
+ audio_format: str | None = None
+ audio_bytes: bytes | None = None
+ #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the
+ #: OpenAI client call through response parsing and local post-process (e.g. audio decode).
+ e2e_latency: float | None = None
+ success: bool = False
+ prompt_tokens: int | None = None
+ cached_tokens: int | None = None
+ multimodal_tokens: dict[str, int] | None = None
+ logprobs: list | None = None
+
+
+@dataclass
+class DiffusionResponse:
+ """Decoded diffusion output from chat completions or offline runner (not raw ``requests``)."""
+
+ text_content: str | None = None
+ images: list[Image.Image] | None = None
+ audios: list[Any] | None = None
+ videos: list[Any] | None = None
+ #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before
+ #: ``chat.completions.create`` through local image / audio decode.
+ e2e_latency: float | None = None
+ success: bool = False
+
+
+@dataclass
+class HttpResponse:
+ """Normalized view of a ``requests`` response from :class:`OnlineOmniClient` HTTP helpers."""
+
+ status_code: int
+ success: bool
+ error_message: str | None = None
+ json_body: dict[str, Any] | list[Any] | None = None
+
+
+@dataclass
+class WebSocketJsonResponse:
+ """First JSON object delivered as a text WebSocket frame (streaming endpoints)."""
+
+ json_body: dict[str, Any]
+
+
+@dataclass
+class OpenPIWebSocketResponse:
+ """Msgpack WebSocket session against ``/v1/realtime/robot/openpi``."""
+
+ server_metadata: dict[str, Any]
+ operation_responses: list[Any]
+ actions: dict[str, np.ndarray] | None = None
+ action_tensors: list[np.ndarray] | None = None
+
+
+def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]:
+ """Build a minimal DROID-style observation payload for the OpenPI robot endpoint."""
+ identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32)
+ identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
+ return {
+ "session_id": session_id,
+ "video": {
+ "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+ "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
+ },
+ "state": {
+ "eef_9d": identity_eef_9d,
+ "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
+ "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
+ },
+ "language": {"annotation.language.language_instruction": [["pick up the object"]]},
+ }
+
+
+DREAMZERO_DEFAULT_PROMPT = (
+ "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+)
+DREAMZERO_ACTION_HORIZON = 24
+DREAMZERO_ACTION_DIM = 8
+DREAMZERO_CAMERA_FILES = {
+ "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
+ "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
+ "observation/wrist_image_left": "wrist_image_left.mp4",
+}
+
+
+def _require_opencv() -> Any:
+ try:
+ import cv2
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("DreamZero OpenPI test dependencies are missing: opencv-python") from exc
+ return cv2
+
+
+def load_dreamzero_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
+ cv2 = _require_opencv()
+ camera_frames: dict[str, np.ndarray] = {}
+ for camera_key, file_name in DREAMZERO_CAMERA_FILES.items():
+ video_path = video_dir / file_name
+ if not video_path.exists():
+ raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
+ cap = cv2.VideoCapture(str(video_path))
+ frames = []
+ while True:
+ ok, frame = cap.read()
+ if not ok:
+ break
+ frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+ cap.release()
+ if not frames:
+ raise RuntimeError(f"No frames loaded from {video_path}")
+ camera_frames[camera_key] = np.stack(frames, axis=0)
+ return camera_frames
+
+
+def build_dreamzero_demo_observations(
+ camera_frames: dict[str, np.ndarray],
+ *,
+ prompt: str,
+ session_id: str,
+ num_chunks: int = 2,
+) -> list[dict[str, Any]]:
+ if num_chunks < 1:
+ raise ValueError("num_chunks must be at least 1")
+
+ relative_offsets = [-23, -16, -8, 0]
+ total_frames = min(frames.shape[0] for frames in camera_frames.values())
+ frame_schedules = [[0]]
+ current_frame = 23
+ for _ in range(num_chunks - 1):
+ indices = [max(current_frame + offset, 0) for offset in relative_offsets]
+ if indices[-1] >= total_frames:
+ break
+ frame_schedules.append(indices)
+ current_frame += DREAMZERO_ACTION_HORIZON
+
+ observations: list[dict[str, Any]] = []
+ for frame_indices in frame_schedules:
+ obs: dict[str, Any] = {}
+ for camera_key, all_frames in camera_frames.items():
+ selected = all_frames[frame_indices]
+ obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
+ obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
+ obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
+ obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
+ obs["prompt"] = prompt
+ obs["session_id"] = session_id
+ observations.append(obs)
+ return observations
+
+
+class OpenPIWebSocketSession:
+ """Persistent msgpack session for ``/v1/realtime/robot/openpi``."""
+
+ DEFAULT_PING_INTERVAL_SECS = 300
+ DEFAULT_PING_TIMEOUT_SECS = 3600
+
+ @staticmethod
+ def _require_dependencies() -> tuple[Any, Any]:
+ try:
+ import websockets.sync.client as websockets_client
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: websockets") from exc
+ try:
+ from openpi_client import msgpack_numpy
+ except ImportError as exc: # pragma: no cover - optional e2e dependency
+ raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: openpi-client") from exc
+ return websockets_client, msgpack_numpy
+
+ def __init__(
+ self,
+ uri: str,
+ *,
+ ping_interval: float = DEFAULT_PING_INTERVAL_SECS,
+ ping_timeout: float = DEFAULT_PING_TIMEOUT_SECS,
+ open_timeout: float = 120.0,
+ close_timeout: float = 120.0,
+ ) -> None:
+ websockets_client, msgpack_numpy = self._require_dependencies()
+ self._msgpack_numpy = msgpack_numpy
+ self._packer = msgpack_numpy.Packer()
+ self._conn = websockets_client.connect(
+ uri,
+ compression=None,
+ max_size=None,
+ ping_interval=ping_interval,
+ ping_timeout=ping_timeout,
+ open_timeout=open_timeout,
+ close_timeout=close_timeout,
+ )
+ server_metadata = msgpack_numpy.unpackb(self._conn.recv())
+ if not isinstance(server_metadata, dict):
+ raise TypeError(f"Expected dict metadata from {uri}, got {type(server_metadata)!r}")
+ self._server_metadata = server_metadata
+
+ def get_server_metadata(self) -> dict[str, Any]:
+ return dict(self._server_metadata)
+
+ def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
+ response = self._send_operation("infer", obs)
+ if not isinstance(response, dict):
+ raise TypeError(f"Expected dict infer response, got {type(response)!r}")
+ return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+
+ def reset(self, reset_info: dict[str, Any] | None = None) -> str:
+ response = self._send_operation("reset", dict(reset_info or {}))
+ return str(response["status"])
+
+ def close(self) -> None:
+ self._conn.close()
+
+ def _send_operation(self, endpoint: str, payload: dict[str, Any]) -> Any:
+ body = dict(payload)
+ body["endpoint"] = endpoint
+ self._conn.send(self._packer.pack(body))
+ raw = self._conn.recv()
+ if isinstance(raw, str):
+ raise RuntimeError(f"OpenPI {endpoint!r} failed: {raw}")
+ decoded = self._msgpack_numpy.unpackb(raw)
+ if isinstance(decoded, dict) and decoded.get("type") == "error":
+ raise RuntimeError(f"OpenPI {endpoint!r} failed: {decoded.get('message')}")
+ if endpoint == "reset":
+ if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
+ raise RuntimeError(f"Unexpected OpenPI reset response: {decoded!r}")
+ return decoded
+ return decoded
+
+
+def _merge_http_expectation_kwargs(
+ base: dict[str, Any] | None,
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+) -> dict[str, Any]:
+ cfg = dict(base or {})
+ if err_code is not None:
+ cfg["err_code"] = err_code
+ if err_message is not None:
+ cfg["err_message"] = err_message
+ return cfg
+
+
+def _merge_ws_expectation_kwargs(
+ base: dict[str, Any] | None,
+ *,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ws_json_type: str | None = None,
+ ws_error_code: str | None = None,
+) -> dict[str, Any]:
+ cfg = dict(base or {})
+ if err_message is not None:
+ cfg["err_message"] = err_message
+ if ws_json_type is not None:
+ cfg["ws_json_type"] = ws_json_type
+ if ws_error_code is not None:
+ cfg["ws_error_code"] = ws_error_code
+ return cfg
+
+
+def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None:
+ jb = resp.json_body
+ want_type = cfg.get("ws_json_type")
+ if want_type is not None:
+ assert jb.get("type") == want_type, (jb, want_type)
+ want_code = cfg.get("ws_error_code")
+ if want_code is not None:
+ assert jb.get("code") == want_code, (jb, want_code)
+ err_message = cfg.get("err_message")
+ if err_message is not None:
+ assert_http_error(resp, err_message=err_message, websocket_json_message=True)
+
+
+def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse:
+ """Concatenate images in order; ``e2e_latency`` is wall-clock of the batch (set by caller) or max of parts."""
+ merged = DiffusionResponse()
+ merged.success = all(p.success for p in parts) and len(parts) > 0
+ imgs: list[Image.Image] = []
+ for p in parts:
+ if p.images:
+ imgs.extend(p.images)
+ merged.images = imgs if imgs else None
+ latencies = [p.e2e_latency for p in parts if p.e2e_latency is not None]
+ merged.e2e_latency = max(latencies) if latencies else None
+ return merged
+
+
+class OnlineOmniClient:
+ def __init__(
+ self,
+ host: str = "127.0.0.1",
+ port: int | None = None,
+ api_key: str = "EMPTY",
+ run_level: str | None = None,
+ *,
+ log_stats: bool = True,
+ ):
+ if port is None:
+ from tests.helpers.runtime import get_open_port
+
+ port = get_open_port()
+ self.base_url = f"http://{host}:{port}"
+ self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
+ self.run_level = run_level
+ self.log_stats = log_stats
+
+ def _print_client_stat(self, message: str) -> None:
+ if self.log_stats:
+ print(message, flush=True)
+
+ def _process_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
+ """Wall clock from *before* ``chat.completions.create`` through stream drain + local decode."""
+ result = OmniResponse()
+ try:
+ text_content = ""
+ audio_data = []
+ for chunk in chat_completion:
+ for choice in chunk.choices:
+ content = getattr(getattr(choice, "delta", None), "content", None)
+ modality = getattr(chunk, "modality", None)
+ if modality == "audio" and content:
+ audio_data.append(content)
+ elif modality == "text" and content:
+ text_content += content
+ # Usage is yielded after the last token
+ if chunk.usage:
+ result.prompt_tokens = chunk.usage.prompt_tokens
+ if details := getattr(chunk.usage, "prompt_tokens_details", None):
+ result.cached_tokens = details.cached_tokens
+ result.multimodal_tokens = getattr(details, "multimodal_tokens", None)
+
+ if audio_data:
+ merged_seg = _merge_base64_audio_to_segment(audio_data)
+ wav_buf = BytesIO()
+ merged_seg.export(wav_buf, format="wav")
+ result.audio_bytes = wav_buf.getvalue()
+ result.text_content = text_content
+ result.audio_data = audio_data
+ result.e2e_latency = time.perf_counter() - wall_start
+ result.success = True
+ except Exception as e:
+ msg = f"Stream processing error: {str(e)}"
+ print(f"Error: {msg}")
+ return result
+
+ def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
+ """Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
+ result = OmniResponse()
+ try:
+ audio_data = None
+ text_content = None
+ for choice in chat_completion.choices:
+ if hasattr(choice.message, "audio") and choice.message.audio is not None:
+ audio_data = choice.message.audio.data
+ if hasattr(choice.message, "content") and choice.message.content is not None:
+ text_content = choice.message.content
+ # Extract cached & prompt token counts for prefix caching tests
+ usage = getattr(chat_completion, "usage", None)
+ if usage:
+ result.prompt_tokens = usage.prompt_tokens
+ if details := getattr(usage, "prompt_tokens_details", None):
+ result.cached_tokens = details.cached_tokens
+ result.multimodal_tokens = getattr(details, "multimodal_tokens", None)
+ if audio_data:
+ result.audio_bytes = base64.b64decode(audio_data)
+ result.text_content = text_content
+ result.e2e_latency = time.perf_counter() - wall_start
+ if chat_completion.choices and chat_completion.choices[0].logprobs is not None:
+ result.logprobs = chat_completion.choices[0].logprobs.content
+ result.success = True
+ except Exception as e:
+ msg = f"Non-stream processing error: {str(e)}"
+ print(f"Error: {msg}")
+ return result
+
+ def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse:
+ """Wall clock from *before* ``chat.completions.create`` through image decode."""
+ result = DiffusionResponse()
+ try:
+ images = []
+ audios = []
+ for choice in chat_completion.choices:
+ content = getattr(choice.message, "content", None)
+ if isinstance(content, list):
+ for item in content:
+ image_url = None
+ if isinstance(item, dict):
+ image_url = item.get("image_url", {}).get("url")
+ else:
+ image_url_obj = getattr(item, "image_url", None)
+ image_url = getattr(image_url_obj, "url", None) if image_url_obj else None
+ if image_url and image_url.startswith("data:image"):
+ b64_data = image_url.split(",", 1)[1]
+ img = Image.open(io.BytesIO(base64.b64decode(b64_data)))
+ img.load()
+ images.append(img)
+
+ # OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`.
+ audio_obj = getattr(choice.message, "audio", None)
+ audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None
+ if audio_b64:
+ audios.append(
+ {
+ "wav_bytes": base64.b64decode(audio_b64),
+ "id": getattr(audio_obj, "id", None),
+ "expires_at": getattr(audio_obj, "expires_at", None),
+ }
+ )
+ result.images = images if images else None
+ result.audios = audios if audios else None
+ result.e2e_latency = time.perf_counter() - wall_start
+ result.success = True
+ except Exception as e:
+ msg = f"Diffusion response processing error: {str(e)}"
+ print(f"Error: {msg}")
+ return result
+
+ def _http_response_from_requests(self, r: requests.Response) -> HttpResponse:
+ payload = _parse_response_json(r)
+ ok = 200 <= r.status_code < 300
+ return HttpResponse(
+ status_code=r.status_code,
+ success=ok,
+ error_message=None if ok else (r.text[:8000] if r.text else None),
+ json_body=payload,
+ )
+
+ def send_health_http_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/health`` (raw ``requests``).
+
+ ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for
+ :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args).
+ """
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0)))
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_models_http_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`)."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = requests.get(
+ self._build_url("/v1/models"),
+ headers={"Accept": "application/json"},
+ timeout=float(cfg.get("timeout", 120.0)),
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_chat_completions_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests)."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_batched_chat_completions_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """Post to batched chat completions."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_completions_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
+ # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_omni_sleep_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_omni_wakeup_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/omni/wakeup``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_voices_list_http_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/v1/audio/voices``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = requests.get(
+ self._build_url("/v1/audio/voices"),
+ headers={"Accept": "application/json"},
+ timeout=float(cfg.get("timeout", 120.0)),
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_voices_create_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_voices_delete_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ name = cfg["name"]
+ timeout = float(cfg.get("timeout", 120.0))
+ path = f"/v1/audio/voices/{quote(str(name), safe='')}"
+ r = requests.delete(
+ self._build_url(path),
+ headers={"Accept": "application/json"},
+ timeout=timeout,
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_speech_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_speech_batch_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/audio/speech/batch``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_audio_generate_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/audio/generate``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_images_generations_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ if cfg.get("err_code") is None:
+ assert resp.success, resp.error_message
+ payload = resp.json_body
+ assert isinstance(payload, dict)
+ assert_images_generations_response(payload, cfg, run_level=self.run_level)
+ return [resp]
+
+ def send_images_edits_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_videos_create_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_videos_sync_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """POST ``/v1/videos/sync``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0)
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_videos_list_http_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/v1/videos`` — optional ``params``, ``timeout``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ r = requests.get(
+ self._build_url("/v1/videos"),
+ params=cfg.get("params"),
+ headers={"Accept": "application/json"},
+ timeout=float(cfg.get("timeout", 120.0)),
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_video_retrieve_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/v1/videos/{video_id}``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ video_id = cfg["video_id"]
+ timeout = float(cfg.get("timeout", 120.0))
+ r = requests.get(
+ self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+ headers={"Accept": "application/json"},
+ timeout=timeout,
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_video_delete_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """DELETE ``/v1/videos/{video_id}``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ video_id = cfg["video_id"]
+ timeout = float(cfg.get("timeout", 120.0))
+ r = requests.delete(
+ self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+ headers={"Accept": "application/json"},
+ timeout=timeout,
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def send_video_content_http_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_code: int | tuple[int, ...] | list[int] | None = None,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ) -> list[HttpResponse]:
+ """GET ``/v1/videos/{video_id}/content``."""
+ cfg = _merge_http_expectation_kwargs(
+ request_config,
+ err_code=err_code,
+ err_message=err_message,
+ )
+ video_id = cfg["video_id"]
+ timeout = float(cfg.get("timeout", 120.0))
+ r = requests.get(
+ self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"),
+ timeout=timeout,
+ )
+ resp = self._http_response_from_requests(r)
+ assert_http_error(
+ resp,
+ err_code=cfg.get("err_code"),
+ err_message=cfg.get("err_message"),
+ )
+ return [resp]
+
+ def _build_ws_url(self, path: str) -> str:
+ """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers."""
+ base = self.base_url.rstrip("/")
+ suffix = "/" + path.lstrip("/")
+ if base.startswith("http://"):
+ return "ws://" + base.removeprefix("http://") + suffix
+ if base.startswith("https://"):
+ return "wss://" + base.removeprefix("https://") + suffix
+ raise ValueError(f"Unsupported base_url for WebSocket: {base!r}")
+
+ def _send_websocket_first_json_request(
+ self,
+ path: str,
+ cfg: dict[str, Any],
+ ) -> list[WebSocketJsonResponse]:
+ """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`.
+
+ ``request_config`` keys:
+
+ - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
+ speaks first, e.g. ``/v1/realtime`` rejection path).
+ - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame
+ (e.g. ``["session.created"]`` on ``/v1/realtime``).
+ - ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
+ - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
+ """
+ send_frames_raw = cfg.get("send_frames")
+ if send_frames_raw is None:
+ frames: list[str] = []
+ elif isinstance(send_frames_raw, str):
+ frames = [send_frames_raw]
+ else:
+ frames = list(send_frames_raw)
+
+ timeout = float(cfg.get("timeout", 120.0))
+ uri = self._build_ws_url(path)
+ skip_types = set(cfg.get("ws_skip_types") or [])
+
+ connect_kw: dict[str, Any] = {}
+ if "ws_max_size" in cfg:
+ connect_kw["max_size"] = cfg["ws_max_size"]
+
+ async def _recv_first_json_object() -> WebSocketJsonResponse:
+ import websockets
+
+ async with websockets.connect(uri, **connect_kw) as ws:
+ for frame in frames:
+ await ws.send(frame)
+ while True:
+ raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
+ if not isinstance(raw, str):
+ raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
+ try:
+ data = json.loads(raw)
+ except json.JSONDecodeError as exc:
+ raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
+ if not isinstance(data, dict):
+ raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
+ if skip_types and data.get("type") in skip_types:
+ continue
+ return WebSocketJsonResponse(json_body=data)
+
+ resp = asyncio.run(_recv_first_json_object())
+ _run_ws_expectations_from_request_config(cfg, resp)
+ return [resp]
+
+ def send_audio_speech_stream_ws_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ws_json_type: str | None = None,
+ ws_error_code: str | None = None,
+ ) -> list[WebSocketJsonResponse]:
+ """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame."""
+ cfg = _merge_ws_expectation_kwargs(
+ request_config,
+ err_message=err_message,
+ ws_json_type=ws_json_type,
+ ws_error_code=ws_error_code,
+ )
+ return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg)
+
+ def send_video_chat_stream_ws_request(
+ self,
+ request_config: dict[str, Any],
+ *,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ws_json_type: str | None = None,
+ ws_error_code: str | None = None,
+ ) -> list[WebSocketJsonResponse]:
+ """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame."""
+ cfg = _merge_ws_expectation_kwargs(
+ request_config,
+ err_message=err_message,
+ ws_json_type=ws_json_type,
+ ws_error_code=ws_error_code,
+ )
+ return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg)
+
+ def send_realtime_ws_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ *,
+ err_message: str | tuple[str, ...] | list[str] | None = None,
+ ws_json_type: str | None = None,
+ ws_error_code: str | None = None,
+ ) -> list[WebSocketJsonResponse]:
+ """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated)."""
+ cfg = _merge_ws_expectation_kwargs(
+ request_config,
+ err_message=err_message,
+ ws_json_type=ws_json_type,
+ ws_error_code=ws_error_code,
+ )
+ return self._send_websocket_first_json_request("/v1/realtime", cfg)
+
+ def send_robot_openpi_ws_request(
+ self,
+ request_config: dict[str, Any] | None = None,
+ ) -> list[OpenPIWebSocketResponse]:
+ """WebSocket ``/v1/realtime/robot/openpi`` — msgpack metadata plus optional infer/reset ops.
+
+ ``request_config`` keys:
+
+ - ``operations``: optional sequence of ``{"endpoint": "infer"|"reset", "payload": {...}}``.
+ - ``run_default_policy_session``: when true and ``operations`` is omitted, run infer then reset
+ using :func:`build_openpi_droid_observation`.
+ - ``run_dreamzero_policy_session``: when true and ``operations`` is omitted, run the DreamZero
+ infer/reset/infer sequence using :func:`build_dreamzero_demo_observations`.
+ - ``video_dir``: required for ``run_dreamzero_policy_session``.
+ - ``prompt``: language instruction for DreamZero observations (optional).
+ - ``session_id``: used by the default policy session observation builder.
+ - ``num_chunks``: DreamZero chunk count (default ``2``).
+ - ``timeout``: seconds to wait for each inbound msgpack frame (default ``120``).
+ - ``ping_interval`` / ``ping_timeout``: forwarded to :func:`websockets.sync.client.connect`.
+ """
+ cfg = dict(request_config or {})
+ operations_cfg = cfg.get("operations")
+ if operations_cfg is not None:
+ operations = [dict(op) for op in operations_cfg]
+ elif cfg.get("run_dreamzero_policy_session"):
+ video_dir = cfg.get("video_dir")
+ if video_dir is None:
+ raise ValueError("run_dreamzero_policy_session requires video_dir")
+ session_id = str(cfg.get("session_id", "dreamzero-smoke"))
+ prompt = str(cfg.get("prompt", DREAMZERO_DEFAULT_PROMPT))
+ num_chunks = int(cfg.get("num_chunks", 2))
+ observations = build_dreamzero_demo_observations(
+ load_dreamzero_camera_frames(Path(video_dir)),
+ prompt=prompt,
+ session_id=session_id,
+ num_chunks=num_chunks,
+ )
+ operations = [{"endpoint": "infer", "payload": obs} for obs in observations]
+ operations.append({"endpoint": "reset", "payload": {}})
+ operations.append({"endpoint": "infer", "payload": observations[0]})
+ elif cfg.get("run_default_policy_session"):
+ session_id = str(cfg.get("session_id", "gr00t-smoke"))
+ operations = [
+ {"endpoint": "infer", "payload": build_openpi_droid_observation(session_id=session_id)},
+ {"endpoint": "reset", "payload": {}},
+ ]
+ else:
+ operations = []
+ timeout = float(cfg.get("timeout", 120.0))
+ ping_interval = float(cfg.get("ping_interval", OpenPIWebSocketSession.DEFAULT_PING_INTERVAL_SECS))
+ ping_timeout = float(cfg.get("ping_timeout", OpenPIWebSocketSession.DEFAULT_PING_TIMEOUT_SECS))
+
+ uri = self._build_ws_url("/v1/realtime/robot/openpi")
+ session = OpenPIWebSocketSession(
+ uri,
+ ping_interval=ping_interval,
+ ping_timeout=ping_timeout,
+ open_timeout=timeout,
+ close_timeout=timeout,
+ )
+ try:
+ operation_responses: list[Any] = []
+ actions: dict[str, np.ndarray] | None = None
+ action_tensors: list[np.ndarray] = []
+ for operation in operations:
+ endpoint = str(operation["endpoint"])
+ payload = dict(operation.get("payload") or {})
+ response = session._send_operation(endpoint, payload)
+ operation_responses.append(response)
+ if endpoint == "infer":
+ if isinstance(response, dict):
+ actions = {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
+ else:
+ action_tensors.append(np.asarray(response, dtype=np.float32))
+ finally:
+ session.close()
+
+ return [
+ OpenPIWebSocketResponse(
+ server_metadata=session.get_server_metadata(),
+ operation_responses=operation_responses,
+ actions=actions,
+ action_tensors=action_tensors or None,
+ )
+ ]
+
+ def send_omni_request(
+ self,
+ request_config: dict[str, Any],
+ request_num: int = 1,
+ *,
+ min_successes: int | None = None,
+ max_concurrency: int | None = None,
+ ) -> list[OmniResponse]:
+ """Chat completions via the OpenAI Python SDK (not raw HTTP).
+
+ ``min_successes`` gates the batch on how many of the ``request_num`` requests pass
+ rather than requiring each one to, and returns only the successes; ``max_concurrency``
+ caps how many are in flight at once. See ``SuccessRateGate``.
+ """
+ responses: list[OmniResponse] = []
+ stream = request_config.get("stream", False)
+ modalities = request_config.get("modalities", ["text", "audio"])
+ extra_body: dict[str, Any] = {}
+ if "speaker" in request_config:
+ extra_body["speaker"] = request_config["speaker"]
+ if request_config.get("use_audio_in_video"):
+ mm = dict(extra_body.get("mm_processor_kwargs") or {})
+ mm["use_audio_in_video"] = True
+ extra_body["mm_processor_kwargs"] = mm
+ if "sampling_params_list" in request_config:
+ extra_body["sampling_params_list"] = request_config["sampling_params_list"]
+ if request_config.get("extra_body"):
+ extra_body.update(request_config["extra_body"])
+
+ create_kwargs: dict[str, Any] = {
+ "model": request_config.get("model"),
+ "messages": request_config.get("messages"),
+ "stream": stream,
+ "modalities": modalities,
+ }
+ if "logprobs" in request_config:
+ create_kwargs["logprobs"] = request_config["logprobs"]
+ if "top_logprobs" in request_config:
+ create_kwargs["top_logprobs"] = request_config["top_logprobs"]
+ if "stream_options" in request_config:
+ create_kwargs["stream_options"] = request_config["stream_options"]
+ if extra_body:
+ create_kwargs["extra_body"] = extra_body
+
+ def _one():
+ wall_start = time.perf_counter()
+ chat_completion = self.client.chat.completions.create(**create_kwargs)
+ return (
+ self._process_stream_omni_response(chat_completion, wall_start=wall_start)
+ if stream
+ else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
+ )
+
+ if min_successes is not None:
+ gate = SuccessRateGate(min_successes=min_successes, concurrency=max_concurrency)
+
+ def _sample() -> OmniResponse:
+ resp = _one()
+ assert_omni_response(resp, request_config, run_level=self.run_level)
+ return resp
+
+ # Not _print_client_stat: the count is the gate's verdict, so it has to survive
+ # log_stats=False.
+ return collect_at_success_rate(
+ _sample,
+ gate,
+ request_num=request_num,
+ report=lambda line: print(f"[omni] {line}", flush=True),
+ )
+
+ if request_num == 1:
+ resp = _one()
+ assert_omni_response(resp, request_config, run_level=self.run_level)
+ if resp.e2e_latency is not None:
+ self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[omni] request#1 completed")
+ responses.append(resp)
+ return responses
+
+ with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+ futures = {executor.submit(_one): i + 1 for i in range(request_num)}
+ for future in concurrent.futures.as_completed(futures):
+ request_idx = futures[future]
+ resp = future.result()
+ assert_omni_response(resp, request_config, run_level=self.run_level)
+ if resp.e2e_latency is not None:
+ self._print_client_stat(f"[omni] request#{request_idx} success in {resp.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat(f"[omni] request#{request_idx} completed")
+ responses.append(resp)
+ return responses
+
+ def _process_stream_audio_speech_response(
+ self, response, *, response_format: str | None = None, wall_start: float
+ ) -> OmniResponse:
+ """
+ Process streaming /v1/audio/speech responses into an OmniResponse.
+
+ This mirrors _process_stream_omni_response but operates on low-level
+ audio bytes. Whisper transcription runs in assert_audio_speech_response
+ when the run_level requires it.
+ """
+ result = OmniResponse()
+
+ try:
+ # Aggregate all audio bytes from the streaming response.
+ data = bytearray()
+
+ # Preferred OpenAI helper.
+ if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")):
+ for chunk in response.iter_bytes():
+ if chunk:
+ data.extend(chunk)
+ else:
+ # Generic iterable-of-bytes fallback (e.g., generator or list of chunks).
+ try:
+ iterator = iter(response)
+ except TypeError:
+ iterator = None
+
+ if iterator is not None:
+ for chunk in iterator:
+ if not chunk:
+ continue
+ if isinstance(chunk, (bytes, bytearray)):
+ data.extend(chunk)
+ elif hasattr(chunk, "data"):
+ data.extend(chunk.data) # type: ignore[arg-type]
+ elif hasattr(chunk, "content"):
+ data.extend(chunk.content) # type: ignore[arg-type]
+ else:
+ raise TypeError(f"Unsupported stream chunk type: {type(chunk)}")
+ else:
+ raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}")
+
+ raw_bytes = bytes(data)
+
+ # Populate OmniResponse.
+ result.audio_bytes = raw_bytes
+ result.e2e_latency = time.perf_counter() - wall_start
+ result.success = True
+ result.audio_format = getattr(response, "response", None)
+ if result.audio_format is not None:
+ result.audio_format = result.audio_format.headers.get("content-type", "")
+
+ except Exception as e:
+ msg = f"Audio speech stream processing error: {str(e)}"
+ print(f"Error: {msg}")
+
+ return result
+
+ def _process_non_stream_audio_speech_response(
+ self, response, *, response_format: str | None = None, wall_start: float
+ ) -> OmniResponse:
+ """
+ Process non-streaming /v1/audio/speech responses into an OmniResponse.
+
+ This mirrors _process_non_stream_omni_response but for the binary
+ audio payload returned by audio.speech.create.
+ """
+ result = OmniResponse()
+
+ try:
+ # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content)
+ if hasattr(response, "read") and callable(getattr(response, "read")):
+ raw_bytes = response.read()
+ elif hasattr(response, "content"):
+ raw_bytes = response.content # type: ignore[assignment]
+ else:
+ raise TypeError(f"Unsupported audio speech response type: {type(response)}")
+
+ result.audio_bytes = raw_bytes
+ result.e2e_latency = time.perf_counter() - wall_start
+ result.success = True
+ result.audio_format = getattr(response, "response", None)
+ if result.audio_format is not None:
+ result.audio_format = result.audio_format.headers.get("content-type", "")
+
+ except Exception as e:
+ msg = f"Audio speech non-stream processing error: {str(e)}"
+ print(f"Error: {msg}")
+
+ return result
+
+ def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+ """
+ Call the /v1/audio/speech endpoint using the same configuration-dict
+ style as send_omni_request, but via the OpenAI Python client's
+ audio.speech APIs.
+
+ Expected keys in request_config:
+ - model: model name/path (required)
+ - input: text to synthesize (required)
+ - response_format: audio format such as "wav" or "pcm" (optional)
+ - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body)
+ - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response``
+ - transcript_expected_text: local expected spoken text; defaults to ``input``
+ - timeout: request timeout in seconds (float, optional, default 120.0)
+ - stream: whether to use streaming API (bool, optional, default False)
+
+ For negative / contract checks (expected 4xx), use
+ :meth:`send_audio_speech_http_request` with ``err_code`` / ``err_message``.
+ """
+ timeout = float(request_config.get("timeout", 120.0))
+
+ model = request_config["model"]
+ text_input = request_config["input"]
+ stream = bool(request_config.get("stream", False))
+ voice = request_config.get("voice", None)
+
+ # Standard OpenAI param: use omit when not provided to keep default behavior.
+ response_format = request_config.get("response_format", omit)
+
+ # Qwen3-TTS custom fields, forwarded via extra_body.
+ extra_body: dict[str, Any] = {}
+ # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
+ for key in (
+ "task_type",
+ "ref_text",
+ "ref_audio",
+ "language",
+ "max_new_tokens",
+ "seed",
+ "instructions",
+ "speed",
+ "stream_format",
+ "x_vector_only_mode",
+ ):
+ if key in request_config:
+ extra_body[key] = request_config[key]
+
+ responses: list[OmniResponse] = []
+
+ speech_fmt: str | None = None if response_format is omit else str(response_format).lower()
+
+ print(f"[audio.speech] start model={model}, stream={stream}, request_num={request_num}, timeout={timeout:.1f}s")
+
+ if request_num == 1:
+ if stream:
+ # Use streaming response helper.
+ wall_start = time.perf_counter()
+ with self.client.audio.speech.with_streaming_response.create(
+ model=model,
+ input=text_input,
+ response_format=response_format,
+ extra_body=extra_body or None,
+ timeout=timeout,
+ voice=voice,
+ ) as resp:
+ omni_resp = self._process_stream_audio_speech_response(
+ resp, response_format=speech_fmt, wall_start=wall_start
+ )
+ else:
+ # Non-streaming response.
+ wall_start = time.perf_counter()
+ resp = self.client.audio.speech.create(
+ model=model,
+ input=text_input,
+ response_format=response_format,
+ extra_body=extra_body or None,
+ timeout=timeout,
+ voice=voice,
+ )
+ omni_resp = self._process_non_stream_audio_speech_response(
+ resp, response_format=speech_fmt, wall_start=wall_start
+ )
+
+ assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
+ if omni_resp.e2e_latency is not None:
+ self._print_client_stat(f"[audio.speech] request#1 success in {omni_resp.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[audio.speech] request#1 completed")
+ responses.append(omni_resp)
+ return responses
+ else:
+ # request_num > 1: concurrent requests (use same params as single-request path)
+
+ if stream:
+
+ def _stream_task(request_idx: int):
+ wall_start = time.perf_counter()
+ with self.client.audio.speech.with_streaming_response.create(
+ model=model,
+ input=text_input,
+ response_format=response_format,
+ extra_body=extra_body or None,
+ timeout=timeout,
+ voice=voice,
+ ) as resp:
+ result = self._process_stream_audio_speech_response(
+ resp, response_format=speech_fmt, wall_start=wall_start
+ )
+ if result.e2e_latency is not None:
+ self._print_client_stat(
+ f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
+ )
+ else:
+ self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
+ return result
+
+ with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+ futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)}
+ for future in concurrent.futures.as_completed(futures):
+ request_idx = futures[future]
+ try:
+ omni_resp = future.result()
+ except Exception as e:
+ print(
+ f"[audio.speech] request#{request_idx} failed "
+ f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
+ )
+ raise
+ assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
+ responses.append(omni_resp)
+ else:
+
+ def _non_stream_task(request_idx: int):
+ wall_start = time.perf_counter()
+ r = self.client.audio.speech.create(
+ model=model,
+ input=text_input,
+ response_format=response_format,
+ extra_body=extra_body or None,
+ timeout=timeout,
+ voice=voice,
+ )
+ result = self._process_non_stream_audio_speech_response(
+ r, response_format=speech_fmt, wall_start=wall_start
+ )
+ if result.e2e_latency is not None:
+ self._print_client_stat(
+ f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
+ )
+ else:
+ self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
+ return result
+
+ with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+ futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)}
+ for future in concurrent.futures.as_completed(futures):
+ request_idx = futures[future]
+ try:
+ omni_resp = future.result()
+ except Exception as e:
+ print(
+ f"[audio.speech] request#{request_idx} failed "
+ f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
+ )
+ raise
+ assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
+ responses.append(omni_resp)
+
+ return responses
+
+ def send_diffusion_request(
+ self, request_config: dict[str, Any] | list[dict[str, Any]], request_num: int = 1
+ ) -> list[DiffusionResponse]:
+ """
+ Send OpenAI requests for diffusion models.
+ If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel
+ (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order.
+
+ Args:
+ request_config: A single request configuration dict, or a list of
+ request configuration dicts (one request per element)
+ request_num: Number of requests to send concurrently, defaults to 1 (single request)
+ Returns:
+ list[DiffusionResponse]: List of DiffusionResponse objects containing the response data
+ """
+ responses: list[DiffusionResponse] = []
+
+ def _create_from_config(cfg: dict[str, Any]) -> tuple[Any, float]:
+ stream = cfg.get("stream", False)
+ if stream:
+ raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
+ modalities = cfg.get("modalities", omit) # Most diffusion models don't require modalities param
+ eb = cfg.get("extra_body")
+ extra = copy.deepcopy(eb) if eb else None
+ wall_start = time.perf_counter()
+ chat_completion = self.client.chat.completions.create(
+ model=cfg.get("model"),
+ messages=cfg.get("messages"),
+ extra_body=extra,
+ modalities=modalities,
+ )
+ return chat_completion, wall_start
+
+ if isinstance(request_config, list):
+ if not request_config:
+ raise ValueError("request_config list must not be empty")
+ if request_num != 1:
+ raise ValueError("request_num is not supported when request_config is a list")
+ with concurrent.futures.ThreadPoolExecutor(max_workers=len(request_config)) as executor:
+ futures = {
+ executor.submit(_create_from_config, cfg): (i + 1, cfg) for i, cfg in enumerate(request_config)
+ }
+ for future in concurrent.futures.as_completed(futures):
+ request_idx, cfg = futures[future]
+ chat_completion, wall_start = future.result()
+ response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
+ assert_diffusion_response(response, cfg, run_level=self.run_level)
+ if response.e2e_latency is not None:
+ self._print_client_stat(
+ f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s"
+ )
+ else:
+ self._print_client_stat(f"[diffusion] request#{request_idx} completed")
+ responses.append(response)
+ return responses
+
+ size_splits = _split_request_config_by_per_output_sizes(request_config)
+ if size_splits is not None:
+ if request_num != 1:
+ raise ValueError(
+ "request_num must be 1 when extra_body height/width are lists (split into concurrent per-size calls)"
+ )
+ t0 = time.perf_counter()
+ with concurrent.futures.ThreadPoolExecutor(max_workers=len(size_splits)) as executor:
+ split_futures = [executor.submit(_create_from_config, sub) for sub in size_splits]
+ chat_completions = [f.result() for f in split_futures]
+ parts = [self._process_diffusion_response(cc, wall_start=ws) for cc, ws in chat_completions]
+ merged = _merge_diffusion_responses(parts)
+ merged.e2e_latency = time.perf_counter() - t0
+ assert_diffusion_response(merged, request_config, run_level=self.run_level)
+ if merged.e2e_latency is not None:
+ self._print_client_stat(f"[diffusion] request#1 success in {merged.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[diffusion] request#1 completed")
+ return [merged]
+
+ if request_num == 1:
+ # Send single request
+ chat_completion, wall_start = _create_from_config(request_config)
+ response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
+ assert_diffusion_response(response, request_config, run_level=self.run_level)
+ if response.e2e_latency is not None:
+ self._print_client_stat(f"[diffusion] request#1 success in {response.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[diffusion] request#1 completed")
+ responses.append(response)
+ return responses
+
+ # Send concurrent requests for the same request_config
+ with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+ idx_by_future = {executor.submit(_create_from_config, request_config): i + 1 for i in range(request_num)}
+ for future in concurrent.futures.as_completed(idx_by_future):
+ request_idx = idx_by_future[future]
+ chat_completion, wall_start = future.result()
+ response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
+ assert_diffusion_response(response, request_config, run_level=self.run_level)
+ if response.e2e_latency is not None:
+ self._print_client_stat(f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat(f"[diffusion] request#{request_idx} completed")
+ responses.append(response)
+ return responses
+
+ def send_video_diffusion_request(
+ self, request_config: dict[str, Any], request_num: int = 1
+ ) -> list[DiffusionResponse]:
+ """
+ Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content.
+
+ For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc.
+ """
+ if request_num != 1:
+ raise NotImplementedError("Concurrent video diffusion requests are not currently implemented")
+
+ form_data = request_config.get("form_data")
+ if not isinstance(form_data, dict):
+ raise ValueError("Video request_config must contain 'form_data'")
+ normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
+ files: dict[str, tuple[str, BytesIO, str]] = {}
+ image_reference = request_config.get("image_reference")
+ video_reference = request_config.get("video_reference")
+ if image_reference and video_reference:
+ raise ValueError("Only one of image_reference or video_reference can be provided")
+ if image_reference:
+ if image_reference.startswith("data:image"):
+ header, encoded = image_reference.split(",", 1)
+ content_type = header.split(";")[0].removeprefix("data:")
+ extension = content_type.split("/")[-1]
+ file_data = base64.b64decode(encoded)
+ files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
+ else:
+ normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
+ if video_reference:
+ if video_reference.startswith("data:video"):
+ header, encoded = video_reference.split(",", 1)
+ content_type = header.split(";")[0].removeprefix("data:")
+ extension = content_type.split("/")[-1]
+ file_data = base64.b64decode(encoded)
+ files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
+ else:
+ normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference})
+
+ result = DiffusionResponse()
+ create_url = self._build_url("/v1/videos")
+ response = requests.post(
+ create_url,
+ data=normalized_form_data,
+ files=files,
+ headers={"Accept": "application/json"},
+ timeout=60,
+ )
+ start_time = time.perf_counter()
+ response.raise_for_status()
+ job_data = response.json()
+ video_id = job_data["id"]
+ self._wait_until_video_completed(video_id)
+ end_time = time.perf_counter()
+ video_content = self._download_video_content(video_id)
+ result.success = True
+ result.videos = [video_content]
+ result.e2e_latency = end_time - start_time
+ assert_diffusion_response(result, request_config, run_level=self.run_level)
+ if result.e2e_latency is not None:
+ self._print_client_stat(f"[diffusion] request#1 success in {result.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[diffusion] request#1 completed")
+ return [result]
+
+ def _post_json_endpoint(
+ self,
+ path: str,
+ request_config: dict[str, Any],
+ *,
+ default_timeout: float,
+ ) -> requests.Response:
+ url = self._build_url(path)
+ timeout = float(request_config.get("timeout", default_timeout))
+ if "raw_body" in request_config:
+ raw = request_config["raw_body"]
+ payload = raw.encode("utf-8") if isinstance(raw, str) else raw
+ return requests.post(
+ url,
+ data=payload,
+ headers={"Content-Type": "application/json", "Accept": "application/json"},
+ timeout=timeout,
+ )
+ if "json" not in request_config:
+ raise ValueError(f"{path} request_config must include 'json' or 'raw_body'")
+ return requests.post(
+ url,
+ json=request_config["json"],
+ headers={"Accept": "application/json", "Content-Type": "application/json"},
+ timeout=timeout,
+ )
+
+ def _post_form_endpoint(
+ self,
+ path: str,
+ request_config: dict[str, Any],
+ *,
+ default_timeout: float = 120.0,
+ ) -> requests.Response:
+ url = self._build_url(path)
+ timeout = float(request_config.get("timeout", default_timeout))
+ data = request_config.get("data")
+ files = request_config.get("files")
+ if data is None and not files:
+ data = {}
+ return requests.post(
+ url,
+ data=data,
+ files=files,
+ headers={"Accept": "application/json"} if not files else {"Accept": "application/json"},
+ timeout=timeout,
+ )
+
+ def send_streaming_video_diffusion_request(
+ self,
+ request_config: dict[str, Any],
+ request_num: int = 1,
+ *,
+ timeout_seconds: float = 600.0,
+ ) -> list[DiffusionResponse]:
+ """
+ Send a native ``/v1/realtime/video`` WebSocket request and return one
+ finalized MP4 artifact assembled from the streamed binary fragments.
+ """
+ if request_num != 1:
+ raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented")
+
+ response = asyncio.run(
+ self._send_streaming_video_diffusion_request_once(
+ request_config,
+ timeout_seconds=timeout_seconds,
+ )
+ )
+ assert_diffusion_response(response, request_config, run_level=self.run_level)
+ if response.e2e_latency is not None:
+ self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s")
+ else:
+ self._print_client_stat("[diffusion.stream] request#1 completed")
+ return [response]
+
+ async def _send_streaming_video_diffusion_request_once(
+ self,
+ request_config: dict[str, Any],
+ *,
+ timeout_seconds: float,
+ ) -> DiffusionResponse:
+ form_data = request_config.get("form_data")
+ if not isinstance(form_data, dict):
+ raise ValueError("Video request_config must contain 'form_data'")
+ payload: dict[str, Any] = {
+ "type": "session.start",
+ **{key: value for key, value in form_data.items() if value is not None},
+ }
+ model = request_config.get("model")
+ if model is not None:
+ payload["model"] = model
+ payload.setdefault("format", "m4s")
+
+ fps = float(payload.get("fps") or 16)
+ stream_format = payload["format"]
+ url = self._build_ws_url("/v1/realtime/video")
+
+ result = DiffusionResponse()
+ chunks: list[bytes] = []
+ start_time = time.perf_counter()
+ deadline = start_time + timeout_seconds
+
+ import websockets
+
+ async with websockets.connect(url, max_size=None) as websocket:
+ await websocket.send(json.dumps(payload))
+
+ while True:
+ remaining = deadline - time.perf_counter()
+ if remaining <= 0:
+ raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s")
+
+ message = await asyncio.wait_for(websocket.recv(), timeout=remaining)
+ if isinstance(message, bytes):
+ chunks.append(message)
+ continue
+
+ msg = json.loads(message)
+ msg_type = msg.get("type")
+ if msg_type == "video.start":
+ stream_format = msg.get("format") or stream_format
+ continue
+ if msg_type == "session.done":
+ break
+ if msg_type == "error":
+ raise RuntimeError(str(msg.get("message", msg)))
+
+ from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
+ from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat
+
+ streamed_bytes = b"".join(chunks)
+ if not streamed_bytes:
+ raise RuntimeError("Streaming video request completed without binary video chunks")
+ result.videos = [
+ finalize_streaming_video_bytes(
+ streamed_bytes,
+ input_format=cast(StreamingVideoFormat, stream_format),
+ fps=fps,
+ )
+ ]
+ result.e2e_latency = time.perf_counter() - start_time
+ result.success = True
+ return result
+
+ def _wait_until_video_completed(
+ self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
+ ) -> None:
+ status_url = self._build_url(f"/v1/videos/{video_id}")
+ deadline = time.monotonic() + timeout_seconds
+ while time.monotonic() < deadline:
+ status_resp = requests.get(status_url, headers={"Accept": "application/json"}, timeout=30)
+ status_resp.raise_for_status()
+ status_data = status_resp.json()
+ current_status = status_data["status"]
+ if current_status == "completed":
+ return
+ if current_status == "failed":
+ error_msg = status_data.get("last_error", "Unknown error")
+ raise RuntimeError(f"Job failed: {error_msg}")
+ time.sleep(poll_interval_seconds)
+ raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s")
+
+ def _download_video_content(self, video_id: str) -> bytes:
+ download_url = self._build_url(f"/v1/videos/{video_id}/content")
+ video_resp = requests.get(download_url, stream=True, timeout=60)
+ video_resp.raise_for_status()
+ video_bytes = BytesIO()
+ for chunk in video_resp.iter_content(chunk_size=8192):
+ if chunk:
+ video_bytes.write(chunk)
+ return video_bytes.getvalue()
+
+ def _build_url(self, path: str) -> str:
+ return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
+
+
+class OfflineOmniClient:
+ def __init__(self, omni_runner: OmniRunner):
+ self.runner = omni_runner
+
+ def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> OmniResponse:
+ result = OmniResponse()
+ try:
+ text_content = None
+ audio_content = None
+ for stage_output in outputs:
+ if getattr(stage_output, "final_output_type", None) == "text":
+ text_content = stage_output.outputs[0].text
+ if getattr(stage_output, "final_output_type", None) == "audio":
+ audio_content = stage_output.outputs[0].multimodal_output["audio"]
+ result.audio_content = audio_content
+ result.text_content = text_content
+ result.success = True
+ except Exception as e:
+ msg = f"Output processing error: {str(e)}"
+ result.success = False
+ print(f"Error: {msg}")
+ return result
+
+ def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse:
+ result = DiffusionResponse()
+ output = outputs[0]
+ if isinstance(output.images[0], list):
+ # Returning frames of images as a video
+ result.videos = output.images
+ else:
+ # Returning actual images
+ result.images = output.images
+ # [TODO] Add audio processing when tests are introduced
+ result.success = True
+ return result
+
+ def send_omni_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
+ if request_config is None:
+ request_config = {}
+ prompts_value = request_config.get("prompts")
+ prompts: str | list[str] = prompts_value if isinstance(prompts_value, (str, list)) else ""
+ videos = request_config.get("videos")
+ images = request_config.get("images")
+ audios = request_config.get("audios")
+ modalities = request_config.get("modalities", ["text", "audio"])
+ outputs = self.runner.generate_multimodal(
+ prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
+ )
+ response = self._process_omni_output(outputs)
+ assert_omni_response(response, request_config, run_level="core_model")
+ return response
+
+ def send_diffusion_request(self, request_config: dict[str, Any]) -> DiffusionResponse:
+ prompt = request_config.get("prompt")
+ if prompt is None:
+ prompts = request_config.get("prompts")
+ if not prompts:
+ raise ValueError("request_config must contain a prompt")
+ if len(prompts) > 1:
+ raise ValueError(
+ "In the current internal data structure, "
+ "only one prompt is supported for diffusion requests. "
+ "Because one prompt can contain multiple images or videos, "
+ "the current internal data structure is ambiguous when multiple prompts are provided."
+ )
+ prompt = prompts[0]
+
+ # Sync previous `extra_body` field with sampling params object
+ sampling_params: OmniDiffusionSamplingParams | None = request_config.get("sampling_params")
+ if sampling_params is None:
+ extra_body = request_config.get("extra_body", {})
+ if extra_body:
+ sampling_params = OmniDiffusionSamplingParams(**extra_body)
+ else:
+ extra_body = asdict(sampling_params)
+ request_config["extra_body"] = extra_body
+ if not extra_body:
+ logger.warning("No sampling params provided in request_config, will skip output assertion")
+
+ negative_prompt = extra_body.get("negative_prompt") or request_config.get("negative_prompt")
+ videos = request_config.get("videos")
+ images = request_config.get("images")
+ audios = request_config.get("audios")
+ # Full dict (e.g. image + mask_image for inpainting) or partial; merged with top-level image/video keys.
+ extra_multi_modal = request_config.get("multi_modal_data")
+ modalities = request_config.get("modalities") # only used by limited models. Do not add default value here
+
+ prompt_object = OmniTextPrompt(prompt=prompt)
+ if negative_prompt:
+ prompt_object["negative_prompt"] = negative_prompt
+ multi_modal: dict = {}
+ if extra_multi_modal is not None:
+ multi_modal.update(dict(extra_multi_modal))
+ if videos is not None:
+ multi_modal["video"] = videos
+ if images is not None:
+ multi_modal["image"] = images
+ if audios is not None:
+ multi_modal["audio"] = audios
+ if multi_modal:
+ prompt_object["multi_modal_data"] = multi_modal
+ if modalities:
+ prompt_object["modalities"] = modalities # pyright: ignore[reportGeneralTypeIssues]
+
+ start_time = time.perf_counter()
+ raw_outputs = self.runner.generate([prompt_object], [sampling_params] if sampling_params else None)
+ end_time = time.perf_counter()
+ if not isinstance(raw_outputs, list):
+ raw_outputs = list(raw_outputs)
+
+ response = self._process_diffusion_output(raw_outputs)
+ response.e2e_latency = end_time - start_time
+ assert_diffusion_response(response, request_config, run_level="core_model")
+ return response
+
+ def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniResponse:
+ """
+ Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response.
+
+ request_config must contain:
+ - 'input' or 'prompts': text to synthesize.
+ Optional keys:
+ - 'voice' -> speaker (CustomVoice)
+ - 'task_type' -> task_type in additional_information (default: "CustomVoice")
+ - 'language' -> language in additional_information (default: "Auto")
+ - 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048)
+ - 'response_format' -> desired audio format (used only for assertion)
+ """
+ input_text = request_config.get("input") or request_config.get("prompts")
+ if input_text is None:
+ raise ValueError("request_config must contain 'input' or 'prompts' for TTS")
+ if isinstance(input_text, list):
+ input_text = input_text[0] if input_text else ""
+
+ mm_processor_kwargs: dict[str, Any] = {}
+ if "voice" in request_config:
+ mm_processor_kwargs["speaker"] = request_config["voice"]
+ if "task_type" in request_config:
+ mm_processor_kwargs["task_type"] = request_config["task_type"]
+ if "ref_audio" in request_config:
+ mm_processor_kwargs["ref_audio"] = request_config["ref_audio"]
+ if "ref_text" in request_config:
+ mm_processor_kwargs["ref_text"] = request_config["ref_text"]
+ if "language" in request_config:
+ mm_processor_kwargs["language"] = request_config["language"]
+ if "max_new_tokens" in request_config:
+ mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"]
+
+ outputs = self.runner.generate_multimodal(
+ prompts=input_text,
+ modalities=["audio"],
+ mm_processor_kwargs=mm_processor_kwargs or None,
+ )
+ mm_out: dict[str, Any] | None = None
+ for stage_out in outputs:
+ if getattr(stage_out, "final_output_type", None) == "audio":
+ mm_out = stage_out.outputs[0].multimodal_output
+ break
+ if mm_out is None:
+ raise AssertionError("No audio output from pipeline")
+
+ audio_data = mm_out.get("audio")
+ if audio_data is None:
+ raise AssertionError("No audio tensor in multimodal output")
+
+ sr_raw = mm_out.get("sr")
+ sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+ sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val)
+ wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
+ wav_buf = io.BytesIO()
+ sf.write(
+ wav_buf,
+ wav_tensor.float().cpu().numpy().reshape(-1),
+ samplerate=sr,
+ format="WAV",
+ subtype="PCM_16",
+ )
+ result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav")
+ assert_audio_speech_response(result, request_config, run_level="core_model")
+ return result
+
+ def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]:
+ return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages)
+
+ def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
+ return self.runner.stop_profile(stages=stages)
+
+
+# Backward-compatible aliases (historical test fixture / type names).
+OpenAIClientHandler = OnlineOmniClient
+OmniRunnerHandler = OfflineOmniClient
+
+
+__all__ = [
+ "DiffusionResponse",
+ "HttpResponse",
+ "WebSocketJsonResponse",
+ "OpenPIWebSocketResponse",
+ "build_openpi_droid_observation",
+ "build_dreamzero_demo_observations",
+ "DREAMZERO_ACTION_DIM",
+ "DREAMZERO_ACTION_HORIZON",
+ "DREAMZERO_CAMERA_FILES",
+ "DREAMZERO_DEFAULT_PROMPT",
+ "load_dreamzero_camera_frames",
+ "OpenPIWebSocketSession",
+ "OmniResponse",
+ "OfflineOmniClient",
+ "OnlineOmniClient",
+ "OpenAIClientHandler",
+ "OmniRunnerHandler",
+]
diff --git a/tests/helpers/env.py b/tests/helpers/env.py
deleted file mode 100644
index 39cec30e7..000000000
--- a/tests/helpers/env.py
+++ /dev/null
@@ -1,297 +0,0 @@
-"""Test environment / lifecycle helpers (device cleanup hooks and memory monitoring for tests).
-
-``vllm_omni.platforms`` is imported only inside functions that need it so importing this module
-at pytest plugin load does not run before session autouse fixtures.
-"""
-
-from __future__ import annotations
-
-import gc
-import os
-import subprocess
-import threading
-import time
-from contextlib import contextmanager
-
-import torch
-
-from vllm_omni.platforms import current_omni_platform
-
-
-def get_physical_device_indices(devices):
- visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
- if visible_devices is None:
- return devices
- visible_indices = [int(x) for x in visible_devices.split(",")]
- index_mapping = {i: physical for i, physical in enumerate(visible_indices)}
- return [index_mapping[i] for i in devices if i in index_mapping]
-
-
-def pick_least_used_device_indices(num_devices: int) -> list[int]:
- """Return physical device indices for the least-used accelerators.
-
- Args:
- num_devices: Number of device indices to return.
-
- Returns:
- Physical device indices sorted from least to most used memory.
- """
- if num_devices <= 0:
- raise ValueError(f"num_devices must be positive, got {num_devices}.")
-
- logical_count = current_omni_platform.get_device_count()
- if logical_count < num_devices:
- raise RuntimeError(f"Need {num_devices} devices, but only {logical_count} are available.")
-
- device_usage: list[tuple[int, int]] = []
- for device_index in range(logical_count):
- with current_omni_platform.device(device_index):
- free_bytes, total_bytes = current_omni_platform.get_device_memory()
- device_usage.append((total_bytes - free_bytes, device_index))
- device_usage.sort()
- logical_indices = [device_index for _, device_index in device_usage[:num_devices]]
- return get_physical_device_indices(logical_indices)
-
-
-def wait_for_gpu_memory_to_clear(
- *,
- devices: list[int],
- threshold_bytes: int | None = None,
- threshold_ratio: float | None = None,
- timeout_s: float = 120,
-) -> None:
- assert threshold_bytes is not None or threshold_ratio is not None
- devices = get_physical_device_indices(devices)
- start_time = time.time()
-
- device_list = ", ".join(str(d) for d in devices)
- if threshold_bytes is not None:
- condition_str = f"Memory usage ≤ {threshold_bytes / 2**30:.2f} GiB"
-
- def is_free(used, total):
- return used <= threshold_bytes / 2**30
- else:
- condition_str = f"Memory usage ratio ≤ {threshold_ratio * 100:.1f}%"
-
- def is_free(used, total):
- return used / total <= threshold_ratio
-
- print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
-
- @contextmanager
- def smi_scope():
- if current_omni_platform.is_rocm():
- from amdsmi import amdsmi_init, amdsmi_shut_down
-
- amdsmi_init()
- try:
- yield
- finally:
- amdsmi_shut_down()
- elif current_omni_platform.is_cuda():
- from vllm.third_party.pynvml import nvmlInit, nvmlShutdown
-
- nvmlInit()
- try:
- yield
- finally:
- nvmlShutdown()
- else:
- yield
-
- def get_mem_gib(device: int) -> tuple[float, float]:
- if current_omni_platform.is_rocm():
- from amdsmi import amdsmi_get_gpu_vram_usage, amdsmi_get_processor_handles
-
- info = amdsmi_get_gpu_vram_usage(amdsmi_get_processor_handles()[device])
- return info["vram_used"] / 2**10, info["vram_total"] / 2**10
- if current_omni_platform.is_npu():
- free_bytes, total_bytes = torch.npu.mem_get_info(device)
- return (total_bytes - free_bytes) / 2**30, total_bytes / 2**30
- from vllm.third_party.pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
-
- info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(device))
- return info.used / 2**30, info.total / 2**30
-
- with smi_scope():
- while True:
- output_raw = {d: get_mem_gib(d) for d in devices}
- output = {
- d: f"{used:.1f}GiB/{total:.1f}GiB ({(used / total) * 100 if total > 0 else 0:.1f}%)"
- for d, (used, total) in output_raw.items()
- }
-
- print("[GPU Memory Status] Current usage:")
- for device_id, mem_info in output.items():
- print(f" GPU {device_id}: {mem_info}")
-
- dur_s = time.time() - start_time
- if all(is_free(used, total) for used, total in output_raw.values()):
- print(f"[GPU Memory Freed] Devices {device_list} meet memory condition")
- print(f" Condition: {condition_str}")
- print(f" Wait time: {dur_s:.1f} seconds ({dur_s / 60:.1f} minutes)")
- break
-
- if dur_s >= timeout_s:
- raise ValueError(
- f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
- f"Condition: {condition_str}\n"
- f"Current status:\n" + "\n".join(f" GPU {d}: {output[d]}" for d in devices)
- )
-
- gc.collect()
- current_omni_platform.empty_cache()
- time.sleep(5)
-
-
-def _run_smi(label: str, cmd: list[str], head_lines: int, timeout: float = 5) -> None:
- print("\n" + "=" * 80)
- print(label)
- print("=" * 80)
- try:
- result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
- if result.returncode == 0 and result.stdout.strip():
- lines = result.stdout.strip().split("\n")
- for line in lines[:head_lines]:
- print(line)
- if len(lines) > head_lines:
- print(f"... (showing first {head_lines} of {len(lines)} lines)")
- else:
- print(f"{cmd[0]} command failed or produced no output")
- except (subprocess.TimeoutExpired, FileNotFoundError):
- print(f"{cmd[0]} not available or timed out")
- except Exception as e:
- print(f"Error running {' '.join(cmd)}: {e}")
-
-
-def _print_gpu_processes() -> None:
- """Print device information via nvidia-smi/npu-smi/amd-smi."""
- from vllm_omni.platforms import current_omni_platform
-
- if current_omni_platform.is_cuda():
- _run_smi("NVIDIA GPU Information (nvidia-smi)", ["nvidia-smi"], 20)
- _run_smi("Detailed GPU Processes (nvidia-smi pmon)", ["nvidia-smi", "pmon", "-c", "1"], 100, timeout=3)
- elif current_omni_platform.is_npu():
- _run_smi("Ascend NPU Information (npu-smi info)", ["npu-smi", "info"], 40)
- elif current_omni_platform.is_rocm():
- _run_smi("AMD GPU Information (amd-smi)", ["amd-smi"], 30)
- _run_smi("Detailed AMD GPU Processes (amd-smi process)", ["amd-smi", "process"], 100, timeout=3)
- else:
- print("\n" + "=" * 80)
- print("WARNING: No supported device platform detected")
- print("=" * 80)
-
- print("\n" + "=" * 80)
- print("System Processes with GPU keywords")
- print("=" * 80)
-
-
-def _cleanup_stale_device_locks() -> None:
- """Remove stale device-initialization lock files whose recorded PID is dead.
-
- Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a
- crashed / killed test run and block subsequent orchestrator startups.
- """
- import glob as _glob
-
- for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"):
- try:
- with open(lock_file) as fh:
- content = fh.read().strip()
- if not content:
- continue
- pid = int(content)
- except (OSError, ValueError):
- continue
-
- try:
- os.kill(pid, 0)
- except ProcessLookupError:
- print(f" Removing stale device lock {lock_file} (PID {pid} is dead)")
- try:
- os.unlink(lock_file)
- except OSError:
- pass
- except PermissionError:
- pass
-
-
-def run_pre_test_cleanup() -> None:
- print("Pre-test GPU status:")
-
- _cleanup_stale_device_locks()
-
- num_gpus = current_omni_platform.device_count()
- if num_gpus > 0:
- try:
- wait_for_gpu_memory_to_clear(
- devices=list(range(num_gpus)),
- threshold_ratio=0.05,
- timeout_s=60,
- )
- except Exception as e:
- print(f"Pre-test cleanup note: {e}")
-
-
-def run_post_test_cleanup() -> None:
- if current_omni_platform.is_available():
- gc.collect()
- current_omni_platform.empty_cache()
-
- print("Post-test GPU status:")
- _print_gpu_processes()
-
-
-class DeviceMemoryMonitor:
- """Poll global device memory usage."""
-
- def __init__(self, device_index: int, interval: float = 0.05):
- self.device_index = device_index
- self.interval = interval
- self._peak_used_mb = 0.0
- self._stop_event = threading.Event()
- self._thread: threading.Thread | None = None
-
- def start(self) -> None:
- from vllm_omni.platforms import current_omni_platform
-
- def monitor_loop() -> None:
- while not self._stop_event.is_set():
- try:
- with current_omni_platform.device(self.device_index):
- free_bytes, total_bytes = current_omni_platform.mem_get_info()
- used_mb = (total_bytes - free_bytes) / (1024**2)
- self._peak_used_mb = max(self._peak_used_mb, used_mb)
- except Exception:
- pass
- time.sleep(self.interval)
-
- self._thread = threading.Thread(target=monitor_loop, daemon=False)
- self._thread.start()
-
- def stop(self) -> None:
- if self._thread is None:
- return
- self._stop_event.set()
- self._thread.join(timeout=2.0)
-
- @property
- def peak_used_mb(self) -> float:
- from vllm_omni.platforms import current_omni_platform
-
- fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2)
- fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2)
- return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
-
- def __del__(self):
- self.stop()
-
-
-__all__ = [
- "DeviceMemoryMonitor",
- "get_physical_device_indices",
- "pick_least_used_device_indices",
- "run_post_test_cleanup",
- "run_pre_test_cleanup",
- "wait_for_gpu_memory_to_clear",
-]
diff --git a/tests/helpers/fixtures/clean.py b/tests/helpers/fixtures/clean.py
new file mode 100644
index 000000000..c07f563bb
--- /dev/null
+++ b/tests/helpers/fixtures/clean.py
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Opt-in cleanup fixtures for GPU memory, pipeline registry, and speaker cache."""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+
+
+@pytest.fixture
+def clean_gpu_memory_between_tests():
+ """Opt-in GPU pre/post hooks for a test (no environment-variable gate).
+
+ Use as a test parameter or ``@pytest.mark.usefixtures("clean_gpu_memory_between_tests")``.
+ """
+ from tests.helpers.clean import cleanup_test_environment
+
+ print("\n=== PRE-TEST DEVICE CLEANUP ===")
+ cleanup_test_environment()
+ yield
+ cleanup_test_environment()
+
+
+@pytest.fixture
+def clean_pipeline_registry():
+ """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it."""
+ snapshot = dict(OMNI_PIPELINES)
+ yield
+ OMNI_PIPELINES.clear()
+ OMNI_PIPELINES.update(snapshot)
+
+
+@pytest.fixture
+def clean_speaker_cache():
+ """Reset the process-wide speaker cache singleton before and after the test."""
+ import vllm_omni.utils.speaker_cache as _sc
+
+ def _reset():
+ with _sc._SINGLETON_LOCK:
+ if _sc._SINGLETON is not None:
+ _sc._SINGLETON.clear()
+ _sc._SINGLETON = None
+
+ _reset()
+ yield
+ _reset()
diff --git a/tests/helpers/fixtures/config.py b/tests/helpers/fixtures/config.py
index 2e309b672..2f46004a6 100644
--- a/tests/helpers/fixtures/config.py
+++ b/tests/helpers/fixtures/config.py
@@ -1,14 +1,47 @@
-"""Fixtures config creation & out of tree registry management."""
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Session-scoped config and environment fixtures."""
+
+from __future__ import annotations
+
+import os
import pytest
-
-from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
+import torch
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
-@pytest.fixture
-def clean_pipeline_registry():
- """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it."""
- snapshot = dict(OMNI_PIPELINES)
+@pytest.fixture(scope="session", autouse=True)
+def default_env():
+ # Keep behavior but avoid import-time side effects (RFC #2299).
+ keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE")
+ previous = {key: os.environ.get(key) for key in keys}
+ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn"
+ os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or (
+ "cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu"
+ )
yield
- OMNI_PIPELINES.clear()
- OMNI_PIPELINES.update(snapshot)
+ for key, value in previous.items():
+ if value is None:
+ os.environ.pop(key, None)
+ else:
+ os.environ[key] = value
+
+
+@pytest.fixture(scope="session", autouse=True)
+def default_vllm_config():
+ """Set a default VllmConfig for the whole test session.
+
+ Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and
+ deferred imports of ``tests.helpers.runtime`` both see the same context.
+ Function-scoped autouse ran too late for ``OmniRunner`` setup and could
+ desynchronize vLLM init vs request preprocessing (e.g. renderer state).
+ """
+ # Use CPU device if no GPU is available (e.g., in CI environments)
+ has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0
+ device = "cuda" if has_gpu else "cpu"
+ device_config = DeviceConfig(device=device)
+
+ with set_current_vllm_config(VllmConfig(device_config=device_config)):
+ yield
diff --git a/tests/helpers/fixtures/env.py b/tests/helpers/fixtures/env.py
deleted file mode 100644
index 2d38e4907..000000000
--- a/tests/helpers/fixtures/env.py
+++ /dev/null
@@ -1,74 +0,0 @@
-import os
-
-import pytest
-import torch
-from vllm.distributed import parallel_state
-
-
-@pytest.fixture(scope="session", autouse=True)
-def default_env():
- # Keep behavior but avoid import-time side effects (RFC #2299).
- keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE")
- previous = {key: os.environ.get(key) for key in keys}
- os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn"
- os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or (
- "cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu"
- )
- yield
- for key, value in previous.items():
- if value is None:
- os.environ.pop(key, None)
- else:
- os.environ[key] = value
-
-
-@pytest.fixture(scope="session")
-def model_prefix() -> str:
- prefix = os.environ.get("MODEL_PREFIX", "")
- return f"{prefix.rstrip('/')}/" if prefix else ""
-
-
-@pytest.fixture
-def clean_gpu_memory_between_tests():
- """Opt-in GPU pre/post hooks for a test (no environment-variable gate).
-
- Use as a test parameter or ``@pytest.mark.usefixtures("clean_gpu_memory_between_tests")``.
- """
- from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
-
- print("\n=== PRE-TEST GPU CLEANUP ===")
- run_pre_test_cleanup()
- yield
- run_post_test_cleanup()
-
-
-@pytest.fixture(scope="session", autouse=True)
-def default_vllm_config():
- """Set a default VllmConfig for the whole test session.
-
- Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and
- deferred imports of ``tests.helpers.runtime`` both see the same context.
- Function-scoped autouse ran too late for ``OmniRunner`` setup and could
- desynchronize vLLM init vs request preprocessing (e.g. renderer state).
- """
- from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
-
- # Use CPU device if no GPU is available (e.g., in CI environments)
- has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0
- device = "cuda" if has_gpu else "cpu"
- device_config = DeviceConfig(device=device)
-
- with set_current_vllm_config(VllmConfig(device_config=device_config)):
- yield
-
-
-@pytest.fixture()
-def init_fake_tp_group(mocker):
- """Provide a fake TP group so vllm linear layers can be instantiated."""
- mock_tp = mocker.MagicMock()
- mock_tp.world_size = 1
- mock_tp.rank_in_group = 0
- old = parallel_state._TP
- parallel_state._TP = mock_tp
- yield
- parallel_state._TP = old
diff --git a/tests/helpers/fixtures/media.py b/tests/helpers/fixtures/media.py
new file mode 100644
index 000000000..c34de093e
--- /dev/null
+++ b/tests/helpers/fixtures/media.py
@@ -0,0 +1,24 @@
+"""Fixtures for the reused Whisper transcription worker."""
+
+import sys
+
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="module")
+def release_audio_transcriber_after_module():
+ """Free the transcription worker's device memory once a test module is done.
+
+ A backstop. The server and runner fixtures already release the worker when
+ each of their instances tears down, which is what keeps Whisper off the
+ device while the next parametrized server initializes. This covers the
+ remaining case: a module that transcribes without using those fixtures.
+
+ Reached through ``sys.modules`` because ``tests/conftest.py`` deliberately
+ keeps ``tests.helpers.media`` lazy: a module that never transcribed should
+ not pay to import it.
+ """
+ yield
+ media = sys.modules.get("tests.helpers.media")
+ if media is not None:
+ media.release_audio_transcriber()
diff --git a/tests/helpers/fixtures/pytest_collection.py b/tests/helpers/fixtures/pytest_collection.py
new file mode 100644
index 000000000..912749242
--- /dev/null
+++ b/tests/helpers/fixtures/pytest_collection.py
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Pytest collection-phase hooks.
+
+All ``pytest_collection_modifyitems`` policies live here so their order is
+explicit. Add a named helper and call it from the hook; do not register a
+second ``pytest_collection_modifyitems`` in another plugin.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+
+def pytest_collection_modifyitems(config, items):
+ _pin_xdist_serial(config, items)
+
+
+def _pin_xdist_serial(config, items):
+ """Pin unmarked tests to one xdist worker.
+
+ Only tests explicitly marked ``pytest.mark.xdist`` run concurrently.
+ Others are pinned to a single xdist worker.
+ """
+ if not config.pluginmanager.hasplugin("xdist"):
+ return
+ for item in items:
+ if not item.get_closest_marker("xdist"):
+ item.add_marker(pytest.mark.xdist_group(name="serial"))
diff --git a/tests/helpers/fixtures/run_args.py b/tests/helpers/fixtures/pytest_run_args.py
similarity index 68%
rename from tests/helpers/fixtures/run_args.py
rename to tests/helpers/fixtures/pytest_run_args.py
index 975584d20..3766e9c4e 100644
--- a/tests/helpers/fixtures/run_args.py
+++ b/tests/helpers/fixtures/pytest_run_args.py
@@ -1,3 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Pytest CLI options and the fixtures that expose them."""
+
+from __future__ import annotations
+
import pytest
diff --git a/tests/helpers/fixtures/runtime.py b/tests/helpers/fixtures/runtime.py
index 3bffa475d..5c4d0c055 100644
--- a/tests/helpers/fixtures/runtime.py
+++ b/tests/helpers/fixtures/runtime.py
@@ -1,10 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Runtime fixtures (OmniRunner / OmniServer). Imports are deferred to fixture time.
Loading ``tests.helpers.runtime`` at plugin import time (before session fixtures)
pulls in vLLM/vllm_omni too early and breaks initialization order vs the legacy
-monolithic conftest. Defer imports until fixtures run so ``default_env`` /
-``default_vllm_config`` run first. Implementation helpers live in
-``tests.helpers.runtime`` (``iter_omni_server`` / ``iter_omni_runner``).
+monolithic conftest. Defer imports until fixtures run so session fixtures in
+``tests.helpers.fixtures.config`` run first. Server/runner helpers live in
+``tests.helpers.runtime``; request clients live in ``tests.helpers.client``.
"""
from __future__ import annotations
@@ -21,19 +24,30 @@ if TYPE_CHECKING:
omni_fixture_lock = threading.Lock()
+@pytest.fixture
+def _normalized_hardware_marks(request: pytest.FixtureRequest):
+ """Dummy carrier for mixed-``num_cards`` ``@hardware_test`` mark normalization.
+
+ Pytest has no API to clone an item and only change marks. ``@hardware_test``
+ therefore parametrizes this fixture (``indirect=True``) so each collected
+ item gets one platform's SKU and ``cards_{n}`` without changing test
+ signatures. The returned platform id is unused by test bodies.
+ """
+ return request.param
+
+
@pytest.fixture(scope="function")
def omni_server_function(
request: pytest.FixtureRequest,
run_level: str,
- model_prefix: str,
) -> Generator[OmniServer, Any, None]:
from tests.helpers.runtime import iter_omni_server
- yield from iter_omni_server(request, run_level, model_prefix, omni_fixture_lock)
+ yield from iter_omni_server(request, run_level, omni_fixture_lock)
@pytest.fixture(scope="module")
-def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
+def omni_server(request: pytest.FixtureRequest, run_level: str) -> Generator[OmniServer, Any, None]:
"""Start vLLM-Omni through the standard or stage-CLI launcher.
The fixture stays module-scoped because multi-stage initialization is costly.
@@ -42,16 +56,16 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
"""
from tests.helpers.runtime import iter_omni_server
- yield from iter_omni_server(request, run_level, model_prefix, omni_fixture_lock)
+ yield from iter_omni_server(request, run_level, omni_fixture_lock)
@pytest.fixture
-def openai_client(request: pytest.FixtureRequest, run_level: str):
+def online_client(request: pytest.FixtureRequest, run_level: str):
"""Resolve ``omni_server`` lazily so parametrized server fixtures work like upstream."""
- from tests.helpers.runtime import OpenAIClientHandler
+ from tests.helpers.client import OnlineOmniClient
server = request.getfixturevalue("omni_server")
- return OpenAIClientHandler(
+ return OnlineOmniClient(
host=server.host,
port=server.port,
api_key="EMPTY",
@@ -61,12 +75,18 @@ def openai_client(request: pytest.FixtureRequest, run_level: str):
@pytest.fixture
-def openai_client_function(request: pytest.FixtureRequest, run_level: str):
+def openai_client(online_client):
+ """Historical alias for :func:`online_client`."""
+ return online_client
+
+
+@pytest.fixture
+def online_client_function(request: pytest.FixtureRequest, run_level: str):
"""Resolve ``omni_server_function`` lazily for function-scoped reliability tests."""
- from tests.helpers.runtime import OpenAIClientHandler
+ from tests.helpers.client import OnlineOmniClient
server = request.getfixturevalue("omni_server_function")
- return OpenAIClientHandler(
+ return OnlineOmniClient(
host=server.host,
port=server.port,
api_key="EMPTY",
@@ -75,10 +95,15 @@ def openai_client_function(request: pytest.FixtureRequest, run_level: str):
)
+@pytest.fixture
+def openai_client_function(online_client_function):
+ """Historical alias for :func:`online_client_function`."""
+ return online_client_function
+
+
@pytest.fixture(scope="function")
def omni_runner_function(
request: pytest.FixtureRequest,
- model_prefix: str,
run_level: str,
) -> Generator[OmniRunner, Any, None]:
"""Function-scoped :class:`~tests.helpers.runtime.OmniRunner` (cf. :func:`omni_server_function`).
@@ -88,30 +113,42 @@ def omni_runner_function(
"""
from tests.helpers.runtime import iter_omni_runner
- yield from iter_omni_runner(request, model_prefix, run_level, omni_fixture_lock)
+ yield from iter_omni_runner(request, run_level, omni_fixture_lock)
@pytest.fixture(scope="module")
-def omni_runner(request: pytest.FixtureRequest, model_prefix: str, run_level: str) -> Generator[OmniRunner, Any, None]:
+def omni_runner(request: pytest.FixtureRequest, run_level: str) -> Generator[OmniRunner, Any, None]:
"""Module-scoped :class:`~tests.helpers.runtime.OmniRunner` (cf. :func:`omni_server`).
Reuses one runner for the whole module to amortize multi-stage init cost.
"""
from tests.helpers.runtime import iter_omni_runner
- yield from iter_omni_runner(request, model_prefix, run_level, omni_fixture_lock)
+ yield from iter_omni_runner(request, run_level, omni_fixture_lock)
@pytest.fixture
-def omni_runner_handler_function(omni_runner_function: OmniRunner):
- """Resolve :class:`~tests.helpers.runtime.OmniRunnerHandler` for :func:`omni_runner_function`."""
- from tests.helpers.runtime import OmniRunnerHandler
+def offline_client_function(omni_runner_function: OmniRunner):
+ """Resolve :class:`~tests.helpers.client.OfflineOmniClient` for :func:`omni_runner_function`."""
+ from tests.helpers.client import OfflineOmniClient
- return OmniRunnerHandler(omni_runner_function)
+ return OfflineOmniClient(omni_runner_function)
@pytest.fixture
-def omni_runner_handler(omni_runner: OmniRunner):
- from tests.helpers.runtime import OmniRunnerHandler
+def omni_runner_handler_function(offline_client_function):
+ """Historical alias for :func:`offline_client_function`."""
+ return offline_client_function
- return OmniRunnerHandler(omni_runner)
+
+@pytest.fixture
+def offline_client(omni_runner: OmniRunner):
+ from tests.helpers.client import OfflineOmniClient
+
+ return OfflineOmniClient(omni_runner)
+
+
+@pytest.fixture
+def omni_runner_handler(offline_client):
+ """Historical alias for :func:`offline_client`."""
+ return offline_client
diff --git a/tests/helpers/fixtures/speaker_cache.py b/tests/helpers/fixtures/speaker_cache.py
deleted file mode 100644
index cf7a62122..000000000
--- a/tests/helpers/fixtures/speaker_cache.py
+++ /dev/null
@@ -1,21 +0,0 @@
-"""Fixtures for the process-wide speaker cache singleton."""
-
-from __future__ import annotations
-
-import pytest
-
-
-@pytest.fixture
-def fresh_speaker_cache():
- """Reset the process-wide speaker cache singleton before and after the test."""
- import vllm_omni.utils.speaker_cache as _sc
-
- def _reset():
- with _sc._SINGLETON_LOCK:
- if _sc._SINGLETON is not None:
- _sc._SINGLETON.clear()
- _sc._SINGLETON = None
-
- _reset()
- yield
- _reset()
diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py
index 676732a8e..18bdbd394 100644
--- a/tests/helpers/mark.py
+++ b/tests/helpers/mark.py
@@ -1,19 +1,33 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …)."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …).
+
+Module import is stdlib-only so ``tools/pre_commit/check_test_marks.py`` can load
+this file without pytest or vllm. Decorator helpers import pytest lazily.
+"""
from __future__ import annotations
import re
from functools import lru_cache
from pathlib import Path
+from typing import TYPE_CHECKING
-import pytest
-from vllm.platforms import current_platform
+if TYPE_CHECKING:
+ import pytest
-# Marker description tag in ``pyproject.toml`` ``tool.pytest.ini_options.markers``.
-# Example: ``"H100: [hardware-resource] Tests that require H100 GPU"``.
+# Marker description tags in ``pyproject.toml`` ``tool.pytest.ini_options.markers``.
+# Example: ``"H100: [hardware-resource] [cuda] Tests that require H100 GPU"``.
_HARDWARE_RESOURCE_MARKER_TAG = "[hardware-resource]"
+_HARDWARE_CARDS_MARKER_TAG = "[hardware-cards]"
+_HARDWARE_PLATFORM_MARKER_TAG = "[hardware-platform]"
+_CI_LEVEL_MARKER_TAG = "[ci-level]"
+# On ``[hardware-platform]`` lines (not SKUs): also attach ``pytest.mark.gpu``.
+_GPU_DEVICE_CLASS_TAG = "[gpu]"
+_CARDS_MARK_NAME_RE = re.compile(r"^cards_(\d+)$")
+# Indirect fixture in ``tests.helpers.fixtures.runtime``: mixed-count ``@hardware_test``
+# parametrizes it so each collected item carries one platform's SKU and ``cards_{n}``.
+HARDWARE_MARK_NORMALIZATION_FIXTURE = "_normalized_hardware_marks"
def _repo_root() -> Path:
@@ -34,23 +48,16 @@ def _load_pyproject() -> dict:
return {}
-@lru_cache(maxsize=1)
-def get_hardware_mark_list() -> frozenset[str]:
- """Return hardware SKU marker names from ``pyproject.toml``.
-
- A marker is treated as a hardware resource marker when its description
- contains the ``[hardware-resource]`` tag, e.g.::
-
- "H100: [hardware-resource] Tests that require H100 GPU"
- """
+@lru_cache(maxsize=16)
+def _marker_names_with_tag(tag: str) -> frozenset[str]:
+ """Return pytest marker names whose description contains ``tag``."""
data = _load_pyproject()
entries = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("markers", [])
if not entries:
- # Fallback when TOML parsing fails: scan tagged marker lines.
text = (_repo_root() / "pyproject.toml").read_text(encoding="utf-8")
return frozenset(
re.findall(
- rf'^\s*"([A-Za-z0-9_]+)\s*:\s*{re.escape(_HARDWARE_RESOURCE_MARKER_TAG)}',
+ rf'^\s*"([A-Za-z0-9_]+)\s*:.*{re.escape(tag)}',
text,
flags=re.M,
)
@@ -59,7 +66,7 @@ def get_hardware_mark_list() -> frozenset[str]:
names: set[str] = set()
for entry in entries:
text = str(entry)
- if _HARDWARE_RESOURCE_MARKER_TAG not in text:
+ if tag not in text:
continue
name = text.split(":", 1)[0].strip()
if name:
@@ -67,138 +74,256 @@ def get_hardware_mark_list() -> frozenset[str]:
return frozenset(names)
-def cuda_marks(*, res: str, num_cards: int):
+@lru_cache(maxsize=1)
+def get_hardware_mark_list() -> frozenset[str]:
+ """Return hardware SKU marker names tagged ``[hardware-resource]``."""
+ return _marker_names_with_tag(_HARDWARE_RESOURCE_MARKER_TAG)
+
+
+@lru_cache(maxsize=1)
+def get_supported_card_counts() -> frozenset[int]:
+ """Return ``n`` for each ``cards_{n}`` marker tagged ``[hardware-cards]``."""
+ counts: set[int] = set()
+ for name in _marker_names_with_tag(_HARDWARE_CARDS_MARKER_TAG):
+ matched = _CARDS_MARK_NAME_RE.fullmatch(name)
+ if matched:
+ counts.add(int(matched.group(1)))
+ return frozenset(counts)
+
+
+@lru_cache(maxsize=1)
+def get_supported_platforms() -> frozenset[str]:
+ """Return platform marker names tagged ``[hardware-platform]``.
+
+ Includes ``cpu`` / ``gpu`` (device-class marks) and accelerator keys.
+ ``hardware_marks`` / ``hardware_test`` ``res`` dicts reject ``cpu`` / ``gpu``.
+ """
+ return _marker_names_with_tag(_HARDWARE_PLATFORM_MARKER_TAG)
+
+
+@lru_cache(maxsize=1)
+def get_level_markers() -> frozenset[str]:
+ """Return CI level marker names tagged ``[ci-level]`` (``core_model``, ``slow``, …)."""
+ return _marker_names_with_tag(_CI_LEVEL_MARKER_TAG)
+
+
+@lru_cache(maxsize=8)
+def _skus_for_platform(platform: str) -> frozenset[str]:
+ """SKU names tagged ``[hardware-resource]`` and ``[platform]`` (e.g. ``[cuda]``)."""
+ return get_hardware_mark_list() & _marker_names_with_tag(f"[{platform}]")
+
+
+def _require_sku(platform: str, res: str) -> None:
+ allowed = _skus_for_platform(platform)
+ if res not in allowed:
+ supported = ", ".join(sorted(allowed)) or f"(none tagged [hardware-resource] [{platform}])"
+ raise ValueError(f"Invalid {platform} resource type: {res}. Supported: {supported}")
+
+
+def _cards_mark(num_cards: int) -> pytest.MarkDecorator:
+ """Return ``cards_{n}`` so collection can filter with ``-m cards_2`` (etc.).
+
+ Multi-card selection is ``not cards_1`` (or an explicit ``cards_2`` / ``cards_4``).
+ """
+ import pytest
+
+ if not isinstance(num_cards, int) or isinstance(num_cards, bool) or num_cards < 1:
+ raise ValueError(f"num_cards must be a positive int, got {num_cards!r}")
+ if num_cards not in get_supported_card_counts():
+ counts = get_supported_card_counts()
+ supported = ", ".join(f"cards_{n}" for n in sorted(counts))
+ supported = supported or "(none tagged [hardware-cards])"
+ raise ValueError(f"num_cards={num_cards} has no registered marker. Supported: {supported}.")
+ return getattr(pytest.mark, f"cards_{num_cards}")
+
+
+def _cuda_marks(*, res: str, num_cards: int):
+ import pytest
+ from vllm.platforms import current_platform
+
+ _require_sku("cuda", res)
+ test_resource = getattr(pytest.mark, res)
test_platform_detail = pytest.mark.cuda
- if res == "L4":
- test_resource = pytest.mark.L4
- elif res == "H100":
- test_resource = pytest.mark.H100
- else:
- raise ValueError(f"Invalid CUDA resource type: {res}. Supported: L4, H100")
- marks = [test_resource, test_platform_detail]
- if num_cards == 1:
+ marks = [test_resource, test_platform_detail, _cards_mark(num_cards)]
+ if num_cards == 1 or not current_platform.is_cuda():
return marks
- test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards)
-
- if not current_platform.is_cuda():
- return marks + [test_distributed]
- test_skipif = pytest.mark.skipif(
- current_platform.device_count() < num_cards,
- reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
- )
- return marks + [test_distributed, test_skipif]
+ return marks + [
+ pytest.mark.skipif(
+ current_platform.device_count() < num_cards,
+ reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
+ )
+ ]
-def rocm_marks(*, res: str, num_cards: int):
- test_platform_detail = pytest.mark.rocm
- if res == "MI325":
- test_resource = pytest.mark.MI325
- else:
- raise ValueError(f"Invalid ROCm resource type: {res}. Supported: MI325")
- marks = [test_resource, test_platform_detail]
- if num_cards == 1:
+def _rocm_marks(*, res: str, num_cards: int):
+ import pytest
+
+ _require_sku("rocm", res)
+ return [getattr(pytest.mark, res), pytest.mark.rocm, _cards_mark(num_cards)]
+
+
+def _xpu_marks(*, res: str, num_cards: int):
+ import pytest
+ from vllm.platforms import current_platform
+
+ _require_sku("xpu", res)
+ marks = [getattr(pytest.mark, res), pytest.mark.xpu, _cards_mark(num_cards)]
+ if num_cards == 1 or not current_platform.is_xpu():
return marks
- test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards)
- return marks + [test_distributed]
+ return marks + [
+ pytest.mark.skipif(
+ current_platform.device_count() < num_cards,
+ reason=f"Need at least {num_cards} XPUs to run the test.",
+ )
+ ]
-def xpu_marks(*, res: str, num_cards: int):
- test_platform_detail = pytest.mark.xpu
- if res == "B60":
- test_resource = pytest.mark.B60
- else:
- raise ValueError(f"Invalid XPU resource type: {res}. Supported: B60")
- marks = [test_resource, test_platform_detail]
- if num_cards == 1:
- return marks
- test_distributed = pytest.mark.distributed_xpu(num_cards=num_cards)
+def _musa_marks(*, res: str, num_cards: int):
+ import pytest
- if not current_platform.is_xpu():
- return marks + [test_distributed]
- test_skipif = pytest.mark.skipif(
- current_platform.device_count() < num_cards,
- reason=f"Need at least {num_cards} XPUs to run the test.",
- )
- return marks + [test_distributed, test_skipif]
+ _require_sku("musa", res)
+ return [getattr(pytest.mark, res), pytest.mark.musa, _cards_mark(num_cards)]
-def musa_marks(*, res: str, num_cards: int):
- test_platform_detail = pytest.mark.musa
- if res == "S5000":
- test_resource = pytest.mark.S5000
- else:
- raise ValueError(f"Invalid MUSA resource type: {res}. Supported: S5000")
- marks = [test_resource, test_platform_detail]
- if num_cards == 1:
- return marks
- test_distributed = pytest.mark.distributed_musa(num_cards=num_cards)
- return marks + [test_distributed]
+@lru_cache(maxsize=1)
+def _gpu_res_platforms() -> frozenset[str]:
+ """Accelerator platforms tagged ``[hardware-platform] [gpu]`` (not the ``gpu`` mark)."""
+ return (get_supported_platforms() & _marker_names_with_tag(_GPU_DEVICE_CLASS_TAG)) - {"gpu"}
-def gpu_marks(*, res: str, num_cards: int):
- test_platform = pytest.mark.gpu
- if res in ("L4", "H100"):
- return [test_platform] + cuda_marks(res=res, num_cards=num_cards)
- if res == "MI325":
- return [test_platform] + rocm_marks(res=res, num_cards=num_cards)
- if res == "B60":
- return [test_platform] + xpu_marks(res=res, num_cards=num_cards)
- if res == "S5000":
- return [test_platform] + musa_marks(res=res, num_cards=num_cards)
- raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325, B60, S5000")
+def _npu_marks(*, res: str, num_cards: int):
+ import pytest
+
+ _require_sku("npu", res)
+ return [pytest.mark.npu, getattr(pytest.mark, res), _cards_mark(num_cards)]
-def npu_marks(*, res: str, num_cards: int):
- test_platform = pytest.mark.npu
- if res == "A2":
- test_resource = pytest.mark.A2
- elif res == "A3":
- test_resource = pytest.mark.A3
- else:
- test_resource = None
- if num_cards == 1:
- return [mark for mark in [test_platform, test_resource] if mark is not None]
- test_distributed = pytest.mark.distributed_npu(num_cards=num_cards)
- return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None]
+def _res_platforms() -> frozenset[str]:
+ """Accelerator keys allowed in ``hardware_test(res=...)`` / ``hardware_marks(res=...)``.
+
+ A ``[hardware-platform]`` name is a ``res`` key only if some SKU is tagged
+ with that platform (``[cuda]``, ``[npu]``, …). ``cpu`` / ``gpu`` have no SKUs.
+ """
+ return frozenset(p for p in get_supported_platforms() if _skus_for_platform(p))
+
+
+def _normalize_num_cards(res: dict[str, str], num_cards: int | dict[str, int]) -> dict[str, int]:
+ allowed = _res_platforms()
+ device_class = get_supported_platforms() - allowed
+ for platform in res:
+ if platform in device_class:
+ raise ValueError(
+ f"{platform!r} is not a res dict key; use pytest.mark.{platform} "
+ f"or an accelerator key ({', '.join(sorted(allowed))})."
+ )
+ if platform not in allowed:
+ supported = ", ".join(sorted(allowed)) or "(none)"
+ raise ValueError(f"Unsupported platform: {platform}. res keys: {supported}.")
+ if isinstance(num_cards, int) and not isinstance(num_cards, bool):
+ return {platform: num_cards for platform in res}
+ if not isinstance(num_cards, dict):
+ raise ValueError(f"num_cards must be a positive int or dict, got {num_cards!r}")
+ num_cards_dict = dict(num_cards)
+ for platform in num_cards_dict:
+ if platform not in res:
+ raise ValueError(f"Platform '{platform}' in num_cards but not in res.")
+ for platform in res:
+ num_cards_dict.setdefault(platform, 1)
+ return num_cards_dict
+
+
+def _marks_for_platform(platform: str, resource: str, num_cards: int) -> list[pytest.MarkDecorator]:
+ import pytest
+
+ builders = {
+ "cuda": _cuda_marks,
+ "rocm": _rocm_marks,
+ "xpu": _xpu_marks,
+ "musa": _musa_marks,
+ "npu": _npu_marks,
+ }
+ builder = builders.get(platform)
+ if builder is None:
+ raise ValueError(f"Unsupported platform: {platform}")
+ marks = builder(res=resource, num_cards=num_cards)
+ if platform in _gpu_res_platforms():
+ return [pytest.mark.gpu] + marks
+ return marks
+
+
+def _skipif_not_platform(platform: str) -> pytest.MarkDecorator:
+ import pytest
+
+ from vllm_omni.platforms import current_omni_platform
+
+ checkers = {
+ "cuda": current_omni_platform.is_cuda,
+ "rocm": current_omni_platform.is_rocm,
+ "xpu": current_omni_platform.is_xpu,
+ "npu": current_omni_platform.is_npu,
+ "musa": current_omni_platform.is_musa,
+ }
+ return pytest.mark.skipif(not checkers[platform](), reason=f"Requires {platform} platform")
+
+
+def _apply_marks(func, marks: list[pytest.MarkDecorator]):
+ for mark in reversed(marks):
+ func = mark(func)
+ return func
def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
- for platform, _ in res.items():
- if platform not in ("cuda", "rocm", "xpu", "npu", "musa"):
- raise ValueError(f"Unsupported platform: {platform}")
- if isinstance(num_cards, int):
- num_cards_dict = {platform: num_cards for platform in res.keys()}
- else:
- num_cards_dict = num_cards
- for platform in num_cards_dict.keys():
- if platform not in res:
- raise ValueError(f"Platform '{platform}' in num_cards but not in res.")
- for platform in res.keys():
- if platform not in num_cards_dict:
- num_cards_dict[platform] = 1
+ """Return marks for a **single** pytest item.
+
+ Different per-platform ``num_cards`` cannot share one item: ``cards_2`` and
+ ``cards_4`` would both match ``-m cards_2`` / ``-m cards_4``. Use
+ ``@hardware_test`` (one collected variant per platform) or call this once
+ per platform and attach each list to its own ``pytest.param``.
+ """
+ num_cards_dict = _normalize_num_cards(res, num_cards)
+ unique_counts = set(num_cards_dict.values())
+ if len(res) > 1 and len(unique_counts) > 1:
+ raise ValueError(
+ "hardware_marks() cannot attach different per-platform num_cards to one pytest item "
+ f"(got {num_cards_dict}). Use @hardware_test, which collects one variant per platform, "
+ "or call hardware_marks() once per platform."
+ )
all_marks: list[pytest.MarkDecorator] = []
for platform, resource in res.items():
- cards = num_cards_dict[platform]
- if platform in ("cuda", "rocm", "xpu"):
- marks = gpu_marks(res=resource, num_cards=cards)
- elif platform == "musa":
- marks = musa_marks(res=resource, num_cards=cards)
- elif platform == "npu":
- marks = npu_marks(res=resource, num_cards=cards)
- else:
- raise ValueError(f"Unsupported platform: {platform}")
- all_marks.extend(marks)
+ all_marks.extend(_marks_for_platform(platform, resource, num_cards_dict[platform]))
return all_marks
def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
- all_marks = hardware_marks(res=res, num_cards=num_cards)
+ """Apply hardware marks; split collected items when card counts differ by platform.
- def wrapper(f):
- func = f
- for mark in reversed(all_marks):
- func = mark(func)
- return func
+ A single item cannot carry both ``cards_4`` (CUDA) and ``cards_2`` (ROCm):
+ ``-m "H100 and cards_4 and cuda"`` would otherwise also match
+ ``-m "H100 and cards_2 and cuda"``. When ``num_cards`` differs across
+ platforms, this decorator parametrizes one variant per platform so each
+ item only has that platform's SKU and ``cards_{n}``.
+ """
+ import pytest
- return wrapper
+ num_cards_dict = _normalize_num_cards(res, num_cards)
+ mixed_counts = len(res) > 1 and len(set(num_cards_dict.values())) > 1
+ if not mixed_counts:
+ all_marks = hardware_marks(res=res, num_cards=num_cards)
+
+ def apply_union_marks(f):
+ return _apply_marks(f, all_marks)
+
+ return apply_union_marks
+
+ params = []
+ for platform, resource in res.items():
+ marks = _marks_for_platform(platform, resource, num_cards_dict[platform])
+ marks.append(_skipif_not_platform(platform))
+ params.append(pytest.param(platform, marks=marks, id=platform))
+
+ def apply_split_params(f):
+ func = pytest.mark.usefixtures(HARDWARE_MARK_NORMALIZATION_FIXTURE)(f)
+ return pytest.mark.parametrize(HARDWARE_MARK_NORMALIZATION_FIXTURE, params, indirect=True)(func)
+
+ return apply_split_params
diff --git a/tests/helpers/media.py b/tests/helpers/media.py
index 8eb82fc1f..623c218ad 100644
--- a/tests/helpers/media.py
+++ b/tests/helpers/media.py
@@ -1,23 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Synthetic media generation and media/text utilities for tests."""
import atexit
import base64
import concurrent.futures
-import gc
import hashlib
import io
import logging
import math
+import mimetypes
import multiprocessing
import os
import random
import re
import subprocess
+import sys
import tempfile
+import threading
import time
+from concurrent.futures.process import BrokenProcessPool
from contextlib import contextmanager
from pathlib import Path
-from typing import Any
+from typing import Any, Literal, TypedDict, overload
import numpy as np
import soundfile as sf
@@ -391,8 +397,16 @@ def generate_synthetic_video(
import cv2
import imageio
+ class _BouncingBall(TypedDict):
+ x: float
+ y: float
+ vx: float
+ vy: float
+ radius: int
+ color_bgr: tuple[int, int, int]
+
num_balls = random.randint(3, 8)
- balls = []
+ balls: list[_BouncingBall] = []
for _ in range(num_balls):
radius = min(width, height) // 8
if radius < 1:
@@ -404,7 +418,16 @@ def generate_synthetic_video(
vx = speed * math.cos(angle)
vy = speed * math.sin(angle)
color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
- balls.append({"x": x, "y": y, "vx": vx, "vy": vy, "radius": radius, "color_bgr": color_bgr})
+ balls.append(
+ {
+ "x": float(x),
+ "y": float(y),
+ "vx": float(vx),
+ "vy": float(vy),
+ "radius": radius,
+ "color_bgr": color_bgr,
+ }
+ )
video_frames = []
for _ in range(num_frames):
@@ -524,27 +547,34 @@ _AUDIO_MIME_BY_SUFFIX = {
}
-def get_asset_path(relative_path: str | os.PathLike) -> Path:
- """Resolve a path under ``tests/assets/`` to its absolute on-disk location."""
- return _TEST_ASSETS_ROOT / Path(relative_path)
+def _asset_mime_type(path: Path) -> str:
+ mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower())
+ if mime is not None:
+ return mime
+ guessed, _ = mimetypes.guess_type(path.name)
+ return guessed or "application/octet-stream"
-def load_test_audio_data_url(relative_path: str | os.PathLike) -> str:
- """Load a vendored test audio file under ``tests/assets/`` as a base64 data URL.
+@overload
+def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[False] = False) -> Path: ...
- Used by tests that need real reference audio (e.g. voice cloning) without
- relying on the server's ability to fetch external URLs at request time.
+
+@overload
+def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[True]) -> str: ...
+
+
+def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: bool = False) -> Path | str:
+ """Resolve a path under ``tests/assets/``.
+
+ When ``as_data_url`` is true, read the file and return a base64 data URL
+ suitable for embedding reference media in API requests without external URLs.
"""
- path = get_asset_path(relative_path)
- mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower(), "application/octet-stream")
+ path = _TEST_ASSETS_ROOT / Path(relative_path)
+ if not as_data_url:
+ return path
+
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
- return f"data:{mime};base64,{encoded}"
-
-
-def decode_b64_image(b64: str):
- img = Image.open(io.BytesIO(base64.b64decode(b64)))
- img.load()
- return img
+ return f"data:{_asset_mime_type(path)};base64,{encoded}"
def concat_audio(audio_val) -> np.ndarray:
@@ -657,6 +687,10 @@ def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
@contextmanager
def _serialize_whisper_model_download(model_size: str = "small"):
"""Serialize Whisper cache writes across processes (Linux/Unix), per model."""
+ if sys.platform == "win32":
+ yield
+ return
+
import fcntl
lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock"
@@ -670,59 +704,152 @@ def _serialize_whisper_model_download(model_size: str = "small"):
f.close()
-def _whisper_transcribe_in_current_process(
- output_path: str, model_size: str = "small", language: str | None = None
-) -> str:
- import whisper
-
+def _select_whisper_device() -> str:
device_index = None
from vllm_omni.platforms import current_omni_platform
if current_omni_platform.is_available():
n = current_omni_platform.get_device_count()
# Single-GPU runners (e.g. the L4 nightly): the model server already
- # occupies device 0. Loading Whisper there, once per concurrent
- # request, competes for VRAM and OOMs. Only borrow an accelerator
- # when a spare device exists; otherwise validate on CPU.
+ # occupies device 0, and a Whisper model resident there competes with
+ # it for VRAM and OOMs. Only borrow an accelerator when a spare device
+ # exists; otherwise validate on CPU.
if n > 1:
device_index = n - 1
- if device_index is not None:
- torch_device = current_omni_platform.get_torch_device(device_index)
- current_omni_platform.set_device(torch_device)
- device = str(torch_device)
- use_accelerator = True
- else:
- use_accelerator = False
- device = "cpu"
+ if device_index is None:
+ return "cpu"
- with _serialize_whisper_model_download(model_size):
- model = whisper.load_model(model_size, device=device)
- try:
- text = model.transcribe(
- output_path,
- temperature=0.0,
- word_timestamps=True,
- condition_on_previous_text=False,
- # None keeps whisper's auto-detection. Do not default this to a
- # language: callers include non-English audio tests.
- language=language,
- )["text"]
- finally:
- del model
- gc.collect()
- if use_accelerator:
- current_omni_platform.synchronize()
- current_omni_platform.empty_cache()
+ torch_device = current_omni_platform.get_torch_device(device_index)
+ current_omni_platform.set_device(torch_device)
+ return str(torch_device)
+
+
+# Populated in the transcription worker, not in the pytest process.
+_WHISPER_MODELS: dict[str, Any] = {}
+
+
+def _get_whisper_model(model_size: str) -> Any:
+ model = _WHISPER_MODELS.get(model_size)
+ if model is None:
+ import whisper
+
+ # The device is picked on first load and the model stays on it for the
+ # worker's lifetime: the current server or runner fixture instance, or
+ # the test module for callers that transcribe without those fixtures.
+ device = _select_whisper_device()
+ with _serialize_whisper_model_download(model_size):
+ model = whisper.load_model(model_size, device=device)
+ _WHISPER_MODELS[model_size] = model
+ return model
+
+
+def _whisper_transcribe_in_current_process(
+ output_path: str, model_size: str = "small", language: str | None = None
+) -> str:
+ model = _get_whisper_model(model_size)
+ text = model.transcribe(
+ output_path,
+ temperature=0.0,
+ word_timestamps=True,
+ condition_on_previous_text=False,
+ # None keeps whisper's auto-detection. Do not default this to a
+ # language: callers include non-English audio tests.
+ language=language,
+ )["text"]
return text or ""
+# Serializes a whole submit->result->cleanup on the parent side, so at most one
+# call is ever in flight on the single-worker pool. The child already runs
+# transcriptions serially, so this costs no throughput; it lets a failed call
+# discard the worker without disrupting another caller's in-flight future.
+_TRANSCRIBER_CALL_LOCK = threading.Lock()
+# Guards the _TRANSCRIBER pointer itself.
+_TRANSCRIBER_LOCK = threading.Lock()
+_TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None
+
+
+def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor:
+ global _TRANSCRIBER
+ with _TRANSCRIBER_LOCK:
+ if _TRANSCRIBER is None:
+ ctx = multiprocessing.get_context("spawn")
+ _TRANSCRIBER = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx)
+ return _TRANSCRIBER
+
+
+def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> None:
+ """Drop ``executor``, but only while it is still the current one.
+
+ Identity-checked so a stale reference can never shut down a newer worker that
+ was installed after ``executor`` was replaced.
+ """
+ global _TRANSCRIBER
+ with _TRANSCRIBER_LOCK:
+ if _TRANSCRIBER is not executor:
+ return
+ _TRANSCRIBER = None
+ # Joining the worker can block; do it outside the lock.
+ executor.shutdown(wait=True)
+
+
+def release_audio_transcriber() -> None:
+ """Shut the transcription worker down, freeing the device memory its models hold.
+
+ Called when a server or runner fixture instance tears down, so that the next
+ one -- including the next parametrization inside the same test module -- does
+ not initialize its model while Whisper still occupies the device. A module
+ teardown fixture repeats it for tests that transcribe without those fixtures.
+
+ Takes the call lock, so it waits for any in-flight transcription to finish
+ rather than shutting the worker down underneath it.
+ """
+ global _TRANSCRIBER
+ with _TRANSCRIBER_CALL_LOCK:
+ with _TRANSCRIBER_LOCK:
+ executor, _TRANSCRIBER = _TRANSCRIBER, None
+ if executor is not None:
+ executor.shutdown(wait=True)
+
+
def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
- """Convert an audio file to text in an isolated subprocess."""
- ctx = multiprocessing.get_context("spawn")
- with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
- future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language)
- return future.result()
+ """Convert an audio file to text in a reused, isolated subprocess.
+
+ The worker outlives the call so its Whisper model is loaded once rather than
+ once per transcription. The call lock serializes callers onto the single
+ worker (the child already transcribes serially, so this costs no throughput),
+ which lets a failed call tear the worker down without racing another caller.
+ The worker caches one model per size it is asked for, so a run that escalates
+ to a stronger ASR keeps both models resident until release.
+
+ Any failure discards the worker, restoring the failure isolation of the old
+ one-process-per-call design: a failure (a ``torch`` OOM, or a
+ ``KeyboardInterrupt``/``SystemExit`` transported through the future or
+ raised while ``result()`` blocks) propagates unchanged after the worker --
+ and its resident model -- is torn down, and a dead worker
+ (``BrokenProcessPool``) is additionally retried once.
+ """
+ with _TRANSCRIBER_CALL_LOCK:
+ for attempt in range(2):
+ executor = _get_transcriber()
+ try:
+ return executor.submit(
+ _whisper_transcribe_in_current_process, output_path, model_size, language
+ ).result()
+ except BrokenProcessPool:
+ _discard_transcriber(executor)
+ if attempt == 1:
+ raise
+ except BaseException:
+ # Any other failure -- a task exception (a torch OOM included), or
+ # a KeyboardInterrupt/SystemExit transported through the future or
+ # interrupting result() -- leaves the worker and its model
+ # resident; drop it so it cannot contaminate later calls, matching
+ # the old per-call teardown. Do not retry.
+ _discard_transcriber(executor)
+ raise
+ raise AssertionError("unreachable")
def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str:
@@ -742,11 +869,10 @@ __all__ = [
"convert_audio_bytes_to_text",
"convert_audio_file_to_text",
"cosine_similarity_text",
- "decode_b64_image",
"generate_synthetic_audio",
"generate_synthetic_image",
"generate_synthetic_video",
"get_asset_path",
- "load_test_audio_data_url",
"preprocess_text",
+ "release_audio_transcriber",
]
diff --git a/tests/helpers/monitor.py b/tests/helpers/monitor.py
new file mode 100644
index 000000000..3402bd186
--- /dev/null
+++ b/tests/helpers/monitor.py
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Device memory monitoring helpers for tests."""
+
+from __future__ import annotations
+
+import threading
+import time
+
+from vllm_omni.platforms import current_omni_platform
+
+
+class DeviceMemoryMonitor:
+ """Poll global device memory usage."""
+
+ def __init__(self, device_index: int, interval: float = 0.05):
+ self.device_index = device_index
+ self.interval = interval
+ self._peak_used_mb = 0.0
+ self._stop_event = threading.Event()
+ self._thread: threading.Thread | None = None
+
+ def start(self) -> None:
+ def monitor_loop() -> None:
+ while not self._stop_event.is_set():
+ try:
+ with current_omni_platform.device(self.device_index):
+ free_bytes, total_bytes = current_omni_platform.mem_get_info()
+ used_mb = (total_bytes - free_bytes) / (1024**2)
+ self._peak_used_mb = max(self._peak_used_mb, used_mb)
+ except Exception:
+ pass
+ time.sleep(self.interval)
+
+ self._thread = threading.Thread(target=monitor_loop, daemon=False)
+ self._thread.start()
+
+ def stop(self) -> None:
+ if self._thread is None:
+ return
+ self._stop_event.set()
+ self._thread.join(timeout=2.0)
+
+ @property
+ def peak_used_mb(self) -> float:
+ fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2)
+ fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2)
+ return max(self._peak_used_mb, fallback_alloc, fallback_reserved)
+
+ def __del__(self):
+ self.stop()
+
+
+__all__ = ["DeviceMemoryMonitor"]
diff --git a/tests/helpers/process.py b/tests/helpers/process.py
deleted file mode 100644
index 094de9652..000000000
--- a/tests/helpers/process.py
+++ /dev/null
@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import functools
-import os
-import signal
-import subprocess
-import sys
-import tempfile
-from collections.abc import Callable
-from contextlib import ExitStack, suppress
-from typing import Any, Literal
-
-import cloudpickle
-from typing_extensions import ParamSpec
-from vllm.platforms import current_platform
-
-_P = ParamSpec("_P")
-
-
-def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, None]:
- """Decorator to fork a new process for each test function."""
-
- @functools.wraps(func)
- def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
- os.setpgrp()
- from _pytest.outcomes import Skipped
-
- with (
- tempfile.NamedTemporaryFile(
- delete=False, mode="w+b", prefix=f"vllm_test_{func.__name__}_{os.getpid()}_", suffix=".exc"
- ) as exc_file,
- ExitStack() as delete_after,
- ):
- exc_file_path = exc_file.name
- delete_after.callback(os.remove, exc_file_path)
- pid = os.fork()
- if pid == 0:
- delete_after.pop_all()
- try:
- func(*args, **kwargs)
- except Skipped as e:
- print(str(e))
- os._exit(0)
- except Exception as e:
- import traceback
-
- tb_string = traceback.format_exc()
- exc_to_serialize: dict[str, Any]
- try:
- exc_to_serialize = {"pickled_exception": e}
- cloudpickle.dumps(exc_to_serialize)
- except (Exception, KeyboardInterrupt):
- exc_to_serialize = {
- "exception_type": type(e).__name__,
- "exception_msg": str(e),
- "traceback": tb_string,
- }
- try:
- with open(exc_file_path, "wb") as f:
- cloudpickle.dump(exc_to_serialize, f)
- except Exception:
- print(tb_string)
- os._exit(1)
- else:
- os._exit(0)
- else:
- pgid = os.getpgid(pid)
- _pid, _exitcode = os.waitpid(pid, 0)
- old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
- os.killpg(pgid, signal.SIGTERM)
- signal.signal(signal.SIGTERM, old_signal_handler)
- if _exitcode != 0:
- exc_info = {}
- if os.path.exists(exc_file_path):
- with suppress(Exception), open(exc_file_path, "rb") as f:
- exc_info = cloudpickle.load(f)
- if (original_exception := exc_info.get("pickled_exception")) is not None:
- assert isinstance(original_exception, Exception)
- raise original_exception
- if (original_tb := exc_info.get("traceback")) is not None:
- raise AssertionError(
- f"Test {func.__name__} failed when called with args {args} and kwargs {kwargs}"
- f" (exit code: {_exitcode}):\n{original_tb}"
- ) from None
- raise AssertionError(
- f"function {func.__name__} failed when called with args {args} and kwargs {kwargs}"
- f" (exit code: {_exitcode})"
- ) from None
-
- return wrapper
-
-
-def spawn_new_process_for_each_test(f: Callable[_P, None]) -> Callable[_P, None]:
- """Decorator to spawn a new process for each test function."""
-
- @functools.wraps(f)
- def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
- if os.environ.get("RUNNING_IN_SUBPROCESS") == "1":
- return f(*args, **kwargs)
-
- import torch.multiprocessing as mp
-
- with suppress(RuntimeError):
- mp.set_start_method("spawn")
-
- module_name = f.__module__
- env = os.environ.copy()
- env["RUNNING_IN_SUBPROCESS"] = "1"
-
- with tempfile.TemporaryDirectory() as tempdir:
- output_filepath = os.path.join(tempdir, "new_process.tmp")
- input_bytes = cloudpickle.dumps((f, output_filepath))
- cmd = [sys.executable, "-m", f"{module_name}"]
- returned = subprocess.run(cmd, input=input_bytes, capture_output=True, env=env)
- try:
- returned.check_returncode()
- except Exception as e:
- raise RuntimeError(f"Error raised in subprocess:\n{returned.stderr.decode()}") from e
-
- return wrapper
-
-
-def create_new_process_for_each_test(
- method: Literal["spawn", "fork"] | None = None,
-) -> Callable[[Callable[_P, None]], Callable[_P, None]]:
- """Creates a decorator that runs each test function in a new process."""
- if method is None:
- use_spawn = current_platform.is_xpu()
- method = "spawn" if use_spawn else "fork"
- assert method in ["spawn", "fork"], "Method must be either 'spawn' or 'fork'"
- if method == "fork":
- return fork_new_process_for_each_test
- return spawn_new_process_for_each_test
diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py
index 82fcc308d..db4cba893 100644
--- a/tests/helpers/runtime.py
+++ b/tests/helpers/runtime.py
@@ -1,13 +1,13 @@
-"""Server/client/runner runtime primitives for tests."""
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Server and offline runner runtime primitives for tests.
+
+Request clients and response types live in ``tests.helpers.client``;
+this module re-exports them for backward-compatible imports.
+"""
-import asyncio
-import base64
-import concurrent.futures
-import copy
import errno
-import gc
-import io
-import json
import os
import socket
import subprocess
@@ -16,126 +16,27 @@ import tempfile
import threading
import time
from collections.abc import Generator
-from dataclasses import asdict, dataclass
-from io import BytesIO
+from contextlib import contextmanager
from pathlib import Path
-from typing import Any, NamedTuple, cast
-from urllib.parse import quote
+from typing import Any, NamedTuple
import numpy as np
import psutil
-import requests
-import soundfile as sf
-import torch
import yaml
-from openai import APIError, OpenAI, omit
-from PIL import Image
-from vllm import TextPrompt, envs
-from vllm.distributed.parallel_state import (
- destroy_distributed_environment,
- destroy_model_parallel,
-)
+from vllm import TextPrompt
from vllm.logger import init_logger
-from tests.helpers.assertions import (
- assert_audio_speech_response,
- assert_diffusion_response,
- assert_http_error,
- assert_images_generations_response,
- assert_omni_response,
-)
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
+from tests.helpers.clean import cleanup_test_environment
from tests.helpers.media import (
- _merge_base64_audio_to_segment,
- decode_b64_image,
+ release_audio_transcriber,
)
+from tests.helpers.stage_config import load_stage_ids, load_stage_replica_counts
from tests.model_tests.diffusion.utils import resolve_tiny_model_path
from vllm_omni.config.stage_config import resolve_deploy_yaml
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
-
-def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
- # Reset environment variable cache
- envs.disable_envs_cache()
-
- # Reset rocm_aiter_ops class variables to match current os.environ.
- # These are class-level attributes that persist across tests and are
- # NOT restored by monkeypatch (which only restores os.environ).
- from vllm_omni.platforms import current_omni_platform
-
- if current_omni_platform.is_rocm():
- from vllm._aiter_ops import rocm_aiter_ops
-
- rocm_aiter_ops.refresh_env_variables()
-
- # Ensure all objects are not frozen before cleanup
- gc.unfreeze()
-
- destroy_model_parallel()
- destroy_distributed_environment()
- if shutdown_ray:
- import ray # Lazy import Ray
-
- ray.shutdown()
- gc.collect()
-
- if not current_omni_platform.is_cpu():
- current_omni_platform.empty_cache()
- try:
- torch._C._host_emptyCache()
- except AttributeError:
- logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
-
-
-def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None:
- try:
- data = r.json()
- if isinstance(data, (dict, list)):
- return data
- except Exception:
- pass
- return None
-
-
-def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
- """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
- eb = cfg.get("extra_body")
- if not eb:
- return None
- h, w = eb.get("height"), eb.get("width")
- if (isinstance(h, (list, tuple)) or isinstance(w, (list, tuple))) and not (
- isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))
- ):
- raise ValueError("extra_body height and width must both be lists or both be scalars")
- if not (isinstance(h, (list, tuple)) and isinstance(w, (list, tuple))):
- return None
- if len(h) != len(w):
- raise ValueError(f"height and width lists must have equal length; got {len(h)=} {len(w)=}")
- n = len(h)
- n_out = eb.get("num_outputs_per_prompt")
- if n_out is not None:
- n_out = int(n_out)
- if n_out != n:
- raise ValueError(
- "When height/width are lists, num_outputs_per_prompt must equal their length; "
- f"got num_outputs_per_prompt={n_out}, len(lists)={n}"
- )
- splits: list[dict[str, Any]] = []
- for i in range(n):
- sub = copy.deepcopy(cfg)
- sub_eb = dict(sub.get("extra_body") or {})
- sub_eb["height"] = int(h[i])
- sub_eb["width"] = int(w[i])
- sub_eb["num_outputs_per_prompt"] = 1
- sub["extra_body"] = sub_eb
- splits.append(sub)
- return splits
-
-
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
PromptImageInput = list[Any] | Any | None
PromptVideoInput = list[Any] | Any | None
@@ -174,16 +75,30 @@ def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int:
) from last_exc
+def get_distributed_init_method(prefix: str = "torch_dist_init_") -> str:
+ """Return a ``file://`` init_method for a ``torch.distributed`` process group.
+
+ Args:
+ prefix: Prefix for the temporary rendezvous filename. Defaults to
+ ``torch_dist_init_``.
+
+ Returns:
+ An init_method string for ``torch.distributed.init_process_group``.
+ """
+ with tempfile.NamedTemporaryFile(prefix=prefix) as f:
+ return f"file://{f.name}"
+
+
def dummy_messages_from_mix_data(
- system_prompt: dict[str, Any] = None,
+ system_prompt: dict[str, Any] | None = None,
video_data_url: Any = None,
audio_data_url: Any = None,
image_data_url: Any = None,
- content_text: str = None,
+ content_text: str | None = None,
):
"""Create messages with video、image、audio data URL for OpenAI API."""
if content_text is not None:
- content = [{"type": "text", "text": content_text}]
+ content: list[dict[str, Any]] = [{"type": "text", "text": content_text}]
else:
content = []
@@ -246,9 +161,7 @@ class OmniServer:
env_dict: dict[str, str] | None = None,
use_omni: bool = True,
) -> None:
- run_pre_test_cleanup()
- run_post_test_cleanup()
- cleanup_dist_env_and_memory()
+ cleanup_test_environment()
self.model = model
args = list(serve_args)
self.serve_args = args
@@ -261,7 +174,6 @@ class OmniServer:
def _start_server(self) -> None:
env = os.environ.copy()
- env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
if self.env_dict is not None:
env.update(self.env_dict)
@@ -322,7 +234,7 @@ class OmniServer:
if proc.status() != psutil.STATUS_ZOMBIE:
return False
try:
- os.waitpid(proc.pid, os.WNOHANG)
+ os.waitpid(proc.pid, getattr(os, "WNOHANG", 0))
except ChildProcessError:
pass
return True
@@ -446,15 +358,19 @@ class OmniServer:
pass
def __enter__(self):
- self._start_server()
+ try:
+ self._start_server()
+ except BaseException:
+ # ``__exit__`` is not invoked when ``__enter__`` raises; roll back
+ # any processes / log handles already launched.
+ self.__exit__(None, None, None)
+ raise
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.proc:
self._kill_process_tree(self.proc.pid)
- run_pre_test_cleanup()
- run_post_test_cleanup()
- cleanup_dist_env_and_memory()
+ cleanup_test_environment()
class OmniServerStageCli(OmniServer):
@@ -473,7 +389,6 @@ class OmniServerStageCli(OmniServer):
super().__init__(model, serve_args or [], port=port, env_dict=env_dict, use_omni=True)
self.stage_config_path = stage_config_path
self.master_port = get_open_port()
- self.visible_device_list = self._load_visible_device_list(env_dict)
resolved_cfg = resolve_deploy_yaml(stage_config_path)
# Dump the resolved deploy config so CI logs show each stage's
# gpu_memory_utilization / max_model_len / max_num_seqs after
@@ -484,113 +399,15 @@ class OmniServerStageCli(OmniServer):
f"{yaml.safe_dump(resolved_cfg, sort_keys=False, default_flow_style=False)}",
flush=True,
)
- self.stage_runtime_devices = self._load_stage_runtime_devices(resolved_cfg)
- self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg)
+ self.stage_ids = stage_ids or load_stage_ids(resolved_cfg)
if 0 not in self.stage_ids:
raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}")
- self.stage_replica_counts = self._load_stage_replica_counts(resolved_cfg)
+ self.stage_replica_counts = load_stage_replica_counts(resolved_cfg)
self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {}
+ self._stage_log_paths: dict[tuple[int, int], Path] = {}
+ self._stage_log_files: dict[tuple[int, int], Any] = {}
self.proc = None
- @staticmethod
- def _stage_entries(cfg: dict) -> list[dict]:
- """Return the list of stage entries from either legacy (``stage_args``)
- or new-schema (``stages``) deploy YAMLs."""
- return cfg.get("stage_args") or cfg.get("stages") or []
-
- @staticmethod
- def _load_stage_ids(resolved_config: dict) -> list[int]:
- stage_ids = [
- stage["stage_id"] for stage in OmniServerStageCli._stage_entries(resolved_config) if "stage_id" in stage
- ]
- if not stage_ids:
- raise ValueError("No stage IDs found in resolved config")
- return stage_ids
-
- @staticmethod
- def _load_stage_runtime_devices(resolved_config: dict) -> dict[int, str]:
- runtime_devices: dict[int, str] = {}
- for stage in OmniServerStageCli._stage_entries(resolved_config):
- stage_id = stage.get("stage_id")
- # New schema: stage.devices is flat at stage level.
- # Legacy schema: stage.runtime.devices is nested.
- devices = stage.get("devices") or stage.get("runtime", {}).get("devices")
- if stage_id is not None and devices:
- runtime_devices[int(stage_id)] = str(devices)
- return runtime_devices
-
- @staticmethod
- def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]:
- replica_counts: dict[int, int] = {}
- for stage in OmniServerStageCli._stage_entries(resolved_config):
- stage_id = stage.get("stage_id")
- if stage_id is None:
- continue
- replica_counts[int(stage_id)] = max(
- 1,
- int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)),
- )
- return replica_counts
-
- @classmethod
- def _parse_device_list(cls, devices: str | int) -> list[str]:
- if isinstance(devices, int):
- if devices < 0:
- raise ValueError("Device IDs must be non-negative integers")
- return [str(devices)]
- return [token.strip() for token in str(devices).split(",") if token.strip()]
-
- @classmethod
- def _load_visible_device_list(cls, env_dict: dict[str, str] | None) -> list[str] | None:
- env = os.environ.copy()
- if env_dict is not None:
- env.update(env_dict)
-
- env_var = getattr(current_omni_platform, "device_control_env_var", None)
- if env_var and env_var in env:
- return [token.strip() for token in env[env_var].split(",") if token.strip()]
- return None
-
- @classmethod
- def _map_stage_devices(cls, stage_id: int, visible_device_list: list[str] | None, devices: str) -> str:
- device_list = cls._parse_device_list(devices)
-
- if visible_device_list is None:
- return ",".join(device_list)
-
- if not all(device.isdigit() for device in device_list):
- raise ValueError("Logical devices must be non-negative integers")
-
- logical_ids = [int(device) for device in device_list]
- if logical_ids and max(logical_ids) >= len(visible_device_list):
- raise ValueError(
- f"Stage {stage_id} has logical IDs {device_list}, one or more of which exceed the number of visible devices"
- )
-
- return ",".join(visible_device_list[idx] for idx in logical_ids)
-
- def _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str:
- replica_count = self.stage_replica_counts.get(stage_id, 1)
- if replica_count == 1:
- return devices
-
- device_list = self._parse_device_list(devices)
- if len(device_list) % replica_count != 0:
- raise ValueError(
- f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); "
- "device count must be divisible by replica count"
- )
- devices_per_replica = len(device_list) // replica_count
- start = replica_id * devices_per_replica
- return ",".join(device_list[start : start + devices_per_replica])
-
- def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None:
- replica_devices = self._devices_for_replica(stage_id, devices, replica_id)
- mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_devices)
- env_var = getattr(current_omni_platform, "device_control_env_var", None)
- if env_var:
- env[env_var] = mapped_devices
-
def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]:
cmd = [
sys.executable,
@@ -621,7 +438,6 @@ class OmniServerStageCli(OmniServer):
def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None:
env = os.environ.copy()
- env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if self.env_dict is not None:
env.update(self.env_dict)
@@ -631,11 +447,9 @@ class OmniServerStageCli(OmniServer):
# debugging "Stage N exited before API server ready" doesn't rely on
# guessing; the file is surfaced in the RuntimeError message.
log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log"
- self._stage_log_paths = getattr(self, "_stage_log_paths", {})
stage_key = (stage_id, replica_id)
self._stage_log_paths[stage_key] = log_path
- log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in __exit__
- self._stage_log_files = getattr(self, "_stage_log_files", {})
+ log_fh = open(log_path, "w", buffering=1) # noqa: SIM115 - closed in ``__exit__``
self._stage_log_files[stage_key] = log_fh
proc = subprocess.Popen(
cmd,
@@ -655,7 +469,7 @@ class OmniServerStageCli(OmniServer):
for (stage_id, replica_id), proc in self.stage_procs.items():
ret = proc.poll()
if ret is not None:
- log_path = getattr(self, "_stage_log_paths", {}).get((stage_id, replica_id))
+ log_path = self._stage_log_paths.get((stage_id, replica_id))
tail = ""
if log_path and log_path.exists():
try:
@@ -712,7 +526,7 @@ class OmniServerStageCli(OmniServer):
captures them post-run. Head covers engine init; tail covers
whatever state the stage was in when it was torn down.
"""
- log_paths = getattr(self, "_stage_log_paths", {}) or {}
+ log_paths = self._stage_log_paths
for stage_id, replica_id in sorted(log_paths):
log_path = log_paths[(stage_id, replica_id)]
if not log_path or not log_path.exists():
@@ -746,1877 +560,30 @@ class OmniServerStageCli(OmniServer):
proc = self.stage_procs[stage_key]
if proc.poll() is None:
self._kill_process_tree(proc.pid)
- run_pre_test_cleanup()
- run_post_test_cleanup()
- cleanup_dist_env_and_memory()
-
-@dataclass
-class OmniResponse:
- """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
-
- text_content: str | None = None
- audio_data: list[str] | None = None
- audio_content: str | None = None
- audio_format: str | None = None
- audio_bytes: bytes | None = None
- #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the
- #: OpenAI client call through response parsing and local post-process (e.g. audio decode).
- e2e_latency: float | None = None
- success: bool = False
- prompt_tokens: int | None = None
- cached_tokens: int | None = None
- logprobs: list | None = None
- #: HTTP status + error text for the error-handling path (e.g. validator
- #: rejections); populated when the OpenAI client raises an APIError.
- status_code: int | None = None
- error_message: str | None = None
-
-
-@dataclass
-class DiffusionResponse:
- """Decoded diffusion output from chat completions or offline runner (not raw ``requests``)."""
-
- text_content: str | None = None
- images: list[Image.Image] | None = None
- audios: list[Any] | None = None
- videos: list[Any] | None = None
- #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before
- #: ``chat.completions.create`` through local image / audio decode.
- e2e_latency: float | None = None
- success: bool = False
-
-
-@dataclass
-class HttpResponse:
- """Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers."""
-
- status_code: int
- success: bool
- error_message: str | None = None
- json_body: dict[str, Any] | list[Any] | None = None
-
-
-@dataclass
-class WebSocketJsonResponse:
- """First JSON object delivered as a text WebSocket frame (streaming endpoints)."""
-
- json_body: dict[str, Any]
-
-
-@dataclass
-class OpenPIWebSocketResponse:
- """Msgpack WebSocket session against ``/v1/realtime/robot/openpi``."""
-
- server_metadata: dict[str, Any]
- operation_responses: list[Any]
- actions: dict[str, np.ndarray] | None = None
- action_tensors: list[np.ndarray] | None = None
-
-
-def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]:
- """Build a minimal DROID-style observation payload for the OpenPI robot endpoint."""
- identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32)
- identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32)
- return {
- "session_id": session_id,
- "video": {
- "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
- "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8),
- },
- "state": {
- "eef_9d": identity_eef_9d,
- "gripper_position": np.zeros((1, 1, 1), dtype=np.float32),
- "joint_position": np.zeros((1, 1, 7), dtype=np.float32),
- },
- "language": {"annotation.language.language_instruction": [["pick up the object"]]},
- }
-
-
-DREAMZERO_DEFAULT_PROMPT = (
- "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-)
-DREAMZERO_ACTION_HORIZON = 24
-DREAMZERO_ACTION_DIM = 8
-DREAMZERO_CAMERA_FILES = {
- "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
- "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
- "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def _require_opencv() -> Any:
- try:
- import cv2
- except ImportError as exc: # pragma: no cover - optional e2e dependency
- raise ModuleNotFoundError("DreamZero OpenPI test dependencies are missing: opencv-python") from exc
- return cv2
-
-
-def load_dreamzero_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
- cv2 = _require_opencv()
- camera_frames: dict[str, np.ndarray] = {}
- for camera_key, file_name in DREAMZERO_CAMERA_FILES.items():
- video_path = video_dir / file_name
- if not video_path.exists():
- raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
- cap = cv2.VideoCapture(str(video_path))
- frames = []
- while True:
- ok, frame = cap.read()
- if not ok:
- break
- frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
- cap.release()
- if not frames:
- raise RuntimeError(f"No frames loaded from {video_path}")
- camera_frames[camera_key] = np.stack(frames, axis=0)
- return camera_frames
-
-
-def build_dreamzero_demo_observations(
- camera_frames: dict[str, np.ndarray],
- *,
- prompt: str,
- session_id: str,
- num_chunks: int = 2,
-) -> list[dict[str, Any]]:
- if num_chunks < 1:
- raise ValueError("num_chunks must be at least 1")
-
- relative_offsets = [-23, -16, -8, 0]
- total_frames = min(frames.shape[0] for frames in camera_frames.values())
- frame_schedules = [[0]]
- current_frame = 23
- for _ in range(num_chunks - 1):
- indices = [max(current_frame + offset, 0) for offset in relative_offsets]
- if indices[-1] >= total_frames:
- break
- frame_schedules.append(indices)
- current_frame += DREAMZERO_ACTION_HORIZON
-
- observations: list[dict[str, Any]] = []
- for frame_indices in frame_schedules:
- obs: dict[str, Any] = {}
- for camera_key, all_frames in camera_frames.items():
- selected = all_frames[frame_indices]
- obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
- obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
- obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
- obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
- obs["prompt"] = prompt
- obs["session_id"] = session_id
- observations.append(obs)
- return observations
-
-
-class OpenPIWebSocketSession:
- """Persistent msgpack session for ``/v1/realtime/robot/openpi``."""
-
- DEFAULT_PING_INTERVAL_SECS = 300
- DEFAULT_PING_TIMEOUT_SECS = 3600
-
- @staticmethod
- def _require_dependencies() -> tuple[Any, Any]:
- try:
- import websockets.sync.client as websockets_client
- except ImportError as exc: # pragma: no cover - optional e2e dependency
- raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: websockets") from exc
- try:
- from openpi_client import msgpack_numpy
- except ImportError as exc: # pragma: no cover - optional e2e dependency
- raise ModuleNotFoundError("GR00T OpenPI test dependencies are missing: openpi-client") from exc
- return websockets_client, msgpack_numpy
-
- def __init__(
- self,
- uri: str,
- *,
- ping_interval: float = DEFAULT_PING_INTERVAL_SECS,
- ping_timeout: float = DEFAULT_PING_TIMEOUT_SECS,
- open_timeout: float = 120.0,
- close_timeout: float = 120.0,
- ) -> None:
- websockets_client, msgpack_numpy = self._require_dependencies()
- self._msgpack_numpy = msgpack_numpy
- self._packer = msgpack_numpy.Packer()
- self._conn = websockets_client.connect(
- uri,
- compression=None,
- max_size=None,
- ping_interval=ping_interval,
- ping_timeout=ping_timeout,
- open_timeout=open_timeout,
- close_timeout=close_timeout,
- )
- server_metadata = msgpack_numpy.unpackb(self._conn.recv())
- if not isinstance(server_metadata, dict):
- raise TypeError(f"Expected dict metadata from {uri}, got {type(server_metadata)!r}")
- self._server_metadata = server_metadata
-
- def get_server_metadata(self) -> dict[str, Any]:
- return dict(self._server_metadata)
-
- def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]:
- response = self._send_operation("infer", obs)
- if not isinstance(response, dict):
- raise TypeError(f"Expected dict infer response, got {type(response)!r}")
- return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
-
- def reset(self, reset_info: dict[str, Any] | None = None) -> str:
- response = self._send_operation("reset", dict(reset_info or {}))
- return str(response["status"])
-
- def close(self) -> None:
- self._conn.close()
-
- def _send_operation(self, endpoint: str, payload: dict[str, Any]) -> Any:
- body = dict(payload)
- body["endpoint"] = endpoint
- self._conn.send(self._packer.pack(body))
- raw = self._conn.recv()
- if isinstance(raw, str):
- raise RuntimeError(f"OpenPI {endpoint!r} failed: {raw}")
- decoded = self._msgpack_numpy.unpackb(raw)
- if isinstance(decoded, dict) and decoded.get("type") == "error":
- raise RuntimeError(f"OpenPI {endpoint!r} failed: {decoded.get('message')}")
- if endpoint == "reset":
- if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
- raise RuntimeError(f"Unexpected OpenPI reset response: {decoded!r}")
- return decoded
- return decoded
-
-
-def _merge_http_expectation_kwargs(
- base: dict[str, Any] | None,
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
-) -> dict[str, Any]:
- cfg = dict(base or {})
- if err_code is not None:
- cfg["err_code"] = err_code
- if err_message is not None:
- cfg["err_message"] = err_message
- return cfg
-
-
-def _merge_ws_expectation_kwargs(
- base: dict[str, Any] | None,
- *,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ws_json_type: str | None = None,
- ws_error_code: str | None = None,
-) -> dict[str, Any]:
- cfg = dict(base or {})
- if err_message is not None:
- cfg["err_message"] = err_message
- if ws_json_type is not None:
- cfg["ws_json_type"] = ws_json_type
- if ws_error_code is not None:
- cfg["ws_error_code"] = ws_error_code
- return cfg
-
-
-def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None:
- jb = resp.json_body
- want_type = cfg.get("ws_json_type")
- if want_type is not None:
- assert jb.get("type") == want_type, (jb, want_type)
- want_code = cfg.get("ws_error_code")
- if want_code is not None:
- assert jb.get("code") == want_code, (jb, want_code)
- err_message = cfg.get("err_message")
- if err_message is not None:
- assert_http_error(resp, err_message=err_message, websocket_json_message=True)
-
-
-def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse:
- """Concatenate images in order; ``e2e_latency`` is wall-clock of the batch (set by caller) or max of parts."""
- merged = DiffusionResponse()
- merged.success = all(p.success for p in parts) and len(parts) > 0
- imgs: list[Image.Image] = []
- for p in parts:
- if p.images:
- imgs.extend(p.images)
- merged.images = imgs if imgs else None
- latencies = [p.e2e_latency for p in parts if p.e2e_latency is not None]
- merged.e2e_latency = max(latencies) if latencies else None
- return merged
-
-
-class OpenAIClientHandler:
- def __init__(
- self,
- host: str = "127.0.0.1",
- port: int = None,
- api_key: str = "EMPTY",
- run_level: str = None,
- *,
- log_stats: bool = True,
- ):
- if port is None:
- port = get_open_port()
- self.base_url = f"http://{host}:{port}"
- self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
- self.run_level = run_level
- self.log_stats = log_stats
-
- def _print_client_stat(self, message: str) -> None:
- if self.log_stats:
- print(message, flush=True)
-
- def _process_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
- """Wall clock from *before* ``chat.completions.create`` through stream drain + local decode."""
- result = OmniResponse()
- try:
- text_content = ""
- audio_data = []
- for chunk in chat_completion:
- for choice in chunk.choices:
- content = getattr(getattr(choice, "delta", None), "content", None)
- modality = getattr(chunk, "modality", None)
- if modality == "audio" and content:
- audio_data.append(content)
- elif modality == "text" and content:
- text_content += content
- # Usage is yielded after the last token
- if chunk.usage:
- result.prompt_tokens = chunk.usage.prompt_tokens
- if details := getattr(chunk.usage, "prompt_tokens_details", None):
- result.cached_tokens = details.cached_tokens
-
- if audio_data:
- merged_seg = _merge_base64_audio_to_segment(audio_data)
- wav_buf = BytesIO()
- merged_seg.export(wav_buf, format="wav")
- result.audio_bytes = wav_buf.getvalue()
- result.text_content = text_content
- result.audio_data = audio_data
- result.e2e_latency = time.perf_counter() - wall_start
- result.success = True
- except Exception as e:
- msg = f"Stream processing error: {str(e)}"
- print(f"Error: {msg}")
- return result
-
- def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
- """Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
- result = OmniResponse()
- try:
- audio_data = None
- text_content = None
- for choice in chat_completion.choices:
- if hasattr(choice.message, "audio") and choice.message.audio is not None:
- audio_data = choice.message.audio.data
- if hasattr(choice.message, "content") and choice.message.content is not None:
- text_content = choice.message.content
- # Extract cached & prompt token counts for prefix caching tests
- usage = getattr(chat_completion, "usage", None)
- if usage:
- result.prompt_tokens = usage.prompt_tokens
- if details := getattr(usage, "prompt_tokens_details", None):
- result.cached_tokens = details.cached_tokens
- if audio_data:
- result.audio_bytes = base64.b64decode(audio_data)
- result.text_content = text_content
- result.e2e_latency = time.perf_counter() - wall_start
- if chat_completion.choices and chat_completion.choices[0].logprobs is not None:
- result.logprobs = chat_completion.choices[0].logprobs.content
- result.success = True
- except Exception as e:
- msg = f"Non-stream processing error: {str(e)}"
- print(f"Error: {msg}")
- return result
-
- def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse:
- """Wall clock from *before* ``chat.completions.create`` through image decode."""
- result = DiffusionResponse()
- try:
- images = []
- audios = []
- for choice in chat_completion.choices:
- content = getattr(choice.message, "content", None)
- if isinstance(content, list):
- for item in content:
- image_url = None
- if isinstance(item, dict):
- image_url = item.get("image_url", {}).get("url")
- else:
- image_url_obj = getattr(item, "image_url", None)
- image_url = getattr(image_url_obj, "url", None) if image_url_obj else None
- if image_url and image_url.startswith("data:image"):
- b64_data = image_url.split(",", 1)[1]
- images.append(decode_b64_image(b64_data))
-
- # OpenAI audio responses (e.g. AudioX text-to-audio) populate `message.audio`.
- audio_obj = getattr(choice.message, "audio", None)
- audio_b64 = getattr(audio_obj, "data", None) if audio_obj is not None else None
- if audio_b64:
- audios.append(
- {
- "wav_bytes": base64.b64decode(audio_b64),
- "id": getattr(audio_obj, "id", None),
- "expires_at": getattr(audio_obj, "expires_at", None),
- }
- )
- result.images = images if images else None
- result.audios = audios if audios else None
- result.e2e_latency = time.perf_counter() - wall_start
- result.success = True
- except Exception as e:
- msg = f"Diffusion response processing error: {str(e)}"
- print(f"Error: {msg}")
- return result
-
- def _http_response_from_requests(self, r: requests.Response) -> HttpResponse:
- payload = _parse_response_json(r)
- ok = 200 <= r.status_code < 300
- return HttpResponse(
- status_code=r.status_code,
- success=ok,
- error_message=None if ok else (r.text[:8000] if r.text else None),
- json_body=payload,
- )
-
- def send_health_http_request(
- self,
- request_config: dict[str, Any] | None = None,
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/health`` (raw ``requests``).
-
- ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for
- :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args).
- """
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0)))
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_models_http_request(
- self,
- request_config: dict[str, Any] | None = None,
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`)."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = requests.get(
- self._build_url("/v1/models"),
- headers={"Accept": "application/json"},
- timeout=float(cfg.get("timeout", 120.0)),
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_chat_completions_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests)."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_batched_chat_completions_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """Post to batched chat completions."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_completions_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
- # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_omni_sleep_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_omni_wakeup_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/omni/wakeup``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_voices_list_http_request(
- self,
- request_config: dict[str, Any] | None = None,
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/v1/audio/voices``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = requests.get(
- self._build_url("/v1/audio/voices"),
- headers={"Accept": "application/json"},
- timeout=float(cfg.get("timeout", 120.0)),
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_voices_create_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_voices_delete_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- name = cfg["name"]
- timeout = float(cfg.get("timeout", 120.0))
- path = f"/v1/audio/voices/{quote(str(name), safe='')}"
- r = requests.delete(
- self._build_url(path),
- headers={"Accept": "application/json"},
- timeout=timeout,
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_speech_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_speech_batch_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/audio/speech/batch``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_audio_generate_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/audio/generate``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_images_generations_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- if cfg.get("err_code") is None:
- assert resp.success, resp.error_message
- payload = resp.json_body
- assert isinstance(payload, dict)
- assert_images_generations_response(payload, cfg, run_level=self.run_level)
- return [resp]
-
- def send_images_edits_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_videos_create_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_videos_sync_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """POST ``/v1/videos/sync``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0)
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_videos_list_http_request(
- self,
- request_config: dict[str, Any] | None = None,
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/v1/videos`` — optional ``params``, ``timeout``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- r = requests.get(
- self._build_url("/v1/videos"),
- params=cfg.get("params"),
- headers={"Accept": "application/json"},
- timeout=float(cfg.get("timeout", 120.0)),
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_video_retrieve_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/v1/videos/{video_id}``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- video_id = cfg["video_id"]
- timeout = float(cfg.get("timeout", 120.0))
- r = requests.get(
- self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
- headers={"Accept": "application/json"},
- timeout=timeout,
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_video_delete_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """DELETE ``/v1/videos/{video_id}``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- video_id = cfg["video_id"]
- timeout = float(cfg.get("timeout", 120.0))
- r = requests.delete(
- self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
- headers={"Accept": "application/json"},
- timeout=timeout,
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def send_video_content_http_request(
- self,
- request_config: dict[str, Any],
- *,
- err_code: int | tuple[int, ...] | list[int] | None = None,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ) -> list[HttpResponse]:
- """GET ``/v1/videos/{video_id}/content``."""
- cfg = _merge_http_expectation_kwargs(
- request_config,
- err_code=err_code,
- err_message=err_message,
- )
- video_id = cfg["video_id"]
- timeout = float(cfg.get("timeout", 120.0))
- r = requests.get(
- self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"),
- timeout=timeout,
- )
- resp = self._http_response_from_requests(r)
- assert_http_error(
- resp,
- err_code=cfg.get("err_code"),
- err_message=cfg.get("err_message"),
- )
- return [resp]
-
- def _build_ws_url(self, path: str) -> str:
- """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers."""
- base = self.base_url.rstrip("/")
- suffix = "/" + path.lstrip("/")
- if base.startswith("http://"):
- return "ws://" + base.removeprefix("http://") + suffix
- if base.startswith("https://"):
- return "wss://" + base.removeprefix("https://") + suffix
- raise ValueError(f"Unsupported base_url for WebSocket: {base!r}")
-
- def _send_websocket_first_json_request(
- self,
- path: str,
- cfg: dict[str, Any],
- ) -> list[WebSocketJsonResponse]:
- """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`.
-
- ``request_config`` keys:
-
- - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
- speaks first, e.g. ``/v1/realtime`` rejection path).
- - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame
- (e.g. ``["session.created"]`` on ``/v1/realtime``).
- - ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
- - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
- """
- send_frames_raw = cfg.get("send_frames")
- if send_frames_raw is None:
- frames: list[str] = []
- elif isinstance(send_frames_raw, str):
- frames = [send_frames_raw]
- else:
- frames = list(send_frames_raw)
-
- timeout = float(cfg.get("timeout", 120.0))
- uri = self._build_ws_url(path)
- skip_types = set(cfg.get("ws_skip_types") or [])
-
- connect_kw: dict[str, Any] = {}
- if "ws_max_size" in cfg:
- connect_kw["max_size"] = cfg["ws_max_size"]
-
- async def _recv_first_json_object() -> WebSocketJsonResponse:
- import websockets
-
- async with websockets.connect(uri, **connect_kw) as ws:
- for frame in frames:
- await ws.send(frame)
- while True:
- raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
- if not isinstance(raw, str):
- raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
+ # Close per-stage log handles; delete temp files unless VLLM_OMNI_KEEP_LOG is set.
+ keep_logs = os.environ.get("VLLM_OMNI_KEEP_LOG", "").lower() in ("1", "true", "yes")
+ log_files = self._stage_log_files
+ log_paths = self._stage_log_paths
+ for stage_key in set(log_files) | set(log_paths):
+ log_fh = log_files.get(stage_key)
+ if log_fh is not None:
+ try:
+ log_fh.close()
+ except Exception:
+ pass
+ if not keep_logs:
+ log_path = log_paths.get(stage_key)
+ if log_path is not None:
try:
- data = json.loads(raw)
- except json.JSONDecodeError as exc:
- raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
- if not isinstance(data, dict):
- raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
- if skip_types and data.get("type") in skip_types:
- continue
- return WebSocketJsonResponse(json_body=data)
-
- resp = asyncio.run(_recv_first_json_object())
- _run_ws_expectations_from_request_config(cfg, resp)
- return [resp]
-
- def send_audio_speech_stream_ws_request(
- self,
- request_config: dict[str, Any],
- *,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ws_json_type: str | None = None,
- ws_error_code: str | None = None,
- ) -> list[WebSocketJsonResponse]:
- """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame."""
- cfg = _merge_ws_expectation_kwargs(
- request_config,
- err_message=err_message,
- ws_json_type=ws_json_type,
- ws_error_code=ws_error_code,
- )
- return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg)
-
- def send_video_chat_stream_ws_request(
- self,
- request_config: dict[str, Any],
- *,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ws_json_type: str | None = None,
- ws_error_code: str | None = None,
- ) -> list[WebSocketJsonResponse]:
- """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame."""
- cfg = _merge_ws_expectation_kwargs(
- request_config,
- err_message=err_message,
- ws_json_type=ws_json_type,
- ws_error_code=ws_error_code,
- )
- return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg)
-
- def send_realtime_ws_request(
- self,
- request_config: dict[str, Any] | None = None,
- *,
- err_message: str | tuple[str, ...] | list[str] | None = None,
- ws_json_type: str | None = None,
- ws_error_code: str | None = None,
- ) -> list[WebSocketJsonResponse]:
- """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated)."""
- cfg = _merge_ws_expectation_kwargs(
- request_config,
- err_message=err_message,
- ws_json_type=ws_json_type,
- ws_error_code=ws_error_code,
- )
- return self._send_websocket_first_json_request("/v1/realtime", cfg)
-
- def send_robot_openpi_ws_request(
- self,
- request_config: dict[str, Any] | None = None,
- ) -> list[OpenPIWebSocketResponse]:
- """WebSocket ``/v1/realtime/robot/openpi`` — msgpack metadata plus optional infer/reset ops.
-
- ``request_config`` keys:
-
- - ``operations``: optional sequence of ``{"endpoint": "infer"|"reset", "payload": {...}}``.
- - ``run_default_policy_session``: when true and ``operations`` is omitted, run infer then reset
- using :func:`build_openpi_droid_observation`.
- - ``run_dreamzero_policy_session``: when true and ``operations`` is omitted, run the DreamZero
- infer/reset/infer sequence using :func:`build_dreamzero_demo_observations`.
- - ``video_dir``: required for ``run_dreamzero_policy_session``.
- - ``prompt``: language instruction for DreamZero observations (optional).
- - ``session_id``: used by the default policy session observation builder.
- - ``num_chunks``: DreamZero chunk count (default ``2``).
- - ``timeout``: seconds to wait for each inbound msgpack frame (default ``120``).
- - ``ping_interval`` / ``ping_timeout``: forwarded to :func:`websockets.sync.client.connect`.
- """
- cfg = dict(request_config or {})
- operations_cfg = cfg.get("operations")
- if operations_cfg is not None:
- operations = [dict(op) for op in operations_cfg]
- elif cfg.get("run_dreamzero_policy_session"):
- video_dir = cfg.get("video_dir")
- if video_dir is None:
- raise ValueError("run_dreamzero_policy_session requires video_dir")
- session_id = str(cfg.get("session_id", "dreamzero-smoke"))
- prompt = str(cfg.get("prompt", DREAMZERO_DEFAULT_PROMPT))
- num_chunks = int(cfg.get("num_chunks", 2))
- observations = build_dreamzero_demo_observations(
- load_dreamzero_camera_frames(Path(video_dir)),
- prompt=prompt,
- session_id=session_id,
- num_chunks=num_chunks,
- )
- operations = [{"endpoint": "infer", "payload": obs} for obs in observations]
- operations.append({"endpoint": "reset", "payload": {}})
- operations.append({"endpoint": "infer", "payload": observations[0]})
- elif cfg.get("run_default_policy_session"):
- session_id = str(cfg.get("session_id", "gr00t-smoke"))
- operations = [
- {"endpoint": "infer", "payload": build_openpi_droid_observation(session_id=session_id)},
- {"endpoint": "reset", "payload": {}},
- ]
- else:
- operations = []
- timeout = float(cfg.get("timeout", 120.0))
- ping_interval = float(cfg.get("ping_interval", OpenPIWebSocketSession.DEFAULT_PING_INTERVAL_SECS))
- ping_timeout = float(cfg.get("ping_timeout", OpenPIWebSocketSession.DEFAULT_PING_TIMEOUT_SECS))
-
- uri = self._build_ws_url("/v1/realtime/robot/openpi")
- session = OpenPIWebSocketSession(
- uri,
- ping_interval=ping_interval,
- ping_timeout=ping_timeout,
- open_timeout=timeout,
- close_timeout=timeout,
- )
- try:
- operation_responses: list[Any] = []
- actions: dict[str, np.ndarray] | None = None
- action_tensors: list[np.ndarray] = []
- for operation in operations:
- endpoint = str(operation["endpoint"])
- payload = dict(operation.get("payload") or {})
- response = session._send_operation(endpoint, payload)
- operation_responses.append(response)
- if endpoint == "infer":
- if isinstance(response, dict):
- actions = {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()}
- else:
- action_tensors.append(np.asarray(response, dtype=np.float32))
- finally:
- session.close()
-
- return [
- OpenPIWebSocketResponse(
- server_metadata=session.get_server_metadata(),
- operation_responses=operation_responses,
- actions=actions,
- action_tensors=action_tensors or None,
- )
- ]
-
- def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
- """Chat completions via the OpenAI Python SDK (not raw HTTP)."""
- responses: list[OmniResponse] = []
- stream = request_config.get("stream", False)
- modalities = request_config.get("modalities", ["text", "audio"])
- extra_body: dict[str, Any] = {}
- if "speaker" in request_config:
- extra_body["speaker"] = request_config["speaker"]
- if request_config.get("use_audio_in_video"):
- mm = dict(extra_body.get("mm_processor_kwargs") or {})
- mm["use_audio_in_video"] = True
- extra_body["mm_processor_kwargs"] = mm
- if "sampling_params_list" in request_config:
- extra_body["sampling_params_list"] = request_config["sampling_params_list"]
- if request_config.get("extra_body"):
- extra_body.update(request_config["extra_body"])
-
- create_kwargs: dict[str, Any] = {
- "model": request_config.get("model"),
- "messages": request_config.get("messages"),
- "stream": stream,
- "modalities": modalities,
- }
- if "logprobs" in request_config:
- create_kwargs["logprobs"] = request_config["logprobs"]
- if "top_logprobs" in request_config:
- create_kwargs["top_logprobs"] = request_config["top_logprobs"]
- if "stream_options" in request_config:
- create_kwargs["stream_options"] = request_config["stream_options"]
- if extra_body:
- create_kwargs["extra_body"] = extra_body
-
- if request_num == 1:
- wall_start = time.perf_counter()
- chat_completion = self.client.chat.completions.create(**create_kwargs)
- resp = (
- self._process_stream_omni_response(chat_completion, wall_start=wall_start)
- if stream
- else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
- )
- assert_omni_response(resp, request_config, run_level=self.run_level)
- if resp.e2e_latency is not None:
- self._print_client_stat(f"[omni] request#1 success in {resp.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[omni] request#1 completed")
- responses.append(resp)
- return responses
-
- def _one():
- wall_start = time.perf_counter()
- chat_completion = self.client.chat.completions.create(**create_kwargs)
- return (
- self._process_stream_omni_response(chat_completion, wall_start=wall_start)
- if stream
- else self._process_non_stream_omni_response(chat_completion, wall_start=wall_start)
- )
-
- with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
- futures = {executor.submit(_one): i + 1 for i in range(request_num)}
- for future in concurrent.futures.as_completed(futures):
- request_idx = futures[future]
- resp = future.result()
- assert_omni_response(resp, request_config, run_level=self.run_level)
- if resp.e2e_latency is not None:
- self._print_client_stat(f"[omni] request#{request_idx} success in {resp.e2e_latency:.3f}s")
- else:
- self._print_client_stat(f"[omni] request#{request_idx} completed")
- responses.append(resp)
- return responses
-
- def _process_stream_audio_speech_response(
- self, response, *, response_format: str | None = None, wall_start: float
- ) -> OmniResponse:
- """
- Process streaming /v1/audio/speech responses into an OmniResponse.
-
- This mirrors _process_stream_omni_response but operates on low-level
- audio bytes. Whisper transcription runs in assert_audio_speech_response
- when the run_level requires it.
- """
- result = OmniResponse()
-
- try:
- # Aggregate all audio bytes from the streaming response.
- data = bytearray()
-
- # Preferred OpenAI helper.
- if hasattr(response, "iter_bytes") and callable(getattr(response, "iter_bytes")):
- for chunk in response.iter_bytes():
- if chunk:
- data.extend(chunk)
- else:
- # Generic iterable-of-bytes fallback (e.g., generator or list of chunks).
- try:
- iterator = iter(response)
- except TypeError:
- iterator = None
-
- if iterator is not None:
- for chunk in iterator:
- if not chunk:
- continue
- if isinstance(chunk, (bytes, bytearray)):
- data.extend(chunk)
- elif hasattr(chunk, "data"):
- data.extend(chunk.data) # type: ignore[arg-type]
- elif hasattr(chunk, "content"):
- data.extend(chunk.content) # type: ignore[arg-type]
- else:
- raise TypeError(f"Unsupported stream chunk type: {type(chunk)}")
- else:
- raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}")
-
- raw_bytes = bytes(data)
-
- # Populate OmniResponse.
- result.audio_bytes = raw_bytes
- result.e2e_latency = time.perf_counter() - wall_start
- result.success = True
- result.audio_format = getattr(response, "response", None)
- if result.audio_format is not None:
- result.audio_format = result.audio_format.headers.get("content-type", "")
-
- except Exception as e:
- msg = f"Audio speech stream processing error: {str(e)}"
- print(f"Error: {msg}")
-
- return result
-
- def _process_non_stream_audio_speech_response(
- self, response, *, response_format: str | None = None, wall_start: float
- ) -> OmniResponse:
- """
- Process non-streaming /v1/audio/speech responses into an OmniResponse.
-
- This mirrors _process_non_stream_omni_response but for the binary
- audio payload returned by audio.speech.create.
- """
- result = OmniResponse()
-
- try:
- # OpenAI non-streaming audio.speech.create returns HttpxBinaryResponseContent (.read() or .content)
- if hasattr(response, "read") and callable(getattr(response, "read")):
- raw_bytes = response.read()
- elif hasattr(response, "content"):
- raw_bytes = response.content # type: ignore[assignment]
- else:
- raise TypeError(f"Unsupported audio speech response type: {type(response)}")
-
- result.audio_bytes = raw_bytes
- result.e2e_latency = time.perf_counter() - wall_start
- result.success = True
- result.audio_format = getattr(response, "response", None)
- if result.audio_format is not None:
- result.audio_format = result.audio_format.headers.get("content-type", "")
-
- except Exception as e:
- msg = f"Audio speech non-stream processing error: {str(e)}"
- print(f"Error: {msg}")
-
- return result
-
- def send_audio_speech_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
- """
- Call the /v1/audio/speech endpoint using the same configuration-dict
- style as send_omni_request, but via the OpenAI Python client's
- audio.speech APIs.
-
- Expected keys in request_config:
- - model: model name/path (required)
- - input: text to synthesize (required)
- - response_format: audio format such as "wav" or "pcm" (optional)
- - task_type, ref_text, ref_audio: TTS-specific extras (optional, passed via extra_body)
- - min_audio_bytes: optional minimum ``len(audio_bytes)`` checked in ``assert_audio_speech_response``
- - status_code: if set, HTTP status is asserted (int or e.g. ``(400, 422)``); uses APIError handling
- - err_message: optional substring(s) to match against error text (``str`` or list/tuple of alternatives;
- see ``assert_audio_speech_response``). If set, uses the same APIError path as ``status_code``.
- - When both ``status_code`` and ``err_message`` are absent (or each is ``None``), the normal request path
- is used (no try/except around ``APIError``).
- - timeout: request timeout in seconds (float, optional, default 120.0)
- - stream: whether to use streaming API (bool, optional, default False)
- """
- timeout = float(request_config.get("timeout", 120.0))
-
- model = request_config["model"]
- text_input = request_config["input"]
- stream = bool(request_config.get("stream", False))
- voice = request_config.get("voice", None)
-
- # Standard OpenAI param: use omit when not provided to keep default behavior.
- response_format = request_config.get("response_format", omit)
-
- # Qwen3-TTS custom fields, forwarded via extra_body.
- extra_body: dict[str, Any] = {}
- # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
- for key in (
- "task_type",
- "ref_text",
- "ref_audio",
- "language",
- "max_new_tokens",
- "seed",
- "instructions",
- "speed",
- "stream_format",
- "x_vector_only_mode",
- ):
- if key in request_config:
- extra_body[key] = request_config[key]
-
- responses: list[OmniResponse] = []
-
- speech_fmt: str | None = None if response_format is omit else str(response_format).lower()
-
- print(f"[audio.speech] start model={model}, stream={stream}, request_num={request_num}, timeout={timeout:.1f}s")
-
- # Error validation path: only when at least one of these is set to a non-``None`` value.
- expect_error_handling = (request_config.get("status_code") is not None) or (
- request_config.get("err_message") is not None
- )
- if expect_error_handling and request_num != 1:
- raise ValueError(
- "request_config error validation (status_code / err_message) is only supported when request_num=1"
- )
-
- if request_num == 1:
- if expect_error_handling:
- # ``status`` and/or ``err_message`` requested: catch APIError (4xx) and (optionally) assert body text;
- # HTTP 200 with JSON error body is handled in ``assert_audio_speech_response``.
- try:
- if stream:
- wall_start = time.perf_counter()
- with self.client.audio.speech.with_streaming_response.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- ) as resp:
- omni_resp = self._process_stream_audio_speech_response(
- resp, response_format=speech_fmt, wall_start=wall_start
- )
- else:
- wall_start = time.perf_counter()
- resp = self.client.audio.speech.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- )
- omni_resp = self._process_non_stream_audio_speech_response(
- resp, response_format=speech_fmt, wall_start=wall_start
- )
- except APIError as e:
- sc = getattr(e, "status_code", None)
- if sc is None:
- raise
- omni_resp = OmniResponse(
- success=False,
- status_code=sc,
- error_message=str(e),
- )
- else:
- if getattr(omni_resp, "status_code", None) is None:
- omni_resp.status_code = 200
- elif stream:
- # Use streaming response helper.
- wall_start = time.perf_counter()
- with self.client.audio.speech.with_streaming_response.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- ) as resp:
- omni_resp = self._process_stream_audio_speech_response(
- resp, response_format=speech_fmt, wall_start=wall_start
- )
- else:
- # Non-streaming response.
- wall_start = time.perf_counter()
- resp = self.client.audio.speech.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- )
- omni_resp = self._process_non_stream_audio_speech_response(
- resp, response_format=speech_fmt, wall_start=wall_start
- )
-
- assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
- if omni_resp.e2e_latency is not None:
- self._print_client_stat(f"[audio.speech] request#1 success in {omni_resp.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[audio.speech] request#1 completed")
- responses.append(omni_resp)
- return responses
- else:
- # request_num > 1: concurrent requests (use same params as single-request path)
-
- if stream:
-
- def _stream_task(request_idx: int):
- wall_start = time.perf_counter()
- with self.client.audio.speech.with_streaming_response.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- ) as resp:
- result = self._process_stream_audio_speech_response(
- resp, response_format=speech_fmt, wall_start=wall_start
- )
- if result.e2e_latency is not None:
- self._print_client_stat(
- f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
- )
- else:
- self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
- return result
-
- with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
- futures = {executor.submit(_stream_task, i + 1): i + 1 for i in range(request_num)}
- for future in concurrent.futures.as_completed(futures):
- request_idx = futures[future]
- try:
- omni_resp = future.result()
- except Exception as e:
- print(
- f"[audio.speech] request#{request_idx} failed "
- f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
- )
- raise
- assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
- responses.append(omni_resp)
- else:
-
- def _non_stream_task(request_idx: int):
- wall_start = time.perf_counter()
- r = self.client.audio.speech.create(
- model=model,
- input=text_input,
- response_format=response_format,
- extra_body=extra_body or None,
- timeout=timeout,
- voice=voice,
- )
- result = self._process_non_stream_audio_speech_response(
- r, response_format=speech_fmt, wall_start=wall_start
- )
- if result.e2e_latency is not None:
- self._print_client_stat(
- f"[audio.speech] request#{request_idx} success in {result.e2e_latency:.3f}s"
- )
- else:
- self._print_client_stat(f"[audio.speech] request#{request_idx} completed")
- return result
-
- with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
- futures = {executor.submit(_non_stream_task, i + 1): i + 1 for i in range(request_num)}
- for future in concurrent.futures.as_completed(futures):
- request_idx = futures[future]
- try:
- omni_resp = future.result()
- except Exception as e:
- print(
- f"[audio.speech] request#{request_idx} failed "
- f"(stream={stream}, timeout={timeout:.1f}s): {e!r}"
- )
- raise
- assert_audio_speech_response(omni_resp, request_config, run_level=self.run_level)
- responses.append(omni_resp)
-
- return responses
-
- def send_diffusion_request(
- self, request_config: dict[str, Any] | list[dict[str, Any]], request_num: int = 1
- ) -> list[DiffusionResponse]:
- """
- Send OpenAI requests for diffusion models.
- If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel
- (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order.
-
- Args:
- request_config: A single request configuration dict, or a list of
- request configuration dicts (one request per element)
- request_num: Number of requests to send concurrently, defaults to 1 (single request)
- Returns:
- list[DiffusionResponse]: List of DiffusionResponse objects containing the response data
- """
- responses: list[DiffusionResponse] = []
-
- def _create_from_config(cfg: dict[str, Any]) -> tuple[Any, float]:
- stream = cfg.get("stream", False)
- if stream:
- raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test")
- modalities = cfg.get("modalities", omit) # Most diffusion models don't require modalities param
- eb = cfg.get("extra_body")
- extra = copy.deepcopy(eb) if eb else None
- wall_start = time.perf_counter()
- chat_completion = self.client.chat.completions.create(
- model=cfg.get("model"),
- messages=cfg.get("messages"),
- extra_body=extra,
- modalities=modalities,
- )
- return chat_completion, wall_start
-
- if isinstance(request_config, list):
- if not request_config:
- raise ValueError("request_config list must not be empty")
- if request_num != 1:
- raise ValueError("request_num is not supported when request_config is a list")
- with concurrent.futures.ThreadPoolExecutor(max_workers=len(request_config)) as executor:
- futures = {
- executor.submit(_create_from_config, cfg): (i + 1, cfg) for i, cfg in enumerate(request_config)
- }
- for future in concurrent.futures.as_completed(futures):
- request_idx, cfg = futures[future]
- chat_completion, wall_start = future.result()
- response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
- assert_diffusion_response(response, cfg, run_level=self.run_level)
- if response.e2e_latency is not None:
- self._print_client_stat(
- f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s"
- )
- else:
- self._print_client_stat(f"[diffusion] request#{request_idx} completed")
- responses.append(response)
- return responses
-
- size_splits = _split_request_config_by_per_output_sizes(request_config)
- if size_splits is not None:
- if request_num != 1:
- raise ValueError(
- "request_num must be 1 when extra_body height/width are lists (split into concurrent per-size calls)"
- )
- t0 = time.perf_counter()
- with concurrent.futures.ThreadPoolExecutor(max_workers=len(size_splits)) as executor:
- futures = [executor.submit(_create_from_config, sub) for sub in size_splits]
- chat_completions = [f.result() for f in futures]
- parts = [self._process_diffusion_response(cc, wall_start=ws) for cc, ws in chat_completions]
- merged = _merge_diffusion_responses(parts)
- merged.e2e_latency = time.perf_counter() - t0
- assert_diffusion_response(merged, request_config, run_level=self.run_level)
- if merged.e2e_latency is not None:
- self._print_client_stat(f"[diffusion] request#1 success in {merged.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[diffusion] request#1 completed")
- return [merged]
-
- if request_num == 1:
- # Send single request
- chat_completion, wall_start = _create_from_config(request_config)
- response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
- assert_diffusion_response(response, request_config, run_level=self.run_level)
- if response.e2e_latency is not None:
- self._print_client_stat(f"[diffusion] request#1 success in {response.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[diffusion] request#1 completed")
- responses.append(response)
- return responses
-
- # Send concurrent requests for the same request_config
- with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
- futures = {executor.submit(_create_from_config, request_config): i + 1 for i in range(request_num)}
- for future in concurrent.futures.as_completed(futures):
- request_idx = futures[future]
- chat_completion, wall_start = future.result()
- response = self._process_diffusion_response(chat_completion, wall_start=wall_start)
- assert_diffusion_response(response, request_config, run_level=self.run_level)
- if response.e2e_latency is not None:
- self._print_client_stat(f"[diffusion] request#{request_idx} success in {response.e2e_latency:.3f}s")
- else:
- self._print_client_stat(f"[diffusion] request#{request_idx} completed")
- responses.append(response)
- return responses
-
- def send_video_diffusion_request(
- self, request_config: dict[str, Any], request_num: int = 1
- ) -> list[DiffusionResponse]:
- """
- Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content.
-
- For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc.
- """
- if request_num != 1:
- raise NotImplementedError("Concurrent video diffusion requests are not currently implemented")
-
- form_data = request_config.get("form_data")
- if not isinstance(form_data, dict):
- raise ValueError("Video request_config must contain 'form_data'")
- normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
- files: dict[str, tuple[str, BytesIO, str]] = {}
- image_reference = request_config.get("image_reference")
- video_reference = request_config.get("video_reference")
- if image_reference and video_reference:
- raise ValueError("Only one of image_reference or video_reference can be provided")
- if image_reference:
- if image_reference.startswith("data:image"):
- header, encoded = image_reference.split(",", 1)
- content_type = header.split(";")[0].removeprefix("data:")
- extension = content_type.split("/")[-1]
- file_data = base64.b64decode(encoded)
- files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
- else:
- normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
- if video_reference:
- if video_reference.startswith("data:video"):
- header, encoded = video_reference.split(",", 1)
- content_type = header.split(";")[0].removeprefix("data:")
- extension = content_type.split("/")[-1]
- file_data = base64.b64decode(encoded)
- files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type)
- else:
- normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference})
-
- result = DiffusionResponse()
- create_url = self._build_url("/v1/videos")
- response = requests.post(
- create_url,
- data=normalized_form_data,
- files=files,
- headers={"Accept": "application/json"},
- timeout=60,
- )
- start_time = time.perf_counter()
- response.raise_for_status()
- job_data = response.json()
- video_id = job_data["id"]
- self._wait_until_video_completed(video_id)
- end_time = time.perf_counter()
- video_content = self._download_video_content(video_id)
- result.success = True
- result.videos = [video_content]
- result.e2e_latency = end_time - start_time
- assert_diffusion_response(result, request_config, run_level=self.run_level)
- if result.e2e_latency is not None:
- self._print_client_stat(f"[diffusion] request#1 success in {result.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[diffusion] request#1 completed")
- return [result]
-
- def _post_json_endpoint(
- self,
- path: str,
- request_config: dict[str, Any],
- *,
- default_timeout: float,
- ) -> requests.Response:
- url = self._build_url(path)
- timeout = float(request_config.get("timeout", default_timeout))
- if "raw_body" in request_config:
- raw = request_config["raw_body"]
- payload = raw.encode("utf-8") if isinstance(raw, str) else raw
- return requests.post(
- url,
- data=payload,
- headers={"Content-Type": "application/json", "Accept": "application/json"},
- timeout=timeout,
- )
- if "json" not in request_config:
- raise ValueError(f"{path} request_config must include 'json' or 'raw_body'")
- return requests.post(
- url,
- json=request_config["json"],
- headers={"Accept": "application/json", "Content-Type": "application/json"},
- timeout=timeout,
- )
-
- def _post_form_endpoint(
- self,
- path: str,
- request_config: dict[str, Any],
- *,
- default_timeout: float = 120.0,
- ) -> requests.Response:
- url = self._build_url(path)
- timeout = float(request_config.get("timeout", default_timeout))
- data = request_config.get("data")
- files = request_config.get("files")
- if data is None and not files:
- data = {}
- return requests.post(
- url,
- data=data,
- files=files,
- headers={"Accept": "application/json"} if not files else {"Accept": "application/json"},
- timeout=timeout,
- )
-
- def send_streaming_video_diffusion_request(
- self,
- request_config: dict[str, Any],
- request_num: int = 1,
- *,
- timeout_seconds: float = 600.0,
- ) -> list[DiffusionResponse]:
- """
- Send a native ``/v1/realtime/video`` WebSocket request and return one
- finalized MP4 artifact assembled from the streamed binary fragments.
- """
- if request_num != 1:
- raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented")
-
- response = asyncio.run(
- self._send_streaming_video_diffusion_request_once(
- request_config,
- timeout_seconds=timeout_seconds,
- )
- )
- assert_diffusion_response(response, request_config, run_level=self.run_level)
- if response.e2e_latency is not None:
- self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s")
- else:
- self._print_client_stat("[diffusion.stream] request#1 completed")
- return [response]
-
- async def _send_streaming_video_diffusion_request_once(
- self,
- request_config: dict[str, Any],
- *,
- timeout_seconds: float,
- ) -> DiffusionResponse:
- form_data = request_config.get("form_data")
- if not isinstance(form_data, dict):
- raise ValueError("Video request_config must contain 'form_data'")
- payload: dict[str, Any] = {
- "type": "session.start",
- **{key: value for key, value in form_data.items() if value is not None},
- }
- model = request_config.get("model")
- if model is not None:
- payload["model"] = model
- payload.setdefault("format", "m4s")
-
- fps = float(payload.get("fps") or 16)
- stream_format = payload["format"]
- url = self._build_ws_url("/v1/realtime/video")
-
- result = DiffusionResponse()
- chunks: list[bytes] = []
- start_time = time.perf_counter()
- deadline = start_time + timeout_seconds
-
- import websockets
-
- async with websockets.connect(url, max_size=None) as websocket:
- await websocket.send(json.dumps(payload))
-
- while True:
- remaining = deadline - time.perf_counter()
- if remaining <= 0:
- raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s")
-
- message = await asyncio.wait_for(websocket.recv(), timeout=remaining)
- if isinstance(message, bytes):
- chunks.append(message)
- continue
-
- msg = json.loads(message)
- msg_type = msg.get("type")
- if msg_type == "video.start":
- stream_format = msg.get("format") or stream_format
- continue
- if msg_type == "session.done":
- break
- if msg_type == "error":
- raise RuntimeError(str(msg.get("message", msg)))
-
- from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes
- from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat
-
- streamed_bytes = b"".join(chunks)
- if not streamed_bytes:
- raise RuntimeError("Streaming video request completed without binary video chunks")
- result.videos = [
- finalize_streaming_video_bytes(
- streamed_bytes,
- input_format=cast(StreamingVideoFormat, stream_format),
- fps=fps,
- )
- ]
- result.e2e_latency = time.perf_counter() - start_time
- result.success = True
- return result
-
- def _wait_until_video_completed(
- self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
- ) -> None:
- status_url = self._build_url(f"/v1/videos/{video_id}")
- deadline = time.monotonic() + timeout_seconds
- while time.monotonic() < deadline:
- status_resp = requests.get(status_url, headers={"Accept": "application/json"}, timeout=30)
- status_resp.raise_for_status()
- status_data = status_resp.json()
- current_status = status_data["status"]
- if current_status == "completed":
- return
- if current_status == "failed":
- error_msg = status_data.get("last_error", "Unknown error")
- raise RuntimeError(f"Job failed: {error_msg}")
- time.sleep(poll_interval_seconds)
- raise TimeoutError(f"Video job {video_id} did not complete within {timeout_seconds}s")
-
- def _download_video_content(self, video_id: str) -> bytes:
- download_url = self._build_url(f"/v1/videos/{video_id}/content")
- video_resp = requests.get(download_url, stream=True, timeout=60)
- video_resp.raise_for_status()
- video_bytes = BytesIO()
- for chunk in video_resp.iter_content(chunk_size=8192):
- if chunk:
- video_bytes.write(chunk)
- return video_bytes.getvalue()
-
- def _build_url(self, path: str) -> str:
- return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
+ Path(log_path).unlink(missing_ok=True)
+ except Exception:
+ pass
+ self._stage_log_files = {}
+ if not keep_logs:
+ self._stage_log_paths = {}
+
+ cleanup_test_environment()
class OmniRunner:
@@ -2625,7 +592,6 @@ class OmniRunner:
model_name: str,
seed: int = 42,
stage_init_timeout: int = 600,
- batch_timeout: int = 10,
# Bumped from 900s -> 1800s to give CI cold-cache loads of large
# diffusion models enough headroom (Buildkite #8418 hit a 6-second
# overrun loading Tongyi-MAI/Z-Image-Turbo: weights alone took 690s,
@@ -2635,30 +601,34 @@ class OmniRunner:
# affects the test runner wrapper.
init_timeout: int = 1800,
log_stats: bool = False,
- stage_configs_path: str | None = None,
+ deploy_config: str | None = None,
**kwargs,
) -> None:
startup_t0 = time.perf_counter()
- cleanup_dist_env_and_memory()
- run_pre_test_cleanup()
- run_post_test_cleanup()
+ cleanup_test_environment()
self.model_name = model_name
self.seed = seed
self._prompt_len_estimate_cache: dict[str, Any] = {}
- from vllm_omni.entrypoints.omni import Omni
+ self.omni: Any = None
+ try:
+ from vllm_omni.entrypoints.omni import Omni
- self.omni = Omni(
- model=model_name,
- log_stats=log_stats,
- stage_init_timeout=stage_init_timeout,
- batch_timeout=batch_timeout,
- init_timeout=init_timeout,
- stage_configs_path=stage_configs_path,
- **kwargs,
- )
- startup_s = time.perf_counter() - startup_t0
- if log_stats:
- print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
+ self.omni = Omni(
+ model=model_name,
+ log_stats=log_stats,
+ stage_init_timeout=stage_init_timeout,
+ init_timeout=init_timeout,
+ deploy_config=deploy_config,
+ **kwargs,
+ )
+ 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.__exit__(None, None, None)
+ raise
def get_default_sampling_params_list(self) -> list[Any]:
if not hasattr(self.omni, "default_sampling_params_list"):
@@ -2833,14 +803,14 @@ class OmniRunner:
f"<|im_start|>user\n{user_content}<|im_end|>\n"
f"<|im_start|>assistant\n"
)
- input_dict: dict[str, Any] = {"prompt": full_prompt}
+ prompt_dict: dict[str, Any] = {"prompt": full_prompt}
if multi_modal_data:
- input_dict["multi_modal_data"] = multi_modal_data
+ prompt_dict["multi_modal_data"] = multi_modal_data
if modalities:
- input_dict["modalities"] = modalities
+ prompt_dict["modalities"] = modalities
if mm_processor_kwargs:
- input_dict["mm_processor_kwargs"] = mm_processor_kwargs
- omni_inputs.append(input_dict)
+ prompt_dict["mm_processor_kwargs"] = mm_processor_kwargs
+ omni_inputs.append(prompt_dict)
return omni_inputs
def generate(
@@ -2919,197 +889,11 @@ class OmniRunner:
return self
def __exit__(self, exc_type, exc_val, exc_tb):
- if hasattr(self.omni, "close"):
- self.omni.close()
+ omni = getattr(self, "omni", None)
+ if omni is not None and hasattr(omni, "close"):
+ omni.close()
self._cleanup_process()
- run_pre_test_cleanup()
- run_post_test_cleanup()
- cleanup_dist_env_and_memory()
-
-
-class OmniRunnerHandler:
- def __init__(self, omni_runner: OmniRunner):
- self.runner = omni_runner
-
- def _process_omni_output(self, outputs: list[OmniRequestOutput]) -> OmniResponse:
- result = OmniResponse()
- try:
- text_content = None
- audio_content = None
- for stage_output in outputs:
- if getattr(stage_output, "final_output_type", None) == "text":
- text_content = stage_output.outputs[0].text
- if getattr(stage_output, "final_output_type", None) == "audio":
- audio_content = stage_output.outputs[0].multimodal_output["audio"]
- result.audio_content = audio_content
- result.text_content = text_content
- result.success = True
- except Exception as e:
- msg = f"Output processing error: {str(e)}"
- result.success = False
- print(f"Error: {msg}")
- return result
-
- def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse:
- result = DiffusionResponse()
- output = outputs[0]
- if isinstance(output.images[0], list):
- # Returning frames of images as a video
- result.videos = output.images
- else:
- # Returning actual images
- result.images = output.images
- # [TODO] Add audio processing when tests are introduced
- result.success = True
- return result
-
- def send_omni_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
- if request_config is None:
- request_config = {}
- prompts = request_config.get("prompts")
- videos = request_config.get("videos")
- images = request_config.get("images")
- audios = request_config.get("audios")
- modalities = request_config.get("modalities", ["text", "audio"])
- outputs = self.runner.generate_multimodal(
- prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
- )
- response = self._process_omni_output(outputs)
- assert_omni_response(response, request_config, run_level="core_model")
- return response
-
- def send_diffusion_request(self, request_config: dict[str, Any]) -> DiffusionResponse:
- prompt = request_config.get("prompt")
- if prompt is None:
- prompts = request_config.get("prompts")
- if not prompts:
- raise ValueError("request_config must contain a prompt")
- if len(prompts) > 1:
- raise ValueError(
- "In the current internal data structure, "
- "only one prompt is supported for diffusion requests. "
- "Because one prompt can contain multiple images or videos, "
- "the current internal data structure is ambiguous when multiple prompts are provided."
- )
- prompt = prompts[0]
-
- # Sync previous `extra_body` field with sampling params object
- sampling_params: OmniDiffusionSamplingParams | None = request_config.get("sampling_params")
- if sampling_params is None:
- extra_body = request_config.get("extra_body", {})
- if extra_body:
- sampling_params = OmniDiffusionSamplingParams(**extra_body)
- else:
- extra_body = asdict(sampling_params)
- request_config["extra_body"] = extra_body
- if not extra_body:
- logger.warning("No sampling params provided in request_config, will skip output assertion")
-
- negative_prompt = extra_body.get("negative_prompt") or request_config.get("negative_prompt")
- videos = request_config.get("videos")
- images = request_config.get("images")
- audios = request_config.get("audios")
- # Full dict (e.g. image + mask_image for inpainting) or partial; merged with top-level image/video keys.
- extra_multi_modal = request_config.get("multi_modal_data")
- modalities = request_config.get("modalities") # only used by limited models. Do not add default value here
-
- prompt_object = OmniTextPrompt(prompt=prompt)
- if negative_prompt:
- prompt_object["negative_prompt"] = negative_prompt
- multi_modal: dict = {}
- if extra_multi_modal is not None:
- multi_modal.update(dict(extra_multi_modal))
- if videos is not None:
- multi_modal["video"] = videos
- if images is not None:
- multi_modal["image"] = images
- if audios is not None:
- multi_modal["audio"] = audios
- if multi_modal:
- prompt_object["multi_modal_data"] = multi_modal
- if modalities:
- prompt_object["modalities"] = modalities # pyright: ignore[reportGeneralTypeIssues]
-
- start_time = time.perf_counter()
- response = self.runner.generate([prompt_object], [sampling_params] if sampling_params else None)
- end_time = time.perf_counter()
-
- response = self._process_diffusion_output(response)
- response.e2e_latency = end_time - start_time
- assert_diffusion_response(response, request_config, run_level="core_model")
- return response
-
- def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniResponse:
- """
- Offline TTS: text -> audio via generate_multimodal, then validate with assert_audio_speech_response.
-
- request_config must contain:
- - 'input' or 'prompts': text to synthesize.
- Optional keys:
- - 'voice' -> speaker (CustomVoice)
- - 'task_type' -> task_type in additional_information (default: "CustomVoice")
- - 'language' -> language in additional_information (default: "Auto")
- - 'max_new_tokens' -> max_new_tokens in additional_information (default: 2048)
- - 'response_format' -> desired audio format (used only for assertion)
- """
- input_text = request_config.get("input") or request_config.get("prompts")
- if input_text is None:
- raise ValueError("request_config must contain 'input' or 'prompts' for TTS")
- if isinstance(input_text, list):
- input_text = input_text[0] if input_text else ""
-
- mm_processor_kwargs: dict[str, Any] = {}
- if "voice" in request_config:
- mm_processor_kwargs["speaker"] = request_config["voice"]
- if "task_type" in request_config:
- mm_processor_kwargs["task_type"] = request_config["task_type"]
- if "ref_audio" in request_config:
- mm_processor_kwargs["ref_audio"] = request_config["ref_audio"]
- if "ref_text" in request_config:
- mm_processor_kwargs["ref_text"] = request_config["ref_text"]
- if "language" in request_config:
- mm_processor_kwargs["language"] = request_config["language"]
- if "max_new_tokens" in request_config:
- mm_processor_kwargs["max_new_tokens"] = request_config["max_new_tokens"]
-
- outputs = self.runner.generate_multimodal(
- prompts=input_text,
- modalities=["audio"],
- mm_processor_kwargs=mm_processor_kwargs or None,
- )
- mm_out: dict[str, Any] | None = None
- for stage_out in outputs:
- if getattr(stage_out, "final_output_type", None) == "audio":
- mm_out = stage_out.outputs[0].multimodal_output
- break
- if mm_out is None:
- raise AssertionError("No audio output from pipeline")
-
- audio_data = mm_out.get("audio")
- if audio_data is None:
- raise AssertionError("No audio tensor in multimodal output")
-
- sr_raw = mm_out.get("sr")
- sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
- sr = int(sr_val.item() if hasattr(sr_val, "item") else sr_val)
- wav_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
- wav_buf = io.BytesIO()
- sf.write(
- wav_buf,
- wav_tensor.float().cpu().numpy().reshape(-1),
- samplerate=sr,
- format="WAV",
- subtype="PCM_16",
- )
- result = OmniResponse(success=True, audio_bytes=wav_buf.getvalue(), audio_format="audio/wav")
- assert_audio_speech_response(result, request_config, run_level="core_model")
- return result
-
- def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]:
- return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages)
-
- def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
- return self.runner.stop_profile(stages=stages)
+ cleanup_test_environment()
# ---------------------------------------------------------------------------
@@ -3118,16 +902,37 @@ class OmniRunnerHandler:
# ---------------------------------------------------------------------------
+@contextmanager
+def _whisper_device_free_around():
+ """Keep the Whisper judge off the accelerator around a server/runner lifecycle.
+
+ Released on entry so the next instance -- e.g. the next parametrization in
+ the same module -- initializes on a clean device, and again on exit
+ (including when a test raises) so it does not linger for the instance after.
+ """
+ release_audio_transcriber()
+ try:
+ yield
+ finally:
+ release_audio_transcriber()
+
+
+def get_model_prefix() -> str:
+ """Return ``MODEL_PREFIX`` with a trailing slash, or ``""`` when unset."""
+ prefix = os.environ.get("MODEL_PREFIX", "")
+ return f"{prefix.rstrip('/')}/" if prefix else ""
+
+
def iter_omni_server(
request: Any,
run_level: str,
- model_prefix: str,
omni_fixture_lock: threading.Lock,
) -> Generator[Any, Any, None]:
"""Start/stop an Omni HTTP server; used by ``omni_server`` / ``omni_server_function`` fixtures."""
from tests.helpers.stage_config import stage_config_path_for_run_level
- with omni_fixture_lock:
+ model_prefix = get_model_prefix()
+ with omni_fixture_lock, _whisper_device_free_around():
params: OmniServerParams = request.param
# For now, when a tiny model is substituted, we preserve the original model
# name via --served-model-name (so that the server still accepts requests with
@@ -3179,12 +984,7 @@ def iter_omni_server(
print("OmniServer stopping...")
else:
if stage_config_path is not None:
- from vllm_omni.entrypoints.utils import is_new_format_deploy_config
-
- if is_new_format_deploy_config(stage_config_path):
- server_args += ["--deploy-config", stage_config_path]
- else:
- server_args += ["--stage-configs-path", stage_config_path]
+ server_args += ["--deploy-config", stage_config_path]
with (
OmniServer(
@@ -3213,14 +1013,14 @@ def iter_omni_server(
def iter_omni_runner(
request: Any,
- model_prefix: str,
run_level: str,
omni_fixture_lock: threading.Lock,
) -> Generator[Any, None, None]:
"""Yield an :class:`OmniRunner`; used by ``omni_runner`` / ``omni_runner_function`` fixtures."""
from tests.helpers.stage_config import stage_config_path_for_run_level
- with omni_fixture_lock:
+ model_prefix = get_model_prefix()
+ with omni_fixture_lock, _whisper_device_free_around():
param = request.param
if not isinstance(param, (tuple, list)) or len(param) not in (2, 3):
raise ValueError(
@@ -3237,7 +1037,7 @@ def iter_omni_runner(
model = model_prefix + model
if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
model = resolve_tiny_model_path(model)
- with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
+ with OmniRunner(model, seed=42, deploy_config=stage_config_path, **extra_omni_kwargs) as runner:
print("OmniRunner started successfully")
yield runner
print("OmniRunner stopping...")
@@ -3245,6 +1045,166 @@ def iter_omni_runner(
print("OmniRunner stopped")
+# ─────────────────────────────────────────────────────────────────────
+# π0 (Pi-Zero) OpenPI websocket policy client
+#
+# Minimal client for the OpenPI realtime robot protocol
+# (``/v1/realtime/robot/openpi``): connect → read handshake metadata → send
+# observation → receive an action chunk. Used by the π0 online-serving e2e
+# (tests/e2e/online_serving/test_pi0_expansion.py). π0's observation is 3 RGB
+# cameras + proprioceptive state + a language prompt, built as synthetic numpy
+# frames here (no video assets).
+# ─────────────────────────────────────────────────────────────────────
+PI0_OPENPI_PATH = "/v1/realtime/robot/openpi"
+PI0_OPENPI_DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+# π0 / pi0_base camera identities (must match the server's image_feature_keys,
+# i.e. vllm_omni/deploy/pi0.yaml, whose image_key_map is empty).
+PI0_CAMERA_KEYS = (
+ "observation.images.base_0_rgb",
+ "observation.images.left_wrist_0_rgb",
+ "observation.images.right_wrist_0_rgb",
+)
+PI0_ACTION_HORIZON = 50
+PI0_ACTION_DIM = 32
+PI0_STATE_DIM = 32
+PI0_IMAGE_SIZE = 224
+
+
+def pi0_openpi_require_dependencies() -> None:
+ """Raise ModuleNotFoundError if the OpenPI websocket client deps are missing."""
+ missing = []
+ try:
+ import websockets.sync.client # noqa: F401
+ except ImportError:
+ missing.append("websockets")
+ try:
+ from openpi_client import msgpack_numpy # noqa: F401
+ except ImportError:
+ missing.append("openpi-client")
+ if missing:
+ raise ModuleNotFoundError(f"π0 OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+def _pi0_decode_action_response(response: bytes | str) -> np.ndarray:
+ from openpi_client import msgpack_numpy
+
+ if isinstance(response, str):
+ raise RuntimeError(f"Inference failed: {response}")
+ decoded = msgpack_numpy.unpackb(response)
+ if isinstance(decoded, dict) and decoded.get("type") == "error":
+ raise RuntimeError(f"Inference failed: {decoded.get('message', decoded)}")
+ return np.asarray(decoded, dtype=np.float32)
+
+
+def pi0_make_dummy_obs(*, prompt: str, session_id: str, image_size: int = PI0_IMAGE_SIZE) -> dict[str, Any]:
+ """A single π0 observation: 3 blank cameras (HWC uint8) + zero state + prompt."""
+ obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in PI0_CAMERA_KEYS}
+ obs["state"] = np.zeros(PI0_STATE_DIM, dtype=np.float32)
+ obs["prompt"] = prompt
+ obs["session_id"] = session_id
+ return obs
+
+
+def pi0_openpi_run_policy_session(
+ *,
+ host: str = "127.0.0.1",
+ port: int = 8000,
+ path: str = PI0_OPENPI_PATH,
+ prompt: str = PI0_OPENPI_DEFAULT_PROMPT,
+ session_id: str | None = None,
+ num_steps: int = 2,
+) -> dict[str, Any]:
+ """Connect, read handshake metadata, send ``num_steps`` observations."""
+ import uuid
+
+ import websockets.sync.client as websockets_client
+ from openpi_client import msgpack_numpy
+
+ session_id = session_id or str(uuid.uuid4())
+ uri = f"ws://{host}:{port}{path}"
+ packer = msgpack_numpy.Packer()
+ conn = websockets_client.connect(uri, compression=None, max_size=None, ping_interval=300, ping_timeout=3600)
+ try:
+ metadata = msgpack_numpy.unpackb(conn.recv())
+ if not isinstance(metadata, dict):
+ raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+ actions = []
+ for _ in range(num_steps):
+ payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id)
+ payload["endpoint"] = "infer"
+ conn.send(packer.pack(payload))
+ actions.append(_pi0_decode_action_response(conn.recv()))
+ return {"metadata": dict(metadata), "actions": actions, "session_id": session_id}
+ finally:
+ conn.close()
+
+
+def pi0_openpi_validate_session_result(
+ result: dict[str, Any],
+ *,
+ expected_action_horizon: int = PI0_ACTION_HORIZON,
+ expected_action_dim: int = PI0_ACTION_DIM,
+) -> None:
+ """Assert the handshake metadata + every returned action chunk for π0."""
+ metadata = result["metadata"]
+ required_keys = (
+ "image_resolution",
+ "needs_wrist_camera",
+ "needs_session_id",
+ "action_space",
+ "action_horizon",
+ "action_dim",
+ )
+ missing = [key for key in required_keys if key not in metadata]
+ if missing:
+ raise AssertionError(f"Missing π0 metadata keys: {missing}")
+
+ if tuple(metadata["image_resolution"]) != (PI0_IMAGE_SIZE, PI0_IMAGE_SIZE):
+ raise AssertionError(f"Unexpected image_resolution: {metadata['image_resolution']!r}")
+ if not metadata["needs_wrist_camera"]:
+ raise AssertionError("π0 test expects needs_wrist_camera=True")
+ if metadata["needs_session_id"]:
+ raise AssertionError("π0 is stateless; needs_session_id must be False")
+ if metadata["action_space"] != "joint_position":
+ raise AssertionError(f"Unexpected action_space: {metadata['action_space']!r}")
+ if int(metadata["action_horizon"]) != expected_action_horizon:
+ raise AssertionError(f"Unexpected action_horizon: {metadata['action_horizon']!r}")
+ if int(metadata["action_dim"]) != expected_action_dim:
+ raise AssertionError(f"Unexpected action_dim: {metadata['action_dim']!r}")
+
+ actions = result["actions"]
+ if not actions:
+ raise AssertionError("No actions returned from the server")
+ for index, action in enumerate(actions):
+ if action.shape != (expected_action_horizon, expected_action_dim):
+ raise AssertionError(
+ f"Action {index} shape mismatch: expected "
+ f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+ )
+ if not np.isfinite(action).all():
+ raise AssertionError(f"Action {index} contains non-finite values")
+
+
+# Re-export client types for backward-compatible imports.
+from tests.helpers.client import ( # noqa: E402
+ DREAMZERO_ACTION_DIM,
+ DREAMZERO_ACTION_HORIZON,
+ DREAMZERO_CAMERA_FILES,
+ DiffusionResponse,
+ HttpResponse,
+ OfflineOmniClient,
+ OmniResponse,
+ OmniRunnerHandler,
+ OnlineOmniClient,
+ OpenAIClientHandler,
+ OpenPIWebSocketResponse,
+ OpenPIWebSocketSession,
+ WebSocketJsonResponse,
+ build_dreamzero_demo_observations,
+ build_openpi_droid_observation,
+ load_dreamzero_camera_frames,
+)
+
__all__ = [
"DiffusionResponse",
"HttpResponse",
@@ -3259,11 +1219,18 @@ __all__ = [
"OpenPIWebSocketSession",
"OmniResponse",
"OmniRunner",
- "OmniRunnerHandler",
+ "OfflineOmniClient",
"OmniServer",
"OmniServerParams",
"OmniServerStageCli",
+ "OnlineOmniClient",
"OpenAIClientHandler",
+ "OmniRunnerHandler",
+ "get_model_prefix",
"get_open_port",
"dummy_messages_from_mix_data",
+ "pi0_openpi_require_dependencies",
+ "pi0_openpi_run_policy_session",
+ "pi0_openpi_validate_session_result",
+ "pi0_make_dummy_obs",
]
diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py
new file mode 100644
index 000000000..f1f2710b5
--- /dev/null
+++ b/tests/helpers/serving_chat.py
@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared scaffolding for OmniOpenAIServingChat unit tests.
+
+Provides factories for a minimal serving-chat instance, chat requests, and
+OmniRequestOutput objects, plus SSE stream collection/parsing utilities.
+"""
+
+import json
+from unittest.mock import MagicMock
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+ ChatCompletionResponseStreamChoice,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.outputs import CompletionOutput, RequestOutput
+
+from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.outputs import OmniRequestOutput
+
+
+def make_text_omni_output(
+ request_id: str = "test-req",
+ text: str = "hello",
+ token_ids: list[int] | None = None,
+ finish_reason: str | None = None,
+ index: int = 0,
+ num_prompt_tokens: int = 3,
+ num_cached_tokens: int | None = None,
+) -> OmniRequestOutput:
+ """Build an OmniRequestOutput wrapping a text RequestOutput."""
+ if token_ids is None:
+ token_ids = [10, 11, 12]
+ res = RequestOutput(
+ request_id=request_id,
+ prompt="test",
+ prompt_token_ids=list(range(num_prompt_tokens)),
+ prompt_logprobs=None,
+ outputs=[
+ CompletionOutput(
+ index=index,
+ text=text,
+ token_ids=token_ids,
+ cumulative_logprob=0.0,
+ logprobs=None,
+ finish_reason=finish_reason,
+ stop_reason=None,
+ )
+ ],
+ finished=finish_reason is not None,
+ num_cached_tokens=num_cached_tokens,
+ )
+ return OmniRequestOutput.from_stage_output(
+ res,
+ request_id=request_id,
+ final_output_type="text",
+ finished=finish_reason is not None,
+ )
+
+
+def mock_audio_choices(index: int = 0, role: str = "assistant"):
+ return [
+ ChatCompletionResponseStreamChoice(
+ index=index,
+ delta=DeltaMessage(role=role, content="dGVzdA=="),
+ logprobs=None,
+ finish_reason="stop",
+ )
+ ]
+
+
+def build_serving_chat(
+ *,
+ engine_client=None,
+ models=None,
+ online_renderer=None,
+ response_role: str = "assistant",
+ request_logger=None,
+ chat_template=None,
+ chat_template_content_format="auto",
+ **init_overrides,
+) -> OmniOpenAIServingChat:
+ """Create a real OmniOpenAIServingChat via its actual __init__.
+
+ engine_client/models/online_renderer default to minimal mocks; pass your
+ own to wire in test-specific behavior (e.g. engine_client.stage_configs).
+ Any other OmniOpenAIServingChat.__init__ keyword — enable_auto_tools,
+ enable_prompt_tokens_details, trust_request_chat_template, etc. — can be
+ passed through **init_overrides instead of set on the instance after
+ construction.
+ """
+ if engine_client is None:
+ engine_client = MagicMock()
+ engine_client.errored = False
+ if models is None:
+ models = OpenAIServingModels(
+ engine_client=engine_client,
+ base_model_paths=[],
+ )
+ if online_renderer is None:
+ online_renderer = MagicMock()
+
+ instance = OmniOpenAIServingChat(
+ engine_client=engine_client,
+ models=models,
+ response_role=response_role,
+ online_renderer=online_renderer,
+ request_logger=request_logger,
+ chat_template=chat_template,
+ chat_template_content_format=chat_template_content_format,
+ **init_overrides,
+ )
+ instance._create_audio_choice = MagicMock(
+ side_effect=lambda omni_res, role, request, stream=False: mock_audio_choices(
+ index=omni_res.outputs[0].index,
+ role=role,
+ )
+ )
+ return instance
+
+
+def make_request(
+ modalities: list[str],
+ n: int = 1,
+ *,
+ stream: bool = True,
+ include_usage: bool = False,
+) -> ChatCompletionRequest:
+ req = ChatCompletionRequest(
+ model="test-model",
+ messages=[{"role": "user", "content": "hello"}],
+ n=n,
+ stream=stream,
+ stream_options={"include_usage": True} if include_usage else None,
+ )
+ req.modalities = modalities # type: ignore[attr-defined]
+ return req
+
+
+def parse_sse_chunks(lines: list[str]) -> list[dict]:
+ """Parse SSE lines into JSON dicts."""
+ prefix = "data: "
+ chunks = []
+ for line in lines:
+ line = line.strip()
+ if not line.startswith(prefix):
+ continue
+ payload = line[len(prefix) :].strip()
+ if payload == "[DONE]":
+ continue
+ chunks.append(json.loads(payload))
+ return chunks
+
+
+async def collect_stream(gen):
+ result = []
+ async for item in gen:
+ result.append(item)
+ return result
diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py
index 8706bb1bc..36d8c406a 100644
--- a/tests/helpers/stage_config.py
+++ b/tests/helpers/stage_config.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Config/message construction helpers used by tests."""
import atexit
@@ -8,11 +11,13 @@ from typing import Any
import yaml
+from vllm_omni.config.stage_config import load_deploy_config
+
def modify_stage_config(
yaml_path: str,
- updates: dict[str, Any] = None,
- deletes: dict[str, Any] = None,
+ updates: dict[str, Any] | None = None,
+ deletes: dict[str, Any] | None = None,
) -> str:
"""
Modify configurations in a YAML file, supporting both top-level and stage-specific modifications,
@@ -23,17 +28,17 @@ def modify_stage_config(
updates: Dictionary containing both top-level and stage-specific modifications to add or update.
Format: {
'async_chunk': True,
- 'stage_args': {
- 0: {'engine_args.max_model_len': 5800},
- 1: {'engine_args.max_num_seqs': 2}
+ 'stages': {
+ 0: {'max_model_len': 5800},
+ 1: {'max_num_seqs': 2}
}
}
deletes: Dictionary containing configurations to delete.
Format: {
'old_config': None, # Delete entire key
- 'stage_args': {
- 0: ['engine_args.old_param'],
- 1: ['runtime.unused_setting']
+ 'stages': {
+ 0: ['old_param'],
+ 1: ['unused_setting']
}
}
@@ -164,24 +169,20 @@ def modify_stage_config(
else:
print(f"Path {path} does not exist")
- _stage_key = "stages" if "stages" in config else "stage_args"
+ stages = get_stage_entries(config)
# Apply deletions first
if deletes:
for key, value in deletes.items():
- if key in ("stage_args", "stages"):
+ if key == "stages":
if value and isinstance(value, dict):
- stage_args = config.get(_stage_key, [])
- if not stage_args:
- raise ValueError("stage_args does not exist in config")
-
for stage_id, delete_paths in value.items():
if not delete_paths:
continue
# Find stage by ID
target_stage = None
- for stage in stage_args:
+ for stage in stages:
if stage.get("stage_id") == int(stage_id):
target_stage = stage
break
@@ -204,28 +205,24 @@ def modify_stage_config(
# Apply updates
if updates:
for key, value in updates.items():
- if key in ("stage_args", "stages"):
+ if key == "stages":
if value and isinstance(value, dict):
- stage_args = config.get(_stage_key, [])
- if not stage_args:
- raise ValueError("stage_args does not exist in config")
-
for stage_id, stage_updates in value.items():
# Find stage by ID
target_stage = None
- for stage in stage_args:
+ for stage in stages:
if stage.get("stage_id") == int(stage_id):
target_stage = stage
break
if target_stage is None:
- available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+ available_ids = [s.get("stage_id") for s in stages if "stage_id" in s]
raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
# Apply updates to this stage
for update_path, val in stage_updates.items():
# Check if this is a simple key (not dot-separated)
- # Example: 'engine_input_source' vs 'engine_args.max_model_len'
+ # Example: 'engine_input_source' vs nested 'default_sampling_params.max_tokens'
if "." not in update_path:
# Direct key assignment (e.g., updating a list value)
target_stage[update_path] = val
@@ -624,6 +621,36 @@ def _materialize_ci_overlay(model_type: str) -> Path:
return out
+def get_stage_entries(cfg: dict[str, Any]) -> list[dict[str, Any]]:
+ """Return the ``stages`` list from a new-schema deploy config dict.
+
+ Raises ``ValueError`` when the top-level ``stages`` key is missing.
+ """
+ stages = cfg.get("stages")
+ if stages is None:
+ raise ValueError("Deploy config missing top-level 'stages'")
+ return list(stages)
+
+
+def load_stage_ids(resolved_config: dict[str, Any]) -> list[int]:
+ """Return ``stage_id`` values from a resolved deploy config dict."""
+ stage_ids = [stage["stage_id"] for stage in get_stage_entries(resolved_config) if "stage_id" in stage]
+ if not stage_ids:
+ raise ValueError("No stage IDs found in resolved config")
+ return stage_ids
+
+
+def load_stage_replica_counts(resolved_config: dict[str, Any]) -> dict[int, int]:
+ """Return ``{stage_id: num_replicas}`` from flat stage entries (default 1)."""
+ replica_counts: dict[int, int] = {}
+ for stage in get_stage_entries(resolved_config):
+ stage_id = stage.get("stage_id")
+ if stage_id is None:
+ continue
+ replica_counts[int(stage_id)] = max(1, int(stage.get("num_replicas", 1)))
+ return replica_counts
+
+
def get_deploy_config_path(rel_path: str) -> str:
"""Resolve a deploy yaml; ``ci/.yaml`` materializes from ``_CI_OVERLAYS``."""
if rel_path.startswith("ci/") and rel_path.endswith(".yaml"):
@@ -638,33 +665,30 @@ def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]:
with open(get_deploy_config_path(rel_path), encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
- stage_key = "stages" if "stages" in cfg else "stage_args"
- for stage in cfg.get(stage_key, []):
+ for stage in get_stage_entries(cfg):
if stage.get("stage_id") == stage_id:
return stage
raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}")
-def _get_config_value_by_path(config_dict: dict, path: str) -> Any:
- """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``)."""
- current: Any = config_dict
- for key in path.split("."):
- if not isinstance(current, dict) or key not in current:
- return None
- current = current[key]
- return current
+def get_deploy_duplex_max_sessions(rel_path: str) -> int:
+ """Return the duplex session capacity a deploy yaml admits.
+
+ Loads through ``load_deploy_config`` so ``base_config`` merging and
+ ``DuplexSessionRuntimeConfig`` defaults stay in lockstep with the server
+ instead of being re-implemented here.
+ """
+ return load_deploy_config(get_deploy_config_path(rel_path)).duplex_session.max_sessions
-def _stage_load_format_paths(stage_config_path: str) -> tuple[str, str, list[int]]:
- """Return ``(stage_key, load_format_field_path, stage_ids)`` for a deploy YAML."""
+def _stage_ids_from_deploy_yaml(stage_config_path: str) -> list[int]:
+ """Return ``stage_id`` values from a new-schema deploy YAML (``stages``)."""
with open(stage_config_path, encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
- new_schema_stages = cfg.get("stages")
- stage_key = "stages" if new_schema_stages is not None else "stage_args"
- load_format_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
- stage_entries = cfg.get(stage_key, [])
- stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
- return stage_key, load_format_path, stage_ids
+ try:
+ return load_stage_ids(cfg)
+ except ValueError as exc:
+ raise ValueError(f"Deploy YAML missing top-level 'stages': {stage_config_path}") from exc
def _add_dummy_load_format(
@@ -674,10 +698,10 @@ def _add_dummy_load_format(
"""For ``core_model`` runs, patch every stage in the deploy YAML to ``load_format: dummy``."""
if run_level != "core_model" or stage_config_path is None:
return stage_config_path
- stage_key, load_format_path, stage_ids = _stage_load_format_paths(stage_config_path)
+ stage_ids = _stage_ids_from_deploy_yaml(stage_config_path)
return modify_stage_config(
stage_config_path,
- updates={stage_key: {stage_id: {load_format_path: "dummy"} for stage_id in stage_ids}},
+ updates={"stages": {stage_id: {"load_format": "dummy"} for stage_id in stage_ids}},
)
@@ -688,25 +712,27 @@ def _delete_dummy_load_format(
"""For ``advanced_model`` / ``full_model``, strip ``load_format: dummy`` so real weights load."""
if run_level not in {"advanced_model", "full_model"} or stage_config_path is None:
return stage_config_path
- stage_key, load_format_path, _stage_ids = _stage_load_format_paths(stage_config_path)
with open(stage_config_path, encoding="utf-8") as f:
cfg = yaml.safe_load(f) or {}
- stage_entries = cfg.get(stage_key, [])
+ try:
+ stages = get_stage_entries(cfg)
+ except ValueError as exc:
+ raise ValueError(f"Deploy YAML missing top-level 'stages': {stage_config_path}") from exc
deletes: dict[int, list[str]] = {}
- for stage in stage_entries:
+ for stage in stages:
stage_id = stage.get("stage_id")
if stage_id is None:
continue
- if _get_config_value_by_path(stage, load_format_path) == "dummy":
- deletes[stage_id] = [load_format_path]
+ if stage.get("load_format") == "dummy":
+ deletes[stage_id] = ["load_format"]
if not deletes:
return stage_config_path
return modify_stage_config(
stage_config_path,
- deletes={stage_key: deletes},
+ deletes={"stages": deletes},
)
@@ -720,6 +746,10 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st
__all__ = [
"get_deploy_config_path",
"get_deploy_config_stage",
+ "get_deploy_duplex_max_sessions",
+ "get_stage_entries",
+ "load_stage_ids",
+ "load_stage_replica_counts",
"modify_stage_config",
"stage_config_path_for_run_level",
]
diff --git a/tests/helpers/tests/test_assertions.py b/tests/helpers/tests/test_assertions.py
index 9cb1560b9..98deaa511 100644
--- a/tests/helpers/tests/test_assertions.py
+++ b/tests/helpers/tests/test_assertions.py
@@ -1,12 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
from tests.helpers import assertions
-from tests.helpers.assertions import _assert_transcript_matches, _resolve_audio_transcript
+from tests.helpers.assertions import (
+ _assert_transcript_matches,
+ _resolve_audio_transcript,
+ assert_audio_speech_response,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -71,6 +75,40 @@ def test_resolve_transcript_leaves_language_unset_by_default(monkeypatch):
assert captured["model_size"] == "small"
+@pytest.mark.parametrize(
+ ("request_config", "expected_text"),
+ [
+ ({"input": "spoken words"}, "spoken words"),
+ (
+ {
+ "input": "<|style:whispering|>spoken words",
+ "transcript_expected_text": "spoken words",
+ },
+ "spoken words",
+ ),
+ ],
+)
+def test_speech_transcript_expected_text(monkeypatch, request_config, expected_text):
+ captured: dict = {}
+
+ monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "spoken words")
+
+ def capture_match(_transcript, _audio_bytes, expected_text, **_kwargs):
+ captured["expected_text"] = expected_text
+
+ monkeypatch.setattr(assertions, "_assert_transcript_matches", capture_match)
+ response = SimpleNamespace(success=True, audio_bytes=b"fake-wav", audio_format="audio/wav")
+ request_config["response_format"] = "wav"
+
+ assert_audio_speech_response(
+ response,
+ request_config,
+ "advanced_model",
+ )
+
+ assert captured["expected_text"] == expected_text
+
+
def test_escalated_transcript_keeps_declared_language(monkeypatch):
# The escalated pass must honour the same language, otherwise a request that
# pins one silently falls back to auto-detection on retry.
diff --git a/tests/helpers/tests/test_mark.py b/tests/helpers/tests/test_mark.py
new file mode 100644
index 000000000..d2fed4b6d
--- /dev/null
+++ b/tests/helpers/tests/test_mark.py
@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from tests.helpers.mark import (
+ _gpu_res_platforms,
+ hardware_marks,
+ hardware_test,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _mark_names(marks) -> set[str]:
+ return {mark.name for mark in marks}
+
+
+def test_gpu_platforms_come_from_pyproject_gpu_tag():
+ assert "npu" not in _gpu_res_platforms()
+ assert "gpu" not in _gpu_res_platforms()
+
+
+def test_default_num_cards_adds_cards_1():
+ names = _mark_names(hardware_marks(res={"cuda": "H100"}))
+ assert "cards_1" in names
+ assert "H100" in names
+ assert "cards_2" not in names
+
+
+def test_h100_two_cards_is_selectable_with_cards_2():
+ names = _mark_names(hardware_marks(res={"cuda": "H100"}, num_cards=2))
+ assert names >= {"H100", "cuda", "gpu", "cards_2"}
+ assert "cards_1" not in names
+ assert "cards_4" not in names
+
+
+def test_h100_four_cards_is_selectable_with_cards_4():
+ names = _mark_names(hardware_marks(res={"cuda": "H100"}, num_cards=4))
+ assert "cards_4" in names
+ assert "cards_2" not in names
+
+
+def test_npu_card_count_mark_is_independent_of_sku():
+ names = _mark_names(hardware_marks(res={"npu": "A2"}, num_cards=4))
+ assert "cards_4" in names
+ assert "npu" in names
+ assert "gpu" not in names
+
+
+def test_hardware_marks_rejects_mixed_counts_on_one_item():
+ with pytest.raises(ValueError, match="different per-platform num_cards"):
+ hardware_marks(
+ res={"cuda": "H100", "rocm": "MI325"},
+ num_cards={"cuda": 4, "rocm": 2},
+ )
+
+
+def test_same_count_multi_platform_stays_one_item():
+ names = _mark_names(hardware_marks(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2))
+ assert names >= {"H100", "cuda", "MI325", "rocm", "cards_2"}
+
+ @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+ def _probe():
+ return None
+
+ names = {mark.name for mark in _probe.pytestmark}
+ assert "parametrize" not in names
+ assert "cards_2" in names
+ assert "H100" in names
+ assert "MI325" in names
+
+
+def test_hardware_test_decorator_applies_cards_mark():
+ @hardware_test(res={"cuda": "H100"}, num_cards=2)
+ def _probe():
+ return None
+
+ names = {mark.name for mark in _probe.pytestmark}
+ assert "cards_2" in names
+ assert "H100" in names
+
+
+def _hardware_platform_params(func):
+ marks = func.pytestmark if isinstance(func.pytestmark, list) else [func.pytestmark]
+ for mark in marks:
+ if mark.name != "parametrize":
+ continue
+ argnames = mark.args[0]
+ if argnames == "_normalized_hardware_marks":
+ return mark.args[1]
+ raise AssertionError("expected @hardware_test to parametrize _normalized_hardware_marks")
+
+
+def test_mixed_card_counts_split_so_cuda_filter_does_not_see_rocm_count():
+ @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+ def _probe():
+ return None
+
+ variants = {param.id: {mark.name for mark in param.marks} for param in _hardware_platform_params(_probe)}
+ cuda_names = variants["cuda"]
+ rocm_names = variants["rocm"]
+
+ assert cuda_names >= {"H100", "cuda", "cards_4"}
+ assert "cards_2" not in cuda_names
+ assert "rocm" not in cuda_names
+ assert "MI325" not in cuda_names
+
+ assert rocm_names >= {"MI325", "rocm", "cards_2"}
+ assert "cards_4" not in rocm_names
+ assert "cuda" not in rocm_names
+ assert "H100" not in rocm_names
+
+
+def test_unsupported_card_count_raises():
+ with pytest.raises(ValueError, match="num_cards=16"):
+ hardware_marks(res={"cuda": "H100"}, num_cards=16)
+
+
+@pytest.mark.parametrize("platform", ["cpu", "gpu"])
+def test_res_rejects_cpu_and_gpu(platform):
+ with pytest.raises(ValueError, match="not a res dict key"):
+ hardware_marks(res={platform: "H100"})
+
+
+def test_npu_rejects_unsupported_sku():
+ with pytest.raises(ValueError, match="Invalid npu resource type"):
+ hardware_marks(res={"npu": "A100"})
+
+
+@pytest.mark.parametrize("sku", ["H200", "B200", "H800"])
+def test_cuda_registered_skus_attach_sku_and_cards(sku):
+ names = _mark_names(hardware_marks(res={"cuda": sku}, num_cards=2))
+ assert names >= {sku, "cuda", "gpu", "cards_2"}
+
+
+def test_cuda_rejects_sku_from_another_platform():
+ with pytest.raises(ValueError, match="Invalid cuda resource type"):
+ hardware_marks(res={"cuda": "MI325"})
diff --git a/tests/helpers/tests/test_media.py b/tests/helpers/tests/test_media.py
index 4a8e5393d..8045f7a69 100644
--- a/tests/helpers/tests/test_media.py
+++ b/tests/helpers/tests/test_media.py
@@ -2,6 +2,8 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
+from concurrent.futures.process import BrokenProcessPool
+from contextlib import nullcontext
from types import SimpleNamespace
import numpy as np
@@ -22,7 +24,11 @@ def _stub_transcribe(monkeypatch) -> dict:
captured.update(kwargs)
return {"text": "London"}
- monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=lambda size, device=None: FakeModel()))
+ monkeypatch.setitem(
+ sys.modules,
+ "whisper",
+ SimpleNamespace(load_model=lambda size, device=None: FakeModel()),
+ )
# Keep the unit test off the accelerator probe so it stays hermetic.
monkeypatch.setitem(
sys.modules,
@@ -50,32 +56,183 @@ def test_transcribe_defaults_to_auto_language(monkeypatch):
assert captured.get("language") is None
+class _FakeExecutor:
+ """Stand-in for ProcessPoolExecutor recording what was submitted and shut down."""
+
+ def __init__(self, outcomes=()):
+ self._outcomes = list(outcomes)
+ self.submitted: list[tuple] = []
+ self.shutdown_calls = 0
+
+ def submit(self, fn, *args):
+ self.submitted.append((fn, args))
+ outcome = self._outcomes.pop(0) if self._outcomes else "London"
+
+ def _result():
+ if isinstance(outcome, BaseException):
+ raise outcome
+ return outcome
+
+ return SimpleNamespace(result=_result)
+
+ def shutdown(self, wait=True):
+ self.shutdown_calls += 1
+
+
+def _patch_executors(monkeypatch, outcomes_per_executor=()) -> list[_FakeExecutor]:
+ """Replace the pool constructor and hand back the fakes it produces, in creation order."""
+ created: list[_FakeExecutor] = []
+
+ def factory(*args, **kwargs):
+ index = len(created)
+ outcomes = outcomes_per_executor[index] if index < len(outcomes_per_executor) else ()
+ executor = _FakeExecutor(outcomes)
+ executor.init_kwargs = kwargs
+ created.append(executor)
+ return executor
+
+ monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", factory)
+ return created
+
+
+@pytest.fixture(autouse=True)
+def _reset_transcriber_singletons():
+ """The worker and its model cache are module-level state; do not leak them across tests."""
+ media.release_audio_transcriber()
+ media._WHISPER_MODELS.clear()
+ yield
+ media.release_audio_transcriber()
+ media._WHISPER_MODELS.clear()
+
+
def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path):
"""Cover the two hops the tests above skip: bytes -> file -> executor.submit.
The real call crosses a spawn ProcessPoolExecutor, so capture what gets
submitted rather than what the worker eventually does.
"""
- submitted: dict = {}
-
- class FakeExecutor:
- def __enter__(self):
- return self
-
- def __exit__(self, *exc):
- return False
-
- def submit(self, fn, *args):
- submitted["fn"] = fn
- submitted["args"] = args
- return SimpleNamespace(result=lambda: "London")
-
- monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", lambda *a, **kw: FakeExecutor())
+ created = _patch_executors(monkeypatch)
wav = tmp_path / "clip.wav"
sf.write(wav, np.zeros(2400, dtype=np.float32), 24000)
assert media.convert_audio_bytes_to_text(wav.read_bytes(), "small", "en") == "London"
- assert submitted["fn"] is media._whisper_transcribe_in_current_process
- assert submitted["args"][1:] == ("small", "en")
+ fn, args = created[0].submitted[0]
+ assert fn is media._whisper_transcribe_in_current_process
+ assert args[1:] == ("small", "en")
+
+
+def test_parent_reuses_one_spawned_worker_across_calls(monkeypatch):
+ """Two calls share one spawned worker.
+
+ Pins the parent-side contract: the pool is constructed once, with
+ ``spawn`` (CUDA is not fork-safe) and ``max_workers=1`` (one resident
+ model). Reuse of the worker's *loaded model* is pinned by the worker-side
+ cache test below, which is where the startup cost actually lives.
+ """
+ created = _patch_executors(monkeypatch)
+
+ media.convert_audio_file_to_text("/tmp/a.wav")
+ media.convert_audio_file_to_text("/tmp/b.wav")
+
+ assert len(created) == 1
+ assert len(created[0].submitted) == 2
+ assert created[0].init_kwargs["max_workers"] == 1
+ assert created[0].init_kwargs["mp_context"].get_start_method() == "spawn"
+
+
+def test_release_forces_a_new_process_pool(monkeypatch):
+ created = _patch_executors(monkeypatch)
+
+ media.convert_audio_file_to_text("/tmp/a.wav")
+ media.release_audio_transcriber()
+ media.convert_audio_file_to_text("/tmp/b.wav")
+
+ assert len(created) == 2
+ assert created[0].shutdown_calls == 1
+
+
+def test_worker_loads_each_model_size_once(monkeypatch):
+ """The worker-side cache is what removes the repeated whisper.load_model cost."""
+ loaded: list[str] = []
+
+ class FakeModel:
+ def transcribe(self, path, **kwargs):
+ return {"text": "London"}
+
+ def load_model(size, device=None):
+ loaded.append(size)
+ return FakeModel()
+
+ monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=load_model))
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm_omni.platforms",
+ SimpleNamespace(current_omni_platform=SimpleNamespace(is_available=lambda: False)),
+ )
+ monkeypatch.setattr(media, "_serialize_whisper_model_download", lambda model_size: nullcontext())
+
+ media._whisper_transcribe_in_current_process("/tmp/a.wav", "small")
+ media._whisper_transcribe_in_current_process("/tmp/b.wav", "small")
+
+ assert loaded == ["small"]
+
+ # The ASR escalation path asks the same worker for a stronger model.
+ media._whisper_transcribe_in_current_process("/tmp/c.wav", "large-v3")
+
+ assert loaded == ["small", "large-v3"]
+
+
+@pytest.mark.parametrize(
+ ("second_outcome", "recovers"),
+ [(["London"], True), ([BrokenProcessPool("second")], False)],
+)
+def test_broken_pool_is_discarded_and_retried_once(monkeypatch, second_outcome, recovers):
+ """A dead worker is discarded and the call retried exactly once.
+
+ The retry either succeeds, or -- if the replacement also dies -- fails with
+ both pools discarded and no executor left installed.
+ """
+ created = _patch_executors(
+ monkeypatch,
+ outcomes_per_executor=([BrokenProcessPool("first")], second_outcome),
+ )
+
+ if recovers:
+ assert media.convert_audio_file_to_text("/tmp/a.wav") == "London"
+ else:
+ with pytest.raises(BrokenProcessPool):
+ media.convert_audio_file_to_text("/tmp/a.wav")
+ assert created[1].shutdown_calls == 1
+ assert media._TRANSCRIBER is None
+
+ assert len(created) == 2 # discarded the first, retried on a fresh one
+ assert created[0].shutdown_calls == 1
+
+
+@pytest.mark.parametrize(
+ "error",
+ [RuntimeError("CUDA out of memory"), KeyboardInterrupt()],
+ ids=["Exception", "BaseException"],
+)
+def test_transcription_failure_propagates_and_discards_the_worker(monkeypatch, error):
+ """Any failure -- an Exception or a BaseException like KeyboardInterrupt --
+ propagates AND tears the worker down.
+
+ It arrives as the task's own exception, not ``BrokenProcessPool``, but the
+ worker -- and its resident model -- is still discarded, restoring the old
+ one-process-per-call isolation. It is not retried; the next call builds a
+ fresh worker.
+ """
+ created = _patch_executors(monkeypatch, outcomes_per_executor=([error],))
+
+ with pytest.raises(type(error)):
+ media.convert_audio_file_to_text("/tmp/a.wav")
+
+ assert len(created) == 1 # not retried
+ assert created[0].shutdown_calls == 1
+ assert media._TRANSCRIBER is None
+
+ assert media.convert_audio_file_to_text("/tmp/b.wav") == "London"
+ assert len(created) == 2 # a fresh worker for the next call
diff --git a/tests/helpers/tests/test_runtime_omni_runner.py b/tests/helpers/tests/test_runtime_omni_runner.py
new file mode 100644
index 000000000..aaada20f2
--- /dev/null
+++ b/tests/helpers/tests/test_runtime_omni_runner.py
@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Tests for OmniRunner startup rollback."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.runtime import OmniRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.MonkeyPatch) -> None:
+ """``__exit__`` is skipped when construction raises; rollback must still run."""
+ cleaned: list[str] = []
+ monkeypatch.setattr(
+ "tests.helpers.runtime.cleanup_test_environment",
+ lambda: cleaned.append("env"),
+ )
+ monkeypatch.setattr(
+ OmniRunner,
+ "_cleanup_process",
+ lambda self: cleaned.append("process"),
+ )
+
+ class _BoomOmni:
+ def __init__(self, *args, **kwargs) -> None:
+ raise RuntimeError("Orchestrator initialization failed")
+
+ monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _BoomOmni)
+
+ with pytest.raises(RuntimeError, match="Orchestrator initialization failed"):
+ with OmniRunner("fake-model"):
+ raise AssertionError("context body must not run after a failed constructor")
+
+ assert cleaned.count("process") == 1
+ # Once at the start of ``__init__``, once from the constructor rollback.
+ assert cleaned.count("env") == 2
diff --git a/tests/helpers/tests/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py
index 5a2ab5d68..564f12cf1 100644
--- a/tests/helpers/tests/test_runtime_stage_cli.py
+++ b/tests/helpers/tests/test_runtime_stage_cli.py
@@ -1,8 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for OmniServerStageCli process planning helpers."""
from __future__ import annotations
from pathlib import Path
+from types import SimpleNamespace
import pytest
@@ -38,23 +42,48 @@ def test_stage_cli_builds_headless_replica_cmd(tmp_path: Path) -> None:
assert cmd[cmd.index("--deploy-config") + 1] == server.stage_config_path
-def test_stage_cli_splits_devices_per_replica(tmp_path: Path) -> None:
+def test_stage_cli_loads_stage_ids_and_replica_counts(tmp_path: Path) -> None:
server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), [])
- assert server._devices_for_replica(1, "1,2,3", 0) == "1"
- assert server._devices_for_replica(1, "1,2,3", 1) == "2"
- assert server._devices_for_replica(1, "1,2,3", 2) == "3"
+ assert server.stage_ids == [0, 1]
+ assert server.stage_replica_counts == {0: 1, 1: 3}
-def test_stage_cli_maps_visible_devices_after_replica_split(tmp_path: Path) -> None:
- server = OmniServerStageCli(
- "fake-model",
- _write_stage_config(tmp_path),
- [],
- env_dict={"CUDA_VISIBLE_DEVICES": "4,5,6,7"},
- )
- env: dict[str, str] = {}
+def test_stage_cli_enter_rolls_back_launched_stages_on_startup_failure(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ """``__exit__`` is skipped when ``__enter__`` raises; rollback must still run."""
+ monkeypatch.setattr("tests.helpers.runtime.time.sleep", lambda *_args, **_kwargs: None)
+ monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
- server._set_stage_device_env(1, env, "1,2,3", replica_id=1)
+ server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), ["--disable-log-stats"])
+ log_path = tmp_path / "stage0.log"
+ log_fh = open(log_path, "w", encoding="utf-8")
+ log_fh.write("partial launch\n")
+ log_fh.flush()
- assert env["CUDA_VISIBLE_DEVICES"] == "6"
+ proc = SimpleNamespace(pid=4242, poll=lambda: None)
+ killed: list[int] = []
+ monkeypatch.setattr(server, "_kill_process_tree", lambda pid: killed.append(pid))
+
+ def _fake_launch(stage_id: int, *, headless: bool, replica_id: int = 0) -> None:
+ stage_key = (stage_id, replica_id)
+ server.stage_procs[stage_key] = proc # type: ignore[assignment]
+ server._stage_log_paths[stage_key] = log_path
+ server._stage_log_files[stage_key] = log_fh
+ if stage_id == 0 and replica_id == 0:
+ server.proc = proc # type: ignore[assignment]
+
+ def _fail_alive() -> None:
+ raise RuntimeError("Stage 0 replica 0 exited with code 1")
+
+ monkeypatch.setattr(server, "_launch_stage", _fake_launch)
+ monkeypatch.setattr(server, "_ensure_stage_processes_alive", _fail_alive)
+
+ with pytest.raises(RuntimeError, match="exited with code 1"):
+ with server:
+ raise AssertionError("context body must not run after a failed enter")
+
+ assert killed == [4242]
+ assert log_fh.closed
+ assert not log_path.exists()
diff --git a/tests/helpers/tests/test_runtime_transcriber_release.py b/tests/helpers/tests/test_runtime_transcriber_release.py
new file mode 100644
index 000000000..5e7f29ebe
--- /dev/null
+++ b/tests/helpers/tests/test_runtime_transcriber_release.py
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""The server/runner fixtures free the Whisper judge around each instance.
+
+iter_omni_server and iter_omni_runner both wrap their body in the same
+_whisper_device_free_around() context manager, so one representative path
+exercises its enter/exit and exception behavior.
+"""
+
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers import runtime
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def events(monkeypatch) -> list[str]:
+ """Record construction and release in order."""
+ calls: list[str] = []
+
+ class FakeServer:
+ def __init__(self, *args, **kwargs):
+ calls.append("construct")
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ return False
+
+ monkeypatch.setattr(runtime, "OmniServer", FakeServer)
+ monkeypatch.setattr(runtime, "release_audio_transcriber", lambda: calls.append("release"))
+ monkeypatch.setattr(
+ "tests.helpers.stage_config.stage_config_path_for_run_level",
+ lambda path, run_level: None,
+ )
+ return calls
+
+
+def _server_generator():
+ request = SimpleNamespace(
+ param=OmniServerParams(model="fake-model"),
+ node=SimpleNamespace(get_closest_marker=lambda name: None),
+ )
+ return runtime.iter_omni_server(request, "core_model", threading.Lock())
+
+
+def test_releases_before_construction_and_after_teardown(events):
+ generator = _server_generator()
+ generator.send(None)
+
+ assert events == ["release", "construct"], "the device must be clear before the model loads"
+
+ with pytest.raises(StopIteration):
+ generator.send(None)
+
+ assert events == ["release", "construct", "release"]
+
+
+def test_failing_test_still_releases(events):
+ """A test that raises must not strand Whisper on the device."""
+ generator = _server_generator()
+ generator.send(None)
+
+ with pytest.raises(RuntimeError, match="assertion blew up"):
+ generator.throw(RuntimeError("assertion blew up"))
+
+ assert events[-1] == "release"
diff --git a/tests/helpers/tests/test_success_rate_helpers.py b/tests/helpers/tests/test_success_rate_helpers.py
new file mode 100644
index 000000000..c1b8f4417
--- /dev/null
+++ b/tests/helpers/tests/test_success_rate_helpers.py
@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers import assertions
+from tests.helpers.assertions import (
+ AUDIO_MISMATCH_MESSAGE,
+ SuccessRateGate,
+ assert_omni_response,
+ collect_at_success_rate,
+ is_quality_failure,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_GATE = SuccessRateGate(min_successes=7, concurrency=2)
+
+
+def _fake_response(text: str) -> SimpleNamespace:
+ return SimpleNamespace(success=True, text_content=text, audio_content=None, audio_bytes=b"fake-wav")
+
+
+def _omni_response_failure(monkeypatch, *, text: str, transcript: str) -> AssertionError:
+ """Drive the real assert_omni_response to failure and hand back what it raised."""
+ monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *a, **kw: transcript)
+ request_config = {"modalities": ["text", "audio"], "key_words": {"text": ["london"]}}
+ with pytest.raises(AssertionError) as excinfo:
+ assert_omni_response(_fake_response(text), request_config, run_level="full_model")
+ return excinfo.value
+
+
+def test_real_audio_mismatch_is_tolerated(monkeypatch):
+ # Ties the classifier to the actual producer: if assert_omni_response's wording drifts,
+ # this fails instead of the gate silently turning every quality failure into a hard one.
+ assert is_quality_failure(_omni_response_failure(monkeypatch, text="London", transcript="런던"))
+
+
+def test_real_keyword_failure_is_not_tolerated(monkeypatch):
+ assert not is_quality_failure(_omni_response_failure(monkeypatch, text="Paris", transcript="London"))
+
+
+@pytest.mark.parametrize("exc", [RuntimeError(AUDIO_MISMATCH_MESSAGE), ValueError("boom")])
+def test_non_assertion_failures_are_not_tolerated(exc):
+ # A transport error carrying the same text is still a serving failure, not a sample.
+ assert not is_quality_failure(exc)
+
+
+class _ScriptedSampler:
+ """One scripted outcome per call: raised if it is an exception, else returned."""
+
+ def __init__(self, outcomes: list) -> None:
+ self._outcomes = outcomes
+ self._lock = threading.Lock()
+ self.calls = 0
+
+ def __call__(self):
+ with self._lock:
+ outcome = self._outcomes[self.calls]
+ self.calls += 1
+ if isinstance(outcome, BaseException):
+ raise outcome
+ return outcome
+
+
+def _mismatch(detail: str) -> AssertionError:
+ return AssertionError(f"{AUDIO_MISMATCH_MESSAGE} (transcript={detail!r})")
+
+
+def test_tolerated_failures_are_counted_and_every_sample_is_sent():
+ sampler = _ScriptedSampler(["ok"] * 7 + [_mismatch("a"), _mismatch("b"), _mismatch("c")])
+
+ responses = collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
+
+ # No early exit once the gate is already satisfied: it would bias the measured rate.
+ assert sampler.calls == 10
+ assert responses == ["ok"] * 7
+
+
+def test_gate_failure_reports_the_count_and_every_failing_sample():
+ outcomes = ["ok"] * 6 + [_mismatch("first"), _mismatch("second"), _mismatch("third"), _mismatch("fourth")]
+
+ with pytest.raises(AssertionError) as excinfo:
+ collect_at_success_rate(_ScriptedSampler(outcomes), _GATE, request_num=10, report=lambda _: None)
+
+ message = str(excinfo.value)
+ assert "6/10" in message and "7/10 gate" in message
+ assert all(detail in message for detail in ("first", "second", "third", "fourth"))
+
+
+def test_serving_failure_is_not_absorbed_into_the_budget():
+ sampler = _ScriptedSampler(["ok", AssertionError("No audio output is generated"), *["ok"] * 8])
+
+ with pytest.raises(AssertionError, match="No audio output"):
+ collect_at_success_rate(sampler, _GATE, request_num=10, report=lambda _: None)
+
+ # Stopped on the serving failure rather than spending the remaining samples.
+ assert sampler.calls < 10
+
+
+def test_unreachable_gate_is_rejected_before_sending():
+ sampler = _ScriptedSampler(["ok"] * 3)
+
+ with pytest.raises(ValueError, match="can never pass"):
+ collect_at_success_rate(sampler, _GATE, request_num=3, report=lambda _: None)
+
+ assert sampler.calls == 0
diff --git a/tests/host_weight_runtime/test_contracts.py b/tests/host_weight_runtime/test_contracts.py
new file mode 100644
index 000000000..5c7036c8e
--- /dev/null
+++ b/tests/host_weight_runtime/test_contracts.py
@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Contract tests for semantic identity and immutable artifact schemas."""
+
+from __future__ import annotations
+
+import ast
+import math
+from pathlib import Path
+from typing import get_type_hints
+
+import pytest
+
+from vllm_omni.host_weight_runtime import (
+ AdaptationIdentity,
+ CanonicalJson,
+ ComponentIdentity,
+ FileManifestEntry,
+ IdentityValidationError,
+ ManifestValidationError,
+ ProducerIdentity,
+ ProductionPolicy,
+ PublicationMarker,
+ RuntimeWeightLayout,
+ TensorKind,
+ TensorManifestEntry,
+ WeightArtifactIdentity,
+ WeightManifest,
+ WeightRepresentation,
+ WeightRestorePlan,
+ WeightRestorer,
+ WeightSourceIdentity,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_SHA_A = "a" * 64
+_SHA_B = "b" * 64
+_SHA_C = "c" * 64
+
+
+def _identity(
+ *,
+ tp_size: int = 1,
+ tp_rank: int = 0,
+ sp_size: int = 1,
+ adaptation: AdaptationIdentity | None = None,
+ source_metadata: CanonicalJson | None = None,
+) -> WeightArtifactIdentity:
+ return WeightArtifactIdentity(
+ schema_version=1,
+ source=WeightSourceIdentity(
+ model_id="test-org/test-model",
+ revision="0123456789abcdef",
+ fingerprint="source-files-sha256",
+ metadata=source_metadata or CanonicalJson.empty(),
+ ),
+ component=ComponentIdentity(name="transformer", ownership="diffusion.dit"),
+ representation=WeightRepresentation(name="runtime-bf16", dtype="torch.bfloat16"),
+ layout=RuntimeWeightLayout(
+ name="final-module-layout",
+ tensor_parallel_size=tp_size,
+ tensor_parallel_rank=tp_rank,
+ sequence_parallel_size=sp_size,
+ sequence_parallel_backend="ulysses" if sp_size > 1 else "none",
+ ),
+ adaptation=adaptation or AdaptationIdentity(),
+ producer=ProducerIdentity(
+ producer_id="test.final-layout",
+ version="1",
+ implementation_fingerprint="implementation-sha256",
+ manifest_schema="test-producer-v1",
+ restorer_schema="test-restorer-v1",
+ ),
+ )
+
+
+def test_canonical_metadata_is_a_stable_immutable_snapshot() -> None:
+ items: list[object] = [1, {"value": 2}]
+ mutable = {"z": items, "a": True}
+ snapshot = CanonicalJson.from_value(mutable)
+ items.append(3)
+
+ assert snapshot.encoded == b'{"a":true,"z":[1,{"value":2}]}'
+ assert snapshot.to_value() == {"a": True, "z": [1, {"value": 2}]}
+ assert CanonicalJson.from_value({"a": True, "z": [1, {"value": 2}]}) == snapshot
+
+
+@pytest.mark.parametrize("value", [math.inf, math.nan, object(), {1: "non-string-key"}])
+def test_canonical_metadata_rejects_process_specific_or_non_json_values(value: object) -> None:
+ with pytest.raises(IdentityValidationError):
+ CanonicalJson.from_value(value)
+
+
+@pytest.mark.parametrize("encoded", [b"not-json", b'{"b":1, "a":2}', b'{"value":NaN}'])
+def test_canonical_metadata_cannot_bypass_stable_encoding(encoded: bytes) -> None:
+ with pytest.raises(IdentityValidationError):
+ CanonicalJson(encoded)
+
+
+def test_identity_round_trip_and_hash_are_deterministic() -> None:
+ first = _identity(source_metadata=CanonicalJson.from_value({"files": ["b", "a"], "size": 3}))
+ second = _identity(source_metadata=CanonicalJson.from_value({"size": 3, "files": ["b", "a"]}))
+
+ assert first == second
+ assert first.key == second.key
+ assert WeightArtifactIdentity.from_dict(first.to_dict()) == first
+ assert len(first.key) == 64
+
+
+def test_identity_separates_semantic_parallel_layout_and_static_adaptation() -> None:
+ base = _identity()
+ tp0 = _identity(tp_size=2, tp_rank=0)
+ tp1 = _identity(tp_size=2, tp_rank=1)
+ sp2 = _identity(sp_size=2)
+ merged_lora = _identity(adaptation=AdaptationIdentity(kind="merged-lora", fingerprint="adapter-sha256"))
+
+ assert len({item.key for item in (base, tp0, tp1, sp2, merged_lora)}) == 5
+ assert "data_parallel" not in base.layout.to_dict()
+ assert "sequence_parallel_rank" not in sp2.layout.to_dict()
+
+
+def test_identity_schema_rejects_unknown_fields() -> None:
+ document = _identity().to_dict()
+ document["transfer_mode"] = "registered"
+
+ with pytest.raises(IdentityValidationError, match="unexpected"):
+ WeightArtifactIdentity.from_dict(document)
+
+
+def test_identity_schema_does_not_coerce_field_types() -> None:
+ document = _identity().to_dict()
+ document["schema_version"] = "1"
+
+ with pytest.raises(IdentityValidationError, match="must be an integer"):
+ WeightArtifactIdentity.from_dict(document)
+
+
+def test_restoration_contract_has_a_validation_only_one_shot_commit_boundary() -> None:
+ assert get_type_hints(WeightRestorePlan.commit)["return"] is type(None)
+ assert "one-shot" in (WeightRestorePlan.__doc__ or "")
+ assert "without mutating" in (WeightRestorer.__doc__ or "")
+
+
+def test_post_load_publication_policy_is_explicitly_enabled() -> None:
+ policy = ProductionPolicy(allow_post_load_publish=True)
+
+ assert policy.allow_local_build
+ assert policy.allow_post_load_publish
+
+
+def test_manifest_and_publication_marker_round_trip_exactly() -> None:
+ identity = _identity()
+ manifest = WeightManifest(
+ schema_version=1,
+ artifact_key=identity.key,
+ identity=identity,
+ producer_schema="test-producer-v1",
+ restorer_schema="test-restorer-v1",
+ tensors=(
+ TensorManifestEntry(
+ name="layer.weight",
+ kind=TensorKind.PARAMETER,
+ role="weight",
+ file_name="weights.safetensors",
+ storage_key="layer.weight",
+ shape=(2, 4),
+ dtype="torch.bfloat16",
+ stride=(4, 1),
+ nbytes=16,
+ sha256=_SHA_A,
+ ),
+ ),
+ files=(FileManifestEntry(file_name="weights.safetensors", size=128, sha256=_SHA_B),),
+ artifact_content_sha256=_SHA_C,
+ format_metadata=CanonicalJson.from_value({"quantization": "none"}),
+ )
+ marker = PublicationMarker(
+ schema_version=1,
+ publication_id="publication-id",
+ manifest_sha256=_SHA_A,
+ artifact_content_sha256=_SHA_C,
+ integrity_mode="full_checksum",
+ )
+
+ assert WeightManifest.from_dict(manifest.to_dict()) == manifest
+ assert PublicationMarker.from_dict(marker.to_dict()) == marker
+
+ document = manifest.to_dict()
+ document["unknown"] = True
+ with pytest.raises(ManifestValidationError, match="unexpected"):
+ WeightManifest.from_dict(document)
+
+ malformed = manifest.to_dict()
+ tensors = malformed["tensors"]
+ assert isinstance(tensors, list) and isinstance(tensors[0], dict)
+ tensors[0]["nbytes"] = "16"
+ with pytest.raises(ManifestValidationError, match="must be an integer"):
+ WeightManifest.from_dict(malformed)
+
+
+@pytest.mark.parametrize("file_name", ["", ".", "..", "../weights.safetensors", "nested/weights.safetensors"])
+def test_manifest_rejects_unsafe_payload_names(file_name: str) -> None:
+ with pytest.raises(ManifestValidationError, match="safe relative basename"):
+ FileManifestEntry(file_name=file_name, size=1, sha256=_SHA_A)
+
+
+def test_package_source_keeps_the_neutral_dependency_boundary() -> None:
+ package_root = Path(__file__).resolve().parents[2] / "vllm_omni" / "host_weight_runtime"
+ assert package_root.is_dir()
+ forbidden = (
+ "vllm_omni.diffusion",
+ "vllm_omni.distributed",
+ "torch.cuda",
+ "torch.distributed",
+ )
+ violations: list[str] = []
+ for source_path in package_root.rglob("*.py"):
+ tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path))
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Import):
+ violations.extend(
+ f"{source_path.name}:{node.lineno}:{alias.name}"
+ for alias in node.names
+ if alias.name.startswith(forbidden)
+ )
+ elif isinstance(node, ast.ImportFrom) and node.level == 0:
+ module = node.module or ""
+ if module.startswith(forbidden) or (
+ module == "torch" and any(alias.name in {"cuda", "distributed"} for alias in node.names)
+ ):
+ violations.append(f"{source_path.name}:{node.lineno}:{module}")
+
+ assert violations == []
diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py
new file mode 100644
index 000000000..7fe716c5d
--- /dev/null
+++ b/tests/host_weight_runtime/test_filesystem_store.py
@@ -0,0 +1,2581 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Lifecycle and concurrency tests for the local filesystem store."""
+
+from __future__ import annotations
+
+import contextlib
+import errno
+import hashlib
+import json
+import multiprocessing as mp
+import os
+import shutil
+import stat
+import tempfile
+import threading
+import time
+import warnings
+from collections.abc import Sequence
+from dataclasses import replace
+from multiprocessing.process import BaseProcess
+from pathlib import Path
+
+import pytest
+import torch
+
+from vllm_omni.host_weight_runtime import (
+ AdaptationIdentity,
+ ArtifactInventoryState,
+ BuildRequest,
+ CapacityPolicy,
+ ComponentIdentity,
+ CoordinationScope,
+ FailureCode,
+ HostWeightError,
+ HostWeightLease,
+ IntegrityPolicy,
+ ProducerIdentity,
+ ProductionMetadata,
+ ProductionSourceMode,
+ ResolutionStage,
+ RuntimeWeightLayout,
+ StorageClass,
+ StorageDomainPolicy,
+ StorageScope,
+ StoreResult,
+ StoreStatus,
+ TensorKind,
+ TensorWriteSpec,
+ ValidationLevel,
+ WeightArtifactIdentity,
+ WeightProductionSpec,
+ WeightRepresentation,
+ WeightSourceIdentity,
+)
+from vllm_omni.host_weight_runtime.filesystem import (
+ FilesystemHostWeightStore,
+ detect_filesystem_type,
+ detect_storage_class,
+)
+from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module
+from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module
+from vllm_omni.host_weight_runtime.filesystem.locks import FileLock
+from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Spawning imports the complete vLLM package and exceeded 10 seconds on a loaded
+# CI worker. Keep process-bootstrap bounds separate from the short store
+# deadlines exercised by these tests.
+_PROCESS_START_TIMEOUT_SECONDS = 60.0
+_PROCESS_COMPLETION_TIMEOUT_SECONDS = 60.0
+_PROCESS_CLEANUP_TIMEOUT_SECONDS = 5.0
+
+
+def _identity(*, tp_rank: int = 0) -> WeightArtifactIdentity:
+ return WeightArtifactIdentity(
+ schema_version=1,
+ source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"),
+ component=ComponentIdentity("transformer", "diffusion.dit"),
+ representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"),
+ layout=RuntimeWeightLayout(
+ "final-module-layout",
+ tensor_parallel_size=2,
+ tensor_parallel_rank=tp_rank,
+ ),
+ adaptation=AdaptationIdentity(),
+ producer=ProducerIdentity(
+ "test.final-layout",
+ "1",
+ "implementation-sha256",
+ "test-producer-v1",
+ "test-restorer-v1",
+ ),
+ )
+
+
+def _weight(offset: int = 0) -> torch.Tensor:
+ return torch.arange(offset, offset + 12, dtype=torch.float32).to(torch.bfloat16).reshape(4, 3)
+
+
+def _specs() -> tuple[TensorWriteSpec, ...]:
+ return (
+ TensorWriteSpec("layer.weight", (4, 3), torch.bfloat16),
+ TensorWriteSpec("layer.scale", (2,), torch.float32, kind=TensorKind.BUFFER, role="scale"),
+ )
+
+
+def _open_fd_targets() -> tuple[str, ...]:
+ targets = []
+ for path in Path("/proc/self/fd").iterdir():
+ with contextlib.suppress(OSError):
+ targets.append(os.readlink(path))
+ return tuple(targets)
+
+
+def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter:
+ staging = tmp_path / "staging"
+ staging.mkdir()
+ return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None)
+
+
+def test_ordered_writer_hashes_payload_during_writes(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ writer = _direct_writer(tmp_path)
+ weight = _weight()
+ scale = torch.tensor([1.0, 2.0])
+
+ def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object:
+ raise AssertionError("ordered payload must not be rescanned")
+
+ monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan)
+ with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+ # The writer's ordered contract follows canonical storage-key order.
+ output.write_tensor("layer.scale", scale)
+ output.write_rows("layer.weight", 0, weight[:2])
+ output.write_rows("layer.weight", 2, weight[2:])
+
+ manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+ payload = writer.staging_dir / "weights.safetensors"
+ assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest()
+ expected_tensor_digests = {
+ "layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(),
+ "layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(),
+ }
+ assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests
+
+
+def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None:
+ writer = _direct_writer(tmp_path)
+ output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+
+ with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"):
+ output.write_tensor("layer.weight", _weight())
+
+ output.abort()
+ assert not list(writer.staging_dir.iterdir())
+
+
+def test_payload_fsync_overlaps_producer_work(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ writer = _direct_writer(tmp_path)
+ started = threading.Event()
+ release = threading.Event()
+ closed = threading.Event()
+ original_fsync = filesystem_writer_module._fsync_payload
+
+ def blocking_fsync(fd: int) -> None:
+ started.set()
+ if not release.wait(timeout=5):
+ raise TimeoutError("test did not release payload fsync")
+ original_fsync(fd)
+
+ monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync)
+ output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True)
+ output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+ output.write_tensor("layer.weight", _weight())
+ close_errors: list[BaseException] = []
+
+ def close_output() -> None:
+ try:
+ output.close()
+ except BaseException as exc:
+ close_errors.append(exc)
+ finally:
+ closed.set()
+
+ close_thread = threading.Thread(target=close_output)
+ close_thread.start()
+ try:
+ assert started.wait(timeout=2)
+ assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it"
+ finally:
+ release.set()
+ close_thread.join(timeout=5)
+
+ assert not close_errors
+ writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+
+def test_payload_fsync_failure_prevents_publication(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ writer = _direct_writer(tmp_path)
+
+ def failing_fsync(fd: int) -> None:
+ try:
+ raise OSError(errno.EIO, "injected fsync failure")
+ finally:
+ os.close(fd)
+
+ monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync)
+ with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output:
+ output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+ output.write_tensor("layer.weight", _weight())
+
+ with pytest.raises(OSError, match="injected fsync failure"):
+ writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+ assert not (writer.staging_dir / "manifest.json").exists()
+ assert not (writer.staging_dir / "READY.json").exists()
+
+
+def test_unordered_payload_fallback_hashes_files_in_parallel(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ writer = _direct_writer(tmp_path)
+ specs = (
+ TensorWriteSpec("a", (2,), torch.float32),
+ TensorWriteSpec("b", (2,), torch.float32),
+ )
+ for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True):
+ with writer.open_tensor_file(name, (spec,)) as output:
+ output.write_tensor(spec.name, torch.tensor([1.0, 2.0]))
+
+ original_scan = filesystem_writer_module._scan_payload
+ barrier = threading.Barrier(2)
+ thread_ids: set[int] = set()
+ thread_ids_lock = threading.Lock()
+
+ def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object:
+ with thread_ids_lock:
+ thread_ids.add(threading.get_ident())
+ barrier.wait(timeout=5)
+ return original_scan(path, file_specs)
+
+ monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan)
+ manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1"))
+
+ assert len(manifest.files) == 2
+ assert len(thread_ids) == 2
+
+
+class FakeProducer:
+ def __init__(
+ self,
+ identity: WeightArtifactIdentity,
+ *,
+ counter_path: Path | None = None,
+ delay_seconds: float = 0.0,
+ started: object | None = None,
+ release: object | None = None,
+ write_mode: str = "rows",
+ weight_offset: int = 0,
+ coordination_scope: CoordinationScope = CoordinationScope.SINGLE_PROCESS,
+ ) -> None:
+ self._spec = WeightProductionSpec(
+ producer_id="test.final-layout",
+ outputs=(identity,),
+ source_mode=ProductionSourceMode.FINALIZED_MODEL,
+ coordination_scope=coordination_scope,
+ )
+ self.counter_path = counter_path
+ self.delay_seconds = delay_seconds
+ self.started = started
+ self.release = release
+ self.write_mode = write_mode
+ self.weight_offset = weight_offset
+
+ @property
+ def spec(self) -> WeightProductionSpec:
+ return self._spec
+
+ def produce(self, writer: object) -> ProductionMetadata:
+ if self.counter_path is not None:
+ fd = os.open(self.counter_path, os.O_CREAT | os.O_APPEND | os.O_WRONLY | os.O_CLOEXEC, 0o600)
+ try:
+ os.write(fd, f"{os.getpid()}\n".encode())
+ os.fsync(fd)
+ finally:
+ os.close(fd)
+ if self.started is not None:
+ self.started.set() # type: ignore[attr-defined]
+ if self.release is not None and not self.release.wait(timeout=10): # type: ignore[attr-defined]
+ raise TimeoutError("test producer was not released")
+ if self.delay_seconds:
+ time.sleep(self.delay_seconds)
+
+ with writer.open_tensor_file("weights.safetensors", _specs()) as output: # type: ignore[attr-defined]
+ if self.write_mode == "rows":
+ output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2])
+ output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:])
+ elif self.write_mode == "complete":
+ output.write_tensor("layer.weight", _weight(self.weight_offset))
+ elif self.write_mode == "incomplete":
+ output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:2])
+ elif self.write_mode == "overlap":
+ output.write_rows("layer.weight", 0, _weight(self.weight_offset)[:3])
+ output.write_rows("layer.weight", 2, _weight(self.weight_offset)[2:])
+ else: # pragma: no cover - test helper guard
+ raise AssertionError(f"unknown write mode {self.write_mode}")
+ output.write_tensor("layer.scale", torch.tensor([1.0, 2.0]))
+ return ProductionMetadata("test-producer-v1", "test-restorer-v1")
+
+
+class CrashProducer(FakeProducer):
+ def produce(self, writer: object) -> ProductionMetadata:
+ output = writer.open_tensor_file("weights.safetensors", _specs()) # type: ignore[attr-defined]
+ output.write_rows("layer.weight", 0, _weight()[:2])
+ os._exit(17)
+
+
+class NonCooperativeStore(FilesystemHostWeightStore):
+ """Test double representing a publisher that does not join the build lock."""
+
+ def _build_lock_path(self, key: str) -> Path:
+ return self.locks_dir / f"{key}.noncooperative-build.lock"
+
+ def _remove_stale_temps_locked(self, _key: str) -> None:
+ return
+
+
+class InitialLookupBarrierStore(FilesystemHostWeightStore):
+ """Make both test workers observe the initial miss before either builds."""
+
+ def __init__(
+ self,
+ domain: StorageDomainPolicy,
+ capacity: CapacityPolicy,
+ integrity: IntegrityPolicy,
+ *,
+ initial_lookup_barrier: object,
+ ) -> None:
+ super().__init__(domain, capacity, integrity)
+ self._initial_lookup_barrier = initial_lookup_barrier
+ self._lookup_count = 0
+
+ def lookup(
+ self,
+ identity: WeightArtifactIdentity,
+ *,
+ validation: ValidationLevel,
+ deadline: float | None = None,
+ ) -> StoreResult:
+ result = super().lookup(identity, validation=validation, deadline=deadline)
+ self._lookup_count += 1
+ if self._lookup_count == 1:
+ self._initial_lookup_barrier.wait(timeout=10) # type: ignore[attr-defined]
+ return result
+
+
+class PausingOpenLeaseStore(FilesystemHostWeightStore):
+ """Pause after opening a lease so another lookup can publish a deny marker."""
+
+ def __init__(
+ self,
+ domain: StorageDomainPolicy,
+ capacity: CapacityPolicy,
+ integrity: IntegrityPolicy,
+ *,
+ lease_opened: threading.Event,
+ resume_lookup: threading.Event,
+ ) -> None:
+ super().__init__(domain, capacity, integrity)
+ self._lease_opened = lease_opened
+ self._resume_lookup = resume_lookup
+
+ def _open_lease(
+ self,
+ identity: WeightArtifactIdentity,
+ entry_path: Path,
+ artifact_lock: FileLock,
+ validation: ValidationLevel,
+ ) -> HostWeightLease:
+ lease = super()._open_lease(identity, entry_path, artifact_lock, validation)
+ self._lease_opened.set()
+ if not self._resume_lookup.wait(5):
+ lease.close()
+ raise TimeoutError("test did not resume paused lookup")
+ return lease
+
+
+def _make_store(
+ root: Path,
+ *,
+ capacity: CapacityPolicy | None = None,
+ integrity: IntegrityPolicy | None = None,
+) -> FilesystemHostWeightStore:
+ root.mkdir(parents=True, exist_ok=True, mode=0o700)
+ domain = StorageDomainPolicy(root=root, storage_class=detect_storage_class(root))
+ return FilesystemHostWeightStore(domain, capacity or CapacityPolicy(), integrity or IntegrityPolicy())
+
+
+def _publish_test_artifact(
+ store: FilesystemHostWeightStore,
+ identity: WeightArtifactIdentity | None = None,
+) -> tuple[WeightArtifactIdentity, Path]:
+ identity = identity or _identity()
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert result.status is StoreStatus.BUILT
+ assert result.lease is not None
+ result.lease.close()
+ return identity, store.artifacts_dir / identity.key
+
+
+def _corrupt_test_payload(artifact: Path) -> None:
+ payload = artifact / "weights.safetensors"
+ os.chmod(artifact, 0o755)
+ os.chmod(payload, 0o644)
+ with payload.open("r+b") as handle:
+ handle.seek(-1, os.SEEK_END)
+ value = handle.read(1)
+ handle.seek(-1, os.SEEK_END)
+ handle.write(bytes([value[0] ^ 0xFF]))
+ handle.flush()
+ os.fsync(handle.fileno())
+ os.chmod(payload, 0o444)
+ os.chmod(artifact, 0o555)
+
+
+def _build_process(root: str, counter: str, initial_lookup_barrier: object, result_queue: object) -> None:
+ try:
+ identity = _identity()
+ root_path = Path(root)
+ root_path.mkdir(parents=True, exist_ok=True, mode=0o700)
+ store = InitialLookupBarrierStore(
+ StorageDomainPolicy(root=root_path, storage_class=detect_storage_class(root_path)),
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ initial_lookup_barrier=initial_lookup_barrier,
+ )
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, counter_path=Path(counter), delay_seconds=0.25),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 15,
+ )
+ fingerprint = result.lease.provenance.backing_fingerprint if result.lease is not None else None
+ if result.lease is not None:
+ result.lease.close()
+ result_queue.put((result.status.value, fingerprint, None)) # type: ignore[attr-defined]
+ except BaseException as exc:
+ result_queue.put(("error", None, repr(exc))) # type: ignore[attr-defined]
+
+
+def _controlled_build_process(root: str, started: object, release: object, result_queue: object) -> None:
+ try:
+ identity = _identity()
+ store = _make_store(Path(root))
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, started=started, release=release),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 15,
+ )
+ if result.lease is not None:
+ result.lease.close()
+ result_queue.put((result.status.value, None)) # type: ignore[attr-defined]
+ except BaseException as exc:
+ result_queue.put(("error", repr(exc))) # type: ignore[attr-defined]
+
+
+def _pausing_weak_lookup_process(root: str, opened: object, resume: object, result_queue: object) -> None:
+ try:
+ policy = StorageDomainPolicy(root=Path(root), storage_class=detect_storage_class(Path(root)))
+ store = PausingOpenLeaseStore(
+ policy,
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ lease_opened=opened, # type: ignore[arg-type]
+ resume_lookup=resume, # type: ignore[arg-type]
+ )
+ result = store.lookup(_identity(), validation=ValidationLevel.MANIFEST_AND_METADATA)
+ if result.lease is not None:
+ result.lease.close()
+ code = result.failure.code.value if result.failure is not None else None
+ result_queue.put((result.status.value, code, None)) # type: ignore[attr-defined]
+ except BaseException as exc:
+ result_queue.put(("error", None, repr(exc))) # type: ignore[attr-defined]
+
+
+def _crash_build_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+ store.get_or_build(
+ BuildRequest(identity),
+ CrashProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 15,
+ )
+
+
+def _crash_cleanup_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+
+ def crash_before_removal(path: Path) -> None:
+ if path.parent == store.quarantine_dir and ".cleanup." in path.name:
+ os._exit(23)
+ FilesystemHostWeightStore._remove_tree(path)
+
+ store._remove_tree = crash_before_removal # type: ignore[method-assign]
+ store.cleanup(identity)
+ raise AssertionError("cleanup did not reach the injected crash point")
+
+
+def _crash_cleanup_before_rename_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+ artifact = store.artifacts_dir / identity.key
+ original_replace = os.replace
+
+ def crash_before_rename(
+ source: str | os.PathLike[str],
+ destination: str | os.PathLike[str],
+ ) -> None:
+ if Path(source) == artifact and ".cleanup." in Path(destination).name:
+ os._exit(31)
+ original_replace(source, destination)
+
+ filesystem_store_module.os.replace = crash_before_rename
+ store.cleanup(identity)
+ raise AssertionError("cleanup did not reach the pre-rename crash point")
+
+
+def _crash_quarantine_after_rename_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+ original_harden = filesystem_store_module._harden_open_directory
+
+ def crash_before_destination_hardening(fd: int, path: Path, mode: int) -> None:
+ if path.parent == store.quarantine_dir:
+ os._exit(37)
+ original_harden(fd, path, mode)
+
+ filesystem_store_module._harden_open_directory = crash_before_destination_hardening
+ with FileLock(store._artifact_lock_path(identity.key), exclusive=True, deadline=None):
+ store._quarantine_locked(identity.key, reason="validation_failure")
+ raise AssertionError("quarantine did not reach the post-rename crash point")
+
+
+def _crash_after_artifacts_fsync_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+ original_fsync_directory = filesystem_store_module._fsync_directory
+
+ def crash_after_artifacts_fsync(path: Path) -> None:
+ original_fsync_directory(path)
+ if path == store.artifacts_dir:
+ os._exit(29)
+
+ filesystem_store_module._fsync_directory = crash_after_artifacts_fsync
+ store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 15,
+ )
+ raise AssertionError("publication did not reach the injected crash point")
+
+
+def _crash_after_publication_hardening_process(root: str) -> None:
+ identity = _identity()
+ store = _make_store(Path(root))
+ artifact = store.artifacts_dir / identity.key
+ original_chmod = os.chmod
+
+ def crash_after_hardening(
+ path: str | os.PathLike[str],
+ mode: int,
+ *args: object,
+ **kwargs: object,
+ ) -> None:
+ original_chmod(path, mode, *args, **kwargs) # type: ignore[arg-type]
+ if Path(path) == artifact and mode == 0o555:
+ os._exit(41)
+
+ filesystem_store_module.os.chmod = crash_after_hardening
+ store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 15,
+ )
+ raise AssertionError("publication did not reach the post-hardening crash point")
+
+
+def _join_processes(
+ processes: Sequence[BaseProcess],
+ timeout: float = _PROCESS_COMPLETION_TIMEOUT_SECONDS,
+) -> None:
+ deadline = time.monotonic() + timeout
+ for process in processes:
+ process.join(max(0.0, deadline - time.monotonic()))
+ timed_out = [process for process in processes if process.is_alive()]
+ if not timed_out:
+ return
+ timed_out_workers = ", ".join(f"{process.name}(pid={process.pid})" for process in timed_out)
+ for process in timed_out:
+ process.terminate()
+ for process in timed_out:
+ process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS)
+ survivors = [process for process in timed_out if process.is_alive()]
+ for process in survivors:
+ process.kill()
+ process.join(_PROCESS_CLEANUP_TIMEOUT_SECONDS)
+ unstoppable = [process for process in survivors if process.is_alive()]
+ assert not unstoppable, "timed-out test worker processes could not be stopped"
+ pytest.fail(f"test worker processes exceeded the {timeout:g}-second group deadline: {timed_out_workers}")
+
+
+def test_build_lookup_share_one_published_backing_and_lease_blocks_cleanup(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.status is StoreStatus.BUILT
+ assert built.lease is not None
+ assert torch.equal(built.lease.tensors["layer.weight"], _weight())
+ assert torch.equal(built.lease.tensors["layer.scale"], torch.tensor([1.0, 2.0]))
+ assert built.lease.mapped_regions
+ assert all(region.address > 0 and region.length > 0 and region.read_only for region in built.lease.mapped_regions)
+ with pytest.raises(TypeError, match="process-local"):
+ built.lease.__reduce_ex__(4)
+
+ second = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert second.status is StoreStatus.HIT
+ assert second.lease is not None
+ assert second.lease.provenance.publication_id == built.lease.provenance.publication_id
+ assert second.lease.provenance.backing_fingerprint == built.lease.provenance.backing_fingerprint
+ payload = store.artifacts_dir / identity.key / "weights.safetensors"
+ assert any(str(payload) in target for target in _open_fd_targets())
+
+ artifact = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert artifact.status is StoreStatus.HIT
+ assert artifact.manifest is not None and artifact.manifest.identity == identity
+ assert artifact.publication_id == built.lease.provenance.publication_id
+ assert artifact.artifact_lock_active
+ domain = store.inspect_domain()
+ assert domain.domain_uuid == built.lease.provenance.domain_uuid
+ assert domain.filesystem_type == detect_filesystem_type(store.root)
+ assert any(
+ item.artifact_key == identity.key and item.state is ArtifactInventoryState.READY and item.size_bytes > 0
+ for item in domain.inventory
+ )
+
+ active = store.cleanup(identity)
+ assert active is not None and active.code is FailureCode.ACTIVE_LEASE
+ second.lease.close()
+ built.lease.close()
+ built.lease.close()
+ assert not built.lease.tensors
+ assert not any(str(payload) in target for target in _open_fd_targets())
+
+ assert store.cleanup(identity) is None
+ assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS
+
+
+def test_publication_hardens_artifact_after_atomic_rename(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert built.status is StoreStatus.BUILT
+ assert built.lease is not None
+ built.lease.close()
+ artifact_dir = store.artifacts_dir / identity.key
+ assert artifact_dir.stat().st_mode & 0o777 == 0o555
+
+
+@pytest.mark.parametrize(
+ ("expected_filesystem", "filesystem_root"),
+ [("overlay", Path("/tmp")), ("tmpfs", Path("/dev/shm"))],
+)
+def test_hardened_artifact_cleanup_on_supported_filesystem(
+ expected_filesystem: str,
+ filesystem_root: Path,
+) -> None:
+ if not filesystem_root.is_dir():
+ pytest.skip(f"{filesystem_root} is unavailable")
+ try:
+ filesystem_type = detect_filesystem_type(filesystem_root)
+ except OSError as exc:
+ pytest.skip(f"cannot inspect {filesystem_root} filesystem type: {exc}")
+ if filesystem_type != expected_filesystem:
+ pytest.skip(f"{filesystem_root} is {filesystem_type!r}, not {expected_filesystem}")
+ if not os.access(filesystem_root, os.W_OK | os.X_OK):
+ pytest.skip(f"{filesystem_root} is not writable")
+
+ with tempfile.TemporaryDirectory(prefix="vllm-omni-hwr-", dir=filesystem_root) as directory:
+ store = _make_store(Path(directory) / "store")
+ identity, artifact = _publish_test_artifact(store)
+ assert artifact.stat().st_mode & 0o777 == 0o555
+
+ assert store.cleanup(identity) is None
+ assert not artifact.exists()
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parametrize("failure_point", ["chmod", "rename"])
+def test_failed_cleanup_move_preserves_ready_immutability(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_point: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ artifact_inode = artifact.stat().st_ino
+
+ if failure_point == "chmod":
+ original_fchmod = os.fchmod
+
+ def fail_make_writable(fd: int, mode: int) -> None:
+ if os.fstat(fd).st_ino == artifact_inode and mode & stat.S_IWUSR:
+ raise OSError(errno.EIO, "injected artifact chmod failure")
+ original_fchmod(fd, mode)
+
+ monkeypatch.setattr(os, "fchmod", fail_make_writable)
+ else:
+ original_replace = os.replace
+
+ def fail_artifact_rename(
+ source: str | os.PathLike[str],
+ destination: str | os.PathLike[str],
+ ) -> None:
+ if Path(source) == artifact:
+ raise OSError(errno.EIO, "injected artifact rename failure")
+ original_replace(source, destination)
+
+ monkeypatch.setattr(os, "replace", fail_artifact_rename)
+
+ failure = store.cleanup(identity)
+
+ assert failure is not None
+ assert failure.code is FailureCode.QUARANTINE_FAILED
+ assert failure.retryable
+ assert artifact.stat().st_mode & 0o777 == 0o555
+ assert not list(store.quarantine_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+def test_move_artifact_one_shot_restore_failure_still_hardens_destination(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ destination = store.quarantine_dir / f"{identity.key}.direct-test"
+ artifact_inode = artifact.stat().st_ino
+ original_fchmod = os.fchmod
+ injected = False
+
+ def fail_first_hardening(fd: int, mode: int) -> None:
+ nonlocal injected
+ if os.fstat(fd).st_ino == artifact_inode and destination.exists() and not mode & 0o222 and not injected:
+ injected = True
+ raise OSError(errno.EIO, "injected artifact hardening failure")
+ original_fchmod(fd, mode)
+
+ monkeypatch.setattr(os, "fchmod", fail_first_hardening)
+ with FileLock(store._artifact_lock_path(identity.key), exclusive=True, deadline=None):
+ filesystem_store_module._move_artifact_locked(artifact, destination)
+
+ assert injected
+ assert not artifact.exists()
+ assert destination.is_dir()
+ assert not destination.stat().st_mode & 0o222
+
+
+@pytest.mark.parametrize("failing_parent", ["artifacts", "quarantine"])
+def test_cleanup_move_fsync_failure_leaves_retryable_immutable_tombstone(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failing_parent: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ original_fsync_directory = filesystem_store_module._fsync_directory
+ failed_path = store.artifacts_dir if failing_parent == "artifacts" else store.quarantine_dir
+
+ def fail_parent_fsync(path: Path) -> None:
+ if path == failed_path:
+ raise OSError(errno.EIO, "injected cleanup fsync failure")
+ original_fsync_directory(path)
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_parent_fsync)
+ failure = store.cleanup(identity)
+
+ assert failure is not None
+ assert failure.code is FailureCode.QUARANTINE_FAILED
+ assert failure.retryable
+ assert not artifact.exists()
+ tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+ assert len(tombstones) == 1
+ assert tombstones[0].stat().st_mode & 0o777 == 0o555
+ assert tombstones[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", original_fsync_directory)
+ assert store.cleanup(identity) is None
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+def test_cleanup_retries_removal_failure_tombstone(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ original_rmtree = shutil.rmtree
+
+ def fail_cleanup_removal(
+ path: str | os.PathLike[str],
+ ) -> None:
+ if Path(path).parent == store.quarantine_dir and ".cleanup." in Path(path).name:
+ raise OSError(errno.EIO, "injected cleanup removal failure")
+ original_rmtree(path)
+
+ monkeypatch.setattr(shutil, "rmtree", fail_cleanup_removal)
+ failure = store.cleanup(identity)
+
+ assert failure is not None
+ assert failure.code is FailureCode.QUARANTINE_FAILED
+ assert failure.retryable
+ assert not artifact.exists()
+ tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+ assert len(tombstones) == 1
+
+ monkeypatch.setattr(shutil, "rmtree", original_rmtree)
+ assert store.cleanup(identity) is None
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parallel
+def test_cleanup_after_move_crash_tombstone_is_removed_by_next_build(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity, artifact = _publish_test_artifact(store)
+ process = mp.get_context("spawn").Process(target=_crash_cleanup_process, args=(str(root),))
+
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 23
+ assert not artifact.exists()
+ tombstones = list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+ assert len(tombstones) == 1
+ assert not tombstones[0].stat().st_mode & 0o222
+
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+
+
+@pytest.mark.parallel
+def test_pre_rename_cleanup_crash_leaves_persistently_denied_artifact(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity, artifact = _publish_test_artifact(store)
+ process = mp.get_context("spawn").Process(
+ target=_crash_cleanup_before_rename_process,
+ args=(str(root),),
+ )
+
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 31
+ assert artifact.is_dir()
+ assert artifact.stat().st_mode & stat.S_IWUSR
+ assert artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ weak_store = _make_store(
+ root,
+ integrity=IntegrityPolicy(require_trusted_permissions=False),
+ )
+ denied = weak_store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)
+ assert denied.status is StoreStatus.INVALID
+ assert denied.failure is not None and denied.failure.code is FailureCode.ARTIFACT_DENIED
+
+ rebuilt = weak_store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert not quarantined[0].stat().st_mode & 0o222
+ assert not quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+
+
+@pytest.mark.parallel
+def test_post_rename_quarantine_crash_is_rehardened_by_next_build(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity, artifact = _publish_test_artifact(store)
+ process = mp.get_context("spawn").Process(
+ target=_crash_quarantine_after_rename_process,
+ args=(str(root),),
+ )
+
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 37
+ assert not artifact.exists()
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert quarantined[0].stat().st_mode & stat.S_IWUSR
+ assert quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert not quarantined[0].stat().st_mode & 0o222
+ assert not quarantined[0].stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+
+
+@pytest.mark.parametrize("failure_point", ["chmod", "artifact_inode_fsync", "artifacts_fsync"])
+def test_post_rename_failure_quarantines_non_authoritative_publication(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_point: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity = _identity()
+ original_fsync_directory = filesystem_store_module._fsync_directory
+ injected = False
+
+ if failure_point == "chmod":
+ original_chmod = os.chmod
+
+ def fail_publication_chmod(
+ path: str | os.PathLike[str],
+ mode: int,
+ *args: object,
+ **kwargs: object,
+ ) -> None:
+ nonlocal injected
+ if Path(path) == store.artifacts_dir / identity.key and mode == 0o555 and not injected:
+ injected = True
+ raise OSError(errno.EIO, "injected publication chmod failure")
+ original_chmod(path, mode, *args, **kwargs) # type: ignore[arg-type]
+
+ monkeypatch.setattr(os, "chmod", fail_publication_chmod)
+ else:
+ failed_path = (
+ store.artifacts_dir / identity.key if failure_point == "artifact_inode_fsync" else store.artifacts_dir
+ )
+
+ def fail_publication_fsync(path: Path) -> None:
+ nonlocal injected
+ if path == failed_path and not injected:
+ injected = True
+ raise OSError(errno.EIO, f"injected {failure_point} publication failure")
+ original_fsync_directory(path)
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_publication_fsync)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert injected
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is ResolutionStage.LIFECYCLE
+ assert result.failure.code is FailureCode.PUBLICATION_FAILED
+ assert not result.failure.retryable
+ assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS
+ quarantined = list(store.quarantine_dir.iterdir())
+ assert len(quarantined) == 1
+ assert ".publication_failure." in quarantined[0].name
+ assert quarantined[0].stat().st_mode & 0o777 == 0o555
+ assert not list(store.tmp_dir.iterdir())
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", original_fsync_directory)
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+def test_uncontained_publication_failure_is_retryable_and_outcome_uncertain(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity = _identity()
+ artifact = store.artifacts_dir / identity.key
+ original_fsync_directory = filesystem_store_module._fsync_directory
+ original_open = os.open
+ publication_failed = False
+
+ def fail_artifact_inode_fsync(path: Path) -> None:
+ nonlocal publication_failed
+ if path == artifact:
+ publication_failed = True
+ raise OSError(errno.EIO, "injected publication inode fsync failure")
+ original_fsync_directory(path)
+
+ def fail_containment_open(
+ path: str | bytes | os.PathLike[str] | os.PathLike[bytes],
+ flags: int,
+ mode: int = 0o777,
+ *,
+ dir_fd: int | None = None,
+ ) -> int:
+ if publication_failed and Path(os.fsdecode(path)) == artifact:
+ raise OSError(errno.EIO, "injected persistent containment failure")
+ return original_open(path, flags, mode, dir_fd=dir_fd)
+
+ with monkeypatch.context() as fault:
+ fault.setattr(filesystem_store_module, "_fsync_directory", fail_artifact_inode_fsync)
+ fault.setattr(os, "open", fail_containment_open)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert publication_failed
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is ResolutionStage.DOMAIN
+ assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert result.failure.retryable
+ details = result.failure.details.to_value()
+ assert isinstance(details, dict)
+ assert details["outcome_uncertain"] is True
+ assert artifact.is_dir()
+ assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ assert not list(store.quarantine_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+@pytest.mark.parallel
+def test_crash_after_publication_hardening_leaves_logically_committed_hit(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity = _identity()
+ process = mp.get_context("spawn").Process(
+ target=_crash_after_publication_hardening_process,
+ args=(str(root),),
+ )
+
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 41
+ artifact = store.artifacts_dir / identity.key
+ assert artifact.is_dir()
+ assert not artifact.stat().st_mode & 0o222
+ assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+ assert not list(store.quarantine_dir.iterdir())
+
+
+@pytest.mark.parallel
+def test_crash_after_final_artifacts_fsync_leaves_authoritative_hit(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity = _identity()
+ process = mp.get_context("spawn").Process(target=_crash_after_artifacts_fsync_process, args=(str(root),))
+
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 29
+ artifact = store.artifacts_dir / identity.key
+ assert artifact.is_dir()
+ assert not artifact.stat().st_mode & 0o222
+ assert not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+ joined = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert joined.status is StoreStatus.HIT
+ assert joined.lease is not None
+ joined.lease.close()
+ assert not list(store.quarantine_dir.iterdir())
+
+
+@pytest.mark.parametrize("lock_kind", ["artifact", "build"])
+def test_post_commit_lock_release_error_keeps_authoritative_outcome(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ lock_kind: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity = _identity()
+ lock_path = (
+ store._artifact_lock_path(identity.key) if lock_kind == "artifact" else store._build_lock_path(identity.key)
+ )
+ original_fsync_directory = filesystem_store_module._fsync_directory
+ original_close = FileLock.close
+ committed = False
+ injected = False
+
+ def observe_publication_commit(path: Path) -> None:
+ nonlocal committed
+ original_fsync_directory(path)
+ if path == store.artifacts_dir:
+ committed = True
+
+ def fail_post_commit_close(lock: FileLock) -> None:
+ nonlocal injected
+ original_close(lock)
+ if committed and lock.path == lock_path and not injected:
+ injected = True
+ raise OSError(errno.EIO, f"injected post-commit {lock_kind} lock release failure")
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", observe_publication_commit)
+ monkeypatch.setattr(FileLock, "close", fail_post_commit_close)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert committed
+ assert injected
+ assert result.status is StoreStatus.BUILT
+ assert result.lease is not None
+ result.lease.close()
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+ assert not list(store.quarantine_dir.iterdir())
+ assert not list(store.tmp_dir.iterdir())
+
+
+def test_joined_build_lock_release_error_keeps_authoritative_outcome(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, _ = _publish_test_artifact(store)
+ build_lock_path = store._build_lock_path(identity.key)
+ producer_count = tmp_path / "producer-count"
+ original_lookup = store.lookup
+ original_close = FileLock.close
+ lookup_count = 0
+ injected = False
+
+ def miss_once(
+ lookup_identity: WeightArtifactIdentity,
+ *,
+ validation: ValidationLevel,
+ deadline: float | None = None,
+ ) -> StoreResult:
+ nonlocal lookup_count
+ lookup_count += 1
+ if lookup_count == 1:
+ return StoreResult(StoreStatus.MISS)
+ return original_lookup(lookup_identity, validation=validation, deadline=deadline)
+
+ def fail_joined_build_lock_close(lock: FileLock) -> None:
+ nonlocal injected
+ original_close(lock)
+ if lock.path == build_lock_path and not injected:
+ injected = True
+ raise OSError(errno.EIO, "injected joined build-lock release failure")
+
+ monkeypatch.setattr(store, "lookup", miss_once)
+ monkeypatch.setattr(FileLock, "close", fail_joined_build_lock_close)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, counter_path=producer_count),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert lookup_count == 2
+ assert injected
+ assert result.status is StoreStatus.JOINED
+ assert result.lease is not None
+ result.lease.close()
+ assert not producer_count.exists()
+
+
+def test_concurrent_lease_close_waits_for_resource_teardown(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ lease = built.lease
+ teardown_started = threading.Event()
+ allow_teardown = threading.Event()
+ second_started = threading.Event()
+ second_returned = threading.Event()
+ errors: list[BaseException] = []
+
+ def blocking_closer() -> None:
+ teardown_started.set()
+ if not allow_teardown.wait(5):
+ raise TimeoutError("test did not release lease teardown")
+
+ def close_lease(*, started: threading.Event | None = None, returned: threading.Event | None = None) -> None:
+ if started is not None:
+ started.set()
+ try:
+ lease.close()
+ except BaseException as exc:
+ errors.append(exc)
+ finally:
+ if returned is not None:
+ returned.set()
+
+ lease._resource_closers = (*lease._resource_closers, blocking_closer)
+ first = threading.Thread(target=close_lease)
+ second = threading.Thread(target=close_lease, kwargs={"started": second_started, "returned": second_returned})
+ first.start()
+ assert teardown_started.wait(5)
+ second.start()
+ assert second_started.wait(5)
+ assert not second_returned.wait(0.2)
+
+ allow_teardown.set()
+ first.join(5)
+ second.join(5)
+ assert not first.is_alive() and not second.is_alive()
+ assert second_returned.is_set()
+ assert not errors
+
+
+@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork-specific flock ownership regression")
+def test_forked_child_close_does_not_release_parent_lease_lock(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore", DeprecationWarning)
+ child_pid = os.fork()
+ if child_pid == 0: # pragma: no cover - assertions run in the parent
+ try:
+ built.lease.close()
+ except BaseException:
+ os._exit(1)
+ os._exit(0)
+ _, child_status = os.waitpid(child_pid, 0)
+ assert os.waitstatus_to_exitcode(child_status) == 0
+
+ active = store.cleanup(identity)
+ assert active is not None and active.code is FailureCode.ACTIVE_LEASE
+ built.lease.close()
+ assert store.cleanup(identity) is None
+
+
+@pytest.mark.parametrize("write_mode", ["incomplete", "overlap"])
+def test_failed_streaming_writes_never_publish_partial_artifacts(tmp_path: Path, write_mode: str) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, write_mode=write_mode),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None and result.failure.code is FailureCode.PRODUCER_FAILED
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+
+@pytest.mark.parametrize(
+ "dtype",
+ [
+ dtype
+ for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz")
+ if (dtype := getattr(torch, name, None)) is not None
+ ],
+)
+def test_available_float8_formats_round_trip_through_streaming_store(tmp_path: Path, dtype: torch.dtype) -> None:
+ identity = replace(
+ _identity(),
+ representation=WeightRepresentation(name="runtime-fp8", dtype=str(dtype)),
+ )
+ expected = torch.zeros((2, 4), dtype=dtype)
+
+ class Float8Producer:
+ spec = WeightProductionSpec(
+ producer_id="test.final-layout",
+ outputs=(identity,),
+ source_mode=ProductionSourceMode.FINALIZED_MODEL,
+ )
+
+ def produce(self, writer: object) -> ProductionMetadata:
+ spec = TensorWriteSpec("weight", tuple(expected.shape), dtype)
+ with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined]
+ output.write_tensor("weight", expected)
+ return ProductionMetadata("test-producer-v1", "test-restorer-v1")
+
+ store = _make_store(tmp_path / "store")
+ result = store.get_or_build(
+ BuildRequest(identity),
+ Float8Producer(),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert result.status is StoreStatus.BUILT
+ assert result.lease is not None
+ actual = result.lease.tensors["weight"]
+ assert actual.dtype is dtype
+ assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+ result.lease.close()
+
+
+def test_lookup_treats_orphan_deny_marker_as_miss(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ deny = store.deny_dir / f"{identity.key}.json"
+ deny.write_text("{}", encoding="utf-8")
+
+ result = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert result.status is StoreStatus.MISS
+
+
+def test_overlapping_weaker_lookup_cannot_return_after_artifact_is_denied(tmp_path: Path) -> None:
+ identity = _identity()
+ root = tmp_path / "store"
+ store = _make_store(root)
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ built.lease.close()
+
+ artifact_dir = store.artifacts_dir / identity.key
+ payload = artifact_dir / "weights.safetensors"
+ os.chmod(artifact_dir, 0o755)
+ os.chmod(payload, 0o644)
+ with payload.open("r+b") as handle:
+ handle.seek(-1, os.SEEK_END)
+ original = handle.read(1)
+ handle.seek(-1, os.SEEK_END)
+ handle.write(bytes([original[0] ^ 0xFF]))
+ handle.flush()
+ os.fsync(handle.fileno())
+ os.chmod(payload, 0o444)
+ os.chmod(artifact_dir, 0o555)
+
+ lease_opened = threading.Event()
+ resume_lookup = threading.Event()
+ overlapping = PausingOpenLeaseStore(
+ StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)),
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ lease_opened=lease_opened,
+ resume_lookup=resume_lookup,
+ )
+ weak_results: list[StoreResult] = []
+
+ def weak_lookup() -> None:
+ weak_results.append(overlapping.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA))
+
+ thread = threading.Thread(target=weak_lookup)
+ thread.start()
+ try:
+ assert lease_opened.wait(5)
+ strong = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert strong.status is StoreStatus.INVALID
+ assert (store.deny_dir / f"{identity.key}.json").is_file()
+ finally:
+ resume_lookup.set()
+ thread.join(5)
+
+ assert not thread.is_alive()
+ assert len(weak_results) == 1
+ weak = weak_results[0]
+ assert weak.status is StoreStatus.INVALID
+ assert weak.failure is not None and weak.failure.code is FailureCode.ARTIFACT_DENIED
+
+
+@pytest.mark.parametrize(
+ ("failure_errno", "expected_stage", "expected_code"),
+ [
+ (errno.ENOSPC, ResolutionStage.CAPACITY, FailureCode.ENOSPC),
+ (errno.EDQUOT, ResolutionStage.CAPACITY, FailureCode.ENOSPC),
+ (errno.EIO, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE),
+ ],
+)
+@pytest.mark.skipif(
+ not Path("/proc/self/fd").is_dir(),
+ reason="fault targeting requires Linux /proc/self/fd",
+)
+def test_deny_write_failure_surfaces_as_retryable_storage_failure(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_errno: int,
+ expected_stage: ResolutionStage,
+ expected_code: FailureCode,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ _corrupt_test_payload(artifact)
+ deny_path = store.deny_dir / f"{identity.key}.json"
+ original_write = os.write
+
+ def fail_deny_write(fd: int, data: bytes) -> int:
+ if Path(os.readlink(f"/proc/self/fd/{fd}")) == deny_path:
+ raise OSError(failure_errno, "injected deny write failure")
+ return original_write(fd, data)
+
+ with monkeypatch.context() as fault:
+ fault.setattr(os, "write", fail_deny_write)
+ failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is expected_stage
+ assert failed.failure.code is expected_code
+ assert failed.failure.retryable
+ assert deny_path.exists()
+ assert not artifact.exists()
+ assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1
+
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert not deny_path.exists()
+
+
+def test_deny_open_failure_is_not_reported_as_successful_invalidation(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ _corrupt_test_payload(artifact)
+ deny_path = store.deny_dir / f"{identity.key}.json"
+ original_open = os.open
+
+ def fail_deny_open(
+ path: str | bytes | os.PathLike[str] | os.PathLike[bytes],
+ flags: int,
+ mode: int = 0o777,
+ *,
+ dir_fd: int | None = None,
+ ) -> int:
+ if Path(os.fsdecode(path)) == deny_path and flags & os.O_CREAT:
+ raise OSError(errno.EIO, "injected deny open failure")
+ return original_open(path, flags, mode, dir_fd=dir_fd)
+
+ with monkeypatch.context() as fault:
+ fault.setattr(os, "open", fail_deny_open)
+ failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is ResolutionStage.DOMAIN
+ assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert failed.failure.retryable
+ assert not deny_path.exists()
+ assert not artifact.exists()
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert not quarantined[0].stat().st_mode & 0o222
+
+ assert store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA).status is StoreStatus.MISS
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert not deny_path.exists()
+
+
+def test_deny_method_eio_synchronously_excludes_artifact_from_weaker_lookup(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ _corrupt_test_payload(artifact)
+
+ def fail_deny(_key: str, _failure: object) -> None:
+ raise OSError(errno.EIO, "injected deny publication failure")
+
+ with monkeypatch.context() as fault:
+ fault.setattr(store, "_write_deny", fail_deny)
+ failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is ResolutionStage.DOMAIN
+ assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert failed.failure.retryable
+ assert "injected deny publication failure" in failed.failure.message
+ assert not artifact.exists()
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA).status is StoreStatus.MISS
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert not quarantined[0].stat().st_mode & 0o222
+
+
+def test_deny_failure_containment_preserves_a_replacement_publication(tmp_path: Path) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ old_inode = (artifact.stat().st_dev, artifact.stat().st_ino)
+ _corrupt_test_payload(artifact)
+ shared_lock = FileLock(
+ store._artifact_lock_path(identity.key),
+ exclusive=False,
+ deadline=time.monotonic() + 10,
+ )
+ rebuilt_publications: list[tuple[str, tuple[int, int]]] = []
+
+ class RebuildOnSharedLockRelease:
+ def close(self) -> None:
+ shared_lock.close()
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt_inode = (artifact.stat().st_dev, artifact.stat().st_ino)
+ rebuilt_publications.append((rebuilt.lease.provenance.publication_id, rebuilt_inode))
+ rebuilt.lease.close()
+
+ containment = store._contain_invalid_after_deny_failure(
+ identity.key,
+ RebuildOnSharedLockRelease(), # type: ignore[arg-type]
+ deadline=time.monotonic() + 10,
+ )
+
+ assert containment["invalidation_marker"] == "durable"
+ assert containment["quarantine"] == "already_replaced"
+ assert len(rebuilt_publications) == 1
+ publication_id, rebuilt_inode = rebuilt_publications[0]
+ assert rebuilt_inode != old_inode
+ assert (artifact.stat().st_dev, artifact.stat().st_ino) == rebuilt_inode
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ assert hit.lease.provenance.publication_id == publication_id
+ hit.lease.close()
+ assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+@pytest.mark.parallel
+def test_deny_failure_blocks_overlapping_weaker_lookup_across_processes(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ ctx = mp.get_context("spawn")
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity, artifact = _publish_test_artifact(store)
+ _corrupt_test_payload(artifact)
+ lease_opened = ctx.Event()
+ resume_lookup = ctx.Event()
+ result_queue = ctx.Queue()
+ process = ctx.Process(
+ target=_pausing_weak_lookup_process,
+ args=(str(root), lease_opened, resume_lookup, result_queue),
+ )
+ strong_results: list[StoreResult] = []
+
+ def fail_deny(_key: str, _failure: object) -> None:
+ raise OSError(errno.EIO, "injected deny publication failure")
+
+ def strong_lookup() -> None:
+ strong_results.append(
+ store.lookup(
+ identity,
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ )
+
+ process.start()
+ strong_thread = threading.Thread(target=strong_lookup)
+ strong_started = False
+ try:
+ assert lease_opened.wait(timeout=_PROCESS_START_TIMEOUT_SECONDS)
+ monkeypatch.setattr(store, "_write_deny", fail_deny)
+ strong_thread.start()
+ strong_started = True
+ marker_deadline = time.monotonic() + 5
+ while not artifact.stat().st_mode & filesystem_store_module._ARTIFACT_INVALIDATION_MODE:
+ assert time.monotonic() < marker_deadline, "invalidation marker was not published"
+ time.sleep(0.01)
+ finally:
+ resume_lookup.set()
+ if strong_started:
+ strong_thread.join(timeout=10)
+ _join_processes([process])
+
+ assert not strong_thread.is_alive()
+ assert process.exitcode == 0
+ assert result_queue.get(timeout=3) == (
+ StoreStatus.INVALID.value,
+ FailureCode.ARTIFACT_DENIED.value,
+ None,
+ )
+ assert len(strong_results) == 1
+ strong = strong_results[0]
+ assert strong.status is StoreStatus.FAILED
+ assert strong.failure is not None and strong.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert not artifact.exists()
+ assert len(list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))) == 1
+
+
+@pytest.mark.parametrize("failure_point", ["entry_open", "ready_read", "payload_hash"])
+def test_operational_lookup_eio_is_retryable_without_denying_valid_artifact(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_point: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+
+ with monkeypatch.context() as fault:
+ if failure_point == "entry_open":
+ original_open = os.open
+
+ def fail_entry_open(
+ path: str | bytes | os.PathLike[str] | os.PathLike[bytes],
+ flags: int,
+ mode: int = 0o777,
+ *,
+ dir_fd: int | None = None,
+ ) -> int:
+ if Path(os.fsdecode(path)) == artifact:
+ raise OSError(errno.EIO, "injected artifact open failure")
+ return original_open(path, flags, mode, dir_fd=dir_fd)
+
+ fault.setattr(os, "open", fail_entry_open)
+ elif failure_point == "ready_read":
+ original_read_json_at = filesystem_store_module._read_json_at
+
+ def fail_ready_read(
+ directory_fd: int,
+ name: str,
+ *,
+ require_read_only: bool = False,
+ ) -> dict[str, object]:
+ if name == "READY.json":
+ raise OSError(errno.EIO, "injected READY metadata read failure")
+ return original_read_json_at(
+ directory_fd,
+ name,
+ require_read_only=require_read_only,
+ )
+
+ fault.setattr(filesystem_store_module, "_read_json_at", fail_ready_read)
+ else:
+
+ def fail_hash(_fd: int, _size: int) -> str:
+ raise OSError(errno.EIO, "injected hash failure")
+
+ fault.setattr(filesystem_store_module, "_sha256_fd", fail_hash)
+ failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is ResolutionStage.DOMAIN
+ assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert failed.failure.retryable
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+@pytest.mark.parametrize("failure_point", ["entry_stat", "deny_stat"])
+@pytest.mark.skipif(
+ not Path("/proc/self/fd").is_dir(),
+ reason="lock descriptor assertions require Linux /proc/self/fd",
+)
+def test_lookup_state_probe_eio_is_typed_and_releases_artifact_lock(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_point: str,
+) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ deny_path = store.deny_dir / f"{identity.key}.json"
+ target = artifact if failure_point == "entry_stat" else deny_path
+ artifact_lock_target = str(store._artifact_lock_path(identity.key))
+ lock_fds_before = _open_fd_targets().count(artifact_lock_target)
+ original_lstat = Path.lstat
+ injected = False
+
+ def fail_state_probe(path: Path) -> os.stat_result:
+ nonlocal injected
+ if path == target and not injected:
+ injected = True
+ raise OSError(errno.EIO, f"injected {failure_point} failure")
+ return original_lstat(path)
+
+ with monkeypatch.context() as fault:
+ fault.setattr(Path, "lstat", fail_state_probe)
+ failed = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert injected
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is ResolutionStage.DOMAIN
+ assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert failed.failure.retryable
+ assert _open_fd_targets().count(artifact_lock_target) == lock_fds_before
+ assert not deny_path.exists()
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+def test_full_checksum_denies_quarantines_and_rebuilds_corrupt_payload(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ rebuild_count = tmp_path / "rebuild-count"
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ built.lease.close()
+
+ artifact_dir = store.artifacts_dir / identity.key
+ _corrupt_test_payload(artifact_dir)
+
+ metadata_only = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)
+ assert metadata_only.status is StoreStatus.HIT
+ assert metadata_only.lease is not None
+ metadata_only.lease.close()
+
+ inspected = store.inspect_artifact(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert inspected.status is StoreStatus.INVALID
+ assert inspected.failure is not None and inspected.failure.code is FailureCode.CHECKSUM_MISMATCH
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+
+ invalid = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert invalid.status is StoreStatus.INVALID
+ assert invalid.failure is not None and invalid.failure.code is FailureCode.CHECKSUM_MISMATCH
+ assert (store.deny_dir / f"{identity.key}.json").is_file()
+
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, counter_path=rebuild_count),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ assert torch.equal(rebuilt.lease.tensors["layer.weight"], _weight())
+ rebuilt.lease.close()
+ assert rebuild_count.read_text(encoding="utf-8").count("\n") == 1
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert not quarantined[0].stat().st_mode & 0o222
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+def test_pre_writable_crash_state_is_quarantined_immutably_and_rebuilt(tmp_path: Path) -> None:
+ store = _make_store(tmp_path / "store")
+ identity, artifact = _publish_test_artifact(store)
+ rebuild_count = tmp_path / "rebuild-count"
+ os.chmod(artifact, 0o755)
+
+ invalid = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+
+ assert invalid.status is StoreStatus.INVALID
+ assert invalid.failure is not None and invalid.failure.code is FailureCode.MANIFEST_CORRUPT
+ assert (store.deny_dir / f"{identity.key}.json").is_file()
+
+ rebuilt = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, counter_path=rebuild_count),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert rebuilt.status is StoreStatus.BUILT
+ assert rebuilt.lease is not None
+ rebuilt.lease.close()
+ assert rebuild_count.read_text(encoding="utf-8").count("\n") == 1
+ assert not (store.artifacts_dir / identity.key).stat().st_mode & 0o222
+ quarantined = list(store.quarantine_dir.glob(f"{identity.key}.validation_failure.*"))
+ assert len(quarantined) == 1
+ assert not quarantined[0].stat().st_mode & 0o222
+ assert not (store.deny_dir / f"{identity.key}.json").exists()
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.quarantine_dir.glob(f"{identity.key}.cleanup.*"))
+
+
+def test_malformed_safetensors_is_a_typed_invalid_artifact(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ built.lease.close()
+
+ artifact_dir = store.artifacts_dir / identity.key
+ payload = artifact_dir / "weights.safetensors"
+ os.chmod(artifact_dir, 0o755)
+ os.chmod(payload, 0o644)
+ with payload.open("r+b") as handle:
+ handle.write(b"\xff" * 8)
+ handle.flush()
+ os.fsync(handle.fileno())
+ os.chmod(payload, 0o444)
+ os.chmod(artifact_dir, 0o555)
+
+ result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)
+
+ assert result.status is StoreStatus.INVALID
+ assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT
+ assert (store.deny_dir / f"{identity.key}.json").is_file()
+
+
+@pytest.mark.parametrize("schema_field", ["producer_schema", "restorer_schema"])
+def test_lookup_rejects_manifest_schema_that_differs_from_identity(
+ tmp_path: Path,
+ schema_field: str,
+) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ built.lease.close()
+
+ artifact_dir = store.artifacts_dir / identity.key
+ manifest_path = artifact_dir / "manifest.json"
+ marker_path = artifact_dir / "READY.json"
+ manifest_document = json.loads(manifest_path.read_bytes())
+ marker_document = json.loads(marker_path.read_bytes())
+ manifest_document[schema_field] = "incompatible-schema"
+ manifest_bytes = json.dumps(
+ manifest_document,
+ ensure_ascii=False,
+ sort_keys=True,
+ separators=(",", ":"),
+ ).encode("utf-8")
+ marker_document["manifest_sha256"] = hashlib.sha256(manifest_bytes).hexdigest()
+ marker_bytes = json.dumps(
+ marker_document,
+ ensure_ascii=False,
+ sort_keys=True,
+ separators=(",", ":"),
+ ).encode("utf-8")
+
+ os.chmod(artifact_dir, 0o755)
+ for path, contents in ((manifest_path, manifest_bytes), (marker_path, marker_bytes)):
+ os.chmod(path, 0o644)
+ path.write_bytes(contents)
+ os.chmod(path, 0o444)
+ os.chmod(artifact_dir, 0o555)
+
+ result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)
+
+ assert result.status is StoreStatus.INVALID
+ assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_INCOMPATIBLE
+ assert (store.deny_dir / f"{identity.key}.json").is_file()
+
+
+def test_noncooperative_conflicting_publication_is_reported_without_replacement(tmp_path: Path) -> None:
+ identity = _identity()
+ root = tmp_path / "store"
+ store = _make_store(root)
+ started = threading.Event()
+ release = threading.Event()
+ thread_result: list[StoreResult] = []
+
+ def build_expected_content() -> None:
+ thread_result.append(
+ store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, started=started, release=release),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ )
+
+ builder = threading.Thread(target=build_expected_content)
+ builder.start()
+ try:
+ assert started.wait(timeout=5)
+ competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy)
+ published = competing.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, weight_offset=100),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert published.status is StoreStatus.BUILT
+ assert published.lease is not None
+ published.lease.close()
+ finally:
+ release.set()
+ builder.join(timeout=10)
+ assert not builder.is_alive()
+ assert len(thread_result) == 1
+ collision = thread_result[0]
+ assert collision.status is StoreStatus.FAILED
+ assert collision.failure is not None and collision.failure.code is FailureCode.IDENTITY_COLLISION
+
+ winner = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert winner.status is StoreStatus.HIT
+ assert winner.lease is not None
+ assert torch.equal(winner.lease.tensors["layer.weight"], _weight(100))
+ winner.lease.close()
+ assert len(list(store.quarantine_dir.iterdir())) == 1
+
+
+def test_same_semantic_competing_publication_remains_authoritative_on_staging_cleanup_eio(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ root = tmp_path / "store"
+ store = _make_store(root)
+ injected = False
+ original_remove_tree = FilesystemHostWeightStore._remove_tree
+
+ class CompetingProducer(FakeProducer):
+ def produce(self, writer: object) -> ProductionMetadata:
+ competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy)
+ published = competing.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert published.status is StoreStatus.BUILT
+ assert published.lease is not None
+ published.lease.close()
+ return super().produce(writer)
+
+ def fail_primary_staging_cleanup_once(path: Path) -> None:
+ nonlocal injected
+ if path.parent == store.tmp_dir and not injected:
+ injected = True
+ raise OSError(errno.EIO, "injected same-semantic staging cleanup failure")
+ original_remove_tree(path)
+
+ monkeypatch.setattr(store, "_remove_tree", fail_primary_staging_cleanup_once)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ CompetingProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert injected
+ assert result.status is StoreStatus.BUILT
+ assert result.lease is not None
+ result.lease.close()
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.deny_dir.iterdir())
+ assert not list(store.quarantine_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+def test_competing_publication_manifest_eio_is_retryable_and_preserves_entry(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ root = tmp_path / "store"
+ store = _make_store(root)
+ artifact = store.artifacts_dir / identity.key
+ manifest_path = artifact / "manifest.json"
+ original_read_json_file = filesystem_store_module._read_json_file
+ injected = False
+
+ class CompetingProducer(FakeProducer):
+ def produce(self, writer: object) -> ProductionMetadata:
+ competing = NonCooperativeStore(store.domain_policy, store.capacity_policy, store.integrity_policy)
+ published = competing.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert published.status is StoreStatus.BUILT
+ assert published.lease is not None
+ published.lease.close()
+ return super().produce(writer)
+
+ def fail_competing_manifest_read_once(path: Path) -> dict[str, object]:
+ nonlocal injected
+ if path == manifest_path and not injected:
+ injected = True
+ raise OSError(errno.EIO, "injected competing manifest read failure")
+ return original_read_json_file(path)
+
+ monkeypatch.setattr(filesystem_store_module, "_read_json_file", fail_competing_manifest_read_once)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ CompetingProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert injected
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is ResolutionStage.DOMAIN
+ assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert result.failure.retryable
+ details = result.failure.details.to_value()
+ assert isinstance(details, dict)
+ assert details["outcome_uncertain"] is True
+ assert details["competing_entry"] == "preserved"
+ assert artifact.is_dir()
+ assert not list(store.quarantine_dir.iterdir())
+ assert not list(store.tmp_dir.iterdir())
+ hit = store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM)
+ assert hit.status is StoreStatus.HIT
+ assert hit.lease is not None
+ hit.lease.close()
+
+
+def test_payload_symlink_is_rejected_without_following_it(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ built = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert built.lease is not None
+ built.lease.close()
+
+ artifact_dir = store.artifacts_dir / identity.key
+ payload = artifact_dir / "weights.safetensors"
+ outside = tmp_path / "outside.safetensors"
+ os.chmod(artifact_dir, 0o755)
+ payload.rename(outside)
+ payload.symlink_to(outside)
+ os.chmod(artifact_dir, 0o555)
+
+ result = store.lookup(identity, validation=ValidationLevel.MANIFEST_AND_METADATA)
+ assert result.status is StoreStatus.INVALID
+ assert result.failure is not None and result.failure.code is FailureCode.MANIFEST_CORRUPT
+
+
+def test_coordination_deadline_does_not_claim_to_cancel_an_active_producer(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ started = time.monotonic()
+
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, delay_seconds=0.15),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 0.02,
+ )
+
+ assert time.monotonic() - started >= 0.1
+ assert result.status is StoreStatus.BUILT
+ assert result.lease is not None
+ result.lease.close()
+
+
+@pytest.mark.parallel
+def test_multiprocess_contention_runs_one_builder_and_joins_same_artifact(tmp_path: Path) -> None:
+ ctx = mp.get_context("spawn")
+ root = tmp_path / "store"
+ _make_store(root)
+ counter = tmp_path / "producer-count"
+ initial_lookup_barrier = ctx.Barrier(2)
+ result_queue = ctx.Queue()
+ processes = [
+ ctx.Process(
+ target=_build_process,
+ args=(str(root), str(counter), initial_lookup_barrier, result_queue),
+ )
+ for _ in range(2)
+ ]
+ for process in processes:
+ process.start()
+ _join_processes(processes)
+
+ results = [result_queue.get(timeout=3) for _ in processes]
+ assert all(process.exitcode == 0 for process in processes)
+ assert sorted(item[0] for item in results) == [StoreStatus.BUILT.value, StoreStatus.JOINED.value]
+ assert all(item[2] is None for item in results)
+ assert len({item[1] for item in results}) == 1
+ assert counter.read_text(encoding="utf-8").count("\n") == 1
+
+
+@pytest.mark.parallel
+def test_reader_sees_miss_during_build_and_waiter_has_bounded_timeout(tmp_path: Path) -> None:
+ ctx = mp.get_context("spawn")
+ root = tmp_path / "store"
+ store = _make_store(root)
+ identity = _identity()
+ started = ctx.Event()
+ release = ctx.Event()
+ result_queue = ctx.Queue()
+ process = ctx.Process(
+ target=_controlled_build_process,
+ args=(str(root), started, release, result_queue),
+ )
+ process.start()
+ try:
+ assert started.wait(timeout=_PROCESS_START_TIMEOUT_SECONDS)
+ assert store.lookup(identity, validation=ValidationLevel.FULL_CHECKSUM).status is StoreStatus.MISS
+ timed_out = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 0.1,
+ )
+ assert timed_out.status is StoreStatus.TIMEOUT
+ assert timed_out.failure is not None and timed_out.failure.code is FailureCode.ACTIVE_BUILD_TIMEOUT
+ finally:
+ release.set()
+ _join_processes([process])
+
+ status, error = result_queue.get(timeout=3)
+ assert process.exitcode == 0
+ assert (status, error) == (StoreStatus.BUILT.value, None)
+
+
+@pytest.mark.parallel
+def test_crashed_builder_leaves_only_invisible_temp_and_next_builder_recovers(tmp_path: Path) -> None:
+ ctx = mp.get_context("spawn")
+ root = tmp_path / "store"
+ store = _make_store(root)
+ process = ctx.Process(target=_crash_build_process, args=(str(root),))
+ process.start()
+ _join_processes([process])
+
+ assert process.exitcode == 17
+ assert list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+ identity = _identity()
+ recovered = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert recovered.status is StoreStatus.BUILT
+ assert recovered.lease is not None
+ recovered.lease.close()
+ assert not list(store.tmp_dir.iterdir())
+
+
+def test_capacity_and_enospc_failures_are_typed_and_clean_staging(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ identity = _identity()
+ limited = _make_store(tmp_path / "limited", capacity=CapacityPolicy(max_artifact_bytes=64))
+ too_large = limited.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert too_large.status is StoreStatus.FAILED
+ assert too_large.failure is not None and too_large.failure.code is FailureCode.ARTIFACT_TOO_LARGE
+ assert not list(limited.tmp_dir.iterdir())
+
+ no_space = _make_store(tmp_path / "enospc")
+
+ def raise_enospc(_fd: int, _offset: int, _length: int) -> None:
+ raise OSError(errno.ENOSPC, "injected no space")
+
+ monkeypatch.setattr(os, "posix_fallocate", raise_enospc)
+ failed = no_space.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None and failed.failure.code is FailureCode.ENOSPC
+ assert not list(no_space.tmp_dir.iterdir())
+ assert not list(no_space.artifacts_dir.iterdir())
+
+
+@pytest.mark.parametrize(
+ ("failure_errno", "expected_stage", "expected_code"),
+ [
+ (errno.ENOSPC, ResolutionStage.CAPACITY, FailureCode.ENOSPC),
+ (errno.EDQUOT, ResolutionStage.CAPACITY, FailureCode.ENOSPC),
+ (errno.EACCES, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE),
+ (errno.EIO, ResolutionStage.DOMAIN, FailureCode.DOMAIN_UNAVAILABLE),
+ ],
+)
+def test_staging_directory_failures_are_typed_and_clean(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ failure_errno: int,
+ expected_stage: ResolutionStage,
+ expected_code: FailureCode,
+) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ original_mkdir = Path.mkdir
+
+ def fail_staging_mkdir(
+ path: Path,
+ mode: int = 0o777,
+ parents: bool = False,
+ exist_ok: bool = False,
+ ) -> None:
+ if path.parent == store.tmp_dir:
+ raise OSError(failure_errno, "injected staging directory failure")
+ original_mkdir(path, mode=mode, parents=parents, exist_ok=exist_ok)
+
+ monkeypatch.setattr(Path, "mkdir", fail_staging_mkdir)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is expected_stage
+ assert result.failure.code is expected_code
+ assert result.failure.retryable
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+
+def test_post_producer_staging_io_failure_is_retryable_domain_failure(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ original_fsync_directory = filesystem_store_module._fsync_directory
+
+ def fail_staging_fsync(path: Path) -> None:
+ if path.parent == store.tmp_dir:
+ raise OSError(errno.EIO, "injected post-producer staging fsync failure")
+ original_fsync_directory(path)
+
+ monkeypatch.setattr(filesystem_store_module, "_fsync_directory", fail_staging_fsync)
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is ResolutionStage.DOMAIN
+ assert result.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert result.failure.retryable
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+
+def test_async_payload_fsync_eio_is_retryable_domain_failure(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+ injection_count = 0
+
+ def fail_payload_fsync(fd: int) -> None:
+ nonlocal injection_count
+ injection_count += 1
+ try:
+ raise OSError(errno.EIO, "injected payload-fsync EIO")
+ finally:
+ os.close(fd)
+
+ with monkeypatch.context() as fault:
+ fault.setattr(filesystem_writer_module, "_fsync_payload", fail_payload_fsync)
+ failed = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert injection_count == 1
+ assert failed.status is StoreStatus.FAILED
+ assert failed.failure is not None
+ assert failed.failure.stage is ResolutionStage.DOMAIN
+ assert failed.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert failed.failure.retryable
+ assert "injected payload-fsync EIO" in failed.failure.message
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+ recovered = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert recovered.status is StoreStatus.BUILT
+ assert recovered.lease is not None
+ recovered.lease.close()
+
+
+def test_producer_raised_eio_remains_nonretryable_producer_failure(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+
+ class FailingProducer(FakeProducer):
+ def produce(self, writer: object) -> ProductionMetadata:
+ del writer
+ raise OSError(errno.EIO, "injected producer failure")
+
+ result = store.get_or_build(
+ BuildRequest(identity),
+ FailingProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+
+ assert result.status is StoreStatus.FAILED
+ assert result.failure is not None
+ assert result.failure.stage is ResolutionStage.PRODUCTION
+ assert result.failure.code is FailureCode.PRODUCER_FAILED
+ assert not result.failure.retryable
+ assert not list(store.tmp_dir.iterdir())
+ assert not list(store.artifacts_dir.iterdir())
+
+
+def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path: Path) -> None:
+ identity = _identity()
+ store_limited = _make_store(tmp_path / "store-limit", capacity=CapacityPolicy(max_store_bytes=1))
+ store_result = store_limited.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert store_result.status is StoreStatus.FAILED
+ assert store_result.failure is not None and store_result.failure.code is FailureCode.STORE_LIMIT_EXCEEDED
+ assert not list(store_limited.artifacts_dir.iterdir())
+
+ free_root = tmp_path / "free-limit"
+ free_root.mkdir(mode=0o700)
+ # Leave enough margin for filesystem accounting to move between the
+ # preflight measurement and the capacity check.
+ unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30)
+ free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve))
+ free_result = free_limited.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert free_result.status is StoreStatus.FAILED
+ assert free_result.failure is not None and free_result.failure.code is FailureCode.INSUFFICIENT_CAPACITY
+ assert not list(free_limited.artifacts_dir.iterdir())
+
+
+def test_domain_capacity_policy_is_authoritative(tmp_path: Path) -> None:
+ root = tmp_path / "store"
+ _make_store(root, capacity=CapacityPolicy(max_store_bytes=4096))
+
+ with pytest.raises(HostWeightError, match="capacity policy differs"):
+ _make_store(root, capacity=CapacityPolicy(max_store_bytes=8192))
+
+
+@pytest.mark.parametrize("filesystem_type", ["nfs4", "cifs", "lustre", "ceph", "unknownfs"])
+def test_remote_and_unknown_filesystems_cannot_back_the_local_store(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ filesystem_type: str,
+) -> None:
+ monkeypatch.setattr(filesystem_store_module, "detect_filesystem_type", lambda _path: filesystem_type)
+ root = tmp_path / "store"
+
+ with pytest.raises(ValueError, match="not node-local|not recognized"):
+ detect_storage_class(root)
+ with pytest.raises(HostWeightError) as error:
+ FilesystemHostWeightStore(
+ StorageDomainPolicy(root=root, storage_class=StorageClass.DISK),
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ )
+
+ assert error.value.failure.code is FailureCode.DOMAIN_POLICY_MISMATCH
+ assert not root.exists()
+
+
+def test_trusted_store_rejects_group_or_world_writable_root(tmp_path: Path) -> None:
+ root = tmp_path / "untrusted"
+ root.mkdir(mode=0o700)
+ root.chmod(0o777)
+
+ with pytest.raises(HostWeightError) as error:
+ FilesystemHostWeightStore(
+ StorageDomainPolicy(root=root, storage_class=detect_storage_class(root)),
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ )
+
+ assert error.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+
+
+def test_unimplemented_coordination_and_fs_verity_return_typed_failures(tmp_path: Path) -> None:
+ identity = _identity()
+ store = _make_store(tmp_path / "store")
+
+ unsupported_producer = store.get_or_build(
+ BuildRequest(identity),
+ FakeProducer(identity, coordination_scope=CoordinationScope.GROUP_COHORT),
+ validation=ValidationLevel.FULL_CHECKSUM,
+ deadline=time.monotonic() + 10,
+ )
+ assert unsupported_producer.status is StoreStatus.FAILED
+ assert unsupported_producer.failure is not None
+ assert unsupported_producer.failure.code is FailureCode.PRODUCER_UNSUPPORTED
+
+ fs_verity = store.lookup(identity, validation=ValidationLevel.FS_VERITY)
+ assert fs_verity.status is StoreStatus.FAILED
+ assert fs_verity.failure is not None and fs_verity.failure.code is FailureCode.UNSUPPORTED
+
+
+def test_unimplemented_numa_enforcement_and_background_scrub_fail_explicitly(tmp_path: Path) -> None:
+ numa_root = tmp_path / "numa"
+ numa_root.mkdir(mode=0o700)
+ with pytest.raises(HostWeightError) as numa_error:
+ FilesystemHostWeightStore(
+ StorageDomainPolicy(
+ root=numa_root,
+ scope=StorageScope.NUMA,
+ storage_class=detect_storage_class(numa_root),
+ numa_prefault=True,
+ ),
+ CapacityPolicy(),
+ IntegrityPolicy(),
+ )
+ assert numa_error.value.failure.code is FailureCode.UNSUPPORTED
+
+ with pytest.raises(HostWeightError) as scrub_error:
+ _make_store(tmp_path / "scrub", integrity=IntegrityPolicy(background_scrub=True))
+ assert scrub_error.value.failure.code is FailureCode.UNSUPPORTED
diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py
new file mode 100644
index 000000000..21bab37de
--- /dev/null
+++ b/tests/host_weight_runtime/test_runtime_resolution.py
@@ -0,0 +1,632 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Runtime resolution and publication-policy tests for Host Weight Runtime."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import cast
+
+import pytest
+import torch
+
+from vllm_omni.host_weight_runtime import (
+ AdaptationIdentity,
+ CanonicalJson,
+ ComponentIdentity,
+ FailureCode,
+ HostWeightFailure,
+ HostWeightLease,
+ HostWeightLeaseCarrier,
+ HostWeightRuntime,
+ HostWeightRuntimeConfig,
+ IntegrityPolicy,
+ LookupPhase,
+ PostLoadPublicationOutcome,
+ PostLoadPublicationReport,
+ ProducerIdentity,
+ ProductionMetadata,
+ ProductionPolicy,
+ ProductionSourceMode,
+ RemoteImportPolicy,
+ RemoteOnMiss,
+ ResolutionOutcome,
+ ResolutionReport,
+ ResolutionStage,
+ RuntimeMode,
+ RuntimeWeightLayout,
+ StorageDomainPolicy,
+ StoreResult,
+ StoreStatus,
+ TensorWriteSpec,
+ ValidationLevel,
+ WeightArtifactIdentity,
+ WeightProductionSpec,
+ WeightRepresentation,
+ WeightSourceIdentity,
+)
+from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _identity() -> WeightArtifactIdentity:
+ return WeightArtifactIdentity(
+ schema_version=1,
+ source=WeightSourceIdentity("test-org/test-model", "0123456789abcdef", "source-files-sha256"),
+ component=ComponentIdentity("transformer", "diffusion.dit"),
+ representation=WeightRepresentation("runtime-bf16", "torch.bfloat16"),
+ layout=RuntimeWeightLayout("final-module-layout"),
+ adaptation=AdaptationIdentity(),
+ producer=ProducerIdentity(
+ "test.final-layout",
+ "1",
+ "implementation-sha256",
+ "test-producer-v1",
+ "test-restorer-v1",
+ ),
+ )
+
+
+class CountingProducer:
+ def __init__(
+ self,
+ identity: WeightArtifactIdentity,
+ *,
+ lookup_phase: LookupPhase = LookupPhase.PRE_LOAD_SAFE,
+ ) -> None:
+ self.calls = 0
+ self._spec = WeightProductionSpec(
+ producer_id="test.final-layout",
+ outputs=(identity,),
+ source_mode=ProductionSourceMode.FINALIZED_MODEL,
+ lookup_phase=lookup_phase,
+ )
+
+ @property
+ def spec(self) -> WeightProductionSpec:
+ return self._spec
+
+ def produce(self, writer: object) -> ProductionMetadata:
+ self.calls += 1
+ spec = TensorWriteSpec("weight", (2, 2), torch.bfloat16)
+ with writer.open_tensor_file("weights.safetensors", (spec,)) as output: # type: ignore[attr-defined]
+ output.write_tensor("weight", torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2))
+ return ProductionMetadata(
+ producer_schema="test-producer-v1",
+ restorer_schema="test-restorer-v1",
+ format_metadata=CanonicalJson.from_value({"source": "test"}),
+ )
+
+
+def _domain(root: Path) -> StorageDomainPolicy:
+ return StorageDomainPolicy(root=root, storage_class=detect_storage_class(root.parent))
+
+
+def test_disabled_mode_does_not_probe_or_create_configured_root(tmp_path: Path) -> None:
+ root = tmp_path / "must-not-exist"
+ runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.DISABLED, domain=_domain(root)))
+
+ resolution = runtime.resolve()
+
+ assert resolution.report.outcome is ResolutionOutcome.CANONICAL_DIRECT
+ assert resolution.report.attempts == ()
+ assert resolution.lease is None
+ assert not root.exists()
+
+
+@pytest.mark.parametrize(
+ ("mode", "expected"),
+ [
+ (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+ (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+ ],
+)
+def test_local_miss_uses_configured_fallback_policy(
+ tmp_path: Path,
+ mode: RuntimeMode,
+ expected: ResolutionOutcome,
+) -> None:
+ runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path / mode.value)))
+
+ resolution = runtime.resolve(_identity())
+
+ assert resolution.report.outcome is expected
+ assert resolution.lease is None
+ assert [attempt.result.value for attempt in resolution.report.attempts] == ["miss", "skipped"]
+ expected_action = "canonical_fallback" if mode is RuntimeMode.PREFERRED else "fail_startup"
+ assert [attempt.action.value for attempt in resolution.report.attempts] == [expected_action, expected_action]
+
+
+def test_local_production_then_exact_warm_hit_emits_one_terminal_report_each(tmp_path: Path) -> None:
+ reports: list[ResolutionReport] = []
+ identity = _identity()
+ producer = CountingProducer(identity)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")),
+ observer=reports.append,
+ )
+
+ cold = runtime.resolve(identity, producer=producer)
+ assert cold.report.outcome is ResolutionOutcome.LOCAL_PRODUCTION
+ assert cold.lease is not None
+ assert torch.equal(
+ cold.lease.tensors["weight"],
+ torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2),
+ )
+ cold.lease.close()
+ assert cold.report.attempts[0].action.value == "try_producer"
+
+ warm = runtime.resolve(identity)
+ assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT
+ assert warm.lease is not None
+ warm.lease.close()
+
+ assert producer.calls == 1
+ assert reports == [cold.report, warm.report]
+ assert all(report.resolution_id for report in reports)
+
+
+def test_lease_carrier_is_single_take_and_process_local(tmp_path: Path) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store"))
+ )
+ resolution = runtime.resolve(identity, producer=producer)
+ assert resolution.lease is not None
+
+ carrier = HostWeightLeaseCarrier(resolution.lease)
+ with pytest.raises(TypeError, match="cannot be serialized"):
+ carrier.__reduce_ex__(4)
+
+ lease = carrier.take()
+ assert lease is resolution.lease
+ with pytest.raises(RuntimeError, match="already taken"):
+ carrier.take()
+ carrier.close() # ownership has moved; the carrier must not close twice.
+ assert not lease.closed
+ lease.close()
+
+
+@pytest.mark.parametrize(
+ ("mode", "expected"),
+ [
+ (RuntimeMode.PREFERRED, ResolutionOutcome.FAILED),
+ (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+ ],
+)
+def test_remote_required_is_explicitly_unsupported_and_skips_local_production(
+ tmp_path: Path,
+ mode: RuntimeMode,
+ expected: ResolutionOutcome,
+) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=mode,
+ domain=_domain(tmp_path / mode.value),
+ remote=RemoteImportPolicy(
+ on_local_miss=RemoteOnMiss.REQUIRE,
+ providers=("future-mooncake",),
+ ),
+ )
+ )
+
+ resolution = runtime.resolve(identity, producer=producer)
+
+ assert resolution.report.outcome is expected
+ assert producer.calls == 0
+ assert [attempt.stage.value for attempt in resolution.report.attempts] == ["lookup", "remote"]
+ assert resolution.report.attempts[-1].failure is not None
+ assert resolution.report.attempts[-1].failure.code.value == "unsupported"
+ assert resolution.report.attempts[-1].action.value == "fail_startup"
+
+
+def test_observer_failure_cannot_change_resolution_outcome(tmp_path: Path) -> None:
+ def broken_observer(_report: object) -> None:
+ raise RuntimeError("metrics backend unavailable")
+
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store")),
+ observer=broken_observer,
+ )
+
+ resolution = runtime.resolve(_identity())
+
+ assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+
+
+def test_preload_resolution_does_not_run_postload_only_producer(tmp_path: Path) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store"))
+ )
+
+ resolution = runtime.resolve(identity, producer=producer)
+
+ assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+ assert producer.calls == 0
+ assert resolution.report.attempts[0].action.value == "canonical_fallback"
+
+
+def test_postload_publication_closes_lease_and_only_warms_future_resolution(
+ tmp_path: Path,
+) -> None:
+ publications: list[PostLoadPublicationReport] = []
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(
+ allow_local_build=False,
+ allow_post_load_publish=True,
+ ),
+ ),
+ publication_observer=publications.append,
+ )
+
+ preload = runtime.resolve(identity, producer=producer)
+ assert preload.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+ assert producer.calls == 0
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+ assert publication.outcome is PostLoadPublicationOutcome.PUBLISHED
+ assert publication.identity_digest == identity.key
+
+ already_present = runtime.publish_after_load(identity, producer=producer)
+ assert already_present.outcome is PostLoadPublicationOutcome.ALREADY_PRESENT
+
+ warm = runtime.resolve(identity)
+ assert warm.report.outcome is ResolutionOutcome.LOCAL_HIT
+ assert warm.lease is not None
+ assert torch.equal(
+ warm.lease.tensors["weight"],
+ torch.arange(4, dtype=torch.float32).to(torch.bfloat16).reshape(2, 2),
+ )
+ warm.lease.close()
+ assert producer.calls == 1
+ assert publications == [publication, already_present]
+ assert isinstance(runtime.store, FilesystemHostWeightStore)
+ assert runtime.store.cleanup(identity) is None
+
+
+def test_postload_publication_policy_skip_does_not_run_producer(tmp_path: Path) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "store"))
+ )
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.POLICY_DISABLED
+ assert producer.calls == 0
+
+
+def test_postload_publication_reports_lease_cleanup_failure(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+ assert runtime.store is not None
+
+ class FailingLease:
+ def close(self) -> None:
+ raise RuntimeError("injected lease cleanup failure")
+
+ def return_failing_lease(*_args: object, **_kwargs: object) -> StoreResult:
+ return StoreResult(StoreStatus.BUILT, lease=cast(HostWeightLease, FailingLease()))
+
+ monkeypatch.setattr(runtime.store, "get_or_build", return_failing_lease)
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.FAILED
+ assert publication.failure is not None
+ assert publication.failure.stage is ResolutionStage.LIFECYCLE
+ assert publication.failure.code is FailureCode.PUBLICATION_FAILED
+ assert publication.failure.details.to_value() == {
+ "exception_type": "RuntimeError",
+ "store_outcome": "built",
+ }
+
+
+def test_disabled_postload_publication_does_not_probe_configured_root(tmp_path: Path) -> None:
+ root = tmp_path / "must-not-exist"
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.DISABLED,
+ domain=_domain(root),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+
+ publication = runtime.publish_after_load()
+
+ assert publication.outcome is PostLoadPublicationOutcome.RUNTIME_DISABLED
+ assert not root.exists()
+
+
+def test_postload_publication_rejects_preload_producer(tmp_path: Path) -> None:
+ identity = _identity()
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+
+ with pytest.raises(ValueError, match="POST_LOAD_ONLY"):
+ runtime.publish_after_load(identity, producer=CountingProducer(identity))
+
+
+def test_postload_publication_failure_is_reported_without_raising(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+ assert runtime.store is not None
+ failure = HostWeightFailure(
+ stage=ResolutionStage.CAPACITY,
+ code=FailureCode.ENOSPC,
+ retryable=True,
+ message="injected post-load publication failure",
+ details=CanonicalJson.empty(),
+ )
+
+ def fail_publication(*_args: object, **_kwargs: object) -> StoreResult:
+ return StoreResult(StoreStatus.FAILED, failure=failure)
+
+ monkeypatch.setattr(runtime.store, "get_or_build", fail_publication)
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.FAILED
+ assert publication.failure is failure
+ assert producer.calls == 0
+
+
+def test_postload_publication_reports_store_initialization_failure(tmp_path: Path) -> None:
+ root = tmp_path / "not-a-directory"
+ root.write_text("occupied", encoding="utf-8")
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(root),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.FAILED
+ assert publication.failure is not None
+ assert publication.failure.code is FailureCode.DOMAIN_UNAVAILABLE
+ assert publication.failure.retryable
+ assert producer.calls == 0
+
+
+def test_postload_publication_maps_joined_and_closes_lease(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+ assert runtime.store is not None
+
+ class ClosingLease:
+ closed = False
+
+ def close(self) -> None:
+ self.closed = True
+
+ lease = ClosingLease()
+
+ def return_joined_lease(*_args: object, **_kwargs: object) -> StoreResult:
+ return StoreResult(StoreStatus.JOINED, lease=cast(HostWeightLease, lease))
+
+ monkeypatch.setattr(runtime.store, "get_or_build", return_joined_lease)
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.JOINED
+ assert lease.closed
+
+
+def test_postload_publication_synthesizes_failure_for_unexpected_store_status(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ producer = CountingProducer(identity, lookup_phase=LookupPhase.POST_LOAD_ONLY)
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ production=ProductionPolicy(allow_post_load_publish=True),
+ )
+ )
+ assert runtime.store is not None
+
+ def return_unexpected_miss(*_args: object, **_kwargs: object) -> StoreResult:
+ return StoreResult(StoreStatus.MISS)
+
+ monkeypatch.setattr(runtime.store, "get_or_build", return_unexpected_miss)
+
+ publication = runtime.publish_after_load(identity, producer=producer)
+
+ assert publication.outcome is PostLoadPublicationOutcome.FAILED
+ assert publication.failure is not None
+ assert publication.failure.stage is ResolutionStage.PRODUCTION
+ assert publication.failure.code is FailureCode.PRODUCER_FAILED
+ assert not publication.failure.retryable
+
+
+@pytest.mark.parametrize(
+ ("mode", "expected"),
+ [
+ (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK),
+ (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED),
+ ],
+)
+def test_retryable_store_initialization_failure_obeys_runtime_mode(
+ tmp_path: Path,
+ mode: RuntimeMode,
+ expected: ResolutionOutcome,
+) -> None:
+ root = tmp_path / "not-a-directory"
+ root.write_text("occupied", encoding="utf-8")
+ runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=mode, domain=_domain(root)))
+
+ resolution = runtime.resolve(_identity())
+
+ assert resolution.report.outcome is expected
+ assert resolution.lease is None
+ assert len(resolution.report.attempts) == 1
+ assert resolution.report.attempts[0].failure is not None
+ assert resolution.report.attempts[0].failure.code.value == "domain_unavailable"
+ assert resolution.report.attempts[0].failure.retryable
+
+
+@pytest.mark.parametrize(
+ ("stage", "code"),
+ [
+ (ResolutionStage.PRODUCTION, FailureCode.PRODUCER_UNSUPPORTED),
+ (ResolutionStage.VALIDATION, FailureCode.IDENTITY_COLLISION),
+ (ResolutionStage.LIFECYCLE, FailureCode.PUBLICATION_FAILED),
+ ],
+)
+def test_nonretryable_post_initialization_failure_cannot_fall_back_in_preferred_mode(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+ stage: ResolutionStage,
+ code: FailureCode,
+) -> None:
+ identity = _identity()
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / code.value))
+ )
+ assert runtime.store is not None
+
+ def fail_production(
+ _request: object,
+ _producer: object,
+ *,
+ validation: object,
+ deadline: float,
+ ) -> StoreResult:
+ del validation, deadline
+ return StoreResult(
+ StoreStatus.FAILED,
+ failure=HostWeightFailure(
+ stage=stage,
+ code=code,
+ retryable=False,
+ message="injected nonretryable production failure",
+ details=CanonicalJson.empty(),
+ ),
+ )
+
+ monkeypatch.setattr(runtime.store, "get_or_build", fail_production)
+ resolution = runtime.resolve(identity, producer=CountingProducer(identity))
+
+ assert resolution.report.outcome is ResolutionOutcome.FAILED
+ assert resolution.report.attempts[-1].action.value == "fail_startup"
+ assert resolution.report.attempts[-1].failure is not None
+ assert resolution.report.attempts[-1].failure.code is code
+
+
+def test_retryable_post_initialization_failure_can_fall_back_in_preferred_mode(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ identity = _identity()
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path / "retryable"))
+ )
+ assert runtime.store is not None
+
+ def fail_production(
+ _request: object,
+ _producer: object,
+ *,
+ validation: object,
+ deadline: float,
+ ) -> StoreResult:
+ del validation, deadline
+ return StoreResult(
+ StoreStatus.FAILED,
+ failure=HostWeightFailure(
+ stage=ResolutionStage.CAPACITY,
+ code=FailureCode.ENOSPC,
+ retryable=True,
+ message="injected recoverable capacity failure",
+ details=CanonicalJson.empty(),
+ ),
+ )
+
+ monkeypatch.setattr(runtime.store, "get_or_build", fail_production)
+ resolution = runtime.resolve(identity, producer=CountingProducer(identity))
+
+ assert resolution.report.outcome is ResolutionOutcome.CANONICAL_FALLBACK
+ assert resolution.report.attempts[-1].action.value == "canonical_fallback"
+
+
+def test_nonretryable_store_configuration_failure_remains_visible_in_preferred_mode(tmp_path: Path) -> None:
+ root = tmp_path / "untrusted"
+ root.mkdir(mode=0o700)
+ root.chmod(0o777)
+ runtime = HostWeightRuntime.from_config(HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(root)))
+
+ resolution = runtime.resolve(_identity())
+
+ assert resolution.report.outcome is ResolutionOutcome.FAILED
+ assert resolution.report.attempts[0].action.value == "fail_startup"
+ assert resolution.report.attempts[0].failure is not None
+ assert not resolution.report.attempts[0].failure.retryable
+
+
+def test_unsupported_integrity_policy_fails_instead_of_disguising_configuration_error(tmp_path: Path) -> None:
+ runtime = HostWeightRuntime.from_config(
+ HostWeightRuntimeConfig(
+ mode=RuntimeMode.PREFERRED,
+ domain=_domain(tmp_path / "store"),
+ integrity=IntegrityPolicy(local_lookup=ValidationLevel.FS_VERITY),
+ )
+ )
+
+ resolution = runtime.resolve(_identity())
+
+ assert resolution.report.outcome is ResolutionOutcome.FAILED
+ assert resolution.report.attempts[0].failure is not None
+ assert resolution.report.attempts[0].failure.code.value == "unsupported"
diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py
new file mode 100644
index 000000000..b46f4080f
--- /dev/null
+++ b/tests/metrics/test_emit_calls.py
@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Verify emit-call wiring in production code paths.
+
+Uses behavior and Prometheus exposition to verify production call semantics.
+"""
+
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+import pytest
+
+from vllm_omni.metrics import OmniPrometheusMetrics
+from vllm_omni.metrics import definitions as defs
+
+if TYPE_CHECKING:
+ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_MODEL = "emit-calls-test"
+
+
+# ---------------------------------------------------------------------------
+# Behavioral pins — verify call semantics with mock OmniPrometheusMetrics.
+# ---------------------------------------------------------------------------
+
+
+def _make_omni_base_with_mock_prom(mocker):
+ """Build a minimal OmniBase shell wired with a mock prom_metrics.
+
+ Mirrors the pattern in ``test_prometheus.py::TestRequestLifecycleGauges``
+ — ``object.__new__`` skips ``__init__`` so we don't need a real engine.
+ """
+ from vllm_omni.entrypoints.omni_base import OmniBase
+
+ obj = object.__new__(OmniBase)
+ obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
+ obj.request_states = {}
+ obj._consumed_metric_messages = {}
+ obj.log_stats = True
+ return obj, obj.prom_metrics
+
+
+class TestFailureCounterWiring:
+ """All failure paths share one request-idempotent counter helper."""
+
+ def test_failure_reason_taxonomy_is_bounded(self) -> None:
+ from vllm_omni.metrics.utils import normalize_failure_reason
+
+ assert normalize_failure_reason("client_abort") == "client_abort"
+ assert normalize_failure_reason("client_disconnect") == "client_disconnect"
+ assert normalize_failure_reason("stage_error") == "stage_error"
+ assert normalize_failure_reason("unknown") == "unknown"
+ assert normalize_failure_reason("timeout for request req-1") == "unknown"
+ assert normalize_failure_reason("") == "unknown"
+ assert normalize_failure_reason(None) == "unknown"
+
+ def test_fire_failure_counter_passes_reason_through(self, mocker) -> None:
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+ obj.request_states["req-1"] = SimpleNamespace(
+ metrics=SimpleNamespace(e2e_done=set()),
+ failure_recorded=False,
+ )
+
+ obj._record_request_failure_once("req-1", reason="client_disconnect")
+
+ prom.request_failed.assert_called_once()
+ prom.inc_requests_failed.assert_called_once_with("client_disconnect")
+
+ def test_fire_failure_counter_normalizes_unrecognized_reason(self, mocker) -> None:
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+ obj.request_states["req-1"] = SimpleNamespace(
+ metrics=SimpleNamespace(e2e_done=set()),
+ failure_recorded=False,
+ )
+
+ obj._record_request_failure_once("req-1", reason="CUDA error for req-1")
+
+ prom.inc_requests_failed.assert_called_once_with("unknown")
+
+ def test_fire_failure_counter_skips_when_request_already_succeeded(self, mocker) -> None:
+ # When the request is in e2e_done (i.e. finalize already fired
+ # request_succeeded), the failure path must NOT double-count.
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+ obj.request_states["req-1"] = SimpleNamespace(
+ metrics=SimpleNamespace(e2e_done={"req-1"}),
+ failure_recorded=False,
+ )
+
+ obj._record_request_failure_once("req-1", reason="client_disconnect")
+
+ prom.request_failed.assert_not_called()
+ prom.inc_requests_failed.assert_not_called()
+
+ def test_fire_failure_counter_skips_when_request_state_missing(self, mocker) -> None:
+ # No request_states entry — already popped by abort path. Fail-safe
+ # must NOT raise.
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+
+ obj._record_request_failure_once("missing-req", reason="oom")
+
+ prom.request_failed.assert_not_called()
+ prom.inc_requests_failed.assert_not_called()
+
+ def test_failure_counter_is_idempotent_across_abort_and_cleanup(self, mocker) -> None:
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+ obj.request_states["req-1"] = SimpleNamespace(
+ metrics=SimpleNamespace(
+ e2e_done=set(),
+ build_and_log_summary=lambda: None,
+ ),
+ failure_recorded=False,
+ )
+
+ obj._record_request_failure_once("req-1", reason="client_disconnect")
+ obj._log_summary_and_cleanup("req-1", reason="stage_error")
+
+ prom.request_failed.assert_called_once()
+ prom.inc_requests_failed.assert_called_once_with("client_disconnect")
+
+ def test_log_summary_and_cleanup_default_reason(self, mocker) -> None:
+ obj, prom = _make_omni_base_with_mock_prom(mocker)
+ obj.request_states["req-1"] = SimpleNamespace(
+ metrics=SimpleNamespace(
+ e2e_done=set(),
+ build_and_log_summary=lambda: None,
+ ),
+ failure_recorded=False,
+ )
+
+ obj._log_summary_and_cleanup("req-1")
+
+ prom.inc_requests_failed.assert_called_once_with("stage_error")
+
+
+class TestEarlyReturnOnLogStatsOff:
+ """When --log-stats is off (default), observe methods must be silent no-ops.
+
+ This is the gating contract — emit sites can call unconditionally and the
+ helper short-circuits.
+ """
+
+ def test_observe_methods_silent_when_log_stats_false(self) -> None:
+ prom = OmniPrometheusMetrics(model_name=_MODEL, log_stats=False)
+ # None of these should raise or write to the registry.
+ prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=1.5)
+ prom.observe_stage_in_queue(stage=0, in_queue_s=0.2)
+ prom.observe_queue_wait(queue_wait_s=0.5)
+ prom.set_stage_waiting_requests(stage=0, n_waiting=3)
+ prom.observe_num_inference_steps(n_steps=20)
+ prom.inc_image_count(n_images=1)
+ prom.observe_image_pixels(n_pixels=512 * 512)
+ prom.set_peak_memory(stage=0, peak_memory_mb=1024.0)
+ prom.inc_requests_failed(reason="oom")
+ prom.observe_kv_wait(connector_type="shm", kv_wait_s=0.01)
+
+
+class TestStageGenerationLabels:
+ def test_stage_gen_time_exposes_bounded_stage_type(self) -> None:
+ from prometheus_client import REGISTRY, generate_latest
+
+ model = _MODEL + "-stage-type"
+ prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+ prom.observe_stage_gen_time(stage=0, stage_type="llm", gen_time_s=0.25)
+ prom.observe_stage_gen_time(stage=1, stage_type="diffusion", gen_time_s=1.5)
+
+ out = generate_latest(REGISTRY).decode()
+ assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="0",stage_type="llm"}} 1.0' in out
+ assert f'{defs.STAGE_GEN_TIME_S}_count{{model_name="{model}",stage="1",stage_type="diffusion"}} 1.0' in out
+
+
+class TestCounterExpositionNames:
+ def test_counter_families_have_exactly_one_total_suffix(self) -> None:
+ from prometheus_client import REGISTRY, generate_latest
+
+ model = _MODEL + "-counter-suffix"
+ prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+ prom.inc_image_count(2)
+ prom.inc_requests_failed("stage_error")
+
+ out = generate_latest(REGISTRY).decode()
+ assert f'{defs.IMAGE_COUNT_METRIC}_total{{model_name="{model}"}} 2.0' in out
+ assert f'{defs.REQUESTS_FAILED}_total{{model_name="{model}",reason="stage_error"}} 1.0' in out
+ assert f"{defs.IMAGE_COUNT_METRIC}_total_total" not in out
+ assert f"{defs.REQUESTS_FAILED}_total_total" not in out
+
+
+class TestQueueWaitExtraction:
+ def test_present_zero_is_a_valid_observation(self) -> None:
+ from vllm_omni.metrics.utils import extract_queue_wait_s
+
+ assert extract_queue_wait_s({"queue_wait_ms": 0.0}) == 0.0
+
+ def test_missing_queue_wait_is_not_synthetic_zero(self) -> None:
+ from vllm_omni.metrics.utils import extract_queue_wait_s
+
+ assert extract_queue_wait_s({"preprocess_ms": 2.0}) is None
+ assert extract_queue_wait_s(None) is None
+
+
+class TestStageInQueueObservation:
+ def test_prometheus_records_valid_zero_wait(self) -> None:
+ from prometheus_client import REGISTRY, generate_latest
+
+ model = _MODEL + "-zero-stage-wait"
+ prom = OmniPrometheusMetrics(model_name=model, log_stats=True)
+
+ prom.observe_stage_in_queue(stage=1, in_queue_s=0.0)
+
+ out = generate_latest(REGISTRY).decode()
+ sample = f'{defs.STAGE_IN_QUEUE_S}_count{{model_name="{model}",stage="1"}} 1.0'
+ assert sample in out
+
+
+class TestStageWorkloadMetricScope:
+ @staticmethod
+ def _observe(mocker, *, stage_type: str, output_unit_type: str):
+ from vllm_omni.metrics.utils import observe_stage_workload_metrics
+
+ prom = mocker.Mock(spec=OmniPrometheusMetrics)
+ stage_metrics = SimpleNamespace(
+ num_inference_steps=50,
+ output_unit_type=output_unit_type,
+ image_pixels=1024 * 1024,
+ output_unit_count=2,
+ )
+ observe_stage_workload_metrics(
+ prom,
+ stage_type=stage_type,
+ stage_metrics=stage_metrics,
+ )
+ return prom
+
+ def test_diffusion_image_observes_steps_and_image_workload(self, mocker) -> None:
+ prom = self._observe(mocker, stage_type="diffusion", output_unit_type="image")
+
+ prom.observe_num_inference_steps.assert_called_once_with(50)
+ prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
+ prom.inc_image_count.assert_called_once_with(2)
+
+ def test_diffusion_video_observes_steps_only(self, mocker) -> None:
+ prom = self._observe(mocker, stage_type="diffusion", output_unit_type="video")
+
+ prom.observe_num_inference_steps.assert_called_once_with(50)
+ prom.observe_image_pixels.assert_not_called()
+ prom.inc_image_count.assert_not_called()
+
+ def test_llm_text_observes_no_workload_metrics(self, mocker) -> None:
+ prom = self._observe(mocker, stage_type="llm", output_unit_type="text")
+
+ prom.observe_num_inference_steps.assert_not_called()
+ prom.observe_image_pixels.assert_not_called()
+ prom.inc_image_count.assert_not_called()
+
+ def test_non_diffusion_image_observes_image_workload_only(self, mocker) -> None:
+ prom = self._observe(mocker, stage_type="llm", output_unit_type="image")
+
+ prom.observe_num_inference_steps.assert_not_called()
+ prom.observe_image_pixels.assert_called_once_with(1024 * 1024)
+ prom.inc_image_count.assert_called_once_with(2)
+
+ @pytest.mark.parametrize("peak_memory_mb", [2048.0, 0.0])
+ def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak_memory_mb: float) -> None:
+ from vllm_omni.entrypoints.omni_base import OmniBase
+
+ obj = object.__new__(OmniBase)
+ obj._enable_ar_profiler = False
+ obj._consumed_metric_messages = {}
+ obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics)
+ obj.mod_metrics = mocker.Mock()
+ obj.engine = SimpleNamespace(
+ get_stage_metadata=lambda _stage_id: SimpleNamespace(
+ stage_type="diffusion",
+ final_output=False,
+ final_output_type="image",
+ )
+ )
+ stage_metrics = SimpleNamespace(
+ stage_gen_time_ms=1000.0,
+ diffusion_metrics={"diffusion_engine_exec_time_s": 0.8, "scheduler_queue_wait_s": 0.1},
+ num_inference_steps=20,
+ output_unit_type="image",
+ image_pixels=1024 * 1024,
+ output_unit_count=1,
+ serving_time_to_first_output_ms=250.0,
+ image_time_to_first_output_ms=250.0,
+ pipeline_timings={},
+ )
+ result = SimpleNamespace(
+ request_id="req-replay",
+ stage_id=1,
+ replica_id=0,
+ stage_submit_ts=10.0,
+ metrics=stage_metrics,
+ engine_outputs=SimpleNamespace(
+ stage_durations={},
+ peak_memory_mb=peak_memory_mb,
+ finished=True,
+ final_output_type="image",
+ ),
+ )
+ aggregator = mocker.Mock()
+ aggregator.stage_events = {}
+ aggregator.stage_first_ts = [None, None]
+ aggregator.stage_last_ts = [None, None]
+
+ for _ in range(2):
+ assert obj._process_single_result(result, 1, aggregator, {}, 0.0, 1) is None
+
+ aggregator.on_stage_metrics.assert_called_once_with(1, "req-replay", stage_metrics, "image")
+ obj.prom_metrics.inc_image_count.assert_called_once_with(1)
+ obj.prom_metrics.observe_image_pixels.assert_called_once_with(1024 * 1024)
+ obj.prom_metrics.observe_num_inference_steps.assert_called_once_with(20)
+ obj.prom_metrics.observe_stage_in_queue.assert_called_once_with(1, 0.1)
+ if peak_memory_mb > 0:
+ obj.prom_metrics.set_peak_memory.assert_called_once_with(1, peak_memory_mb)
+ else:
+ obj.prom_metrics.set_peak_memory.assert_not_called()
+ obj.mod_metrics.observe_image_ttfp.assert_called_once_with("1", "0", 0.25)
+
+
+class TestStageWaitingAggregation:
+ @staticmethod
+ def _make_orchestrator(mocker):
+ from vllm_omni.engine.orchestrator import Orchestrator
+
+ orchestrator = object.__new__(Orchestrator)
+ orchestrator._prom_metrics = mocker.Mock()
+ orchestrator._stage_replica_waiting = {}
+ return orchestrator
+
+ def test_multiple_replicas_are_summed(self, mocker) -> None:
+ orchestrator = self._make_orchestrator(mocker)
+
+ orchestrator._update_stage_replica_waiting(1, 0, 3)
+ orchestrator._update_stage_replica_waiting(1, 1, 5)
+
+ orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 8)
+
+ def test_zero_update_and_stage_isolation(self, mocker) -> None:
+ orchestrator = self._make_orchestrator(mocker)
+
+ orchestrator._update_stage_replica_waiting(1, 0, 3)
+ orchestrator._update_stage_replica_waiting(2, 0, 7)
+ orchestrator._update_stage_replica_waiting(1, 0, 0)
+
+ orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 0)
+ assert orchestrator._stage_replica_waiting[(2, 0)] == 7
+
+ def test_dead_replica_snapshot_is_removed(self, mocker) -> None:
+ orchestrator = self._make_orchestrator(mocker)
+ orchestrator._update_stage_replica_waiting(1, 0, 3)
+ orchestrator._update_stage_replica_waiting(1, 1, 5)
+
+ orchestrator._remove_stage_replica_waiting(1, 1)
+
+ orchestrator._prom_metrics.set_stage_waiting_requests.assert_called_with(1, 3)
+
+
+# ---------------------------------------------------------------------------
+# Behavioral pin: call sites now scope workload observations by stage/output
+# type. Keep the Prometheus wrapper's <= 0 guard as a second line of defense.
+# ---------------------------------------------------------------------------
+
+
+class TestZeroGuardContract:
+ """Zero-valued observations must not bump histogram ``_count``."""
+
+ def test_zero_image_pixels_not_observed(self) -> None:
+ from prometheus_client import REGISTRY, generate_latest
+
+ prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-guard")
+ prom.observe_image_pixels(n_pixels=0)
+ out = generate_latest(REGISTRY).decode()
+ # ``.labels()`` in OmniPrometheusMetrics.__init__ already creates the
+ # child sample, so the ``_count`` line exists with value 0.0 even
+ # without observations. The guard prevents ``_count`` from bumping
+ # to 1.0 — assert the value stays at 0.
+ needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-zero-guard"}}'
+ count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+ assert count_lines, "expected image_pixels_count line to exist after OmniPrometheusMetrics construction"
+ assert float(count_lines[0].split()[-1]) == 0.0, (
+ "zero-pixel observation leaked to registry — guard should early-return on <= 0"
+ )
+
+ def test_zero_num_inference_steps_not_observed(self) -> None:
+ from prometheus_client import REGISTRY, generate_latest
+
+ prom = OmniPrometheusMetrics(model_name=_MODEL + "-zero-steps")
+ prom.observe_num_inference_steps(n_steps=0)
+ out = generate_latest(REGISTRY).decode()
+ needle = f'vllm_omni:num_inference_steps_count{{model_name="{_MODEL}-zero-steps"}}'
+ count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+ assert count_lines, "expected num_inference_steps_count line to exist after construction"
+ assert float(count_lines[0].split()[-1]) == 0.0, (
+ "zero-step observation leaked to registry — guard should early-return on <= 0"
+ )
+
+ def test_positive_image_pixels_observed(self) -> None:
+ # Sanity check: positive values DO get observed.
+ from prometheus_client import REGISTRY, generate_latest
+
+ prom = OmniPrometheusMetrics(model_name=_MODEL + "-pos-guard")
+ prom.observe_image_pixels(n_pixels=512)
+ out = generate_latest(REGISTRY).decode()
+ needle = f'vllm_omni:image_pixels_count{{model_name="{_MODEL}-pos-guard"}}'
+ count_lines = [ln for ln in out.splitlines() if ln.startswith(needle)]
+ assert count_lines, "expected image_pixels_count line after positive observation"
+ assert float(count_lines[0].split()[-1]) == 1.0, "positive-pixel observation did not increment count to 1"
+
+
+# ---------------------------------------------------------------------------
+# KV-wait emit wiring — scheduler ENTER/EXIT lifecycle + orchestrator dispatch.
+# ---------------------------------------------------------------------------
+
+
+def _make_scheduler_shell() -> OmniARScheduler:
+ """Minimal OmniARScheduler shell — skips upstream __init__ via object.__new__."""
+ from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+ obj = object.__new__(OmniARScheduler)
+ obj.requests = {}
+ obj.running = []
+ obj.waiting = []
+ obj._kv_wait_start_ts = {}
+ obj._omni_kv_config = None
+ return obj
+
+
+class TestKvWaitSchedulerEmit:
+ """Scheduler-side half of kv_wait_s: pops start ts, carries the wait across."""
+
+ def test_skips_when_no_start_ts_recorded(self) -> None:
+ sched = _make_scheduler_shell()
+ outputs: dict[int, list] = {}
+
+ sched._emit_kv_wait_output(outputs, "req-no-wait", req=SimpleNamespace(client_index=0))
+
+ assert outputs == {}, "emit must not append output when no start ts recorded"
+ assert "req-no-wait" not in sched._kv_wait_start_ts
+
+ def test_emits_wait_duration_and_connector_type(self) -> None:
+ sched = _make_scheduler_shell()
+
+ sched._kv_wait_start_ts["req-1"] = time.monotonic() - 0.25
+ outputs: dict[int, list] = {}
+ live_req = SimpleNamespace(client_index=3)
+
+ sched._emit_kv_wait_output(outputs, "req-1", req=live_req)
+
+ assert "req-1" not in sched._kv_wait_start_ts, "start ts must be popped after emit"
+ assert 3 in outputs and len(outputs[3]) == 1
+ params = outputs[3][0].kv_transfer_params
+ assert "kv_wait_s" in params and "connector_type" in params
+ assert 0.0 < params["kv_wait_s"] < 1.0
+ assert params["connector_type"] == "unknown"
+
+ def test_connector_type_resolved_from_omni_kv_config(self) -> None:
+ sched = _make_scheduler_shell()
+ sched._omni_kv_config = {"connector_config": {"type": "SharedMemoryConnector"}}
+
+ sched._kv_wait_start_ts["req-3"] = time.monotonic() - 0.01
+ outputs: dict[int, list] = {}
+
+ sched._emit_kv_wait_output(outputs, "req-3", req=SimpleNamespace(client_index=0))
+
+ eco = outputs[0][0]
+ assert eco.kv_transfer_params["connector_type"] == "SharedMemoryConnector"
+
+
+class TestKvWaitTerminalCleanup:
+ def test_finish_requests_clears_requested_wait_only(self) -> None:
+ from vllm.v1.request import RequestStatus
+
+ scheduler = _make_scheduler_shell()
+ scheduler._kv_wait_start_ts = {"req-1": 10.0, "req-2": 20.0}
+
+ finished = scheduler.finish_requests(
+ (request_id for request_id in ["req-1"]),
+ RequestStatus.FINISHED_ABORTED,
+ )
+
+ assert finished == []
+ assert "req-1" not in scheduler._kv_wait_start_ts
+ assert scheduler._kv_wait_start_ts["req-2"] == 20.0
+
+ def test_cleanup_is_idempotent_when_timestamp_is_already_gone(self) -> None:
+ from vllm.v1.request import RequestStatus
+
+ scheduler = _make_scheduler_shell()
+
+ scheduler.finish_requests("req-already-emitted", RequestStatus.FINISHED_ABORTED)
+
+ assert scheduler._kv_wait_start_ts == {}
diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py
index 315594233..694808a31 100644
--- a/tests/metrics/test_modality.py
+++ b/tests/metrics/test_modality.py
@@ -1,5 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
+from types import SimpleNamespace
+
import pytest
from prometheus_client import REGISTRY, generate_latest
@@ -10,6 +15,7 @@ from vllm_omni.metrics.modality import (
observe_audio_streaming_finalize,
observe_modality_at_finalize,
)
+from vllm_omni.metrics.stats import OrchestratorAggregator
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -46,6 +52,11 @@ _EXPECTED_FAMILIES = [
defs.DIFFUSION_EXEC_PER_STEP_S,
defs.DIFFUSION_PREPROCESS_S,
defs.DIFFUSION_POSTPROCESS_S,
+ defs.VAE_DECODE_S,
+ defs.DIFFUSION_FORWARD_S,
+ defs.DIFFUSION_KV_LOAD_S,
+ defs.IMAGE_TTFP_S,
+ defs.DENOISE_STEP_LATENCY_S,
]
@@ -63,6 +74,11 @@ class TestRegistration:
mod.observe_diffusion_exec_per_step("s", "r", 0.01)
mod.observe_diffusion_preprocess("s", "r", 0.01)
mod.observe_diffusion_postprocess("s", "r", 0.2)
+ mod.observe_vae_decode("s", "r", 0.3)
+ mod.observe_diffusion_forward("s", "r", 0.8)
+ mod.observe_diffusion_kv_load("s", "r", 0.04)
+ mod.observe_image_ttfp("s", "r", 2.5)
+ mod.observe_denoise_step_latency("s", "r", 0.05)
out = generate_latest(REGISTRY).decode()
for name in _EXPECTED_FAMILIES:
@@ -247,6 +263,21 @@ class _StubModMetrics:
def observe_diffusion_postprocess(self, s, r, v):
self.calls.append(("observe_diffusion_postprocess", s, r, v))
+ def observe_vae_decode(self, s, r, seconds):
+ self.calls.append(("observe_vae_decode", s, r, seconds))
+
+ def observe_diffusion_forward(self, s, r, seconds):
+ self.calls.append(("observe_diffusion_forward", s, r, seconds))
+
+ def observe_diffusion_kv_load(self, s, r, seconds):
+ self.calls.append(("observe_diffusion_kv_load", s, r, seconds))
+
+ def observe_image_ttfp(self, s, r, seconds):
+ self.calls.append(("observe_image_ttfp", s, r, seconds))
+
+ def observe_denoise_step_latency(self, s, r, seconds):
+ self.calls.append(("observe_denoise_step_latency", s, r, seconds))
+
class _Bag:
def __init__(self, **kw):
@@ -306,11 +337,13 @@ class TestObserveModalityAtFinalize:
def test_diffusion_path_full(self):
stub = _StubModMetrics()
stage_metrics = _Bag(
+ output_unit_type="image",
stage_gen_time_ms=1000.0,
diffusion_metrics={
"preprocess_time_s": 0.05,
"diffusion_engine_exec_time_s": 1.2,
"postprocess_time_s": 0.3,
+ "forward_time_s": 1.0,
"image_num": 1,
"resolution": 1024,
"num_inference_steps": 50,
@@ -328,6 +361,71 @@ class TestObserveModalityAtFinalize:
assert ("observe_diffusion_exec", "2", "0", 1.2) in stub.calls
assert ("observe_diffusion_postprocess", "2", "0", 0.3) in stub.calls
assert ("observe_diffusion_exec_per_step", "2", "0", pytest.approx(0.024)) in stub.calls
+ assert ("observe_denoise_step_latency", "2", "0", pytest.approx(0.02)) in stub.calls
+
+ def test_denoise_step_latency_requires_forward_time(self):
+ stub = _StubModMetrics()
+ stage_metrics = _Bag(
+ output_unit_type="image",
+ num_inference_steps=20,
+ diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
+ )
+ observe_modality_at_finalize(
+ stub,
+ output_type="image",
+ stage_id=2,
+ replica_id=0,
+ stage_metrics=stage_metrics,
+ engine_outputs=_Bag(),
+ )
+ assert not any(c[0] == "observe_denoise_step_latency" for c in stub.calls)
+
+ def test_stage_snapshot_keeps_forward_based_denoise_latency(self):
+ aggregator = OrchestratorAggregator(
+ num_stages=1,
+ log_stats=False,
+ wall_start_ts=0.0,
+ final_stage_id_for_e2e=0,
+ )
+ aggregator.diffusion_metrics["req-image"]["forward_time_s"] = 0.8
+ stage_metrics = _Bag(
+ stage_id=None,
+ request_id=None,
+ final_output_type=None,
+ num_tokens_in=0,
+ num_tokens_out=0,
+ stage_gen_time_ms=1200.0,
+ rx_transfer_bytes=0,
+ rx_decode_time_ms=0.0,
+ rx_in_flight_time_ms=0.0,
+ stage_stats=_Bag(),
+ diffusion_metrics=None,
+ num_inference_steps=20,
+ output_unit_type="image",
+ denoise_step_latency_ms=0.0,
+ )
+
+ stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
+
+ assert stats.denoise_step_latency_ms == pytest.approx(40.0)
+
+ def test_stage_snapshot_falls_back_to_stage_time_without_profiler(self):
+ aggregator = object.__new__(OrchestratorAggregator)
+ stage_metrics = _Bag(
+ stage_id=None,
+ request_id=None,
+ final_output_type=None,
+ stage_gen_time_ms=1200.0,
+ diffusion_metrics=None,
+ num_inference_steps=20,
+ output_unit_type="image",
+ denoise_step_latency_ms=0.0,
+ )
+
+ stats = aggregator._as_stage_request_stats(0, "req-image", stage_metrics, "image")
+
+ assert stats.diffusion_metrics is None
+ assert stats.denoise_step_latency_ms == pytest.approx(60.0)
def test_diffusion_per_step_skipped_without_num_inference_steps(self):
stub = _StubModMetrics()
@@ -422,6 +520,80 @@ class TestObserveModalityAtFinalize:
assert stub.calls == []
+class TestDiffusionUnitConversion:
+ def test_accumulator_converts_all_engine_milliseconds_to_seconds(self):
+ aggregator = OrchestratorAggregator(
+ num_stages=1,
+ log_stats=False,
+ wall_start_ts=0.0,
+ final_stage_id_for_e2e=0,
+ )
+ engine_output = SimpleNamespace(
+ metrics={
+ "preprocess_time_ms": 50.0,
+ "diffusion_engine_exec_time_ms": 1500.0,
+ "postprocess_time_ms": 20.0,
+ "vae_decode_time_ms": 300.0,
+ "forward_time_ms": 800.0,
+ "scheduler_queue_wait_ms": 125.0,
+ "kv_recv_time_ms": 40.0,
+ }
+ )
+
+ aggregator.accumulate_diffusion_metrics("diffusion", "req-units", engine_output)
+
+ assert dict(aggregator.diffusion_metrics["req-units"]) == {
+ "preprocess_time_s": pytest.approx(0.05),
+ "diffusion_engine_exec_time_s": pytest.approx(1.5),
+ "postprocess_time_s": pytest.approx(0.02),
+ "vae_decode_time_s": pytest.approx(0.3),
+ "forward_time_s": pytest.approx(0.8),
+ "scheduler_queue_wait_s": pytest.approx(0.125),
+ "kv_recv_time_s": pytest.approx(0.04),
+ }
+
+
+class TestObserveDiffusionFinalize:
+ @pytest.mark.parametrize(
+ ("metric_key", "call_name", "value"),
+ [
+ ("vae_decode_time_s", "observe_vae_decode", 0.3),
+ ("forward_time_s", "observe_diffusion_forward", 0.8),
+ ("kv_recv_time_s", "observe_diffusion_kv_load", 0.04),
+ ],
+ )
+ def test_emits_optional_breakdown_metric(self, metric_key, call_name, value):
+ stub = _StubModMetrics()
+ observe_modality_at_finalize(
+ stub,
+ output_type="image",
+ stage_id=2,
+ replica_id=0,
+ stage_metrics=_Bag(diffusion_metrics={metric_key: value}),
+ engine_outputs=_Bag(),
+ )
+
+ assert (call_name, "2", "0", value) in stub.calls
+
+ def test_audio_diffusion_stage_fires_both_paths(self):
+ stub = _StubModMetrics()
+ observe_modality_at_finalize(
+ stub,
+ output_type="audio",
+ stage_id=1,
+ replica_id=0,
+ stage_metrics=_Bag(
+ stage_gen_time_ms=1500.0,
+ audio_generated_frames=24000,
+ diffusion_metrics={"diffusion_engine_exec_time_s": 1.2},
+ ),
+ engine_outputs=_Bag(multimodal_output={"audio_sample_rate": 24000}),
+ )
+
+ assert ("observe_diffusion_exec", "1", "0", 1.2) in stub.calls
+ assert ("observe_audio_duration", "1", "0", 1.0) in stub.calls
+
+
class TestObserveAudioFirstPacket:
def test_observes_with_valid_inputs(self):
stub = _StubModMetrics()
@@ -541,3 +713,39 @@ class TestBucketSelection:
# SECONDS_FAST_BUCKETS includes le=0.001 which is absent from SECONDS_BUCKETS.
fast_marker = f'{defs.AUDIO_UNDERRUN_S}_bucket{{le="0.001"'
assert fast_marker in out, "audio_underrun_s should use SECONDS_FAST_BUCKETS containing le=0.001"
+
+ def test_vae_decode_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+ # VAE decode can stretch into seconds for high-resolution latents;
+ # stays on the wide bucket so the 300 s ceiling covers video VAE.
+ stage, replica = "vae_decode_sec", "0"
+ mod.observe_vae_decode(stage, replica, 1.5)
+ out = generate_latest(REGISTRY).decode()
+ sec_marker = f'{defs.VAE_DECODE_S}_bucket{{le="300.0"'
+ assert sec_marker in out, "vae_decode_s should use SECONDS_BUCKETS containing le=300.0"
+
+ def test_diffusion_forward_uses_seconds_buckets(self, mod: OmniModalityMetrics) -> None:
+ # Denoise loop can exceed the 60s fast-bucket ceiling for video gen.
+ stage, replica = "diff_forward_sec", "0"
+ mod.observe_diffusion_forward(stage, replica, 2.0)
+ out = generate_latest(REGISTRY).decode()
+ sec_marker = f'{defs.DIFFUSION_FORWARD_S}_bucket{{le="300.0"'
+ assert sec_marker in out, "diffusion_forward_s should use SECONDS_BUCKETS containing le=300.0"
+
+ def test_denoise_step_latency_uses_seconds_fast_buckets(self, mod: OmniModalityMetrics) -> None:
+ # Per-step denoise is sub-second; stays on the fast bucket for ms-level
+ # resolution.
+ stage, replica = "denoise_fast", "0"
+ mod.observe_denoise_step_latency(stage, replica, 0.075)
+ out = generate_latest(REGISTRY).decode()
+ fast_marker = f'{defs.DENOISE_STEP_LATENCY_S}_bucket{{le="0.001"'
+ assert fast_marker in out, "denoise_step_latency_s should use SECONDS_FAST_BUCKETS containing le=0.001"
+
+
+class TestDenoiseStepLatencyGuard:
+ @pytest.mark.parametrize(("value", "expected_count"), [(0.0, None), (-0.1, None), (0.05, 1.0)])
+ def test_nonpositive_values_are_skipped(self, mod: OmniModalityMetrics, value, expected_count) -> None:
+ stage, replica = f"denoise_{value}", "0"
+ mod.observe_denoise_step_latency(stage, replica, value)
+ out = generate_latest(REGISTRY).decode()
+ prefix = f'{defs.DENOISE_STEP_LATENCY_S}_count{{model_name="{_MODEL}",replica="{replica}",stage="{stage}"}}'
+ assert _sample_value(out, prefix) == expected_count
diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py
index f01933dbc..e45365cf1 100644
--- a/tests/metrics/test_prometheus.py
+++ b/tests/metrics/test_prometheus.py
@@ -198,3 +198,32 @@ class TestRequestLifecycleGauges:
out = generate_latest(registry).decode()
assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-2"}') == 1.0
assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0
+
+ def test_engine_queued_requests_count_as_waiting(self, registry: CollectorRegistry) -> None:
+ from types import SimpleNamespace
+
+ from vllm_omni.entrypoints.omni_base import OmniBase
+ from vllm_omni.metrics.prometheus import OmniPrometheusMetrics, OmniRequestCounter
+
+ obj = object.__new__(OmniBase)
+ obj.engine = SimpleNamespace(
+ _running_counter=OmniRequestCounter(),
+ _engines_waiting_counter=OmniRequestCounter(),
+ )
+ obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-3")
+ obj.request_states = {}
+ obj.log_stats = False
+
+ # Three requests dispatched to a stage engine whose scheduler admits
+ # one and queues two (e.g. a diffusion engine with
+ # max_num_running_reqs=1). The orchestrator reports the queued count
+ # via _engines_waiting_counter; they must show as waiting, not running.
+ for _ in range(3):
+ obj.engine._running_counter.increment()
+ obj.engine._engines_waiting_counter.value = 2
+
+ obj._publish_request_gauges(3)
+
+ out = generate_latest(registry).decode()
+ assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-3"}') == 1.0
+ assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-3"}') == 2.0
diff --git a/tests/model_executor/model_loader/test_resolve_model_path.py b/tests/model_executor/model_loader/test_resolve_model_path.py
new file mode 100644
index 000000000..996313b06
--- /dev/null
+++ b/tests/model_executor/model_loader/test_resolve_model_path.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for resolve_model_to_local_path."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import huggingface_hub
+import pytest
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+from vllm_omni.model_executor.model_loader import weight_utils
+from vllm_omni.model_executor.model_loader.weight_utils import resolve_model_to_local_path
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Both resolution paths are patched in every test, so nothing is ever fetched
+REPO_ID = "test-org/test-model"
+RESOLVED = "/cache/snapshots/deadbeef"
+
+
+@pytest.fixture
+def paths(mocker):
+ """The resolver's two ways out, both patched."""
+ return SimpleNamespace(
+ cache_only=mocker.patch.object(huggingface_hub, "snapshot_download", return_value=RESOLVED),
+ download=mocker.patch.object(weight_utils, "download_weights_from_hf_specific", return_value=RESOLVED),
+ )
+
+
+@pytest.mark.parametrize("allow_download", [False, True])
+def test_local_directory_short_circuits(tmp_path, paths, allow_download):
+ """A local dir resolves to itself without either path running."""
+ assert resolve_model_to_local_path(str(tmp_path), allow_download=allow_download) == str(tmp_path)
+
+ paths.cache_only.assert_not_called()
+ paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "allow_patterns", "cache_dir"),
+ [
+ ({}, None, None),
+ ({"allow_patterns": ["subdir/*"], "cache_dir": "/cache"}, ["subdir/*"], "/cache"),
+ ],
+)
+def test_repo_id_is_cache_only_by_default(paths, kwargs, allow_patterns, cache_dir):
+ """Without `allow_download` nothing is fetched, and narrowing is forwarded."""
+ assert resolve_model_to_local_path(REPO_ID, **kwargs) == RESOLVED
+
+ paths.cache_only.assert_called_once_with(
+ REPO_ID,
+ cache_dir=cache_dir,
+ allow_patterns=allow_patterns,
+ local_files_only=True,
+ )
+ paths.download.assert_not_called()
+
+
+def test_cache_miss_propagates(paths):
+ """A failed resolution raises instead of degrading into the raw repo id."""
+ paths.cache_only.side_effect = LocalEntryNotFoundError("cold cache")
+
+ with pytest.raises(LocalEntryNotFoundError):
+ resolve_model_to_local_path(REPO_ID)
+
+ paths.download.assert_not_called()
+
+
+@pytest.mark.parametrize(
+ ("kwargs", "allow_patterns"),
+ [
+ ({}, ["*"]),
+ ({"allow_patterns": ["subdir/model*.safetensors"]}, ["subdir/model*.safetensors"]),
+ ],
+)
+def test_download_routes_through_the_locked_helper(paths, kwargs, allow_patterns):
+ """Downloads go through the vLLM helper holding the node-wide lock.
+
+ That helper is also what honors `VLLM_USE_MODELSCOPE`;
+ calling `snapshot_download` directly would drop both.
+ """
+ assert resolve_model_to_local_path(REPO_ID, allow_download=True, cache_dir="/cache", **kwargs) == RESOLVED
+
+ paths.download.assert_called_once_with(
+ model_name_or_path=REPO_ID,
+ cache_dir="/cache",
+ allow_patterns=allow_patterns,
+ require_all=True,
+ )
+ paths.cache_only.assert_not_called()
diff --git a/tests/model_executor/models/common/__init__.py b/tests/model_executor/models/common/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py
new file mode 100644
index 000000000..e202321dc
--- /dev/null
+++ b/tests/model_executor/models/common/test_fused_adaptive_group_norm_silu.py
@@ -0,0 +1,292 @@
+"""Unit tests for fused_adaptive_group_norm_silu operator.
+
+Tests numeric correctness against PyTorch native implementation:
+ F.silu(F.group_norm(x, num_groups, weight, bias, eps) * (1 + scale) + shift)
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.common.ops import fused_adaptive_group_norm_silu
+
+# Skip tests if CUDA not available
+pytestmark = [
+ pytest.mark.core_model,
+ pytest.mark.cuda,
+ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"),
+]
+
+
+def _reference(x, weight, bias, scale, shift, num_groups, eps):
+ """AdaGN reference, computed in fp32.
+
+ As with fused_group_norm_silu, a same-dtype reference is unusable: the
+ GroupNorm affine can cancel to near zero, and rounding its O(1) terms to
+ bf16 wipes out the remaining significant digits.
+ """
+ B, C = x.shape[:2]
+ normed = F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)
+ scale_b = scale.float().view(B, C, *([1] * (x.ndim - 2)))
+ shift_b = shift.float().view(B, C, *([1] * (x.ndim - 2)))
+ return F.silu(normed * (1 + scale_b) + shift_b)
+
+
+def _tolerance(dtype):
+ if dtype == torch.float32:
+ return 1e-5, 1e-5
+ if dtype == torch.float16:
+ return 2e-3, 2e-3
+ return 1e-2, 1e-2 # bfloat16
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("channels", [32, 64, 128])
+@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32)])
+@pytest.mark.parametrize("num_groups", [8, 32])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_fused_adaptive_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype):
+ """Test numeric correctness against PyTorch native ops."""
+ if channels % num_groups != 0:
+ pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}")
+
+ H, W = spatial_size
+ eps = 1e-6
+ torch.manual_seed(42)
+ kw = dict(device="cuda", dtype=dtype)
+
+ x = torch.randn(batch_size, channels, H, W, **kw)
+ weight = torch.randn(channels, **kw)
+ bias = torch.randn(channels, **kw)
+ scale = torch.randn(batch_size, channels, **kw)
+ shift = torch.randn(batch_size, channels, **kw)
+
+ ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype)
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps)
+
+ rtol, atol = _tolerance(dtype)
+ torch.testing.assert_close(
+ fused_out,
+ ref_out,
+ rtol=rtol,
+ atol=atol,
+ msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}",
+ )
+
+
+def test_broadcast_scale_shift_shapes():
+ """scale/shift may arrive as (B, C) or already broadcast to (B, C, 1, 1)."""
+ B, C, H, W = 2, 64, 16, 16
+ kw = dict(device="cuda", dtype=torch.float32)
+ x = torch.randn(B, C, H, W, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ flat = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+ expanded = fused_adaptive_group_norm_silu(x, weight, bias, scale.view(B, C, 1, 1), shift.view(B, C, 1, 1), 32, 1e-6)
+ torch.testing.assert_close(flat, expanded, rtol=0, atol=0)
+
+
+def test_output_dtype_matches_eager():
+ """Output dtype must match what eager GroupNorm would produce."""
+ for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+ B, C = 2, 64
+ kw = dict(device="cuda", dtype=dtype)
+ x = torch.randn(B, C, 16, 16, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ eager_dtype = F.silu(F.group_norm(x, 32, weight, bias, 1e-6)).dtype
+ out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+
+ assert out.dtype == eager_dtype, (
+ f"Output dtype {out.dtype} != eager dtype {eager_dtype} for input dtype {dtype}"
+ )
+ assert out.shape == x.shape
+
+
+@pytest.mark.parametrize(
+ "autocast_dtype, input_dtype",
+ [
+ (torch.float16, torch.float16),
+ (torch.bfloat16, torch.bfloat16),
+ (torch.float16, torch.float32),
+ (torch.bfloat16, torch.float32),
+ ],
+)
+def test_autocast_matches_eager(autocast_dtype, input_dtype):
+ """Inside autocast the fused op must stay a drop-in replacement.
+
+ Also pins the fp32-activation corner so the inlined ``is_autocast_enabled
+ -> fp32`` branch cannot silently downcast an already-fp32 activation.
+ """
+ B, C = 2, 64
+ x = torch.randn(B, C, 16, 16, device="cuda", dtype=input_dtype)
+ weight = torch.randn(C, device="cuda", dtype=torch.float32)
+ bias = torch.randn(C, device="cuda", dtype=torch.float32)
+ scale = torch.randn(B, C, device="cuda", dtype=torch.float32)
+ shift = torch.randn(B, C, device="cuda", dtype=torch.float32)
+
+ with torch.autocast("cuda", dtype=autocast_dtype):
+ eager_out = F.silu(
+ F.group_norm(x, 32, weight, bias, 1e-6) * (1 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1)
+ )
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+
+ assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast"
+ assert fused_out.dtype == eager_out.dtype, (
+ f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} "
+ f"under autocast({autocast_dtype}) for input {input_dtype}"
+ )
+ torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5)
+
+
+def test_invalid_channels_not_divisible():
+ """Test that assertion fires when channels not divisible by num_groups."""
+ B, C = 1, 33
+ kw = dict(device="cuda", dtype=torch.float32)
+ x = torch.randn(B, C, 16, 16, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ with pytest.raises(AssertionError, match="must be divisible by num_groups"):
+ fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 8, 8), (2, 64, 128)])
+def test_non_2d_spatial_matches_eager(dtype, shape):
+ """Any spatial rank must work; the wrapper collapses the spatial axes.
+
+ GroupNorm reduces over the whole channel group and every spatial position,
+ so flattening is exact rather than an approximation.
+ """
+ torch.manual_seed(0)
+ B, C = shape[:2]
+ kw = dict(device="cuda", dtype=dtype)
+ x = torch.randn(*shape, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+ ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-6).to(dtype)
+
+ assert fused_out.shape == x.shape
+ rtol, atol = _tolerance(dtype)
+ torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 3])
+def test_scale_shift_from_chunk_are_not_contiguous(batch_size):
+ """Regression: scale/shift arrive as halves of a torch.chunk.
+
+ ResBlock builds them as ``torch.chunk(emb_out, 2, dim=1)`` where emb_out is
+ (B, 2C, 1, 1). For B > 1 each half is a strided view, so the wrapper must
+ ``reshape`` (which copies) instead of ``view`` (which raises).
+ """
+ C, H, W = 64, 16, 16
+ kw = dict(device="cuda", dtype=torch.float32)
+ x = torch.randn(batch_size, C, H, W, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+
+ emb_out = torch.randn(batch_size, 2 * C, 1, 1, **kw)
+ scale, shift = torch.chunk(emb_out, 2, dim=1)
+ if batch_size > 1:
+ assert not scale.is_contiguous(), "test premise: chunk half is strided"
+
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+ ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale) + shift)
+
+ torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5)
+
+
+def test_non_contiguous_input_matches_eager():
+ """A channels-last / permuted activation must still be handled correctly."""
+ B, C, H, W = 2, 64, 16, 16
+ kw = dict(device="cuda", dtype=torch.float32)
+ x = torch.randn(B, H, W, C, **kw).permute(0, 3, 1, 2)
+ assert not x.is_contiguous(), "test premise: input is strided"
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-6)
+ ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6) * (1.0 + scale.view(B, C, 1, 1)) + shift.view(B, C, 1, 1))
+
+ torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize("spatial", [32, 64])
+def test_dit_sized_activation(spatial):
+ """Shapes the HunyuanImage3 DiT actually sees.
+
+ The earlier kernel walked the spatial axis one position at a time with
+ ``spatial_size`` as a ``tl.constexpr``; at 64x64 that is 4096 serial steps
+ per pass. This test pins the realistic size so a regression to that layout
+ shows up as a timeout rather than silently costing throughput.
+ """
+ B, C = 2, 256
+ kw = dict(device="cuda", dtype=torch.bfloat16)
+ x = torch.randn(B, C, spatial, spatial, **kw)
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+ scale = torch.randn(B, C, **kw)
+ shift = torch.randn(B, C, **kw)
+
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, 32, 1e-5)
+ ref_out = _reference(x, weight, bias, scale, shift, 32, 1e-5).to(torch.bfloat16)
+
+ rtol, atol = _tolerance(torch.bfloat16)
+ torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_large_offset_small_variance(dtype):
+ """A large constant offset with a small variance must stay stable.
+
+ Mirrors the fused_group_norm_silu case: ``x ~ 10000 ± 0.1`` has mean ~1e4
+ and variance ~1e-2, so ``E[x^2] - E[x]^2`` cancels catastrophically in fp32
+ (rounding the variance to ~8.0). The Welford reduction keeps it accurate.
+ fp32 tolerance is 0.1 (not 1e-6) because ``x``'s fp32 representation noise
+ (~0.001 ULP at 1e4) is amplified by ``rstd ~ 10`` to ~0.01-0.02 — the same
+ floor PyTorch's own reference hits; the old naive implementation would be
+ off by ~1.0 here. fp16/bf16 quantize the perturbation away, so for those
+ dtypes the meaningful assertion is just "no NaN".
+ """
+ torch.manual_seed(0)
+ B, C, H, W = 2, 64, 16, 16
+ num_groups = 32
+ eps = 1e-6
+ device = torch.device("cuda")
+
+ x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype)
+ x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1
+ weight = torch.randn(C, device=device, dtype=dtype)
+ bias = torch.randn(C, device=device, dtype=dtype)
+ scale = torch.randn(B, C, device=device, dtype=dtype)
+ shift = torch.randn(B, C, device=device, dtype=dtype)
+
+ fused_out = fused_adaptive_group_norm_silu(x, weight, bias, scale, shift, num_groups, eps)
+ ref_out = _reference(x, weight, bias, scale, shift, num_groups, eps).to(dtype)
+
+ assert not torch.isnan(fused_out).any(), "fused output must not contain NaN"
+
+ if dtype == torch.float32:
+ # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0.
+ torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1)
+ # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close,
+ # the no-NaN assertion above is the regression guard.
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/model_executor/models/common/test_fused_group_norm_silu.py b/tests/model_executor/models/common/test_fused_group_norm_silu.py
new file mode 100644
index 000000000..3e9d055af
--- /dev/null
+++ b/tests/model_executor/models/common/test_fused_group_norm_silu.py
@@ -0,0 +1,338 @@
+"""Unit tests for fused_group_norm_silu operator.
+
+Tests numeric correctness against PyTorch native implementation:
+ F.silu(F.group_norm(x, num_groups, weight, bias, eps))
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.common.ops import fused_group_norm_silu
+
+# Skip tests if CUDA not available
+pytestmark = [
+ pytest.mark.core_model,
+ pytest.mark.cuda,
+ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for Triton kernels"),
+]
+
+
+def _tolerance(dtype):
+ """Tolerance against an fp32 reference, rounded once to ``dtype``."""
+ if dtype == torch.float32:
+ return 1e-5, 1e-6
+ if dtype == torch.float16:
+ return 2e-3, 2e-3
+ return 1e-2, 1e-2 # bfloat16
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 4])
+@pytest.mark.parametrize("channels", [32, 64, 128])
+@pytest.mark.parametrize("spatial_size", [(16, 16), (32, 32), (64, 64)])
+@pytest.mark.parametrize("num_groups", [8, 16, 32])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_fused_group_norm_silu_correctness(batch_size, channels, spatial_size, num_groups, dtype):
+ """Test numeric correctness against PyTorch native ops."""
+ # Skip invalid configurations
+ if channels % num_groups != 0:
+ pytest.skip(f"channels {channels} not divisible by num_groups {num_groups}")
+
+ H, W = spatial_size
+ device = torch.device("cuda")
+ eps = 1e-6
+
+ # Create inputs
+ torch.manual_seed(42)
+ x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype)
+ weight = torch.randn(channels, device=device, dtype=dtype)
+ bias = torch.randn(channels, device=device, dtype=dtype)
+
+ # Reference implementation, computed in fp32.
+ #
+ # A same-dtype reference is NOT usable here: GroupNorm's affine step
+ # ``(x - mean) * rstd * weight + bias`` adds two O(1) terms that can cancel
+ # down to ~1e-3. Rounding those terms to bf16 (ULP ~8e-3 near 1.0) destroys
+ # every remaining significant digit -- measured cases flip sign, e.g. a true
+ # +0.006136 comes out as -0.000471. The low-precision "reference" is then
+ # far less accurate than the kernel it is supposed to validate: the fused
+ # output matches the correctly-rounded fp32 result for 99.997% of elements,
+ # a bf16 reference for only 65.2%.
+ ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype)
+
+ # Fused implementation
+ fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps)
+
+ # Check correctness
+ # Tolerances are set from the measured worst case across this parameter
+ # sweep, leaving roughly 2x headroom for each dtype.
+ if dtype == torch.float32:
+ rtol, atol = 1e-5, 1e-6
+ elif dtype == torch.float16:
+ rtol, atol = 2e-3, 2e-3
+ else: # bfloat16
+ rtol, atol = 1e-2, 1e-2
+
+ torch.testing.assert_close(
+ fused_out,
+ ref_out,
+ rtol=rtol,
+ atol=atol,
+ msg=f"Mismatch for batch={batch_size}, C={channels}, HW={spatial_size}, groups={num_groups}, dtype={dtype}",
+ )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_fp32_accumulation_precision(dtype):
+ """Test that fp32 accumulation matches PyTorch's behavior.
+
+ GroupNorm uses fp32 for mean/variance computation regardless of input dtype.
+ Our kernel should match this behavior.
+ """
+ batch_size, channels, H, W = 2, 64, 32, 32
+ num_groups = 32
+ eps = 1e-6
+ device = torch.device("cuda")
+
+ # Create inputs with values that expose precision differences
+ torch.manual_seed(123)
+ x = torch.randn(batch_size, channels, H, W, device=device, dtype=dtype) * 100
+ weight = torch.ones(channels, device=device, dtype=dtype)
+ bias = torch.zeros(channels, device=device, dtype=dtype)
+
+ # Reference: fp32 accumulation, then a single round back to ``dtype``.
+ # (Same discipline as the main correctness test -- a same-dtype reference
+ # wouldn't isolate the accumulation path this test is named after.)
+ ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype)
+
+ # Fused
+ fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps)
+
+ # Should match closely even with large values
+ if dtype == torch.float32:
+ rtol, atol = 1e-5, 1e-5
+ else:
+ rtol, atol = 1e-2, 1e-2
+
+ torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol)
+
+
+def test_edge_case_zero_input():
+ """Test with zero input."""
+ x = torch.zeros(1, 32, 16, 16, device="cuda")
+ weight = torch.ones(32, device="cuda")
+ bias = torch.zeros(32, device="cuda")
+
+ ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6))
+ fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6)
+
+ torch.testing.assert_close(fused_out, ref_out, rtol=1e-5, atol=1e-6)
+
+
+def test_edge_case_large_eps():
+ """Test with large epsilon value."""
+ x = torch.randn(1, 32, 16, 16, device="cuda")
+ weight = torch.randn(32, device="cuda")
+ bias = torch.randn(32, device="cuda")
+ eps = 1e-2 # Larger than typical
+
+ ref_out = F.silu(F.group_norm(x, 32, weight, bias, eps))
+ fused_out = fused_group_norm_silu(x, weight, bias, 32, eps)
+
+ torch.testing.assert_close(fused_out, ref_out, rtol=1e-4, atol=1e-5)
+
+
+def test_output_dtype_matches_eager():
+ """Output dtype must match what eager GroupNorm+SiLU would produce.
+
+ Outside autocast that means the input dtype; inside autocast it means fp32,
+ because group_norm is on autocast's fp32 cast policy. HunyuanImage3's VAE
+ runs fp32 weights under fp16 autocast, so getting this wrong would silently
+ downcast the post-norm activation.
+ """
+ for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+ x = torch.randn(2, 64, 32, 32, device="cuda", dtype=dtype)
+ weight = torch.randn(64, device="cuda", dtype=dtype)
+ bias = torch.randn(64, device="cuda", dtype=dtype)
+
+ eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6))
+ out = fused_group_norm_silu(x, weight, bias, 32, 1e-6)
+
+ assert out.dtype == eager_out.dtype, (
+ f"Output dtype {out.dtype} != eager dtype {eager_out.dtype} for input dtype {dtype}"
+ )
+ assert out.shape == x.shape, f"Output shape {out.shape} != input shape {x.shape}"
+
+
+@pytest.mark.parametrize(
+ "autocast_dtype, input_dtype",
+ [
+ (torch.float16, torch.float16),
+ (torch.bfloat16, torch.bfloat16),
+ (torch.float16, torch.float32),
+ (torch.bfloat16, torch.float32),
+ ],
+)
+def test_autocast_matches_eager(autocast_dtype, input_dtype):
+ """Inside autocast the fused op must stay a drop-in replacement.
+
+ Mirrors the HunyuanImage3 VAE setup: fp32 module weights, activations
+ arriving as fp16/bf16 from an upstream conv, all under autocast. Also pins
+ the fp32-activation corner (some VAE stages) so the inlined
+ ``is_autocast_enabled -> fp32`` branch cannot silently downcast it.
+ """
+ x = torch.randn(2, 64, 32, 32, device="cuda", dtype=input_dtype)
+ weight = torch.randn(64, device="cuda", dtype=torch.float32)
+ bias = torch.randn(64, device="cuda", dtype=torch.float32)
+
+ with torch.autocast("cuda", dtype=autocast_dtype):
+ eager_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6))
+ fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6)
+
+ assert eager_out.dtype == torch.float32, "sanity check: eager group_norm is expected to return fp32 under autocast"
+ assert fused_out.dtype == eager_out.dtype, (
+ f"Output dtype {fused_out.dtype} != eager dtype {eager_out.dtype} "
+ f"under autocast({autocast_dtype}) for input {input_dtype}"
+ )
+ torch.testing.assert_close(fused_out, eager_out, rtol=1e-5, atol=1e-5)
+
+
+def test_invalid_channels_not_divisible():
+ """Test that assertion fires when channels not divisible by num_groups."""
+ x = torch.randn(1, 33, 16, 16, device="cuda") # 33 not divisible by 32
+ weight = torch.randn(33, device="cuda")
+ bias = torch.randn(33, device="cuda")
+
+ with pytest.raises(AssertionError, match="must be divisible by num_groups"):
+ fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6)
+
+
+def test_invalid_weight_shape():
+ """Test that assertion fires with wrong weight shape."""
+ x = torch.randn(1, 64, 16, 16, device="cuda")
+ weight = torch.randn(32, device="cuda") # Wrong: should be 64
+ bias = torch.randn(64, device="cuda")
+
+ with pytest.raises(AssertionError, match="doesn't match channels"):
+ fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("shape", [(2, 64, 3, 16, 16), (1, 32, 1, 8, 8), (2, 64, 128)])
+def test_non_2d_spatial_matches_eager(dtype, shape):
+ """Spatial ranks other than 2 must work: the 3D VAE feeds (N, C, T, H, W).
+
+ GroupNorm reduces over the whole channel group and every spatial position,
+ so flattening the spatial axes inside the wrapper has to be exact.
+ """
+ torch.manual_seed(0)
+ C = shape[1]
+ x = torch.randn(*shape, device="cuda", dtype=dtype)
+ weight = torch.randn(C, device="cuda", dtype=dtype)
+ bias = torch.randn(C, device="cuda", dtype=dtype)
+
+ fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6)
+ ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype)
+
+ assert fused_out.shape == x.shape
+ rtol, atol = _tolerance(dtype)
+ torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_permuted_input_matches_eager(dtype):
+ """A permuted (channel-stride-1) activation must still be correct.
+
+ This is not a hypothetical layout: HunyuanImage3's ``UNetUp`` calls this op
+ on the output of ``rearrange(x, "b (h w) c -> b c h w")``. An earlier
+ version of the kernel indexed such inputs through their strides, which was
+ correct but uncoalesced enough to run at 0.44x of eager, so the wrapper now
+ materializes them first.
+ """
+ torch.manual_seed(0)
+ B, C, H, W = 2, 64, 16, 16
+ kw = dict(device="cuda", dtype=dtype)
+ x = torch.randn(B, H * W, C, **kw).permute(0, 2, 1).reshape(B, C, H, W)
+ assert not x.is_contiguous(), "test premise: input is strided"
+
+ weight = torch.randn(C, **kw)
+ bias = torch.randn(C, **kw)
+
+ fused_out = fused_group_norm_silu(x, weight, bias, num_groups=32, eps=1e-6)
+ ref_out = F.silu(F.group_norm(x.float(), 32, weight.float(), bias.float(), 1e-6)).to(dtype)
+
+ rtol, atol = _tolerance(dtype)
+ torch.testing.assert_close(fused_out, ref_out, rtol=rtol, atol=atol)
+
+
+@pytest.mark.parametrize("batch_size", [1, 4])
+@pytest.mark.parametrize("channels", [64, 128])
+def test_hunyuan_vae_config(batch_size, channels):
+ """Test with HunyuanImage3 VAE actual configuration.
+
+ HunyuanImage3 ResnetBlock uses:
+ - num_groups=32
+ - eps=1e-6
+ - Typical spatial sizes: 16x16 to 128x128 (after tiling)
+ """
+ spatial_sizes = [(16, 16), (32, 32), (64, 64), (128, 128)]
+
+ for H, W in spatial_sizes:
+ x = torch.randn(batch_size, channels, H, W, device="cuda", dtype=torch.float32)
+ weight = torch.randn(channels, device="cuda")
+ bias = torch.randn(channels, device="cuda")
+
+ ref_out = F.silu(F.group_norm(x, 32, weight, bias, 1e-6))
+ fused_out = fused_group_norm_silu(x, weight, bias, 32, 1e-6)
+
+ torch.testing.assert_close(
+ fused_out,
+ ref_out,
+ rtol=1e-5,
+ atol=1e-6,
+ msg=f"HunyuanVAE config: batch={batch_size}, C={channels}, HW=({H},{W})",
+ )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+def test_large_offset_small_variance(dtype):
+ """A large constant offset with a small variance must stay stable.
+
+ ``x ~ 10000 ± 0.1`` has mean ~1e4 and variance ~1e-2, so the naive
+ ``E[x^2] - E[x]^2`` subtracts two ~1e8 quantities and rounds the variance to
+ ~8.0 (the ULP of 1e8) — catastrophic cancellation, producing garbage
+ normalization. The Welford reduction keeps the variance accurate.
+
+ Why fp32 tolerance is 0.1, not the usual 1e-6: at this offset, ``x`` is stored
+ in fp32 with ULP ~0.001, so ``x - mean`` carries ~0.001 of representation
+ noise, and ``rstd ~ 10`` amplifies it to ~0.01-0.02. Even PyTorch's own
+ GroupNorm reference hits this floor (the old naive implementation would be
+ off by ~1.0 here). fp16/bf16 quantize the 0.1 perturbation away entirely, so
+ for those dtypes the meaningful assertion is just "no NaN".
+ """
+ torch.manual_seed(0)
+ B, C, H, W = 2, 64, 16, 16
+ num_groups = 32
+ eps = 1e-6
+ device = torch.device("cuda")
+
+ x = torch.full((B, C, H, W), 10000.0, device=device, dtype=dtype)
+ x = x + torch.randn(B, C, H, W, device=device, dtype=dtype) * 0.1
+ weight = torch.randn(C, device=device, dtype=dtype)
+ bias = torch.randn(C, device=device, dtype=dtype)
+
+ fused_out = fused_group_norm_silu(x, weight, bias, num_groups, eps)
+ # fp32 reference = PyTorch's Welford path, the ground truth.
+ ref_out = F.silu(F.group_norm(x.float(), num_groups, weight.float(), bias.float(), eps)).to(dtype)
+
+ assert not torch.isnan(fused_out).any(), "fused output must not contain NaN"
+
+ if dtype == torch.float32:
+ # ~0.02 is the representation-noise floor at offset 1e4; old impl ~1.0.
+ torch.testing.assert_close(fused_out, ref_out, rtol=0.05, atol=0.1)
+ # fp16/bf16: perturbation is swallowed by the dtype — no meaningful close,
+ # the no-NaN assertion above is the regression guard.
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
index 2a4cb2f1e..906b27919 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for CosyVoice3 components."""
from types import SimpleNamespace
@@ -9,10 +9,41 @@ import torch
import torch.nn as nn
from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
+ CausalHiFTGenerator,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@pytest.fixture
+def causal_hift():
+ return CausalHiFTGenerator(
+ base_channels=32,
+ upsample_rates=[2, 2],
+ upsample_kernel_sizes=[4, 4],
+ source_resblock_kernel_sizes=[3, 3],
+ source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]],
+ )
+
+
+def test_causal_hift_moves_stft_window_with_model(causal_hift):
+ assert causal_hift.get_buffer("stft_window") is causal_hift.stft_window
+ assert "stft_window" not in causal_hift.state_dict()
+ causal_hift.to(dtype=torch.float64)
+ assert causal_hift.stft_window.dtype == torch.float64
+
+
+def test_causal_hift_stft_moves_window_to_input_device(causal_hift):
+ waveform = torch.empty((1, 64), device="meta")
+
+ real, imag = causal_hift._stft(waveform)
+
+ assert real.device == waveform.device
+ assert imag.device == waveform.device
+ assert causal_hift.stft_window.device == waveform.device
+
+
class TestPreLookaheadLayer:
"""Tests for PreLookaheadLayer."""
@@ -250,6 +281,153 @@ class TestCFM:
assert out.shape == mu.shape
+ @pytest.mark.core_model
+ @pytest.mark.cpu
+ def test_trt_estimator_uses_stream_dependencies(self, monkeypatch):
+ from omegaconf import DictConfig
+
+ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ConditionalCFM
+
+ class FakeStream:
+ def __init__(self, name, state):
+ self.name = name
+ self.state = state
+ self.cuda_stream = hash(name)
+ self.synchronize_calls = 0
+ self.waited_on = []
+
+ def synchronize(self):
+ self.synchronize_calls += 1
+
+ def wait_stream(self, stream):
+ self.waited_on.append(stream)
+
+ class FakeStreamContext:
+ def __init__(self, stream):
+ self.stream = stream
+
+ def __enter__(self):
+ self.previous = self.stream.state.current
+ self.stream.state.current = self.stream
+ return self.stream
+
+ def __exit__(self, exc_type, exc, traceback):
+ self.stream.state.current = self.previous
+
+ state = SimpleNamespace(current=None)
+ caller_stream = FakeStream("caller", state)
+ estimator_stream = FakeStream("estimator", state)
+ state.current = caller_stream
+ monkeypatch.setattr(torch.cuda, "current_stream", lambda *args, **kwargs: state.current)
+ stream_contexts = []
+
+ def stream_context(stream):
+ stream_contexts.append(stream)
+ return FakeStreamContext(stream)
+
+ monkeypatch.setattr(torch.cuda, "stream", stream_context)
+
+ class FakeContext:
+ def __init__(self):
+ self.execute_stream = None
+
+ def set_input_shape(self, name, shape):
+ pass
+
+ def set_tensor_address(self, name, address):
+ pass
+
+ def execute_async_v3(self, stream):
+ self.execute_stream = stream
+ return True
+
+ class FakeEngine:
+ @staticmethod
+ def get_tensor_name(index):
+ return f"tensor_{index}"
+
+ context = FakeContext()
+
+ class FakeEstimatorPool:
+ io_dtype = torch.float32
+
+ def __init__(self):
+ self.released = []
+
+ def acquire_estimator(self):
+ return [context, estimator_stream], FakeEngine()
+
+ def release_estimator(self, released_context, released_stream):
+ self.released.append((released_context, released_stream))
+
+ estimator_pool = FakeEstimatorPool()
+ cfm = ConditionalCFM(
+ in_channels=80,
+ cfm_params=DictConfig(
+ {
+ "sigma_min": 1e-6,
+ "solver": "euler",
+ "t_scheduler": "cosine",
+ "training_cfg_rate": 0.2,
+ "inference_cfg_rate": 0.7,
+ }
+ ),
+ n_spks=1,
+ spk_emb_dim=80,
+ estimator=estimator_pool,
+ )
+
+ x = torch.randn(2, 80, 4)
+ mask = torch.ones(2, 1, 4)
+ mu = torch.randn(2, 80, 4)
+ timestep = torch.randn(2)
+ speakers = torch.randn(2, 80)
+ condition = torch.randn(2, 80, 4)
+
+ output = cfm.forward_estimator(x, mask, mu, timestep, speakers, condition)
+
+ assert output.shape == x.shape
+ assert caller_stream.synchronize_calls == 0
+ assert estimator_stream.synchronize_calls == 0
+ assert estimator_stream.waited_on == [caller_stream]
+ assert caller_stream.waited_on == [estimator_stream]
+ assert stream_contexts == [estimator_stream]
+ assert context.execute_stream == estimator_stream.cuda_stream
+ assert estimator_pool.released == [(context, estimator_stream)]
+
+ @pytest.mark.core_model
+ @pytest.mark.cpu
+ def test_trt_context_pool_stores_cuda_stream(self, monkeypatch):
+ from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import TrtContextWrapper
+
+ execution_context = object()
+
+ class FakeEngine:
+ @staticmethod
+ def create_execution_context():
+ return execution_context
+
+ cuda_stream = object()
+ monkeypatch.setattr(torch.cuda, "Stream", lambda device: cuda_stream)
+
+ def reject_stream_context(stream):
+ pytest.fail("TrtContextWrapper must store the CUDA stream, not a StreamContext")
+
+ monkeypatch.setattr(torch.cuda, "stream", reject_stream_context)
+
+ engine = FakeEngine()
+ wrapper = TrtContextWrapper(engine, device="cuda:0")
+ [context, stream], acquired_engine = wrapper.acquire_estimator()
+
+ assert context is execution_context
+ assert stream is cuda_stream
+ assert acquired_engine is engine
+
+ wrapper.release_estimator(context, stream)
+ [reused_context, reused_stream], _ = wrapper.acquire_estimator()
+ assert reused_context is context
+ assert reused_stream is stream
+
class TestSDPAFallback:
"""Test SDPA fallback for float32 inputs."""
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py
new file mode 100644
index 000000000..7a5492d55
--- /dev/null
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_config.py
@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+
+from vllm_omni.config.model import OmniModelArchConfigConvertor
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_attention_geometry_matches_qwen_backbone() -> None:
+ config = CosyVoice3Config()
+ model_arch_config = OmniModelArchConfigConvertor(config, config.get_text_config()).convert()
+
+ assert model_arch_config.total_num_attention_heads == 14
+ assert model_arch_config.total_num_kv_heads == 2
+ assert model_arch_config.head_size == 64
diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
index a2c3142c7..e16927096 100644
--- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
+++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -401,6 +401,52 @@ def test_sample_tolerates_padded_rows_without_history():
assert out.sampled_token_ids.shape == (2, 1)
+def test_sample_excludes_non_finite_logits():
+ model = _make_talker_model()
+ metadata = _make_sampling_metadata(output_token_ids=[[]])
+ metadata.temperature.fill_(0.5)
+ logits = torch.tensor([[float("nan"), 1.0, float("inf"), float("-inf")]], dtype=torch.bfloat16)
+
+ out = model.sample(logits, metadata)
+
+ assert out is not None
+ assert out.sampled_token_ids.tolist() == [[1]]
+
+
+def test_sample_preserves_allowed_token_mask_with_invalid_logits():
+ model = _make_talker_model()
+ metadata = _make_sampling_metadata(output_token_ids=[[]])
+ metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+ logits = torch.tensor([[float("nan"), 1.0, float("inf")]], dtype=torch.float32)
+
+ out = model.sample(logits, metadata)
+
+ assert out is not None
+ assert out.sampled_token_ids.tolist() == [[1]]
+
+
+def test_sample_rejects_rows_without_finite_logits():
+ model = _make_talker_model()
+ metadata = _make_sampling_metadata(output_token_ids=[[]])
+ metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+ logits = torch.tensor([[0.0, float("nan"), 0.0]], dtype=torch.float32)
+
+ with pytest.raises(ValueError, match="no finite logits"):
+ model.sample(logits, metadata)
+
+
+def test_sample_keeps_only_finite_token_after_ras_rejection():
+ model = _make_talker_model()
+ metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+ metadata.allowed_token_ids_mask = torch.tensor([[True, False, True]])
+ logits = torch.tensor([[0.0, 1.0, 0.0]], dtype=torch.float32)
+
+ out = model.sample(logits, metadata)
+
+ assert out is not None
+ assert out.sampled_token_ids.tolist() == [[1]]
+
+
def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
expected = SamplerOutput(
@@ -423,9 +469,14 @@ def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
_, GPUARModelRunner = _cosyvoice3_model_and_runner()
runner = object.__new__(GPUARModelRunner)
runner.input_batch = _DummyInputBatch()
+
+ def model_sample(logits, sampling_metadata):
+ calls.append(logits.clone())
+ return expected
+
runner.model = SimpleNamespace(
prefer_model_sampler=True,
- sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
+ sample=model_sample,
)
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
@@ -459,12 +510,14 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
_, GPUARModelRunner = _cosyvoice3_model_and_runner()
runner = object.__new__(GPUARModelRunner)
runner.input_batch = _DummyInputBatch()
+
+ def model_sample(logits, sampling_metadata):
+ seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+ return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+
runner.model = SimpleNamespace(
prefer_model_sampler=True,
- sample=lambda logits, sampling_metadata: (
- seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
- or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
- ),
+ sample=model_sample,
)
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
@@ -504,12 +557,14 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
_, GPUARModelRunner = _cosyvoice3_model_and_runner()
runner = object.__new__(GPUARModelRunner)
runner.input_batch = _DummyInputBatch()
+
+ def model_sample(logits, sampling_metadata):
+ seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+ return SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+
runner.model = SimpleNamespace(
prefer_model_sampler=True,
- sample=lambda logits, sampling_metadata: (
- seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
- or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
- ),
+ sample=model_sample,
)
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
diff --git a/tests/model_executor/models/dots_tts/__init__.py b/tests/model_executor/models/dots_tts/__init__.py
new file mode 100644
index 000000000..208f01a7c
--- /dev/null
+++ b/tests/model_executor/models/dots_tts/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/tests/model_executor/models/dots_tts/test_talker.py b/tests/model_executor/models/dots_tts/test_talker.py
new file mode 100644
index 000000000..ed3c6bcaa
--- /dev/null
+++ b/tests/model_executor/models/dots_tts/test_talker.py
@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for the dots.tts talker's stop-signal pairing and
+per-request state lifecycle.
+
+Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes
+the methods under test read) — no weights, no GPU, no engine. Mirrors
+tests/model_executor/models/voxcpm2/test_talker_state_eviction.py.
+"""
+
+from __future__ import annotations
+
+import functools
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+torch = pytest.importorskip("torch")
+
+_VOCAB_SIZE = 151672
+
+
+@functools.lru_cache(maxsize=1)
+def _dots_tts_talker_mod():
+ """Defer talker import (pulls vLLM model_executor) until first use."""
+ from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import (
+ DotsTTSForConditionalGeneration,
+ _RequestState,
+ )
+
+ return DotsTTSForConditionalGeneration, _RequestState
+
+
+def _make_bare_talker():
+ DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod()
+ talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration)
+ talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE)
+ talker._active_states = {}
+ talker._pending_requests = []
+ talker._results_queue = []
+ talker._audio_queue = []
+ talker._deferred_cleanup_ids = set()
+ talker._beta_trace = False
+ return talker
+
+
+def _add_state(talker, req_id: str, *, is_stopping: bool = False):
+ _, _RequestState = _dots_tts_talker_mod()
+ state = _RequestState(request_id=req_id)
+ state.prefill_completed = True
+ state.is_stopping = is_stopping
+ talker._active_states[req_id] = state
+ return state
+
+
+def _hidden(bsz: int) -> torch.Tensor:
+ return torch.zeros(bsz, 1536, dtype=torch.float32)
+
+
+class TestComputeLogitsPairing:
+ """compute_logits pairs _results_queue entries with batch rows
+ positionally, so every request must push exactly one entry per step —
+ including prefills, which push a (req_id, None) placeholder."""
+
+ def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None:
+ talker = _make_bare_talker()
+ _add_state(talker, "prefill-req")
+ decode_state = _add_state(talker, "decode-req", is_stopping=True)
+ decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]])
+
+ # Row 0 is the prefill (None placeholder), row 1 the stopping decode.
+ talker._results_queue = [
+ ("prefill-req", None),
+ ("decode-req", torch.tensor([[0.1, 0.9]])),
+ ]
+ logits = talker.compute_logits(_hidden(2))
+
+ assert logits.shape == (2, _VOCAB_SIZE)
+ # Prefill row: forced continue.
+ assert logits[0, 0].item() == 1.0
+ assert logits[0, 1].item() == float("-inf")
+ # Stopping decode row: stop slot wins.
+ assert logits[1, 1].item() == 1.0
+ assert logits[1, 0].item() == 0.0
+ # Everything else stays masked.
+ assert logits[0, 2].item() == float("-inf")
+ assert logits[1, 2].item() == float("-inf")
+ # Queue drained.
+ assert talker._results_queue == []
+
+ def test_below_threshold_stop_logits_force_continue(self) -> None:
+ """Raw (continue, stop) softmax must not reach the sampler: a
+ prob_stop in (0.5, 0.8] would flip greedy argmax even though
+ is_stopping stayed False (the 0.8-threshold contract)."""
+ talker = _make_bare_talker()
+ _add_state(talker, "req", is_stopping=False)
+ talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))]
+
+ logits = talker.compute_logits(_hidden(1))
+
+ assert logits[0, 0].item() == 1.0
+ assert logits[0, 1].item() == float("-inf")
+
+ def test_empty_queue_defaults_to_all_continue(self) -> None:
+ talker = _make_bare_talker()
+ logits = talker.compute_logits(_hidden(3))
+
+ assert torch.all(logits[:, 0] == 1.0)
+ assert torch.all(logits[:, 1] == float("-inf"))
+
+
+class TestStateEviction:
+ """on_requests_finished fires BEFORE forward(), so eviction must be
+ deferred to _flush_deferred_cleanup at the end of the step — the
+ finishing request's audio still drains through the current forward."""
+
+ def test_finished_request_survives_until_flush(self) -> None:
+ talker = _make_bare_talker()
+ _add_state(talker, "req-a")
+ _add_state(talker, "req-b")
+
+ talker.on_requests_finished(["req-a"])
+ assert "req-a" in talker._active_states # not evicted yet
+
+ talker._flush_deferred_cleanup()
+ assert "req-a" not in talker._active_states
+ assert "req-b" in talker._active_states
+ assert talker._deferred_cleanup_ids == set()
+
+ def test_unknown_request_id_is_ignored(self) -> None:
+ talker = _make_bare_talker()
+ talker.on_requests_finished(["never-seen"])
+ talker._flush_deferred_cleanup()
+ assert talker._active_states == {}
diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
index 6d81f300a..f5a7094ef 100755
--- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
+++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py
@@ -1,4 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Deterministic unit tests for higgs-audio v3.
These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules),
@@ -343,7 +344,22 @@ class TestSamplerMethods:
return t
@staticmethod
- def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95):
+ def _make_state_tracking_talker(num_rows):
+ from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+ talker_cls = mod.HiggsAudioV3TalkerForConditionalGeneration
+ talker = talker_cls.__new__(talker_cls)
+ torch.nn.Module.__init__(talker)
+ talker.num_codebooks = 8
+ talker._decode_last_codes = torch.zeros(num_rows, 8, dtype=torch.long)
+ talker._decode_has_codes = torch.zeros(num_rows, dtype=torch.bool)
+ talker._decode_delay_count = torch.zeros(num_rows, dtype=torch.long)
+ talker._decode_eoc_countdown = torch.full((num_rows,), -1, dtype=torch.long)
+ talker._decode_generation_done = torch.zeros(num_rows, dtype=torch.bool)
+ return talker
+
+ @staticmethod
+ def _sampling_metadata(temperature=1.0, top_k=50, top_p=0.95, generators=None):
return type(
"SamplingMetadata",
(),
@@ -352,9 +368,16 @@ class TestSamplerMethods:
"top_k": top_k,
"top_p": top_p,
"all_greedy": temperature is None,
+ "generators": generators or {},
},
)()
+ @staticmethod
+ def _seeded_generator(seed):
+ generator = torch.Generator(device="cpu")
+ generator.manual_seed(seed)
+ return generator
+
def test_sample_respects_mask(self):
"""Tokens masked to -inf must never be sampled."""
t = self._make_batched_sampler_talker(num_rows=1)
@@ -442,6 +465,51 @@ class TestSamplerMethods:
assert torch.count_nonzero(seen[0][0]).item() == 1
assert torch.count_nonzero(seen[0][2]).item() > 1
+ def test_seeded_sampling_is_batch_independent(self):
+ t = self._make_minimal_talker()
+ num_codebooks = 8
+ torch.manual_seed(123)
+ first_logits = torch.randn(num_codebooks, 64)
+ second_logits = torch.randn(num_codebooks, 64)
+ filler_logits = torch.randn(num_codebooks, 64)
+ reference_generator = self._seeded_generator(42)
+ batched_generators = {
+ 1: self._seeded_generator(42),
+ 2: self._seeded_generator(42),
+ }
+
+ reference_first = t._sample_audio_codes(
+ first_logits,
+ self._sampling_metadata(generators={0: reference_generator}),
+ num_codebooks=num_codebooks,
+ )
+ batched_first = t._sample_audio_codes(
+ torch.cat([filler_logits, first_logits, first_logits]),
+ self._sampling_metadata(generators=batched_generators),
+ num_codebooks=num_codebooks,
+ )
+ assert torch.equal(reference_first, batched_first[num_codebooks : 2 * num_codebooks])
+ assert torch.equal(reference_first, batched_first[2 * num_codebooks :])
+
+ reference_second = t._sample_audio_codes(
+ second_logits,
+ self._sampling_metadata(generators={0: reference_generator}),
+ num_codebooks=num_codebooks,
+ )
+ condensed_second = t._sample_audio_codes(
+ torch.cat([second_logits, second_logits]),
+ self._sampling_metadata(
+ generators={
+ 0: batched_generators[1],
+ 1: batched_generators[2],
+ }
+ ),
+ num_codebooks=num_codebooks,
+ )
+
+ assert torch.equal(reference_second, condensed_second[:num_codebooks])
+ assert torch.equal(reference_second, condensed_second[num_codebooks:])
+
def test_delay_masking_forces_boc_during_delay(self):
"""During delay phase, codebooks beyond delay_count must have only BOC allowed."""
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
@@ -678,6 +746,58 @@ class TestSamplerMethods:
assert not t._decode_generation_done[2:5].any()
assert t._decode_eoc_countdown[2:5].eq(-1).all()
+ def test_decode_state_follows_request_after_batch_condensation(self):
+ talker = self._make_state_tracking_talker(3)
+
+ talker._sync_decode_state_with_batch(["A", "B", "C"])
+ c_codes = torch.arange(300, 308)
+ talker._decode_last_codes[:] = torch.stack((torch.arange(100, 108), torch.zeros(8, dtype=torch.long), c_codes))
+ talker._decode_has_codes[:] = torch.tensor([True, False, True])
+ talker._decode_delay_count[:] = torch.tensor([3, 0, 8])
+
+ talker.on_requests_finished({"B"})
+ talker._sync_decode_state_with_batch(["A", "C"])
+
+ assert talker._decode_has_codes[1].item() is True
+ assert talker._decode_delay_count[1].item() == 8
+ assert torch.equal(talker._decode_last_codes[1], c_codes)
+
+ def test_finished_request_state_is_cleared_before_pool_slot_reuse(self):
+ talker = self._make_state_tracking_talker(1)
+
+ talker._sync_decode_state_with_batch(["finished"])
+ talker._decode_has_codes[0] = True
+ talker._decode_last_codes[0] = 7
+ talker._decode_delay_count[0] = 5
+ talker._decode_eoc_countdown[0] = 0
+ talker._decode_generation_done[0] = True
+ talker.on_requests_finished({"finished"})
+ talker._sync_decode_state_with_batch(["new"])
+
+ assert talker._decode_has_codes[0].item() is False
+ assert talker._decode_last_codes[0].eq(0).all()
+ assert talker._decode_delay_count[0].item() == 0
+ assert talker._decode_eoc_countdown[0].item() == -1
+ assert talker._decode_generation_done[0].item() is False
+
+ def test_finished_request_state_is_cleared_when_request_id_is_reused(self):
+ talker = self._make_state_tracking_talker(1)
+
+ talker._sync_decode_state_with_batch(["request"])
+ talker._decode_has_codes[0] = True
+ talker._decode_last_codes[0] = 7
+ talker._decode_delay_count[0] = 5
+ talker._decode_eoc_countdown[0] = 0
+ talker._decode_generation_done[0] = True
+ talker.on_requests_finished({"request"})
+ talker._sync_decode_state_with_batch(["request"])
+
+ assert talker._decode_has_codes[0].item() is False
+ assert talker._decode_last_codes[0].eq(0).all()
+ assert talker._decode_delay_count[0].item() == 0
+ assert talker._decode_eoc_countdown[0].item() == -1
+ assert talker._decode_generation_done[0].item() is False
+
def test_mixed_batch_prefill_mask_targets_request_rows(self):
"""Mixed prefill/decode must reset only prefill request rows."""
t = self._make_batched_sampler_talker(num_rows=3)
@@ -1326,7 +1446,9 @@ class TestRegistry:
assert "higgs_multimodal_qwen3" in OMNI_PIPELINES
- def test_deploy_yaml_exists(self):
+ def test_deploy_yaml_uses_native_flashinfer(self):
+ from vllm_omni.config.stage_config import load_deploy_config
+
for name in (
"higgs_multimodal_qwen3.yaml",
"higgs_multimodal_qwen3_high_throughput.yaml",
@@ -1334,6 +1456,10 @@ class TestRegistry:
):
yaml_path = get_deploy_config_path(name)
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
+ deploy_config = load_deploy_config(yaml_path)
+ stage_0 = next(stage for stage in deploy_config.stages if stage.stage_id == 0)
+ assert stage_0.engine_extras["attention_backend"] == "FLASHINFER"
+ assert stage_0.engine_extras["attention_config"] == {"use_trtllm_attention": False}
# ---- AC-3: Prompt Builder ----
diff --git a/tests/test_indextts2_optional_dependencies.py b/tests/model_executor/models/indextts2/test_optional_dependencies.py
similarity index 99%
rename from tests/test_indextts2_optional_dependencies.py
rename to tests/model_executor/models/indextts2/test_optional_dependencies.py
index cc3c46b63..24c6131fd 100644
--- a/tests/test_indextts2_optional_dependencies.py
+++ b/tests/model_executor/models/indextts2/test_optional_dependencies.py
@@ -14,7 +14,7 @@ from packaging.requirements import Requirement
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-ROOT = Path(__file__).parents[1]
+ROOT = Path(__file__).parents[4]
TEXT_PROCESSING_PATH = ROOT / "vllm_omni" / "model_executor" / "models" / "indextts2" / "text_processing_v2_5.py"
INDEXTTS2_FRONTEND_PACKAGES = {
"fugashi",
diff --git a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py
index 7a0ad68de..21c8dfd14 100644
--- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py
+++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py
@@ -7,8 +7,14 @@ import pytest
import torch
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.model_executor.models.common.ming.aggregator import Aggregator
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
-from vllm_omni.model_executor.models.common.ming.fm import Solver
+from vllm_omni.model_executor.models.common.ming.cfm_cudagraph import CFMGraphExecutor, CFMSampler
+from vllm_omni.model_executor.models.common.ming.cfm_dit import DiT
+from vllm_omni.model_executor.models.common.ming.cfm_head import FlowLoss
+from vllm_omni.model_executor.models.common.ming.cfm_solver import Solver, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.speaker_extractor import SpeakerEmbeddingExtractor
+from vllm_omni.model_executor.models.common.ming.spk_embedding import SpkembExtractor
from vllm_omni.model_executor.models.ming_tts.constants import (
AGGREGATOR_HIDDEN_SIZE,
HISTORY_PATCH_SIZE,
@@ -24,7 +30,6 @@ from vllm_omni.model_executor.models.ming_tts.constants import (
SPEAKER_EMBEDDING_DIM,
VAE_PATCH_SIZE,
)
-from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import _resolve_speaker_embeddings
from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
@@ -44,7 +49,87 @@ def test_ming_tts_audio_vae_uses_common_config():
def test_ming_tts_cfm_solver_uses_common_implementation():
"""Ming dense imports the shared solver implementation directly."""
- assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
+ assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.cfm_solver"
+
+
+@pytest.mark.parametrize(
+ ("symbol", "expected_module"),
+ [
+ (Aggregator, "vllm_omni.model_executor.models.common.ming.aggregator"),
+ (DiT, "vllm_omni.model_executor.models.common.ming.cfm_dit"),
+ (FlowLoss, "vllm_omni.model_executor.models.common.ming.cfm_head"),
+ (CFMGraphExecutor, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+ (CFMSampler, "vllm_omni.model_executor.models.common.ming.cfm_cudagraph"),
+ (SpeakerEmbeddingExtractor, "vllm_omni.model_executor.models.common.ming.speaker_extractor"),
+ (SpkembExtractor, "vllm_omni.model_executor.models.common.ming.spk_embedding"),
+ ],
+)
+def test_ming_tts_reusable_modules_live_in_common(symbol, expected_module):
+ """The pieces a second Ming-family TTS model reuses live under common/ming.
+
+ Guards the promotion: if an implementation drifts back into `ming_tts`
+ (or into another model package), a second model would have to reach across
+ packages to get it.
+ """
+ assert symbol.__module__ == expected_module
+
+
+def test_ming_flash_talker_uses_the_shared_aggregator():
+ """Ming flash no longer carries its own copy of the aggregator.
+
+ Both Ming packages ran near-identical aggregators one directory apart,
+ which is the duplication this module layout exists to remove.
+ """
+ from vllm_omni.model_executor.models.ming_flash_omni import talker_module
+
+ assert talker_module.Aggregator is Aggregator
+ assert not any(
+ isinstance(obj, type) and obj.__module__ == talker_module.__name__ and name == "Aggregator"
+ for name, obj in vars(talker_module).items()
+ )
+
+
+def test_shared_aggregator_matches_the_ming_flash_construction():
+ """The shared aggregator reproduces Ming flash's former default geometry.
+
+ `talker_module.Aggregator` defaulted to `in_channels=64` while the
+ shared class defaults to 4, so the call site now passes `latent_dim`
+ explicitly. This pins the resulting shapes rather than the default.
+ """
+ latent_dim, llm_hidden = 64, 32
+ aggregator = Aggregator(
+ in_channels=latent_dim,
+ llm_input_dim=llm_hidden,
+ hidden_size=16,
+ depth=1,
+ num_heads=2,
+ )
+
+ assert aggregator.x_embedder.in_features == latent_dim
+ assert aggregator.final_layer.linear.out_features == llm_hidden
+
+ out = aggregator(torch.randn(1, 3, latent_dim))
+ # Keeps only the CLS projection: [B, T, D] -> [B, 1, llm_hidden].
+ assert out.shape == (1, 1, llm_hidden)
+
+
+def test_shared_aggregator_rejects_mismatched_latent_width():
+ """Wrong-width input fails loudly instead of silently projecting incorrect."""
+ aggregator = Aggregator(in_channels=64, llm_input_dim=32, hidden_size=16, depth=1, num_heads=2)
+
+ with pytest.raises(ValueError, match="feature dim mismatch"):
+ aggregator(torch.randn(1, 3, 4))
+
+
+def test_cfm_solver_owns_the_timestep_schedule():
+ """get_epss_timesteps sits with the solver, not with the DiT blocks."""
+ from vllm_omni.model_executor.models.common.ming import dit_blocks
+
+ assert not hasattr(dit_blocks, "get_epss_timesteps")
+
+ t = get_epss_timesteps(10, device="cpu", dtype=torch.float32)
+ assert t.shape == (11,)
+ assert float(t[0]) == 0.0 and float(t[-1]) == 1.0
def test_ming_tts_flowloss_preserves_checkpoint_prefix():
diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py
new file mode 100644
index 000000000..af71c07ca
--- /dev/null
+++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_capture_gating.py
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""What is allowed to keep the CFM wrapper from capturing a graph.
+
+These run without a GPU: ``_capture`` is mocked, so nothing here touches a
+CUDA context. The point is the decision made *before* capture, not the capture.
+"""
+
+from typing import NamedTuple
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
+ CFMGraphWrapper,
+ _format_memory_delta,
+ _memory_snapshot,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _InputSignature(NamedTuple):
+ """The three properties ``replay`` reads from an input before capturing.
+
+ Everything ``replay`` does ahead of the capture decision -- the device
+ check and the cache key -- comes from these, and ``_capture`` is mocked
+ out, so the tests need no CUDA tensors and no GPU.
+ """
+
+ shape: tuple[int, ...]
+ dtype: torch.dtype
+ device: torch.device
+
+
+def _cfm_input_signatures(chunk_size: int) -> tuple[_InputSignature, ...]:
+ device = torch.device("cuda:0")
+ shapes = ((2, 16, chunk_size), (2, 1, 8), (2, 2, 8, 2), (2, 2, 1, 4, 8), (2, 2, 8, 2), (2, 2, 1, 12, 8))
+ return tuple(_InputSignature(shape, torch.float32, device) for shape in shapes)
+
+
+def _wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 4) -> CFMGraphWrapper:
+ monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+ wrapper = object.__new__(CFMGraphWrapper)
+ wrapper.max_graphs = max_graphs
+ wrapper.enabled = True
+ wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
+ wrapper.device = torch.device("cuda:0")
+ wrapper._cache = {}
+ wrapper._unsupported = set()
+ wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
+ wrapper._capture = Mock(return_value=None)
+ return wrapper
+
+
+def test_capture_is_not_gated_on_free_device_memory(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Whether to capture must not be decided from free device memory.
+
+ A capture is served by the caching allocator out of its private graph
+ pool, and the allocator falls back on memory it has already reserved. Free
+ device memory is therefore neither an upper nor a lower bound on what a
+ capture can draw, and a threshold on it decides the wrong question -- one
+ whose answer also silently varies with card size and with whatever else
+ shares the device. Anything consulting it here would fail this test.
+ """
+ import vllm_omni.platforms as platforms_module
+
+ def _saturated(device: torch.device | None = None) -> int:
+ raise AssertionError("capture must not depend on device-level free memory")
+
+ # Setting the attribute on the module shadows the lazy ``__getattr__``, so
+ # this neither builds nor requires a real platform singleton.
+ monkeypatch.setattr(
+ platforms_module,
+ "current_omni_platform",
+ Mock(get_free_memory=_saturated),
+ raising=False,
+ )
+
+ wrapper = _wrapper(monkeypatch)
+ wrapper.replay(*_cfm_input_signatures(10))
+
+ wrapper._capture.assert_called_once()
+
+
+def test_memory_reporting_is_inert_off_cuda() -> None:
+ """Capture logging must not break the platforms that have no such counters."""
+ assert _memory_snapshot(torch.device("cpu")) is None
+ assert _format_memory_delta(None, (1, 2)) == ""
+ assert _format_memory_delta((1, 2), None) == ""
+ assert "allocated" in _format_memory_delta((0, 0), (1 << 20, 2 << 20))
diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
index ba3c7292b..a0e7ed4db 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py
@@ -1,3 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import subprocess
+import sys
+from contextlib import contextmanager
from pathlib import Path
from types import SimpleNamespace
@@ -5,8 +11,12 @@ import pytest
import torch
import torch.nn as nn
+import vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav as batched_token2wav_module
from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
BatchedToken2Wav,
+ _token2wav_sdpa_context,
+ plan_token2wav_encode_slices,
+ relpos_encode_token_budget,
)
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
MiniCPMO45Code2Wav,
@@ -50,6 +60,9 @@ class _FakeEstimator(nn.Module):
self.blocks = [_FakeBlock()]
self.cfg_batches: list[int] = []
self.speaker_order: list[list[float]] = []
+ self.attention_cache_dtypes: list[torch.dtype | None] = []
+ self.register_buffer("att_cache_buffer", torch.ones(1), persistent=False)
+ self.register_buffer("cnn_cache_buffer", torch.ones(1), persistent=False)
def t_embedder(self, time):
return time[:, None]
@@ -64,7 +77,8 @@ class _FakeEstimator(nn.Module):
cnn_out,
att_out,
):
- del time, mask, cnn_cache, att_cache
+ del time, mask, cnn_cache
+ self.attention_cache_dtypes.append(att_cache.dtype if isinstance(att_cache, torch.Tensor) else None)
self.cfg_batches.append(inputs.shape[0])
self.speaker_order.append(inputs[:, 2, 0].tolist())
marker = inputs[:, 1, 0]
@@ -98,16 +112,17 @@ class _FakeHiFT(nn.Module):
super().__init__()
self.calls: list[int] = []
- def forward(self, mel, source):
+ def forward(self, batch, device):
+ mel = batch["speech_feat"].transpose(1, 2).to(device)
+ return self.inference(mel, mel.new_zeros((mel.shape[0], 1, 0)))
+
+ def inference(self, mel, source):
del source
self.calls.append(mel.shape[0])
speech = mel[:, 0].repeat_interleave(3, dim=1)
generated_source = speech[:, None]
return speech, generated_source
- def inference(self, mel, source):
- return self.forward(mel, source)
-
class _FakeToken2Wav:
def __init__(self):
@@ -138,13 +153,14 @@ class _FakeToken2Wav:
raise AssertionError("sequential __call__ fallback must never be called")
-def _config(minimum: int = 1):
+def _config(minimum: int = 1, initial: int = 0):
return SimpleNamespace(
model_config=SimpleNamespace(
model="/fake/model",
stage_connector_config={
"extra": {
"code2wav_min_batch_size": minimum,
+ "code2wav_initial_batch_size": initial,
"prompt_cache_id": "shared",
"prompt_wav": "/fake/prompt.wav",
}
@@ -153,14 +169,45 @@ def _config(minimum: int = 1):
)
-def _model():
+def _model(initial: int = 0, minimum: int = 1):
token2wav = _FakeToken2Wav()
backend = BatchedToken2Wav(token2wav)
- model = MiniCPMO45Code2Wav(vllm_config=_config())
+ _enable_fake_ragged_kernel(backend)
+ model = MiniCPMO45Code2Wav(vllm_config=_config(minimum=minimum, initial=initial))
model.backend = backend
return model, token2wav
+def _enable_fake_ragged_kernel(adapter: BatchedToken2Wav) -> None:
+ """Keep fake-model tests focused on ragged grouping and row mapping."""
+ estimator = adapter.flow.decoder.estimator
+ estimator.in_proj = nn.Identity()
+
+ def fake_ragged_kernel(
+ estimator,
+ estimator_input,
+ time_embedding,
+ attn_mask,
+ cnn_cache,
+ att_cache,
+ cnn_cache_buffer,
+ att_cache_buffer,
+ valid_lengths,
+ ):
+ del valid_lengths
+ return estimator.blocks_forward_chunk(
+ estimator_input,
+ time_embedding,
+ attn_mask,
+ cnn_cache,
+ att_cache,
+ cnn_cache_buffer,
+ att_cache_buffer,
+ )
+
+ adapter._blocks_forward_chunk_ragged = fake_ragged_kernel # type: ignore[method-assign]
+
+
def _info(
request_id: str,
chunk_seq: int,
@@ -217,6 +264,437 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches():
assert cache1[0, 0, 0, 0, 0].item() == 20
+def test_relpos_encode_token_budget_leaves_room_for_upsample_and_cache():
+ # The NPU crash was 6968 vs 985: 3333 tokens * 2 + cache overflowed max_pos=5000.
+ assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3) <= 1024
+ assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3, cap=3000) == 2346
+ assert plan_token2wav_encode_slices(3333, max_frames=1024, min_nonfinal=4, last_chunk=True) == [
+ (0, 1024),
+ (1024, 2048),
+ (2048, 3072),
+ (3072, 3333),
+ ]
+ assert plan_token2wav_encode_slices(50, max_frames=40, min_nonfinal=20, last_chunk=False) == [
+ (0, 30),
+ (30, 50),
+ ]
+
+
+def test_decode_batch_splits_overlong_prefill_inside_relpos_budget():
+ token2wav = _FakeToken2Wav()
+ adapter = BatchedToken2Wav(token2wav)
+ prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ states = adapter.setup_batch(prompt, 1)
+ setup_encodes = len(token2wav.flow.encoder.calls)
+ adapter._max_encode_token_frames = lambda _states: 4
+
+ audios, _ = adapter.decode_batch(
+ torch.arange(10, dtype=torch.long).reshape(1, -1),
+ prompt,
+ states,
+ last_chunk=True,
+ )
+
+ assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1]
+ assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True]
+ assert len(audios) == 1
+ assert audios[0].numel() > 0
+
+
+def test_ragged_overlong_prefill_uses_relpos_safe_exact_slices():
+ token2wav = _FakeToken2Wav()
+ token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3)
+ adapter = BatchedToken2Wav(token2wav)
+ _enable_fake_ragged_kernel(adapter)
+ prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ states = adapter.setup_batch(prompt, 2)
+ setup_encodes = len(token2wav.flow.encoder.calls)
+ adapter._max_encode_token_frames = lambda _states: 4
+
+ tokens = [
+ torch.arange(10, dtype=torch.long),
+ torch.arange(4, dtype=torch.long),
+ ]
+ last_chunks = [True, False]
+ audios, next_states = adapter.decode_ragged_batch(
+ tokens,
+ prompt,
+ states,
+ last_chunks=last_chunks,
+ )
+
+ assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1, 1]
+ assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True, False]
+ assert len(audios) == len(next_states) == 2
+ assert all(audio.numel() > 0 for audio in audios)
+
+ for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+ reference_token2wav = _FakeToken2Wav()
+ reference_token2wav.flow.encoder.pre_lookahead_layer = SimpleNamespace(pre_lookahead_len=3)
+ reference_adapter = BatchedToken2Wav(reference_token2wav)
+ reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ reference_states = reference_adapter.setup_batch(reference_prompt, 1)
+ reference_adapter._max_encode_token_frames = lambda _states: 4
+ reference_audio, reference_state = reference_adapter.decode_batch(
+ row_tokens.unsqueeze(0),
+ reference_prompt,
+ reference_states,
+ last_chunk=last_chunk,
+ )
+ torch.testing.assert_close(audios[row], reference_audio[0])
+ for cache_name, expected in reference_state[0].flow_cache.items():
+ torch.testing.assert_close(
+ next_states[row].flow_cache[cache_name],
+ expected,
+ )
+ for cache_name, expected in reference_state[0].hift_cache.items():
+ torch.testing.assert_close(
+ next_states[row].hift_cache[cache_name],
+ expected,
+ )
+
+
+def test_ragged_diffusion_batches_eight_rows_and_preserves_exact_results():
+ token2wav = _FakeToken2Wav()
+ adapter = BatchedToken2Wav(token2wav)
+ _enable_fake_ragged_kernel(adapter)
+ prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ states = adapter.setup_batch(prompt, 8)
+ token2wav.flow.encoder.calls.clear()
+ token2wav.flow.decoder.estimator.cfg_batches.clear()
+ token2wav.hift.calls.clear()
+ tokens = [
+ *(torch.tensor([10 + row, 11 + row, 12 + row]) for row in range(6)),
+ torch.tensor([20, 21]),
+ torch.tensor([30, 31]),
+ ]
+ last_chunks = [False] * 6 + [True, True]
+
+ audios, next_states = adapter.decode_ragged_batch(
+ tokens,
+ prompt,
+ states,
+ last_chunks=last_chunks,
+ )
+
+ assert token2wav.flow.encoder.calls == [6, 2]
+ assert token2wav.flow.decoder.estimator.cfg_batches == [16, 16]
+ assert token2wav.hift.calls == [6, 2]
+ assert len(audios) == len(next_states) == 8
+
+ for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)):
+ reference_adapter = BatchedToken2Wav(_FakeToken2Wav())
+ reference_prompt = reference_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ reference_states = reference_adapter.setup_batch(reference_prompt, 1)
+ reference_audio, reference_state = reference_adapter.decode_batch(
+ row_tokens.unsqueeze(0),
+ reference_prompt,
+ reference_states,
+ last_chunk=last_chunk,
+ )
+ torch.testing.assert_close(audios[row], reference_audio[0])
+ for cache_name, expected in reference_state[0].flow_cache.items():
+ torch.testing.assert_close(
+ next_states[row].flow_cache[cache_name],
+ expected,
+ )
+ for cache_name, expected in reference_state[0].hift_cache.items():
+ torch.testing.assert_close(
+ next_states[row].hift_cache[cache_name],
+ expected,
+ )
+
+
+def test_npu_patch_preserves_ragged_estimator_contract():
+ script = r"""
+import torch
+
+from tests.model_executor.models.minicpmo_4_5.test_code2wav_batching import (
+ _FakeToken2Wav,
+ _enable_fake_ragged_kernel,
+)
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav
+import vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav as npu_patch
+
+
+class _FailGraphRunner:
+ def run(self, *args, **kwargs):
+ raise AssertionError("ragged decode must bypass the exact-shape NPU Graph")
+
+
+npu_patch.apply_minicpmo_4_5_code2wav_patch()
+adapter = BatchedToken2Wav(_FakeToken2Wav())
+_enable_fake_ragged_kernel(adapter)
+prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+states = adapter.setup_batch(prompt, 2)
+npu_patch._backend_graph_runners[adapter] = _FailGraphRunner()
+audios, next_states = adapter.decode_ragged_batch(
+ [
+ torch.tensor([10, 11, 12]),
+ torch.tensor([20, 21]),
+ ],
+ prompt,
+ states,
+ last_chunks=[False, True],
+)
+assert len(audios) == len(next_states) == 2
+"""
+ completed = subprocess.run(
+ [sys.executable, "-c", script],
+ cwd=Path(__file__).resolve().parents[4],
+ capture_output=True,
+ text=True,
+ )
+
+ assert completed.returncode == 0, completed.stderr
+
+
+def test_decode_cfm_enters_platform_sdpa_context(monkeypatch):
+ entered: list[str] = []
+
+ @contextmanager
+ def recording_context():
+ entered.append("enter")
+ yield
+ entered.append("exit")
+
+ monkeypatch.setattr(
+ batched_token2wav_module,
+ "_token2wav_sdpa_context",
+ lambda _device: recording_context(),
+ )
+ adapter = BatchedToken2Wav(_FakeToken2Wav())
+ _enable_fake_ragged_kernel(adapter)
+ adapter._decode_cfm(
+ torch.ones((2, 1, 2)),
+ torch.ones((2, 1)),
+ torch.zeros((2, 1, 2)),
+ cnn_cache=None,
+ att_cache=None,
+ valid_lengths=[2, 1],
+ )
+
+ assert entered == ["enter", "exit"]
+
+
+def test_ragged_decode_bypasses_exact_shape_accelerators():
+ adapter = BatchedToken2Wav(_FakeToken2Wav())
+ _enable_fake_ragged_kernel(adapter)
+
+ class ForbiddenAccelerator:
+ def replay(self, *args):
+ pytest.fail("ragged decode must not use the exact-shape CFM graph")
+
+ def step(self, **kwargs):
+ pytest.fail("ragged decode must not use the exact-shape TensorRT stepper")
+
+ accelerator = ForbiddenAccelerator()
+ adapter._cfm_graph_wrapper = accelerator
+ adapter._trt_stepper = accelerator
+ adapter._decode_cfm(
+ torch.ones((2, 1, 2)),
+ torch.ones((2, 1)),
+ torch.zeros((2, 1, 2)),
+ cnn_cache=None,
+ att_cache=None,
+ valid_lengths=[2, 1],
+ )
+
+
+def test_ragged_decode_fails_when_estimator_has_no_ragged_kernel():
+ adapter = BatchedToken2Wav(_FakeToken2Wav())
+
+ with pytest.raises(
+ RuntimeError,
+ match=r'"reason":"ragged_kernel_unavailable"',
+ ):
+ adapter._decode_cfm(
+ torch.ones((2, 1, 2)),
+ torch.ones((2, 1)),
+ torch.zeros((2, 1, 2)),
+ cnn_cache=None,
+ att_cache=None,
+ valid_lengths=[2, 1],
+ )
+
+
+def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype():
+ adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True)
+ seen_dtypes: list[torch.dtype] = []
+
+ class RecordingGraph:
+ def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out):
+ del time_embedding, cnn_cache
+ seen_dtypes.append(att_cache.dtype)
+ return estimator_input[:, :1], cnn_out, att_out
+
+ adapter._cfm_graph_wrapper = RecordingGraph()
+ adapter._estimator_step(
+ adapter.flow.decoder.estimator,
+ x=torch.ones((2, 1, 2)),
+ mu=torch.ones((2, 1, 2)),
+ time=torch.zeros(2),
+ speakers=torch.ones((2, 1)),
+ cond=torch.zeros((2, 1, 2)),
+ cnn_cache=torch.zeros((1, 2, 2, 1)),
+ att_cache=torch.zeros((1, 2, 1, 1, 2), dtype=torch.bfloat16),
+ )
+
+ assert seen_dtypes == [torch.float32]
+
+
+def test_npu_sdpa_context_uses_platform_patch(monkeypatch):
+ entered: list[str] = []
+
+ @contextmanager
+ def recording_context():
+ entered.append("enter")
+ yield
+ entered.append("exit")
+
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm_omni.platforms.npu.models.step_audio2_token2wav",
+ SimpleNamespace(npu_token2wav_sdpa_context=recording_context),
+ )
+ with _token2wav_sdpa_context(SimpleNamespace(type="npu")):
+ assert entered == ["enter"]
+
+ assert entered == ["enter", "exit"]
+
+
+def test_ragged_outputs_fail_closed_on_middle_hole():
+ adapter = BatchedToken2Wav(_FakeToken2Wav())
+ prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ states = adapter.setup_batch(prompt, 3)
+ audios: list[torch.Tensor | None] = [
+ torch.zeros(1),
+ None,
+ torch.ones(1),
+ ]
+ next_states = [states[0], None, states[2]]
+
+ with pytest.raises(
+ RuntimeError,
+ match=r'"reason":"incomplete_ragged_output","audio_rows":\[1\],"state_rows":\[1\]',
+ ):
+ adapter._require_complete_ragged_outputs(audios, next_states)
+
+
+@pytest.mark.parametrize(
+ "bfloat16_attention_cache",
+ [False, True],
+ ids=["fp32-cache", "bf16-cache"],
+)
+def test_ragged_dit_two_step_decode_matches_per_row_exact_decode(
+ bfloat16_attention_cache: bool,
+):
+ from cosyvoice2.flow.decoder_dit import DiT
+
+ torch.manual_seed(17)
+ template = DiT(
+ in_channels=4,
+ out_channels=1,
+ depth=2,
+ num_heads=2,
+ head_dim=2,
+ hidden_size=4,
+ ).eval()
+ template_state = {name: value.detach().clone() for name, value in template.state_dict().items()}
+ rand_noise = torch.randn(1, 1, 32)
+
+ def make_adapter() -> BatchedToken2Wav:
+ token2wav = _FakeToken2Wav()
+ estimator = DiT(
+ in_channels=4,
+ out_channels=1,
+ depth=2,
+ num_heads=2,
+ head_dim=2,
+ hidden_size=4,
+ ).eval()
+ estimator.load_state_dict(template_state)
+ token2wav.flow.decoder.estimator = estimator
+ token2wav.flow.decoder.rand_noise = rand_noise.clone()
+ return BatchedToken2Wav(
+ token2wav,
+ bfloat16_attention_cache=bfloat16_attention_cache,
+ )
+
+ ragged_adapter = make_adapter()
+ ragged_prompt = ragged_adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ ragged_states = ragged_adapter.setup_batch(ragged_prompt, 2)
+
+ exact_adapters = [make_adapter(), make_adapter()]
+ exact_prompts = [adapter.prepare_prompt("shared", "/fake/prompt.wav") for adapter in exact_adapters]
+ exact_states = [
+ adapter.setup_batch(prompt, 1) for adapter, prompt in zip(exact_adapters, exact_prompts, strict=True)
+ ]
+ tolerance = (2e-2, 2e-3) if bfloat16_attention_cache else (1e-5, 1e-6)
+ steps = [
+ (
+ [torch.tensor([10, 11, 12]), torch.tensor([20, 21, 22])],
+ [False, True],
+ ),
+ (
+ [torch.tensor([13, 14]), torch.tensor([22, 23, 24])],
+ [True, True],
+ ),
+ ]
+ first_step_history_width: int | None = None
+
+ for step, (tokens, last_chunks) in enumerate(steps):
+ if step == 1:
+ assert first_step_history_width is not None
+ assert all(
+ int(state.flow_cache["estimator_att_cache"].shape[4]) == first_step_history_width
+ for state in ragged_states
+ )
+ ragged_audios, ragged_states = ragged_adapter.decode_ragged_batch(
+ tokens,
+ ragged_prompt,
+ ragged_states,
+ last_chunks=last_chunks,
+ )
+ for row, (adapter, prompt) in enumerate(zip(exact_adapters, exact_prompts, strict=True)):
+ exact_audios, exact_states[row] = adapter.decode_batch(
+ tokens[row].unsqueeze(0),
+ prompt,
+ exact_states[row],
+ last_chunk=last_chunks[row],
+ )
+ torch.testing.assert_close(
+ ragged_audios[row],
+ exact_audios[0],
+ rtol=tolerance[0],
+ atol=tolerance[1],
+ )
+ for cache_name, expected in exact_states[row][0].flow_cache.items():
+ actual = ragged_states[row].flow_cache[cache_name]
+ assert actual.dtype == expected.dtype
+ torch.testing.assert_close(
+ actual,
+ expected,
+ rtol=tolerance[0],
+ atol=tolerance[1],
+ )
+ for cache_name, expected in exact_states[row][0].hift_cache.items():
+ torch.testing.assert_close(
+ ragged_states[row].hift_cache[cache_name],
+ expected,
+ rtol=tolerance[0],
+ atol=tolerance[1],
+ )
+
+ estimator_att = ragged_states[0].flow_cache["estimator_att_cache"]
+ assert estimator_att.shape[4] > 0
+ if step == 0:
+ first_step_history_width = int(estimator_att.shape[4])
+ else:
+ assert first_step_history_width is not None
+ assert int(estimator_att.shape[4]) > first_step_history_width
+
+
def test_fade_in_out_limits_overlap_to_available_previous_audio():
speech = torch.arange(6, dtype=torch.float32).reshape(1, -1)
previous = torch.full((1, 3), 2.0)
@@ -256,6 +734,69 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows():
)
+def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep():
+ token2wav = _FakeToken2Wav()
+ adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True)
+ prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+ states = adapter.setup_batch(prompt, 2)
+
+ assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states)
+ assert (
+ states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr()
+ )
+
+ stacked = adapter._stack_flow_cache(states)
+ assert stacked["estimator_att_cache"].dtype == torch.bfloat16
+ restored = adapter._split_flow_cache(stacked, 2)
+ for original, round_tripped in zip(states, restored, strict=True):
+ assert round_tripped["estimator_att_cache"].dtype == torch.bfloat16
+ torch.testing.assert_close(
+ round_tripped["estimator_att_cache"],
+ original.flow_cache["estimator_att_cache"],
+ )
+
+ _, next_states = adapter.decode_batch(
+ torch.tensor([[10, 11], [20, 21]]),
+ prompt,
+ states,
+ last_chunk=False,
+ )
+ assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states)
+ assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32]
+
+
+def test_bfloat16_estimator_attention_cache_materializes_each_timestep():
+ adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=True)
+
+ _, estimator_cnn, estimator_att = adapter._decode_cfm(
+ torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]),
+ torch.ones((2, 1)),
+ torch.zeros((2, 1, 2)),
+ cnn_cache=None,
+ att_cache=None,
+ )
+
+ assert estimator_cnn.dtype == torch.float32
+ assert estimator_att.dtype == torch.bfloat16
+ assert estimator_att.shape[:3] == (2, 1, 4)
+
+
+def test_float32_estimator_attention_cache_materializes_each_timestep():
+ adapter = BatchedToken2Wav(_FakeToken2Wav())
+
+ _, estimator_cnn, estimator_att = adapter._decode_cfm(
+ torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]]),
+ torch.ones((2, 1)),
+ torch.zeros((2, 1, 2)),
+ cnn_cache=None,
+ att_cache=None,
+ )
+
+ assert estimator_cnn.dtype == torch.float32
+ assert estimator_att.dtype == torch.float32
+ assert estimator_att.shape[:3] == (2, 1, 4)
+
+
def test_model_preserves_output_slots_and_prefers_runtime_codes():
model, token2wav = _model()
output = _forward(
@@ -347,6 +888,28 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio():
assert "duplex" in model._states
+def test_initial_empty_segment_markers_respect_initial_batch_limit():
+ model, token2wav = _model(initial=2)
+ names = ["a", "b", "c", "d", "e"]
+ boundaries = []
+ for name in names:
+ boundary = _info(name, 0, [])
+ boundary["meta"].update(
+ {
+ "code_flat_numel": 0,
+ "tts_is_last_chunk": True,
+ "turn_end": False,
+ }
+ )
+ boundaries.append(boundary)
+
+ output = _forward(model, boundaries)
+
+ assert token2wav.flow.encoder.calls == [2, 2, 1]
+ assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names)
+ assert set(model._states) == set(names)
+
+
def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch):
monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
model, _ = _model()
@@ -498,6 +1061,15 @@ def test_empty_final_ignores_generation_scheduler_placeholder_token():
# The runner injects the engine request id on every step (GPU
# _preprocess, NPU _gather_runtime_additional_information)...
{"request_id": "a", "meta": {"request_id": "a"}},
+ # Runtime snapshot bookkeeping alone is not a Talker payload.
+ {
+ "request_id": "a",
+ "meta": {
+ "request_id": "a",
+ "num_processed_tokens": 0,
+ "resumable": True,
+ },
+ },
# ...but a pre-warm step can also reach the model with nothing at all.
{},
],
@@ -599,11 +1171,56 @@ def test_singleton_and_mixed_shape_buckets_use_same_batched_backend_without_fall
assert token2wav.hift.calls[-2:] == [1, 1]
+def test_initial_batch_limit_allows_later_full_batch():
+ model, token2wav = _model(initial=4)
+ names = [f"request-{index}" for index in range(8)]
+
+ for chunk_seq in (0, 1):
+ infos = [_info(name, chunk_seq, [1, 2]) for name in names]
+ _forward(model, infos)
+ assert token2wav.hift.calls[-2:] == [4, 4]
+
+ token2wav.hift.calls.clear()
+ _forward(model, [_info(name, 2, [1, 2]) for name in names])
+ assert token2wav.hift.calls == [8]
+
+
+def test_initial_batch_partition_preserves_minimum():
+ model, token2wav = _model(initial=4, minimum=3)
+ names = [f"request-{index}" for index in range(6)]
+
+ _forward(model, [_info(name, 0, [1, 2]) for name in names])
+
+ assert token2wav.hift.calls == [3, 3]
+
+
+def test_initial_batch_partition_rejects_impossible_remainder():
+ model, _ = _model(initial=4, minimum=3)
+ names = [f"request-{index}" for index in range(5)]
+
+ with pytest.raises(RuntimeError, match="initial_batch_partition_below_minimum"):
+ _forward(model, [_info(name, 0, [1, 2]) for name in names])
+
+
+@pytest.mark.parametrize("chunk_seqs", [[0, 2, 2, 2], [0, 0, 0, 2]])
+def test_initial_steady_partition_rejects_undersized_wave(chunk_seqs):
+ model, _ = _model(initial=4, minimum=3)
+ bucket = [SimpleNamespace(chunk_seq=chunk_seq) for chunk_seq in chunk_seqs]
+
+ with pytest.raises(RuntimeError, match="decode_wave_below_minimum"):
+ list(model._iter_decode_batches([bucket]))
+
+
def test_backend_failure_does_not_commit_any_request_state(monkeypatch):
model, _ = _model()
_forward(
model,
- [_info(name, 0, [index + 1, index + 2]) for index, name in enumerate(("a", "b", "c", "d"))],
+ [
+ _info("a", 0, [1, 2], cache_epoch=0),
+ _info("b", 0, [2, 3], cache_epoch=0),
+ _info("c", 0, [3, 4], cache_epoch=1),
+ _info("d", 0, [4, 5], cache_epoch=1),
+ ],
)
before = dict(model._states)
original = model.backend.decode_batch
@@ -623,8 +1240,8 @@ def test_backend_failure_does_not_commit_any_request_state(monkeypatch):
[
_info("a", 1, [5, 6]),
_info("b", 1, [7, 8]),
- _info("c", 1, [9, 10, 11]),
- _info("d", 1, [12, 13, 14]),
+ _info("c", 1, [9, 10], cache_epoch=1),
+ _info("d", 1, [12, 13], cache_epoch=1),
],
)
assert call_count == 2
diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
index 7ad35b86d..26daeff7b 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-from functools import lru_cache
+import sys
from types import SimpleNamespace
from unittest.mock import Mock
@@ -10,6 +10,7 @@ import torch
import torch.nn as nn
from vllm.platforms import current_platform
+import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
HiFTGenerator,
)
@@ -248,50 +249,203 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
- wrapper._capture_graph.cache_clear()
+ wrapper._flush()
+
+
+def test_importing_wrapper_does_not_resolve_platform() -> None:
+ """Importing this module must not build the OmniPlatform singleton.
+
+ The NPU platform's ``__init__`` patches Code2Wav, which imports
+ ``batched_token2wav`` -> ``cuda_graph_wrapper``. Resolving
+ ``current_omni_platform`` at module scope re-enters
+ ``platforms.__getattr__`` while the singleton is still under construction
+ and the import graph deadlocks. The module reaches for nothing on
+ ``vllm_omni.platforms`` today; this keeps it that way.
+ """
+ import importlib
+
+ import vllm_omni.platforms as platforms_module
+
+ name = "vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper"
+ importlib.import_module(name)
+ saved = platforms_module._current_omni_platform
+ try:
+ del sys.modules[name]
+ platforms_module._current_omni_platform = None
+ importlib.import_module(name)
+ assert platforms_module._current_omni_platform is None
+ finally:
+ platforms_module._current_omni_platform = saved
def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(CFMGraphWrapper)
wrapper.max_graphs = max_graphs
+ wrapper.enabled = True
wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
wrapper.device = torch.device("cuda")
-
- @lru_cache(maxsize=max_graphs)
- def _capture_graph(key: tuple):
- return wrapper._do_capture(key)
-
- wrapper._capture_graph = _capture_graph
- wrapper._do_capture = Mock(return_value=None)
+ wrapper._cache = {}
+ wrapper._unsupported = set()
+ wrapper._stats = {"calls": 0, "hits": 0, "captures": 0, "flushes": 0, "eager": 0}
+ wrapper._capture = Mock(return_value=None)
return wrapper
def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _cfm_mock_wrapper(monkeypatch)
- wrapper._capture_graph = Mock(return_value=None)
inputs = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs)
- wrapper._capture_graph.assert_called_once()
+ wrapper._capture.assert_called_once()
wrapper.graph_fn.assert_called_once()
-def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_cfm_returning_no_entry_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+ """A capture that yields no entry must still serve the request eagerly.
+
+ ``_capture`` is mocked here, so this says nothing about ``_disable``; see
+ ``test_cfm_capture_failure_disables_further_capture`` for that.
+ """
wrapper = _cfm_mock_wrapper(monkeypatch)
- wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
inputs = _cfm_inputs(2, 10, 0)
result = wrapper.replay(*inputs)
wrapper.graph_fn.assert_called_once()
assert result[0] is wrapper.graph_fn.return_value
+ assert wrapper._stats["eager"] == 1
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
- """Verify that LRU cache evicts oldest entries when full."""
+def test_cfm_capture_failure_disables_further_capture(monkeypatch: pytest.MonkeyPatch) -> None:
+ """A real capture failure must stop the wrapper capturing for good.
+
+ A failed capture can leave the capture stream current and the allocator
+ still routing into the graph pool, so the next shape must not try again.
+ """
+ pool = torch.cuda.graph_pool_handle()
+ monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+ torch.manual_seed(0)
+ estimator = _MiniDiT().eval().cuda()
+ wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
+
+ def _explode(*args: object, **kwargs: object) -> None:
+ raise RuntimeError("capture failed")
+
+ monkeypatch.setattr(torch.cuda, "graph", _explode)
+
+ with torch.inference_mode():
+ wrapper.replay(*_cfm_inputs(2, 10, 0))
+ assert wrapper.enabled is False
+ assert wrapper._cache == {}
+
+ captures_after_failure = wrapper._stats["captures"]
+ wrapper.replay(*_cfm_inputs(2, 12, 0))
+
+ assert wrapper._stats["captures"] == captures_after_failure
+ assert wrapper._stats["eager"] == 2
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.MonkeyPatch) -> None:
+ """A key that cannot round-trip is a property of one shape, not of the GPU.
+
+ It must not disable the wrapper or retire the generation the way a capture
+ failure does.
+ """
+ pool = torch.cuda.graph_pool_handle()
+ monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+ torch.manual_seed(0)
+ estimator = _MiniDiT().eval().cuda()
+ wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
+
+ unbuildable_width = 10
+ real_tensors_from_key = wrapper_module._tensors_from_key
+
+ def _reject_one_shape(key: tuple) -> tuple:
+ if key[1][0][2] == unbuildable_width:
+ raise KeyError("torch.int64")
+ return real_tensors_from_key(key)
+
+ monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape)
+
+ with torch.inference_mode():
+ wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
+ assert wrapper.enabled is True
+ assert wrapper._stats["captures"] == 0
+
+ # a capturable shape still gets a graph
+ wrapper.replay(*_cfm_inputs(2, 12, 0))
+ assert wrapper._stats["captures"] == 1
+
+ # and the rejected shape stays eager without retrying the capture
+ wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
+ assert wrapper._stats["captures"] == 1
+ assert wrapper._stats["eager"] == 2
+
+ wrapper._flush()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_hift_replay_survives_a_cfm_generation_flush(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Retiring a CFM generation must not disturb the vocoder.
+
+ Both wrappers capture into ``get_global_graph_pool()``, so the HiFT graphs
+ are exactly the live graphs a CFM flush could strand.
+ """
+ pool = torch.cuda.graph_pool_handle()
+ monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+ torch.manual_seed(0)
+ hift = _small_hift()
+ token2wav = SimpleNamespace(
+ hift=hift,
+ flow=SimpleNamespace(
+ encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)),
+ token_mel_ratio=2,
+ ),
+ mel_cache_len=2,
+ source_cache_len=960,
+ )
+ hift_wrapper = HiFTGraphWrapper(
+ token2wav,
+ connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3},
+ capture_batch_sizes=[1],
+ )
+ hift_wrapper.capture()
+
+ estimator = _MiniDiT().eval().cuda()
+ cfm = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
+
+ speech_feat = torch.randn(1, 80, 4, device="cuda")
+ cache_source = torch.zeros(1, 1, 0, device="cuda")
+
+ with torch.inference_mode():
+ expected_speech, expected_source = hift.inference(speech_feat, cache_source)
+
+ for width in (10, 12, 14, 16):
+ cfm.replay(*_cfm_inputs(2, width, 0))
+ assert cfm._stats["flushes"] >= 1, "cache never overflowed; the test proves nothing"
+
+ actual_speech, actual_source = hift_wrapper.replay(speech_feat, cache_source)
+
+ torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5)
+ torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5)
+
+ cfm._flush()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_cache_flushes_whole_generation(monkeypatch: pytest.MonkeyPatch) -> None:
+ """A full cache is retired all at once, never one graph at a time.
+
+ Destroying a single graph while its peers stay live strands them, so the
+ cache must never hold a graph that outlived one of its generation-mates.
+ """
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
@@ -302,26 +456,32 @@ def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
with torch.inference_mode():
inputs_a = _cfm_inputs(2, 10, 0)
wrapper.replay(*inputs_a)
- assert wrapper._capture_graph.cache_info().currsize == 1
+ assert len(wrapper._cache) == 1
inputs_b = _cfm_inputs(2, 12, 0)
wrapper.replay(*inputs_b)
- assert wrapper._capture_graph.cache_info().currsize == 2
+ assert len(wrapper._cache) == 2
+ # A hit must not grow the cache or trigger a flush.
wrapper.replay(*inputs_a)
- assert wrapper._capture_graph.cache_info().currsize == 2
+ assert len(wrapper._cache) == 2
+ assert wrapper._stats["flushes"] == 0
+ assert wrapper._stats["hits"] == 1
+ # The third distinct shape flushes the generation, then captures alone.
inputs_c = _cfm_inputs(2, 14, 0)
wrapper.replay(*inputs_c)
- assert wrapper._capture_graph.cache_info().currsize == 2
+ assert wrapper._stats["flushes"] == 1
+ assert len(wrapper._cache) == 1
- wrapper.replay(*inputs_b)
- hits_before = wrapper._capture_graph.cache_info().hits
- wrapper.replay(*inputs_b)
- hits_after = wrapper._capture_graph.cache_info().hits
- assert hits_after == hits_before + 1
+ # The flushed shapes are gone, so they capture again rather than hit.
+ hits_before = wrapper._stats["hits"]
+ wrapper.replay(*inputs_a)
+ assert wrapper._stats["hits"] == hits_before
+ wrapper.replay(*inputs_a)
+ assert wrapper._stats["hits"] == hits_before + 1
- wrapper._capture_graph.cache_clear()
+ wrapper._flush()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py
index 9c4b50967..3fd4867cb 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the MiniCPM-o 4.5 stage 0 -> stage 1 bridge.
Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2tts``:
@@ -25,6 +25,7 @@ import torch
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
+from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import llm2tts
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -313,6 +314,44 @@ class TestPromptAndMultiModal:
)
assert out[0]["multi_modal_data"] == mm
+ def test_reference_audio_is_added_to_stage_handoff(self) -> None:
+ hidden_states = torch.zeros((2, _HIDDEN_DIM))
+ waveform = torch.tensor([0.25, -0.5], dtype=torch.float32)
+
+ stage_outputs = llm2tts(
+ [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+ prompt={"multi_modal_data": {"audio": (waveform, 16000)}},
+ )
+
+ intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+ assert intermediate_buffer["codes"]["ref"] == [0.25, -0.5]
+ assert intermediate_buffer["meta"]["ref_audio_sr"] == 16000
+
+ def test_serving_reference_audio_is_added_to_stage_handoff(self) -> None:
+ hidden_states = torch.zeros((2, _HIDDEN_DIM))
+ waveform = torch.tensor([0.125, -0.25], dtype=torch.float32)
+
+ stage_outputs = llm2tts(
+ [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+ prompt={MINICPMO45_REFERENCE_AUDIO_KEY: (waveform, 24000)},
+ )
+
+ intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+ assert intermediate_buffer["codes"]["ref"] == [0.125, -0.25]
+ assert intermediate_buffer["meta"]["ref_audio_sr"] == 24000
+
+ def test_missing_reference_audio_does_not_add_stage_handoff_fields(self) -> None:
+ hidden_states = torch.zeros((2, _HIDDEN_DIM))
+
+ stage_outputs = llm2tts(
+ [_make_thinker_output(prompt_token_ids=[10], output_token_ids=[20], hidden_states=hidden_states)],
+ prompt={"multi_modal_data": {}},
+ )
+
+ intermediate_buffer = stage_outputs[0]["model_intermediate_buffer"]
+ assert "codes" not in intermediate_buffer
+ assert "ref_audio_sr" not in intermediate_buffer.get("meta", {})
+
def test_internal_streaming_context_is_accepted(self) -> None:
hidden = torch.zeros((2, _HIDDEN_DIM))
out = llm2tts(
diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py
index 148f28b16..a05e4076f 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the MiniCPM-o 4.5 pipeline registration.
Covers:
@@ -19,11 +19,15 @@ from pathlib import Path
from types import SimpleNamespace
import pytest
+from transformers import PretrainedConfig
+from vllm.config import ModelConfig
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_config import (
PipelineConfig,
StageExecutionType,
+ StagePipelineConfig,
+ _apply_platform_overrides,
load_deploy_config,
merge_pipeline_deploy,
)
@@ -62,7 +66,9 @@ class TestRegistryDeclaration:
assert not hasattr(pipeline, "max_native_duplex_sessions")
def test_ordinary_pipeline_defaults_to_no_duplex_control(self) -> None:
- pipeline = PipelineConfig(model_type="ordinary")
+ pipeline = PipelineConfig(
+ model_type="ordinary", stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
+ )
assert pipeline.duplex_control_enabled is False
assert pipeline.duplex_serving_adapter is None
assert not hasattr(pipeline, "max_native_duplex_sessions")
@@ -77,10 +83,6 @@ class TestPipelineTopology:
assert len(pipeline.stages) == 3
assert [s.stage_id for s in pipeline.stages] == [0, 1, 2]
- def test_topology_validates(self, pipeline: PipelineConfig) -> None:
- # ``validate`` returns a list of structural errors; empty == valid.
- assert pipeline.validate() == []
-
def test_thinker_stage(self, pipeline: PipelineConfig) -> None:
thinker = pipeline.get_stage(0)
assert thinker is not None
@@ -104,6 +106,10 @@ class TestPipelineTopology:
assert talker.engine_output_type == "latent"
# scope KV cache / mrope sizing to talker sub-config
assert talker.hf_config_name == "tts_config"
+ assert talker.sampling_constraints == {
+ "detokenize": False,
+ "stop_token_ids": [6561],
+ }
assert talker.custom_process_next_stage_input_func == (
"vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload"
)
@@ -162,8 +168,11 @@ class TestDeployTopology:
assert connector["extra"]["connector_get_max_wait"] == 300
expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload"
assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor)
+ if filename != "minicpmo_4_5.yaml":
+ assert "hf_overrides" not in stages[1].yaml_engine_args
if filename == "minicpmo_4_5.yaml":
assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4]
+ assert stages[1].yaml_engine_args["hf_overrides"] == {"tts_config": {"attention_type": "sliding_recompute"}}
memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages]
assert memory_utilizations == [
0.55,
@@ -184,6 +193,42 @@ class TestDeployTopology:
0.35,
]
+ def test_duplex_talker_hf_override_resolves_nested_attention_policy(self) -> None:
+ deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml")
+ stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
+ overrides = stages[1].yaml_engine_args["hf_overrides"]
+ hf_config = PretrainedConfig()
+ hf_config.tts_config = PretrainedConfig(attention_type="full_attention")
+
+ model_config = object.__new__(ModelConfig)
+ model_config._apply_dict_overrides(hf_config, overrides)
+
+ assert hf_config.tts_config.attention_type == "sliding_recompute"
+
+ @pytest.mark.parametrize(
+ "filename",
+ [
+ "minicpmo_4_5.yaml",
+ "minicpmo_4_5_2gpu.yaml",
+ "minicpmo_4_5_3gpu.yaml",
+ ],
+ )
+ def test_npu_graph_configuration_is_stage_scoped(self, filename: str) -> None:
+ deploy = load_deploy_config(_DEPLOY_DIR / filename)
+ connector_extra = deploy.connectors["connector_of_shared_memory"]["extra"]
+ assert "code2wav_enable_npu_graph" not in connector_extra
+ assert "code2wav_max_npu_graphs" not in connector_extra
+
+ deploy = _apply_platform_overrides(deploy, platform="npu")
+ stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
+ assert stages[0].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+ assert stages[1].yaml_engine_args["compilation_config"]["cudagraph_mode"] == "PIECEWISE"
+ assert stages[2].yaml_engine_args["enforce_eager"] is True
+ assert stages[2].yaml_engine_args["additional_config"] == {
+ "code2wav_enable_npu_graph": True,
+ "code2wav_max_npu_graphs": 32,
+ }
+
def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None:
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
talker = pipeline.get_stage(1)
@@ -200,6 +245,14 @@ class TestDeployTopology:
"vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_token_only"
)
+ def test_no_async_chunk_selects_full_payload_processor(self) -> None:
+ deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml")
+ deploy.async_chunk = False
+ stages = merge_pipeline_deploy(OMNI_PIPELINES[_PIPELINE_KEY], deploy)
+
+ assert stages[1].yaml_engine_args["async_chunk"] is False
+ assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(".tts2code2wav_full_payload")
+
def test_code2wav_model_is_lazily_registered() -> None:
assert _OMNI_MODELS["MiniCPMO45Code2Wav"] == (
diff --git a/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py
new file mode 100644
index 000000000..e867e0cab
--- /dev/null
+++ b/tests/model_executor/models/minicpmo_4_5/test_streaming_audio_cache.py
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniCPM-o 4.5 streaming audio cache compatibility."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from transformers.models.whisper.modeling_whisper import WhisperConfig
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+ WHISPER_ATTENTION_CLASSES,
+ MiniCPMWhisperEncoder,
+ MiniCPMWhisperEncoderLayer,
+ _get_audio_cache_length,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _LegacyWhisperAttention(nn.Module):
+ def __init__(self, **kwargs) -> None:
+ super().__init__()
+ self.seen_cache = None
+
+ def forward(self, hidden_states, past_key_value=None, **kwargs):
+ self.seen_cache = past_key_value
+ return hidden_states, None, past_key_value
+
+
+class _CurrentWhisperAttention(nn.Module):
+ def __init__(self, **kwargs) -> None:
+ super().__init__()
+ self.seen_cache = None
+
+ def forward(self, hidden_states, past_key_values=None, **kwargs):
+ self.seen_cache = past_key_values
+ return hidden_states, None
+
+
+@pytest.mark.parametrize(
+ ("implementation", "attention_cls"),
+ [
+ ("test_legacy_cache", _LegacyWhisperAttention),
+ ("test_current_cache", _CurrentWhisperAttention),
+ ],
+)
+def test_whisper_attention_preserves_streaming_cache(
+ monkeypatch: pytest.MonkeyPatch,
+ implementation: str,
+ attention_cls: type[nn.Module],
+) -> None:
+ monkeypatch.setitem(WHISPER_ATTENTION_CLASSES, implementation, attention_cls)
+ config = WhisperConfig(
+ d_model=4,
+ encoder_attention_heads=1,
+ encoder_ffn_dim=8,
+ encoder_layers=1,
+ )
+ config._attn_implementation = implementation
+ layer = MiniCPMWhisperEncoderLayer(config, layer_idx=0)
+ cache = object()
+
+ outputs = layer(
+ torch.randn(1, 2, config.d_model),
+ attention_mask=None,
+ layer_head_mask=None,
+ past_key_values=cache,
+ use_cache=True,
+ )
+
+ assert layer.self_attn.seen_cache is cache
+ assert outputs[-1] is cache
+
+
+def test_audio_cache_length_uses_current_cache_api() -> None:
+ dynamic_cache = SimpleNamespace(get_seq_length=lambda: 7)
+ encoder_decoder_cache = SimpleNamespace(self_attention_cache=dynamic_cache)
+
+ assert _get_audio_cache_length(encoder_decoder_cache) == 7
+
+
+def test_audio_cache_length_supports_legacy_cache() -> None:
+ key = torch.zeros(1, 1, 5, 2)
+ value = torch.zeros_like(key)
+
+ assert _get_audio_cache_length(((key, value),)) == 5
+
+
+def _tiny_encoder_config(attn_implementation: str) -> WhisperConfig:
+ config = WhisperConfig(
+ num_mel_bins=16,
+ d_model=32,
+ encoder_layers=2,
+ encoder_attention_heads=2,
+ encoder_ffn_dim=64,
+ max_source_positions=100,
+ dropout=0.0,
+ attention_dropout=0.0,
+ )
+ config._attn_implementation = attn_implementation
+ return config
+
+
+@pytest.mark.parametrize("attn_implementation", ["eager", "sdpa"])
+def test_encoder_reuses_cache_across_streaming_chunks(attn_implementation: str) -> None:
+ """Encode two chunks against the installed Whisper attention and cache classes.
+
+ Unlike the stub-based tests above, this drives the real transformers cache
+ API that ``get_audio_embedding_streaming`` relies on, so it fails whenever a
+ removed cache method or renamed attention kwarg breaks streaming reuse.
+ """
+ torch.manual_seed(0)
+ config = _tiny_encoder_config(attn_implementation)
+ encoder = MiniCPMWhisperEncoder(config).eval()
+
+ first_frames, second_frames = 20, 14
+ first_len = (first_frames - 1) // 2 + 1
+ second_len = (second_frames - 1) // 2 + 1
+ first_mel = torch.randn(1, config.num_mel_bins, first_frames)
+ second_mel = torch.randn(1, config.num_mel_bins, second_frames)
+
+ with torch.no_grad():
+ first = encoder(
+ first_mel,
+ attention_mask=torch.zeros(1, 1, first_len, first_len),
+ past_key_values=None,
+ use_cache=True,
+ output_hidden_states=True,
+ )
+ cache = first.past_key_values
+ assert cache is not None
+ assert _get_audio_cache_length(cache) == first_len
+
+ second = encoder(
+ second_mel,
+ attention_mask=torch.zeros(1, 1, second_len, first_len + second_len),
+ past_key_values=cache,
+ use_cache=True,
+ output_hidden_states=True,
+ )
+ assert _get_audio_cache_length(second.past_key_values) == first_len + second_len
+
+ stateless = encoder(
+ second_mel,
+ attention_mask=torch.zeros(1, 1, second_len, second_len),
+ past_key_values=None,
+ use_cache=True,
+ output_hidden_states=True,
+ )
+
+ assert torch.isfinite(second.last_hidden_state).all()
+ assert not torch.allclose(second.last_hidden_state, stateless.last_hidden_state)
diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py
index 8a7593623..95d845686 100644
--- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py
+++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py
@@ -1,9 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Request-alignment tests for MiniCPM-o 4.5's native Talker."""
from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any
+
import pytest
import torch
import torch.nn as nn
@@ -11,35 +14,72 @@ import torch.nn as nn
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
+ ConditionalChatTTSConfig,
+)
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
+ _CODEC_PENALTY_WINDOW,
+ _DUPLEX_CODEC_TOKENS_PER_CHUNK,
+ _OFFLINE_CODEC_MAX_NEW_TOKENS,
_REPETITION_PENALTY_CHUNK_SIZE,
MiniCPMO45OmniTTSForConditionalGeneration,
_apply_batched_repetition_penalty,
- _max_audio_tokens,
+ _native_duplex_chunk_budget,
_restore_weight_norm_weight,
+ blank_scheduler_prompt_for_penalties,
)
from vllm_omni.utils.mm_outputs import to_payload_element
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+@dataclass
+class _PenaltyMetadata:
+ """Stand-in for the penalty-relevant slice of vLLM's SamplingMetadata."""
+
+ prompt_token_ids: torch.Tensor | None
+
+
+@dataclass
+class _SamplerOutput:
+ """Stand-in for vLLM's SamplerOutput."""
+
+ sampled_token_ids: torch.Tensor
+
+
+@dataclass
+class _CodecSamplingMetadata:
+ """Stand-in for the codec-penalty slice of vLLM's SamplingMetadata."""
+
+ repetition_penalties: torch.Tensor
+ prompt_token_ids: torch.Tensor | None = None
+ no_penalties: bool = False
+
+
class _FakeNativeTalker(nn.Module):
has_preprocess = True
def __init__(self) -> None:
super().__init__()
- self.forward_kwargs = None
+ self.forward_kwargs: dict[str, Any] | None = None
- def forward(self, **kwargs):
+ def forward(self, **kwargs: Any) -> torch.Tensor:
self.forward_kwargs = kwargs
return torch.ones(2, 4)
+def test_tts_config_exposes_codec_vocab_to_vllm() -> None:
+ cfg = ConditionalChatTTSConfig(num_audio_tokens=6562)
+ assert cfg.vocab_size == 6562
+ assert cfg.eos_token_id == 6561
+
+
def test_wrapper_always_delegates_talker_to_native_ar_path() -> None:
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
nn.Module.__init__(model)
model.model_stage = "tts"
- model.talker = _FakeNativeTalker()
+ talker = _FakeNativeTalker()
+ model.talker = talker
output = model(
input_ids=torch.tensor([1, 2]),
@@ -48,19 +88,26 @@ def test_wrapper_always_delegates_talker_to_native_ar_path() -> None:
)
assert output.shape == (2, 4)
- assert model.talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req"
+ assert talker.forward_kwargs is not None
+ assert talker.forward_kwargs["model_intermediate_buffer"][0]["request_id"] == "req"
def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
nn.Module.__init__(talker)
talker._num_audio_tokens = 8
- talker._batch_stop_logits = None
- talker._request_generators = {}
+ talker._codec_eos_id = 7
+ talker._force_eos_rows = None
+ talker._mask_eos_rows = None
+ talker._pending_force_eos_rows = None
+ talker._penalty_histories = None
talker._request_audio_states = {}
+ talker._request_condition_states = {}
talker._deferred_cleanup_ids = set()
- talker._codec_min_tokens = 50
- talker._codec_seed = 42
+ talker._tts_config = ConditionalChatTTSConfig()
+ talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)])
+ with torch.no_grad():
+ talker.head_code[0].weight.copy_(torch.eye(8, 2))
return talker
@@ -90,46 +137,6 @@ def _reference_repetition_penalty(
return torch.where(logits < 0, logits * alpha, logits / alpha)
-def _make_sampling_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
- talker = _make_talker()
- talker._codec_temperature = 0.8
- talker._codec_top_k = 5
- talker._codec_top_p = 0.85
- talker._codec_repetition_penalty = 1.05
- talker._request_audio_states = {
- "req-a": {"min_tokens": 2},
- "req-b": {"min_tokens": 0},
- }
- talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)])
- with torch.no_grad():
- talker.head_code[0].weight.copy_(
- torch.tensor(
- [
- [0.1, -0.3],
- [0.2, 0.4],
- [-0.5, 0.2],
- [0.7, -0.1],
- [-0.2, 0.8],
- [0.6, 0.3],
- [-0.4, -0.7],
- [0.3, 0.5],
- ]
- )
- )
- return talker
-
-
-@pytest.mark.parametrize(
- ("condition_tokens", "expected"),
- [(3, 64), (100, 1000), (1000, 2048)],
-)
-def test_audio_token_limit_scales_with_condition_length(
- condition_tokens: int,
- expected: int,
-) -> None:
- assert _max_audio_tokens(condition_tokens) == expected
-
-
def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> None:
generator = torch.Generator().manual_seed(42)
weight_v = torch.randn(8, 16, generator=generator, dtype=torch.bfloat16)
@@ -220,133 +227,454 @@ def test_batched_repetition_penalty_matches_rows_across_chunks(mocker) -> None:
]
-def test_batched_codec_sampling_is_request_order_independent() -> None:
- talker = _make_sampling_talker()
- histories = [torch.tensor([1, 1, 2]), torch.tensor([3, 4, 4])]
- hidden = torch.tensor([[0.25, -0.75], [-0.5, 0.5]])
+def test_scheduler_prompt_is_fully_blanked_for_penalties() -> None:
+ # The handoff path fills the prompt with 0s and the non-handoff path passes
+ # thinker token ids; neither is codec history, so every position is padded.
+ prompt = torch.tensor([[0, 0, 0], [12, 0, 340]])
- forward_order = talker._sample_audio_codes(
- hidden,
- histories,
- ["req-a", "req-b"],
- [0, 0],
- )
- talker._request_generators = {}
- reverse_order = talker._sample_audio_codes(
- hidden.flip(0),
- list(reversed(histories)),
- ["req-b", "req-a"],
- [0, 0],
+ blanked = blank_scheduler_prompt_for_penalties(prompt, vocab_size=8)
+
+ assert blanked.tolist() == [[8, 8, 8], [8, 8, 8]]
+ assert prompt.tolist() == [[0, 0, 0], [12, 0, 340]]
+
+
+def test_blanked_prompt_leaves_every_codec_logit_unpenalized() -> None:
+ from vllm.model_executor.layers.utils import apply_penalties
+
+ vocab_size = 8
+ logits = torch.arange(1.0, vocab_size + 1.0).unsqueeze(0)
+ prompt = blank_scheduler_prompt_for_penalties(torch.tensor([[0, 0, 3]]), vocab_size)
+
+ penalized = apply_penalties(
+ logits.clone(),
+ prompt,
+ torch.full((1, 1), vocab_size, dtype=torch.long),
+ torch.zeros(1),
+ torch.zeros(1),
+ torch.full((1,), 1.02),
)
- assert torch.equal(forward_order[0], reverse_order[1])
- assert torch.equal(forward_order[1], reverse_order[0])
+ torch.testing.assert_close(penalized, logits)
-def test_batched_codec_sampling_preserves_request_rng_after_compaction() -> None:
- talker = _make_sampling_talker()
- history_a = torch.tensor([1, 1, 2])
- history_b = torch.tensor([3, 4, 4])
- hidden_a = torch.tensor([[0.25, -0.75]])
- hidden_b = torch.tensor([[-0.5, 0.5]])
-
- first_batch = talker._sample_audio_codes(
- torch.cat([hidden_a, hidden_b]),
- [history_a, history_b],
- ["req-a", "req-b"],
- [0, 0],
- )
- compacted_b = talker._sample_audio_codes(
- hidden_b,
- [history_b],
- ["req-b"],
- [1],
- )
-
- talker._request_generators = {}
- standalone_b_first = talker._sample_audio_codes(
- hidden_b,
- [history_b],
- ["req-b"],
- [0],
- )[0]
- standalone_b_second = talker._sample_audio_codes(
- hidden_b,
- [history_b],
- ["req-b"],
- [1],
- )[0]
-
- assert torch.equal(first_batch[1], standalone_b_first)
- assert torch.equal(compacted_b[0], standalone_b_second)
-
-
-def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None:
+def test_talker_sample_blanks_penalty_prompt_without_touching_runner_state(mocker) -> None:
talker = _make_talker()
- seen: list[tuple[str, list[float], list[int]]] = []
+ metadata = _PenaltyMetadata(prompt_token_ids=torch.tensor([[0, 2, 0]]))
+ captured: dict[str, torch.Tensor] = {}
- def sample(hidden, histories, request_ids, steps):
- assert steps == [0, 0]
- seen.extend(
- (request_id, hidden[row].reshape(-1).tolist(), histories[row].tolist())
- for row, request_id in enumerate(request_ids)
+ class _Recorder:
+ def __call__(self, logits, sampling_metadata):
+ prompt = sampling_metadata.prompt_token_ids
+ assert prompt is not None
+ captured["prompt"] = prompt
+ return None
+
+ mocker.patch(
+ "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler",
+ return_value=_Recorder(),
+ )
+ talker.sample(torch.zeros(1, 8), metadata)
+
+ assert captured["prompt"].tolist() == [[8, 8, 8]]
+ assert metadata.prompt_token_ids is not None
+ assert metadata.prompt_token_ids.tolist() == [[0, 2, 0]]
+
+
+def test_compute_logits_returns_codec_vocab() -> None:
+ talker = _make_talker()
+ hidden = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+ logits = talker.compute_logits(hidden)
+
+ assert logits.shape == (2, 8)
+ assert torch.allclose(logits[0], talker.head_code[0](hidden[0]))
+
+
+def test_compute_logits_forces_eos_for_empty_speech() -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-empty"] = {"finished": True}
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0
+ assert output.multimodal_outputs["meta"]["finished"][0].item() is True
+ assert int(logits.argmax(dim=-1)) == 7
+
+
+def _capture_sampler_logits(mocker) -> dict[str, Any]:
+ captured: dict[str, Any] = {}
+
+ def _record(logits, sampling_metadata):
+ captured["logits"] = logits.clone()
+ captured["metadata"] = sampling_metadata
+ return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True))
+
+ mocker.patch(
+ "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler",
+ return_value=_record,
+ )
+ return captured
+
+
+def test_sample_scores_the_upstream_windowed_codec_penalty(mocker) -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-pen"] = {"finished": False, "step": 0}
+ # ``head_code`` is eye(8, 2) here, so only codes 0 and 1 carry a non-zero
+ # logit and are the only ones a penalty can visibly move.
+ hidden = torch.tensor([[1.0, -2.0]])
+ raw = talker.head_code[0](hidden).float()
+
+ output = talker.make_omni_output(
+ hidden,
+ model_intermediate_buffer=[{"request_id": "req-pen", "codes": {"audio": torch.tensor([[1]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+ captured = _capture_sampler_logits(mocker)
+ talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05])))
+
+ expected = _reference_repetition_penalty(raw, torch.tensor([1]), penalty=1.05, window_size=_CODEC_PENALTY_WINDOW)
+ torch.testing.assert_close(captured["logits"], expected)
+ # A negative logit is pushed further down, so the row really did change.
+ assert captured["logits"][0, 1].item() < raw[0, 1].item()
+ # Scored once: the Sampler's own whole-stream pass is neutralized.
+ assert captured["metadata"].repetition_penalties.tolist() == [1.0]
+ # Consumed once, so a later step cannot rescore a stale row.
+ assert talker._penalty_histories is None
+
+
+def test_codec_penalty_forgets_codes_older_than_the_upstream_window(mocker) -> None:
+ # vLLM's sampler penalty spans the whole stream, so every code ever sampled
+ # stays taxed and a long tail drifts toward never-sampled codes. Upstream
+ # scores only the last ``past_window`` frames.
+ talker = _make_talker()
+ talker._request_audio_states["req-window"] = {"finished": False, "step": 0}
+ hidden = torch.tensor([[1.0, -2.0]])
+ raw = talker.head_code[0](hidden).float()
+ captured = _capture_sampler_logits(mocker)
+
+ def emit(code: int) -> torch.Tensor:
+ output = talker.make_omni_output(
+ hidden,
+ model_intermediate_buffer=[{"request_id": "req-window", "codes": {"audio": torch.tensor([[code]])}}],
+ request_token_spans=[(0, 1)],
)
- return torch.tensor([2, 3])
+ logits = talker.compute_logits(output.text_hidden_states)
+ talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.05])))
+ return captured["logits"]
- batched_sample = mocker.patch.object(talker, "_sample_audio_codes", side_effect=sample)
+ assert emit(1)[0, 1].item() < raw[0, 1].item()
+ # Push code 1 out of the window with a full window of code 0.
+ for _ in range(_CODEC_PENALTY_WINDOW):
+ scored = emit(0)
+
+ assert talker._request_audio_states["req-window"]["recent_codes"] == [0] * _CODEC_PENALTY_WINDOW
+ torch.testing.assert_close(scored[0, 1], raw[0, 1])
+
+
+def test_sample_honours_a_request_that_disabled_penalties(mocker) -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-none"] = {"finished": False, "step": 0}
+ hidden = torch.tensor([[1.0, -2.0]])
+ raw = talker.head_code[0](hidden).float()
+
+ output = talker.make_omni_output(
+ hidden,
+ model_intermediate_buffer=[{"request_id": "req-none", "codes": {"audio": torch.tensor([[1]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+ captured = _capture_sampler_logits(mocker)
+ talker.sample(logits, _CodecSamplingMetadata(repetition_penalties=torch.tensor([1.0]), no_penalties=True))
+
+ torch.testing.assert_close(captured["logits"], raw)
+
+
+def test_sample_restores_forced_eos_that_min_tokens_masked_away(mocker) -> None:
+ # ``min_tokens`` masks every stage stop_token_id, including the codec EOS
+ # compute_logits just forced, so the sampler returns an arbitrary id for a
+ # request the model already finished.
+ talker = _make_talker()
+ talker._request_audio_states["req-empty"] = {"finished": True}
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+ logits[0, talker._codec_eos_id] = float("-inf")
+
+ class _MinTokensSampler:
+ def __call__(self, logits, sampling_metadata):
+ del sampling_metadata
+ return _SamplerOutput(sampled_token_ids=logits.argmax(dim=-1, keepdim=True))
+
+ mocker.patch(
+ "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler",
+ return_value=_MinTokensSampler(),
+ )
+ sampled = talker.sample(logits, _PenaltyMetadata(prompt_token_ids=None))
+
+ assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id]]
+
+
+def test_sample_leaves_unforced_rows_to_the_sampler(mocker) -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-a"] = {"finished": True}
+ talker._request_audio_states["req-b"] = {"finished": False, "step": 0}
+
+ output = talker.make_omni_output(
+ torch.ones(2, 2),
+ model_intermediate_buffer=[
+ {"request_id": "req-a", "codes": {"audio": torch.empty(0)}},
+ {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}},
+ ],
+ request_token_spans=[(0, 1), (1, 2)],
+ )
+ talker.compute_logits(output.text_hidden_states)
+
+ mocker.patch(
+ "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler",
+ return_value=lambda logits, sampling_metadata: _SamplerOutput(sampled_token_ids=torch.tensor([[0], [4]])),
+ )
+ sampled = talker.sample(torch.zeros(2, 8), _PenaltyMetadata(prompt_token_ids=None))
+
+ assert sampled.sampled_token_ids.tolist() == [[talker._codec_eos_id], [4]]
+ # A single compute_logits decision must not leak into the next sample().
+ assert talker._pending_force_eos_rows is None
+
+
+def test_native_duplex_chunk_budget_masks_eos_until_generate_chunk_is_full() -> None:
+ assert _native_duplex_chunk_budget({}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, _DUPLEX_CODEC_TOKENS_PER_CHUNK)
+ assert _native_duplex_chunk_budget({"turn_start": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0)
+ assert _native_duplex_chunk_budget({"turn_end": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0)
+
+ talker = _make_talker()
+ talker._request_audio_states["req-chunk"] = {
+ "finished": False,
+ "step": 0,
+ "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK,
+ "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK,
+ }
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-chunk", "codes": {"audio": torch.empty(0)}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert logits[0, 7].item() == float("-inf")
+ assert torch.isfinite(logits[0, 0])
+ assert output.multimodal_outputs["meta"]["finished"][0].item() is False
+
+
+def test_native_duplex_chunk_budget_forces_eos_after_twenty_five_frames() -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-full"] = {
+ "finished": False,
+ "step": 24,
+ "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK,
+ "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK,
+ }
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-full", "codes": {"audio": torch.tensor([[3]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert talker._request_audio_states["req-full"]["step"] == 25
+ assert talker._request_audio_states["req-full"]["finished"] is True
+ assert int(logits.argmax(dim=-1)) == 7
+ assert output.multimodal_outputs["meta"]["finished"][0].item() is True
+
+
+def test_offline_prefill_caps_codec_budget_to_the_remaining_talker_context() -> None:
+ talker = _make_talker()
+ talker.emb_text = nn.Embedding(1, 2)
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+ talker._text_eos_id = 0
+ talker._tts_bos_id = 0
+ # Long enough that the context, not upstream's max_new_token, is the binding
+ # limit; the Talker must never plan past its own position embeddings.
+ prompt_len = int(talker._tts_config.max_position_embeddings) - _OFFLINE_CODEC_MAX_NEW_TOKENS + 100
+
+ _, _, updates = talker.preprocess(
+ torch.zeros(2, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ _omni_prompt_len=prompt_len,
+ request_id="req-offline-cap",
+ tts_token_ids=torch.empty(0, dtype=torch.long),
+ tts_hidden_states=torch.empty(0, 2),
+ )
+
+ remaining = int(talker._tts_config.max_position_embeddings) - prompt_len
+ assert remaining < _OFFLINE_CODEC_MAX_NEW_TOKENS
+ assert updates["audio_state"]["max_tokens"] == remaining
+ assert updates["audio_state"]["min_tokens"] is None
+
+
+def test_talker_masks_eos_until_recorded_min_tokens() -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-early"] = {
+ "finished": False,
+ "step": 0,
+ "max_tokens": 8,
+ "min_tokens": 4,
+ }
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-early", "codes": {"audio": torch.tensor([[3]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert talker._request_audio_states["req-early"]["step"] == 1
+ assert logits[0, 7].item() == float("-inf")
+ assert torch.isfinite(logits[0, 0])
+ assert output.multimodal_outputs["meta"]["finished"][0].item() is False
+
+
+def test_offline_talker_forces_eos_at_max_new_token() -> None:
+ talker = _make_talker()
+ max_tokens = 8
+ talker._request_audio_states["req-cap"] = {
+ "finished": False,
+ "step": max_tokens - 2,
+ "max_tokens": max_tokens,
+ }
+
+ output = talker.make_omni_output(
+ torch.ones(1, 2),
+ model_intermediate_buffer=[{"request_id": "req-cap", "codes": {"audio": torch.tensor([[3]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert talker._request_audio_states["req-cap"]["step"] == max_tokens - 1
+ assert talker._request_audio_states["req-cap"]["finished"] is True
+ assert int(logits.argmax(dim=-1)) == 7
+ assert output.multimodal_outputs["meta"]["finished"][0].item() is True
+
+
+def test_talker_without_a_recorded_budget_neither_forces_nor_masks_eos() -> None:
+ # A state rebuilt from a transported ``audio_state`` can lack budget keys;
+ # such rows must fall through to the raw codec head.
+ talker = _make_talker()
+ talker._request_audio_states["req-nobudget"] = {"finished": False, "step": 0}
+
+ output = talker.make_omni_output(
+ torch.tensor([[1.0, 0.0]]),
+ model_intermediate_buffer=[{"request_id": "req-nobudget", "codes": {"audio": torch.tensor([[3]])}}],
+ request_token_spans=[(0, 1)],
+ )
+ logits = talker.compute_logits(output.text_hidden_states)
+
+ assert talker._request_audio_states["req-nobudget"]["step"] == 1
+ assert torch.allclose(logits[0], talker.head_code[0](output.text_hidden_states[0]).float())
+
+
+def test_offline_prefill_records_upstream_max_new_token_budget() -> None:
+ talker = _make_talker()
+ talker.emb_text = nn.Embedding(1, 2)
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+ talker._text_eos_id = 0
+ talker._tts_bos_id = 0
+
+ _, _, updates = talker.preprocess(
+ torch.zeros(2, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ request_id="req-offline",
+ tts_token_ids=torch.empty(0, dtype=torch.long),
+ tts_hidden_states=torch.empty(0, 2),
+ )
+
+ # A 2-token boundary prompt leaves far more context than upstream's
+ # max_new_token, so MiniCPMTTS.generate's own ceiling is what binds.
+ assert int(talker._tts_config.max_position_embeddings) - 2 > _OFFLINE_CODEC_MAX_NEW_TOKENS
+ assert updates["audio_state"]["max_tokens"] == _OFFLINE_CODEC_MAX_NEW_TOKENS
+ assert updates["audio_state"]["min_tokens"] is None
+
+
+def test_native_duplex_prefill_records_generate_chunk_budget() -> None:
+ talker = _make_talker()
+ talker.emb_text = nn.Embedding(1, 2)
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+ talker._text_eos_id = 0
+ talker._tts_bos_id = 0
+
+ _, _, mid_turn = talker.preprocess(
+ torch.zeros(2, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ request_id="req-mid",
+ native_duplex=True,
+ tts_token_ids=torch.empty(0, dtype=torch.long),
+ tts_hidden_states=torch.empty(0, 2),
+ )
+ _, _, boundary = talker.preprocess(
+ torch.zeros(2, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ request_id="req-end",
+ native_duplex=True,
+ meta={"turn_end": True},
+ tts_token_ids=torch.empty(0, dtype=torch.long),
+ tts_hidden_states=torch.empty(0, 2),
+ )
+
+ assert mid_turn["audio_state"]["max_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK
+ assert mid_turn["audio_state"]["min_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK
+ assert boundary["audio_state"]["min_tokens"] == 0
+
+
+def test_make_omni_output_packs_the_previous_codec_id() -> None:
+ talker = _make_talker()
infos = [
- {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}},
- {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}},
+ {"request_id": "req-a", "codes": {"audio": torch.tensor([[2]])}},
+ {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}},
]
output = talker.make_omni_output(
- torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]),
+ torch.ones(2, 2),
model_intermediate_buffer=infos,
- request_token_spans=[(0, 2), (2, 3)],
+ request_token_spans=[(0, 1), (1, 2)],
)
- assert seen == [
- ("req-a", [2.0, 0.0], [1]),
- ("req-b", [3.0, 0.0], []),
- ]
- batched_sample.assert_called_once()
- assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2]
- assert infos[1]["audio_codes"]["accumulated"].tolist() == [3]
- assert set(output.multimodal_outputs) == {"codes", "meta"}
- assert "model_outputs" not in output.multimodal_outputs
- assert "sr" not in output.multimodal_outputs
assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]]
assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]]
assert _routed(output, 0)["meta"]["finished"].item() is False
- assert set(output.multimodal_outputs["meta"]) == {"finished"}
- assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0]
-def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
+def test_talker_projects_request_aligned_duplex_metadata() -> None:
talker = _make_talker()
- mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2, 2]))
infos = [
{
"request_id": "req-a",
"native_duplex": True,
"duplex": {"epoch": 3, "turn_id": 7},
"ids": {"tts": [41]},
- "meta": {
- "native_duplex_segment_text": "first",
- "turn_eos_token_id": 99,
- },
- "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
+ "meta": {"native_duplex_segment_text": "first", "turn_eos_token_id": 99},
+ "codes": {"audio": torch.empty(0)},
},
{
"request_id": "req-b",
"native_duplex": True,
"duplex": {"epoch": 4, "turn_id": 8},
"ids": {"tts": [42, 99]},
- "meta": {
- "native_duplex_segment_text": "second",
- "turn_eos_token_id": 99,
- },
- "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
+ "meta": {"native_duplex_segment_text": "second", "turn_eos_token_id": 99},
+ "codes": {"audio": torch.empty(0)},
},
]
@@ -360,7 +688,6 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
assert [value.item() for value in meta["native_duplex"]] == [True, True]
assert [value.item() for value in meta["duplex_epoch"]] == [3, 4]
assert [value.item() for value in meta["duplex_turn_id"]] == [7, 8]
- assert "native_duplex_segment_text" not in meta
assert [bytes(value.tolist()).decode("utf-8") for value in meta["llm_output_text_utf8"]] == [
"first",
"second",
@@ -368,14 +695,9 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
assert [value.item() for value in meta["turn_end"]] == [False, True]
-def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
+def test_talker_rejects_native_duplex_without_fence_identity() -> None:
talker = _make_talker()
- mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2]))
- info = {
- "request_id": "req-missing-fence",
- "native_duplex": True,
- "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
- }
+ info = {"request_id": "req-missing-fence", "native_duplex": True, "codes": {"audio": torch.empty(0)}}
with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"):
talker.make_omni_output(
@@ -385,113 +707,35 @@ def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
)
-def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None:
+def test_decode_preprocess_embeds_the_sampled_codec_id() -> None:
talker = _make_talker()
- sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2]))
- infos = [
- {
- "request_id": "req-prefill",
- "audio_state": {"step": 0},
- "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
- },
- {
- "request_id": "req-decode",
- "audio_state": {"step": 4},
- "audio_codes": {"accumulated": torch.tensor([1])},
- },
- ]
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)])
+ with torch.no_grad():
+ talker.emb_code[0].weight.copy_(torch.arange(32, dtype=torch.float32).reshape(8, 4))
- output = talker.make_omni_output(
- torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]),
- model_intermediate_buffer=infos,
- request_token_spans=[(0, 2), (2, 3)],
- request_sample_eligible=[False, True],
+ _, embeds, updates = talker.preprocess(
+ torch.tensor([3]),
+ None,
+ request_id="req-decode",
+ audio_state={"finished": False},
)
- sample.assert_called_once()
- assert sample.call_args.args[2] == ["req-decode"]
- assert infos[0]["audio_state"]["step"] == 0
- assert infos[0]["audio_codes"]["accumulated"].numel() == 0
- assert infos[1]["audio_state"]["step"] == 5
- assert _routed(output, 0)["codes"]["audio"].shape == (0, 1)
- assert _routed(output, 1)["codes"]["audio"].tolist() == [[2]]
+ assert torch.equal(embeds, talker.emb_code[0](torch.tensor([3])))
+ assert updates["codes"]["audio"].tolist() == [[3]]
-def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None:
+def test_decode_preprocess_drops_codec_eos() -> None:
talker = _make_talker()
- sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([7]))
- info = {
- "request_id": "req-stop",
- "audio_state": {"step": 3},
- "audio_codes": {"accumulated": torch.tensor([4, 5])},
- }
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)])
- first = talker.make_omni_output(
- torch.ones(1, 2),
- model_intermediate_buffer=[info],
- request_token_spans=[(0, 1)],
- )
- first_logits = talker.compute_logits(first.text_hidden_states)
- second = talker.make_omni_output(
- torch.ones(1, 2),
- model_intermediate_buffer=[info],
- request_token_spans=[(0, 1)],
+ _, _, updates = talker.preprocess(
+ torch.tensor([7]),
+ None,
+ request_id="req-eos",
+ audio_state={"finished": False},
)
- sample.assert_called_once()
- assert info["audio_codes"]["accumulated"].tolist() == [4, 5]
- assert first.multimodal_outputs["codes"]["audio"][0].shape == (0, 1)
- assert first.multimodal_outputs["meta"]["finished"][0].item() is True
- assert second.multimodal_outputs["meta"]["finished"][0].item() is False
- assert first_logits.argmax(dim=-1).tolist() == [1]
- assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1]
-
-
-def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None:
- talker = _make_talker()
- mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3]))
- info = {
- "request_id": "req-limit",
- "audio_state": {"step": 1, "max_tokens": 2},
- "audio_codes": {"accumulated": torch.tensor([4, 5])},
- }
-
- output = talker.make_omni_output(
- torch.ones(1, 2),
- model_intermediate_buffer=[info],
- request_token_spans=[(0, 1)],
- )
-
- # MiniCPMTTS.generate_chunk samples once at the max-token boundary to
- # advance RNG state, but the sampled code is not fed into KV or returned.
- assert info["audio_codes"]["accumulated"].tolist() == [4, 5]
- assert output.multimodal_outputs["codes"]["audio"][0].shape == (0, 1)
- assert output.multimodal_outputs["meta"]["finished"][0].item() is True
- assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1]
-
-
-def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None:
- talker = _make_talker()
- mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3]))
- first_info = {
- "request_id": "req-local-state",
- "audio_state": {"step": 1, "max_tokens": 3},
- "audio_codes": {"accumulated": torch.tensor([4])},
- }
-
- talker.make_omni_output(
- torch.ones(1, 2),
- model_intermediate_buffer=[first_info],
- request_token_spans=[(0, 1)],
- )
- second = talker.make_omni_output(
- torch.ones(1, 2),
- model_intermediate_buffer=[{"request_id": "req-local-state"}],
- request_token_spans=[(0, 1)],
- )
-
- assert second.multimodal_outputs["meta"]["finished"][0].item() is True
- assert talker._request_audio_states["req-local-state"]["step"] == 3
+ assert updates["codes"]["audio"].numel() == 0
def test_missing_conditioning_fails_clearly() -> None:
@@ -524,18 +768,19 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None:
assert torch.equal(embeds, talker.emb_text(torch.tensor([5, 6])))
assert updates["audio_state"]["finished"] is True
+ assert updates["codes"]["audio"].numel() == 0
- # Stage 1's sampling min_tokens keeps scheduling decode steps until the stop
- # token becomes eligible, and those steps have no previous code to embed.
- _, decode_embeds, _ = talker.preprocess(
+ # min_tokens can keep scheduling leftover decode rows after an empty
+ # speech segment. Those rows have no previous codec id to embed.
+ _, decode_embeds, decode_updates = talker.preprocess(
torch.zeros(1, dtype=torch.long),
None,
request_id="req-empty",
audio_state=updates["audio_state"],
- audio_codes=updates["audio_codes"],
)
assert decode_embeds.shape == (1, 4)
+ assert decode_updates["codes"]["audio"].numel() == 0
def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None:
@@ -544,7 +789,7 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non
condition = torch.arange(18, dtype=torch.float32).reshape(9, 2)
mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition)
- _, embeds, _ = talker.preprocess(
+ _, embeds, updates = talker.preprocess(
torch.zeros(9, dtype=torch.long),
None,
_omni_is_prefill=True,
@@ -556,46 +801,9 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non
)
assert torch.equal(embeds, condition)
+ assert updates["codes"]["audio"].numel() == 0
state = talker._request_audio_states["req-chunked-prefill"]
- assert state["min_tokens"] == 50
- assert state["max_tokens"] == 64
-
-
-@pytest.mark.parametrize(
- ("meta", "expected_min_tokens"),
- [
- ({"turn_start": True}, 0),
- ({}, 26),
- ({"turn_end": True}, 0),
- ],
-)
-def test_native_duplex_prefill_uses_official_chunk_limits(
- mocker,
- meta,
- expected_min_tokens,
-) -> None:
- talker = _make_talker()
- talker.emb_text = nn.Embedding(1, 2)
- mocker.patch.object(
- talker,
- "_build_condition_embeddings",
- return_value=torch.ones(3, 2),
- )
-
- talker.preprocess(
- torch.zeros(3, dtype=torch.long),
- None,
- _omni_is_prefill=True,
- request_id="req-duplex-chunk",
- native_duplex=True,
- meta=meta,
- tts_token_ids=torch.tensor([1]),
- tts_hidden_states=torch.ones(1, 2),
- )
-
- state = talker._request_audio_states["req-duplex-chunk"]
- assert state["min_tokens"] == expected_min_tokens
- assert state["max_tokens"] == 26
+ assert state["finished"] is False
def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None:
@@ -626,13 +834,164 @@ def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None:
assert condition.shape[0] == token_ids.shape[0] + 1
-def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None:
+def test_native_duplex_rollover_matches_official_sliding_recompute(mocker) -> None:
talker = _make_talker()
- talker._request_generators["req-done"] = torch.Generator()
- talker._request_audio_states["req-done"] = {"step": 1}
+ talker.emb_text = nn.Embedding(1, 2)
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+ with torch.no_grad():
+ talker.emb_code[0].weight.copy_(torch.arange(16, dtype=torch.float32).reshape(8, 2))
+
+ previous_condition = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+ current_condition = torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]])
+ build_condition = mocker.patch.object(
+ talker,
+ "_build_condition_embeddings",
+ side_effect=[previous_condition, current_condition, current_condition],
+ )
+
+ talker.preprocess(
+ torch.zeros(2, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ _omni_prompt_len=2,
+ request_id="req-window",
+ native_duplex=True,
+ meta={"turn_start": True, "streaming_condition_seq": 7},
+ tts_token_ids=torch.tensor([1]),
+ tts_hidden_states=torch.ones(1, 2),
+ )
+
+ older_codes = [3, 4, 5]
+ talker._request_condition_states["req-window"]["base_recent_codes"] = tuple(older_codes)
+ previous_codes = [1, 2]
+ recompute_meta = {
+ "streaming_condition_seq": 8,
+ "streaming_prompt_recompute": True,
+ }
+ expected = torch.cat(
+ [
+ previous_condition,
+ talker.emb_code[0](torch.tensor(previous_codes)),
+ current_condition,
+ ],
+ dim=0,
+ )
+ _, full_embeds, _ = talker.preprocess(
+ torch.zeros(expected.shape[0], dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ _omni_prompt_len=expected.shape[0],
+ request_id="req-window",
+ native_duplex=True,
+ ids={"streaming_prompt_previous_codes": previous_codes},
+ meta=recompute_meta,
+ tts_token_ids=torch.tensor([2]),
+ tts_hidden_states=torch.ones(1, 2),
+ )
+
+ assert torch.equal(full_embeds, expected)
+ assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes
+
+ # Chunked-prefill/preemption retries reuse the immutable recipe for the
+ # same condition instead of advancing the one-chunk window twice.
+ _, retry_slice, _ = talker.preprocess(
+ torch.zeros(3, dtype=torch.long),
+ None,
+ _omni_is_prefill=True,
+ _omni_num_computed_tokens=2,
+ _omni_prompt_len=expected.shape[0],
+ request_id="req-window",
+ native_duplex=True,
+ ids={"streaming_prompt_previous_codes": previous_codes},
+ meta=recompute_meta,
+ tts_token_ids=torch.tensor([2]),
+ tts_hidden_states=torch.ones(1, 2),
+ )
+ assert torch.equal(retry_slice, expected[2:5])
+ assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes
+ assert build_condition.call_count == 3
+
+
+def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> None:
+ talker = _make_talker()
+ condition = torch.ones(2, 2)
+ mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition)
+ common = {
+ "_omni_is_prefill": True,
+ "_omni_prompt_len": 2,
+ "request_id": "req-history",
+ "native_duplex": True,
+ "tts_token_ids": torch.tensor([1]),
+ "tts_hidden_states": torch.ones(1, 2),
+ }
+
+ talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 0}, **common)
+ with pytest.raises(ValueError, match="advanced without its streaming recompute marker"):
+ talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 1}, **common)
+
+
+@pytest.mark.parametrize(
+ ("state", "previous_codes", "error"),
+ [
+ (None, [1], "missing the previous Talker condition"),
+ (
+ {"condition_seq": 0, "condition": torch.ones(2, 2), "base_recent_codes": ()},
+ [7],
+ "invalid or terminal token",
+ ),
+ ],
+)
+def test_native_duplex_rollover_rejects_unverifiable_history(state, previous_codes, error) -> None:
+ talker = _make_talker()
+ talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+ if state is not None:
+ talker._request_condition_states["req-invalid-window"] = state
+
+ with pytest.raises(ValueError, match=error):
+ talker._build_streaming_recompute_embeddings(
+ torch.ones(3, 2),
+ request_id="req-invalid-window",
+ info_dict={"ids": {"streaming_prompt_previous_codes": previous_codes}},
+ meta={
+ "streaming_condition_seq": 1,
+ "streaming_prompt_recompute": True,
+ },
+ )
+
+
+def test_native_duplex_turn_start_resets_rollover_history() -> None:
+ talker = _make_talker()
+ talker._request_condition_states["req-turn"] = {
+ "condition_seq": 4,
+ "condition": torch.full((2, 2), 9.0),
+ "active_embeddings": torch.full((5, 2), 9.0),
+ }
+ current = torch.ones(3, 2)
+
+ result = talker._build_streaming_recompute_embeddings(
+ current,
+ request_id="req-turn",
+ info_dict={},
+ meta={"turn_start": True, "streaming_condition_seq": 5},
+ )
+
+ assert torch.equal(result, current)
+ state = talker._request_condition_states["req-turn"]
+ assert state["condition_seq"] == 5
+ assert torch.equal(state["condition"], current)
+ assert "active_embeddings" not in state
+
+
+def test_request_cleanup_evicts_decode_state() -> None:
+ talker = _make_talker()
+ talker._request_audio_states["req-done"] = {"finished": False}
+ talker._request_condition_states["req-done"] = {
+ "condition_seq": 1,
+ "condition": torch.ones(2, 2),
+ }
talker.on_requests_finished(["req-done"])
talker._flush_deferred_cleanup()
- assert "req-done" not in talker._request_generators
assert "req-done" not in talker._request_audio_states
+ assert "req-done" not in talker._request_condition_states
diff --git a/tests/model_executor/models/moss_tts/test_compact_topk.py b/tests/model_executor/models/moss_tts/test_compact_topk.py
new file mode 100644
index 000000000..e2b06c824
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_compact_topk.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU regression tests for MOSS-TTS compact top-k sampling."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts import modeling_moss_tts_local
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_compact_topk_only_samples_retained_candidates() -> None:
+ logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(7))
+
+ sampled = modeling_moss_tts_local._sample_token(
+ logits,
+ temperature=1.7,
+ top_k=25,
+ top_p=0.8,
+ do_sample=True,
+ generator=torch.Generator().manual_seed(42),
+ )
+
+ retained = logits.topk(25, dim=-1).indices
+ assert torch.all((retained == sampled.unsqueeze(-1)).any(dim=-1))
+
+
+def test_compact_topk_tiny_nucleus_selects_argmax() -> None:
+ logits = torch.randn(32, 1024, generator=torch.Generator().manual_seed(17))
+
+ sampled = modeling_moss_tts_local._sample_token(
+ logits,
+ temperature=1.7,
+ top_k=25,
+ top_p=1e-9,
+ do_sample=True,
+ generator=torch.Generator().manual_seed(42),
+ )
+
+ torch.testing.assert_close(sampled, logits.argmax(dim=-1))
diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py
new file mode 100644
index 000000000..90c386ee0
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU parity tests for MOSS-TTS Local Depth RoPE caching."""
+
+import pytest
+import torch
+from transformers import GPT2Config
+
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import (
+ MossTTSLocalDepthTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def _test_config() -> GPT2Config:
+ return GPT2Config(
+ n_embd=32,
+ n_head=4,
+ n_inner=64,
+ layer_norm_epsilon=1e-6,
+ rope_base=10000.0,
+ )
+
+
+def test_rope_cache_matches_reference_and_reuses_storage() -> None:
+ model = MossTTSLocalDepthTransformer(_test_config()).eval()
+ attn = model.h[0].attn
+ device = torch.device("cpu")
+
+ attn.prepare_rope_cache(12, device, torch.float32)
+ cache_ptr = attn._rope_cos_cache.data_ptr()
+ cos, sin = attn._rope_cos_sin(4, device, torch.float32)
+
+ positions = torch.arange(4, dtype=torch.float32)
+ freqs = torch.einsum("s,d->sd", positions, attn.inv_freq.float())
+ expected_cos = freqs.cos().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+ expected_sin = freqs.sin().repeat_interleave(2, dim=-1).view(1, 4, 1, attn.head_dim)
+ torch.testing.assert_close(cos, expected_cos)
+ torch.testing.assert_close(sin, expected_sin)
+
+ # Smaller slices retain the fixed allocation used by CUDA Graph.
+ attn.prepare_rope_cache(6, device, torch.float32)
+ assert attn._rope_cos_cache.data_ptr() == cache_ptr
diff --git a/tests/model_executor/models/moss_tts/test_moss_fused_load.py b/tests/model_executor/models/moss_tts/test_moss_fused_load.py
new file mode 100644
index 000000000..60039c693
--- /dev/null
+++ b/tests/model_executor/models/moss_tts/test_moss_fused_load.py
@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression tests for MOSS-TTS-Realtime fused code-predictor loading.
+
+The MOSS-TTS-Realtime depth transformer shares ``CodePredictorBaseModel``, whose
+q/k/v and gate/up projections are fused. Before the fix, the talker loaded those
+shards with a per-tensor ``default_weight_loader``, which silently dropped every
+q/k/v/gate/up shard and left the fused params randomly initialized -- no
+exception, just a corrupted model.
+
+These tests assert the property that failure mode violates: after
+``load_weights``, every fused parameter is both reported loaded and numerically
+equal to the concatenation of the shards that fed it.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import (
+ MossTTSLocalTransformerConfig,
+)
+from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import (
+ MossTTSRealtimeLocalTransformer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _tiny_cfg():
+ return MossTTSLocalTransformerConfig(
+ head_dim=8,
+ hidden_size=32,
+ num_hidden_layers=2,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ intermediate_size=64,
+ max_position_embeddings=8,
+ )
+
+
+def _hf_shards(cfg, *, drop: set[str] | None = None):
+ """Synthetic HF-style checkpoint tensors for the depth transformer."""
+ drop = drop or set()
+ h, inter = cfg.hidden_size, cfg.intermediate_size
+ q_out = cfg.num_attention_heads * cfg.head_dim
+ kv_out = cfg.num_key_value_heads * cfg.head_dim
+ out = {}
+ for layer in range(cfg.num_hidden_layers):
+ p = f"model.layers.{layer}"
+ for nm, rows in (("q_proj", q_out), ("k_proj", kv_out), ("v_proj", kv_out)):
+ key = f"{p}.self_attn.{nm}.weight"
+ if nm not in drop:
+ out[key] = torch.randn(rows, h)
+ out[f"{p}.self_attn.o_proj.weight"] = torch.randn(h, q_out)
+ for nm in ("gate_proj", "up_proj"):
+ key = f"{p}.mlp.{nm}.weight"
+ if nm not in drop:
+ out[key] = torch.randn(inter, h)
+ out[f"{p}.mlp.down_proj.weight"] = torch.randn(h, inter)
+ out[f"{p}.input_layernorm.weight"] = torch.randn(h)
+ out[f"{p}.post_attention_layernorm.weight"] = torch.randn(h)
+ for nm in ("q_norm", "k_norm"):
+ out[f"{p}.self_attn.{nm}.weight"] = torch.randn(cfg.head_dim)
+ out["model.norm.weight"] = torch.randn(h)
+ return out
+
+
+def test_every_fused_param_is_loaded_and_correctly_packed():
+ cfg = _tiny_cfg()
+ lt = MossTTSRealtimeLocalTransformer(cfg)
+ shards = _hf_shards(cfg)
+
+ loaded = lt.load_weights(iter([(k, v) for k, v in shards.items()]))
+
+ params = dict(lt.named_parameters())
+ fused = [n for n in params if ".self_attn.qkv_proj." in n or ".mlp.gate_up_proj." in n]
+ assert fused, "model declares no fused params -- test would be vacuous"
+
+ # 1. Reported as loaded (this is what the old path silently failed).
+ unreported = [n for n in fused if n not in loaded]
+ assert not unreported, f"fused params not reported loaded: {unreported}"
+
+ # 2. Numerically equal to the concatenation of their shards, in order.
+ for layer in range(cfg.num_hidden_layers):
+ p = f"model.layers.{layer}"
+ want_qkv = torch.cat([shards[f"{p}.self_attn.{n}.weight"] for n in ("q_proj", "k_proj", "v_proj")], dim=0)
+ got_qkv = params[f"{p}.self_attn.qkv_proj.weight"].detach()
+ assert torch.equal(got_qkv, want_qkv), f"layer {layer}: qkv packing mismatch"
+
+ want_gu = torch.cat([shards[f"{p}.mlp.{n}.weight"] for n in ("gate_proj", "up_proj")], dim=0)
+ got_gu = params[f"{p}.mlp.gate_up_proj.weight"].detach()
+ assert torch.equal(got_gu, want_gu), f"layer {layer}: gate_up packing mismatch"
+
+
+def test_incomplete_shard_group_raises_instead_of_corrupting():
+ """A checkpoint missing v_proj must fail loudly, not load a random qkv."""
+ cfg = _tiny_cfg()
+ lt = MossTTSRealtimeLocalTransformer(cfg)
+ shards = _hf_shards(cfg, drop={"v_proj"})
+
+ with pytest.raises(RuntimeError, match="incomplete fused shards"):
+ lt.load_weights(iter([(k, v) for k, v in shards.items()]))
+
+
+def test_absent_fused_group_raises():
+ """No q/k/v at all must also raise, via the missing-fused-param guard."""
+ cfg = _tiny_cfg()
+ lt = MossTTSRealtimeLocalTransformer(cfg)
+ shards = _hf_shards(cfg, drop={"q_proj", "k_proj", "v_proj"})
+
+ with pytest.raises(RuntimeError, match="missing fused parameters"):
+ lt.load_weights(iter([(k, v) for k, v in shards.items()]))
+
+
+def test_returned_names_keep_the_model_prefix():
+ """Returned names must be ``model.*`` for talker re-prefixing."""
+ cfg = _tiny_cfg()
+ lt = MossTTSRealtimeLocalTransformer(cfg)
+ loaded = lt.load_weights(iter([(k, v) for k, v in _hf_shards(cfg).items()]))
+ assert loaded, "nothing reported loaded"
+ assert all(n.startswith("model.") for n in loaded), sorted(n for n in loaded if not n.startswith("model."))
diff --git a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py
index 35c309979..b5866781c 100644
--- a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py
+++ b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py
@@ -23,6 +23,11 @@ class _RecordingLogitsProcessor:
return torch.zeros((1, 4))
+class _EmbeddingBackbone(nn.Module):
+ def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return input_ids.to(dtype=torch.float32).unsqueeze(-1).repeat(1, 4)
+
+
def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None:
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
MossTTSDelayTalkerForGeneration,
@@ -41,3 +46,74 @@ def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> N
assert logits.shape == (1, 4)
assert model.logits_processor.args == (model.text_lm_head, hidden_states)
assert model.logits_processor.kwargs == {}
+
+
+def test_moss_tts_local_keeps_fixed_shape_rows_and_routes_stop_logits() -> None:
+ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+ MossTTSLocalTalkerForGeneration,
+ )
+
+ model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+ nn.Module.__init__(model)
+ model.audio_pad_token_id = 1024
+ model.audio_assistant_slot_token_id = 2
+ model.im_end_token_id = 3
+ model.text_vocab_size = 8
+ model.n_vq = 4
+ model._batch_state = None
+ model._batch_state_spans = None
+ model._batch_should_continue = None
+
+ valid = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
+ stopped = torch.full((1, 4), model.audio_pad_token_id, dtype=torch.long)
+ infos = [
+ {"audio_state": {}, "audio_codes": {"current": valid}},
+ {"audio_state": {}, "audio_codes": {"current": stopped}},
+ ]
+ spans = [(0, 3), (3, 4)]
+
+ output = model.make_omni_output(
+ torch.randn(4, 8),
+ model_intermediate_buffer=infos,
+ request_token_spans=spans,
+ )
+
+ audio = output.multimodal_outputs["codes"]["audio"]
+ assert len(audio) == 2
+ torch.testing.assert_close(audio[0], valid)
+ # Keeping this row avoids CUDA boolean-index shape discovery; the
+ # talker2codec raw async-chunk processor drops all-pad rows on CPU.
+ torch.testing.assert_close(audio[1], stopped)
+ assert model._batch_should_continue.tolist() == [True, False]
+
+ logits = model.compute_logits(torch.randn(4, 8))
+
+ assert logits is not None
+ assert logits.argmax(dim=-1).tolist() == [2, 2, 2, 3]
+
+
+def test_moss_tts_local_single_token_prefill_does_not_advance_audio() -> None:
+ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
+ MossTTSLocalTalkerForGeneration,
+ )
+
+ model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
+ nn.Module.__init__(model)
+ model.model = _EmbeddingBackbone()
+ model.n_vq = 4
+
+ input_ids = torch.tensor([7], dtype=torch.long)
+ _, embeds, update = model.preprocess(
+ input_ids,
+ None,
+ _omni_is_prefill=True,
+ audio_state={"is_stopping": True},
+ ref_offset=5,
+ )
+
+ torch.testing.assert_close(embeds, torch.full((1, 4), 7.0))
+ assert update == {
+ "audio_state": {"is_stopping": False},
+ "ref_offset": 6,
+ }
+ assert "mtp_inputs" not in update
diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py
new file mode 100644
index 000000000..382e2f4d1
--- /dev/null
+++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_quantization.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for nested Qwen3-Omni quantization config mapping."""
+
+from dataclasses import dataclass
+
+import pytest
+
+from vllm_omni.model_executor.models.qwen3_omni.quantization import (
+ _OUTER_MAPPING_APPLIED,
+ apply_nested_quant_config_mapping,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _QuantConfig:
+ packed_modules_mapping: dict[str, list[str]]
+
+
+@dataclass
+class _NestedModel:
+ quant_config: _QuantConfig | None
+
+
+@dataclass
+class _NestedModelWithPackedModules(_NestedModel):
+ packed_modules_mapping: dict[str, list[str]]
+
+
+def _mapped_quant_config(**packed_modules_mapping):
+ quant_config = _QuantConfig(packed_modules_mapping=packed_modules_mapping)
+ setattr(quant_config, _OUTER_MAPPING_APPLIED, True)
+ return quant_config
+
+
+def test_nested_mapping_ignores_missing_quant_config():
+ apply_nested_quant_config_mapping(_NestedModel(quant_config=None))
+
+
+def test_nested_mapping_allows_model_without_packed_modules_mapping():
+ quant_config = _mapped_quant_config(existing=["module"])
+
+ apply_nested_quant_config_mapping(_NestedModel(quant_config=quant_config))
+
+ assert quant_config.packed_modules_mapping == {"existing": ["module"]}
+
+
+def test_nested_mapping_merges_packed_modules_mapping():
+ quant_config = _mapped_quant_config(existing=["module"])
+ model = _NestedModelWithPackedModules(
+ quant_config=quant_config,
+ packed_modules_mapping={"qkv_proj": ["q_proj", "k_proj", "v_proj"]},
+ )
+
+ apply_nested_quant_config_mapping(model)
+
+ assert quant_config.packed_modules_mapping == {
+ "existing": ["module"],
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
+ }
diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
index 20f8eb2f0..619f2169a 100644
--- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
+++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
@@ -697,3 +697,481 @@ class TestGumbelMaxSampling:
empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
# L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
assert (empirical - target).abs().sum().item() < 0.05
+
+
+class TestCodePredictorFusedProjections:
+ """Reference tests for the fused qkv_proj / gate_up_proj paths.
+
+ These CPU tests confirm that the packed projections implement the same
+ linear algebra as the unfused reference under this setup. GPU executions
+ may still diverge at the bit level when kernel choices change.
+ """
+
+ def test_fused_qkv_matches_separate(self, loaded_target_classes) -> None:
+ torch.manual_seed(0)
+ hidden = 32
+ num_heads, num_kv_heads, head_dim = 4, 2, 8
+ q_size = num_heads * head_dim
+ kv_size = num_kv_heads * head_dim
+ x = torch.randn(2, 5, hidden)
+
+ q = torch.nn.Linear(hidden, q_size, bias=False)
+ k = torch.nn.Linear(hidden, kv_size, bias=False)
+ v = torch.nn.Linear(hidden, kv_size, bias=False)
+ fused = torch.nn.Linear(hidden, q_size + 2 * kv_size, bias=False)
+ with torch.no_grad():
+ fused.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
+
+ ref_q, ref_k, ref_v = q(x), k(x), v(x)
+ out_q, out_k, out_v = fused(x).split([q_size, kv_size, kv_size], dim=-1)
+ assert torch.equal(ref_q, out_q)
+ assert torch.equal(ref_k, out_k)
+ assert torch.equal(ref_v, out_v)
+
+ def test_attention_split_redensifies_kv_slices(self, loaded_target_classes) -> None:
+ cp_config, _ = _make_tiny_config(loaded_target_classes)
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ attn = common_mod.CodePredictorAttention(cp_config)
+
+ qkv = torch.randn(2, 5, attn._q_size + 2 * attn._kv_size)
+ q_ref, k_ref, v_ref = qkv.split([attn._q_size, attn._kv_size, attn._kv_size], dim=-1)
+ q_raw, k_raw, v_raw = attn._split_qkv(qkv)
+
+ assert torch.equal(q_raw, q_ref)
+ assert torch.equal(k_raw, k_ref)
+ assert torch.equal(v_raw, v_ref)
+ assert k_raw.is_contiguous()
+ assert v_raw.is_contiguous()
+
+ def test_fused_gate_up_matches_separate(self, loaded_target_classes) -> None:
+ torch.manual_seed(1)
+ hidden, intermediate = 32, 64
+ x = torch.randn(2, 5, hidden)
+
+ g = torch.nn.Linear(hidden, intermediate, bias=False)
+ u = torch.nn.Linear(hidden, intermediate, bias=False)
+ fused = torch.nn.Linear(hidden, 2 * intermediate, bias=False)
+ with torch.no_grad():
+ fused.weight.copy_(torch.cat([g.weight, u.weight], dim=0))
+
+ ref_g, ref_u = g(x), u(x)
+ out_g, out_u = fused(x).split([intermediate, intermediate], dim=-1)
+ assert torch.equal(ref_g, out_g)
+ assert torch.equal(ref_u, out_u)
+
+ def test_load_weights_repacks_hf_shards(self, mocker: MockerFixture, loaded_target_classes) -> None:
+ """``load_weights`` must transparently re-pack HF q/k/v and gate/up
+ shards into the fused ``qkv_proj`` and ``gate_up_proj`` weights."""
+ _, _, _, code_predictor_model, _ = loaded_target_classes
+ cp_config, _ = _make_tiny_config(loaded_target_classes)
+ model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size)
+
+ # The test harness mocks ``default_weight_loader`` as a no-op; install a
+ # real copier so this test verifies both the returned names and the
+ # packed parameter values/order.
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ # Build a synthetic "HF checkpoint" that exposes separate q/k/v
+ # and gate/up shards, mirroring what HuggingFace ships.
+ param_dict = dict(model.named_parameters(remove_duplicate=False))
+ weights: list[tuple[str, torch.Tensor]] = []
+ expected_fused: dict[str, torch.Tensor] = {}
+ torch.manual_seed(42)
+ for name, param in param_dict.items():
+ if name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = name[: -len(".qkv_proj.weight")]
+ num_heads = cp_config.num_attention_heads
+ num_kv_heads = cp_config.num_key_value_heads
+ head_dim = cp_config.head_dim
+ q_size = num_heads * head_dim
+ kv_size = num_kv_heads * head_dim
+ hidden = cp_config.hidden_size
+ q_w = torch.randn(q_size, hidden)
+ k_w = torch.randn(kv_size, hidden)
+ v_w = torch.randn(kv_size, hidden)
+ weights.append((f"{prefix}.q_proj.weight", q_w))
+ weights.append((f"{prefix}.k_proj.weight", k_w))
+ weights.append((f"{prefix}.v_proj.weight", v_w))
+ expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0)
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = name[: -len(".gate_up_proj.weight")]
+ inter = cp_config.intermediate_size
+ hidden = cp_config.hidden_size
+ gate_w = torch.randn(inter, hidden)
+ up_w = torch.randn(inter, hidden)
+ weights.append((f"{prefix}.gate_proj.weight", gate_w))
+ weights.append((f"{prefix}.up_proj.weight", up_w))
+ expected_fused[name] = torch.cat([gate_w, up_w], dim=0)
+ else:
+ # Pass-through parameters (norms, embeddings, o_proj, down_proj).
+ weights.append((name, torch.randn_like(param)))
+
+ loaded = model.load_weights(weights)
+ after = dict(model.named_parameters(remove_duplicate=False))
+ assert expected_fused, "test config must contain at least one fused projection layer"
+ for fname, expected in expected_fused.items():
+ assert fname in loaded, f"{fname} was not assembled by load_weights"
+ assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order"
+
+ def test_load_weights_repacks_bias_shards(self, mocker: MockerFixture, loaded_target_classes) -> None:
+ """When ``attention_bias=True``, the q/k/v ``.bias`` shards must be
+ concatenated into ``qkv_proj.bias`` instead of being dropped (which
+ would leave the fused bias randomly initialized)."""
+ _, _, _, code_predictor_model, _ = loaded_target_classes
+ cp_config, _ = _make_tiny_config(loaded_target_classes)
+ cp_config.attention_bias = True
+ model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size)
+
+ # The test harness mocks ``default_weight_loader`` as a no-op; install a
+ # real copier so we can assert the packed values land on the parameter.
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ param_dict = dict(model.named_parameters(remove_duplicate=False))
+ weights: list[tuple[str, torch.Tensor]] = []
+ expected_bias: dict[str, torch.Tensor] = {}
+ torch.manual_seed(11)
+ for name, param in param_dict.items():
+ if name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = name[: -len(".qkv_proj.weight")]
+ q_size = cp_config.num_attention_heads * cp_config.head_dim
+ kv_size = cp_config.num_key_value_heads * cp_config.head_dim
+ hidden = cp_config.hidden_size
+ q_w, k_w, v_w = (torch.randn(s, hidden) for s in (q_size, kv_size, kv_size))
+ q_b, k_b, v_b = (torch.randn(s) for s in (q_size, kv_size, kv_size))
+ weights.append((f"{prefix}.q_proj.weight", q_w))
+ weights.append((f"{prefix}.k_proj.weight", k_w))
+ weights.append((f"{prefix}.v_proj.weight", v_w))
+ weights.append((f"{prefix}.q_proj.bias", q_b))
+ weights.append((f"{prefix}.k_proj.bias", k_b))
+ weights.append((f"{prefix}.v_proj.bias", v_b))
+ expected_bias[f"{prefix}.qkv_proj.bias"] = torch.cat([q_b, k_b, v_b], dim=0)
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = name[: -len(".gate_up_proj.weight")]
+ inter, hidden = cp_config.intermediate_size, cp_config.hidden_size
+ weights.append((f"{prefix}.gate_proj.weight", torch.randn(inter, hidden)))
+ weights.append((f"{prefix}.up_proj.weight", torch.randn(inter, hidden)))
+ elif name.endswith(".qkv_proj.bias"):
+ # Fused target is assembled from the q/k/v bias shards above; a
+ # real HF checkpoint never ships a qkv_proj.bias tensor directly.
+ continue
+ else:
+ weights.append((name, torch.randn_like(param)))
+
+ loaded = model.load_weights(weights)
+ after = dict(model.named_parameters())
+ assert expected_bias, "test config must contain at least one qkv layer"
+ for bname, exp in expected_bias.items():
+ assert bname in loaded, f"{bname} bias was not assembled by load_weights"
+ assert torch.equal(after[bname], exp), f"{bname} does not match the packed q/k/v bias"
+
+ def test_wrapper_load_weights_handles_model_prefixed_shards(
+ self,
+ mocker: MockerFixture,
+ loaded_target_classes,
+ ) -> None:
+ _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+ cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+ vllm_config = _make_vllm_config(mocker)
+ predictor = code_predictor_wrapper(
+ vllm_config=vllm_config,
+ config=cp_config,
+ talker_config=talker_config,
+ )
+
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ param_dict = dict(predictor.named_parameters(remove_duplicate=False))
+ weights: list[tuple[str, torch.Tensor]] = []
+ expected_fused: dict[str, torch.Tensor] = {}
+ torch.manual_seed(7)
+ for name, param in param_dict.items():
+ if name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = name[: -len(".qkv_proj.weight")]
+ q_size = cp_config.num_attention_heads * cp_config.head_dim
+ kv_size = cp_config.num_key_value_heads * cp_config.head_dim
+ hidden = cp_config.hidden_size
+ q_w = torch.randn(q_size, hidden)
+ k_w = torch.randn(kv_size, hidden)
+ v_w = torch.randn(kv_size, hidden)
+ weights.append((f"{prefix}.q_proj.weight", q_w))
+ weights.append((f"{prefix}.k_proj.weight", k_w))
+ weights.append((f"{prefix}.v_proj.weight", v_w))
+ expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0)
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = name[: -len(".gate_up_proj.weight")]
+ inter = cp_config.intermediate_size
+ hidden = cp_config.hidden_size
+ gate_w = torch.randn(inter, hidden)
+ up_w = torch.randn(inter, hidden)
+ weights.append((f"{prefix}.gate_proj.weight", gate_w))
+ weights.append((f"{prefix}.up_proj.weight", up_w))
+ expected_fused[name] = torch.cat([gate_w, up_w], dim=0)
+ else:
+ weights.append((name, torch.randn_like(param)))
+
+ loaded = predictor.load_weights(weights)
+ after = dict(predictor.named_parameters(remove_duplicate=False))
+ assert expected_fused, "test config must contain at least one fused projection layer"
+ for fname, expected in expected_fused.items():
+ assert fname in loaded, f"{fname} was not loaded through the outer wrapper"
+ assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order"
+
+ def test_wrapper_load_weights_handles_legacy_unprefixed_model_shards(
+ self,
+ mocker: MockerFixture,
+ loaded_target_classes,
+ ) -> None:
+ """Qwen3-Omni checkpoints may expose code_predictor body tensors as
+ ``layers.*`` / ``codec_embedding.*`` instead of ``model.*``. The
+ wrapper must still route those shards through the fused inner loader."""
+ _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+ cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+ vllm_config = _make_vllm_config(mocker)
+ predictor = code_predictor_wrapper(
+ vllm_config=vllm_config,
+ config=cp_config,
+ talker_config=talker_config,
+ )
+
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ param_dict = dict(predictor.named_parameters(remove_duplicate=False))
+ weights: list[tuple[str, torch.Tensor]] = []
+ expected_fused: dict[str, torch.Tensor] = {}
+ torch.manual_seed(17)
+ for name, param in param_dict.items():
+ if name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = name[: -len(".qkv_proj.weight")]
+ legacy_prefix = prefix.removeprefix("model.")
+ q_size = cp_config.num_attention_heads * cp_config.head_dim
+ kv_size = cp_config.num_key_value_heads * cp_config.head_dim
+ hidden = cp_config.hidden_size
+ q_w = torch.randn(q_size, hidden)
+ k_w = torch.randn(kv_size, hidden)
+ v_w = torch.randn(kv_size, hidden)
+ weights.append((f"{legacy_prefix}.q_proj.weight", q_w))
+ weights.append((f"{legacy_prefix}.k_proj.weight", k_w))
+ weights.append((f"{legacy_prefix}.v_proj.weight", v_w))
+ expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0)
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = name[: -len(".gate_up_proj.weight")]
+ legacy_prefix = prefix.removeprefix("model.")
+ inter = cp_config.intermediate_size
+ hidden = cp_config.hidden_size
+ gate_w = torch.randn(inter, hidden)
+ up_w = torch.randn(inter, hidden)
+ weights.append((f"{legacy_prefix}.gate_proj.weight", gate_w))
+ weights.append((f"{legacy_prefix}.up_proj.weight", up_w))
+ expected_fused[name] = torch.cat([gate_w, up_w], dim=0)
+ elif name.startswith("model."):
+ weights.append((name.removeprefix("model."), torch.randn_like(param)))
+ else:
+ weights.append((name, torch.randn_like(param)))
+
+ loaded = predictor.load_weights(weights)
+ after = dict(predictor.named_parameters(remove_duplicate=False))
+ assert expected_fused, "test config must contain at least one fused projection layer"
+ for fname, expected in expected_fused.items():
+ assert fname in loaded, f"{fname} was not loaded from legacy unprefixed shards"
+ assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order"
+
+ def test_wrapper_load_weights_handles_nested_code_predictor_prefixed_shards(
+ self,
+ mocker: MockerFixture,
+ loaded_target_classes,
+ ) -> None:
+ """Realtime Qwen3-Omni loaders may still prepend ``talker.code_predictor.``
+ before the inner ``model.*`` body; the wrapper must normalize that layer
+ too before delegating to the fused inner loader."""
+ _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+ cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+ vllm_config = _make_vllm_config(mocker)
+ predictor = code_predictor_wrapper(
+ vllm_config=vllm_config,
+ config=cp_config,
+ talker_config=talker_config,
+ prefix="talker.code_predictor",
+ )
+
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ param_dict = dict(predictor.named_parameters(remove_duplicate=False))
+ weights: list[tuple[str, torch.Tensor]] = []
+ expected_fused: dict[str, torch.Tensor] = {}
+ torch.manual_seed(23)
+ wrapper_prefix = "talker.code_predictor."
+ for name, param in param_dict.items():
+ wrapped_name = f"{wrapper_prefix}{name}"
+ if name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = wrapped_name[: -len(".qkv_proj.weight")]
+ q_size = cp_config.num_attention_heads * cp_config.head_dim
+ kv_size = cp_config.num_key_value_heads * cp_config.head_dim
+ hidden = cp_config.hidden_size
+ q_w = torch.randn(q_size, hidden)
+ k_w = torch.randn(kv_size, hidden)
+ v_w = torch.randn(kv_size, hidden)
+ weights.append((f"{prefix}.q_proj.weight", q_w))
+ weights.append((f"{prefix}.k_proj.weight", k_w))
+ weights.append((f"{prefix}.v_proj.weight", v_w))
+ expected_fused[name] = torch.cat([q_w, k_w, v_w], dim=0)
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = wrapped_name[: -len(".gate_up_proj.weight")]
+ inter = cp_config.intermediate_size
+ hidden = cp_config.hidden_size
+ gate_w = torch.randn(inter, hidden)
+ up_w = torch.randn(inter, hidden)
+ weights.append((f"{prefix}.gate_proj.weight", gate_w))
+ weights.append((f"{prefix}.up_proj.weight", up_w))
+ expected_fused[name] = torch.cat([gate_w, up_w], dim=0)
+ else:
+ weights.append((wrapped_name, torch.randn_like(param)))
+
+ loaded = predictor.load_weights(weights)
+ after = dict(predictor.named_parameters(remove_duplicate=False))
+ assert expected_fused, "test config must contain at least one fused projection layer"
+ for fname, expected in expected_fused.items():
+ assert fname in loaded, f"{fname} was not loaded from nested code_predictor shards"
+ assert torch.equal(after[fname], expected), f"{fname} does not match the packed shard order"
+
+ def test_wrapper_load_weights_accepts_incremental_non_layer_shard(
+ self,
+ mocker: MockerFixture,
+ loaded_target_classes,
+ ) -> None:
+ """Qwen3-Omni splits code-predictor weights across checkpoint shards.
+
+ Its first shard contains the transformer layers, while the next shard
+ re-enters the wrapper with only codec embeddings and LM heads. Loading
+ that second shard must not require every fused layer to appear again.
+ """
+ _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+ cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+ vllm_config = _make_vllm_config(mocker)
+ predictor = code_predictor_wrapper(
+ vllm_config=vllm_config,
+ config=cp_config,
+ talker_config=talker_config,
+ )
+
+ common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+ mocker.patch.object(
+ common_mod,
+ "default_weight_loader",
+ lambda param, weight: param.data.copy_(weight),
+ )
+
+ first_shard: list[tuple[str, torch.Tensor]] = []
+ second_shard: list[tuple[str, torch.Tensor]] = []
+ for name, param in predictor.named_parameters(remove_duplicate=False):
+ if name.startswith("model.codec_embedding.") or name.startswith("lm_head."):
+ second_shard.append((name, torch.randn_like(param)))
+ elif name.endswith(".self_attn.qkv_proj.weight"):
+ prefix = name[: -len(".qkv_proj.weight")]
+ q_size = cp_config.num_attention_heads * cp_config.head_dim
+ kv_size = cp_config.num_key_value_heads * cp_config.head_dim
+ hidden = cp_config.hidden_size
+ first_shard.extend(
+ [
+ (f"{prefix}.q_proj.weight", torch.randn(q_size, hidden)),
+ (f"{prefix}.k_proj.weight", torch.randn(kv_size, hidden)),
+ (f"{prefix}.v_proj.weight", torch.randn(kv_size, hidden)),
+ ]
+ )
+ elif name.endswith(".mlp.gate_up_proj.weight"):
+ prefix = name[: -len(".gate_up_proj.weight")]
+ inter = cp_config.intermediate_size
+ hidden = cp_config.hidden_size
+ first_shard.extend(
+ [
+ (f"{prefix}.gate_proj.weight", torch.randn(inter, hidden)),
+ (f"{prefix}.up_proj.weight", torch.randn(inter, hidden)),
+ ]
+ )
+ else:
+ first_shard.append((name, torch.randn_like(param)))
+
+ assert first_shard and second_shard
+ common_mod.CodePredictorWrapper.load_weights(predictor, first_shard)
+ loaded_second = common_mod.CodePredictorWrapper.load_weights(predictor, second_shard)
+
+ assert {name for name, _ in second_shard} <= loaded_second
+
+ def test_load_weights_raises_on_incomplete_shards(self, loaded_target_classes) -> None:
+ """Missing one of q/k/v (or gate/up) must surface a clear error.
+
+ ``CodePredictorBaseModel`` is the inner transformer (its parameter
+ paths look like ``layers..self_attn.qkv_proj.weight`` -- without
+ the ``model.`` prefix that the outer ``CodePredictorWrapper`` strips
+ before forwarding). The test mimics that contract.
+ """
+ _, _, _, code_predictor_model, _ = loaded_target_classes
+ cp_config, _ = _make_tiny_config(loaded_target_classes)
+ model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size)
+
+ # Find any actual self_attn layer prefix, to make the test robust
+ # to renaming of the inner model's parameter scheme.
+ attn_prefix = next(
+ (
+ name[: -len(".qkv_proj.weight")]
+ for name in dict(model.named_parameters()).keys()
+ if name.endswith(".self_attn.qkv_proj.weight")
+ ),
+ None,
+ )
+ assert attn_prefix is not None, "test config must contain at least one self_attn layer"
+
+ num_heads = cp_config.num_attention_heads
+ num_kv_heads = cp_config.num_key_value_heads
+ head_dim = cp_config.head_dim
+ hidden = cp_config.hidden_size
+
+ # Provide only q + k (no v). The fused param exists in the model,
+ # so load_weights must refuse to silently leave it uninitialized.
+ bad_weights: list[tuple[str, torch.Tensor]] = [
+ (f"{attn_prefix}.q_proj.weight", torch.randn(num_heads * head_dim, hidden)),
+ (f"{attn_prefix}.k_proj.weight", torch.randn(num_kv_heads * head_dim, hidden)),
+ # v_proj missing on purpose.
+ ]
+ with pytest.raises(RuntimeError, match="incomplete fused shards"):
+ model.load_weights(bad_weights)
+
+ def test_load_weights_raises_when_fused_shards_absent(self, loaded_target_classes) -> None:
+ """Entirely absent q/k/v or gate/up groups must not leave fused params
+ randomly initialized."""
+ _, _, _, code_predictor_model, _ = loaded_target_classes
+ cp_config, _ = _make_tiny_config(loaded_target_classes)
+ model = code_predictor_model(cp_config, embedding_dim=cp_config.hidden_size)
+
+ weights = [
+ (name, torch.randn_like(param))
+ for name, param in model.named_parameters(remove_duplicate=False)
+ if not (name.endswith(".qkv_proj.weight") or name.endswith(".gate_up_proj.weight"))
+ ]
+
+ with pytest.raises(RuntimeError, match="missing fused parameters"):
+ model.load_weights(weights)
diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py
new file mode 100644
index 000000000..3552c5356
--- /dev/null
+++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py
@@ -0,0 +1,322 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048).
+
+``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary
+for the first ``silence_ban_frames`` decode frames of a request, and only for
+requests running Base voice-clone in x-vector-only mode. ICL onsets are
+legitimate audio and must not be masked.
+
+Two per-request gates guard the mask, and both index by batch position:
+
+* the decode step, from ``sampling_metadata.output_token_ids[i]`` -- the
+ request's generated-token history, so the window closes on its own and a
+ preempted request cannot re-enter it;
+* the generation mode, recorded at prefill in ``preprocess_batch`` and looked
+ up through the batch req_id order the runner supplies via
+ ``set_batch_req_ids``.
+
+These tests drive the methods directly against a talker built with ``__new__``:
+no config, no checkpoint, no device. Only the attributes the masking path reads
+are populated, and ``logits_processor`` is stubbed to return a fixed tensor.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import resolve_x_vector_only
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+ Qwen3TTSTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_VOCAB = 32
+_SILENCE_IDS = (3, 7, 11)
+
+
+def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None):
+ """Talker stub exposing only what ``compute_logits`` touches.
+
+ ``modes`` is one entry per batch row: ``True`` for x-vector-only, ``False``
+ for ICL, ``None`` for a request whose mode was never recorded. Defaults to
+ every row in scope, so a test that says nothing about mode exercises the
+ step gate alone.
+ """
+ talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+
+ silence_mask = torch.zeros((_VOCAB,), dtype=torch.bool)
+ for token in _SILENCE_IDS:
+ silence_mask[token] = True
+
+ if modes is None:
+ modes = [True] * logits.shape[0]
+ req_ids = [f"req{i}" for i in range(len(modes))]
+
+ object.__setattr__(talker, "_silence_ban_frames", ban_frames)
+ object.__setattr__(talker, "_silence_mask", silence_mask)
+ object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool))
+ object.__setattr__(talker, "_batch_req_ids", req_ids)
+ object.__setattr__(
+ talker,
+ "_req_x_vector_only",
+ {rid: mode for rid, mode in zip(req_ids, modes, strict=True) if mode is not None},
+ )
+ object.__setattr__(talker, "lm_head", None)
+ object.__setattr__(talker, "logits_processor", lambda _head, _hidden: logits.clone())
+ return talker
+
+
+def _metadata(*steps: int):
+ """sampling_metadata whose output_token_ids encode one step per request."""
+ return SimpleNamespace(output_token_ids=[[0] * step for step in steps])
+
+
+def _banned(row: torch.Tensor) -> set[int]:
+ return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()}
+
+
+# -------------------- the step gate --------------------
+
+
+def test_ban_disabled_leaves_logits_unchanged() -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=0, logits=logits)
+ out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0))
+ assert _banned(out[0]) == set(), "ban_frames=0 must be a no-op"
+
+
+@pytest.mark.parametrize("step", [0, 1, 2])
+def test_early_steps_mask_only_the_silence_vocabulary(step: int) -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step))
+ assert _banned(out[0]) == set(_SILENCE_IDS), f"step {step} must mask exactly the silence ids"
+
+
+@pytest.mark.parametrize("step", [3, 4, 50])
+def test_steps_past_the_window_are_untouched(step: int) -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((1, 4)), _metadata(step))
+ assert _banned(out[0]) == set(), f"step {step} is past the window; nothing may be masked"
+
+
+def test_mixed_step_batch_masks_only_the_early_request() -> None:
+ logits = torch.zeros((2, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 3))
+ assert _banned(out[0]) == set(_SILENCE_IDS), "the step-0 request must be masked"
+ assert _banned(out[1]) == set(), "the step-3 request must not be masked"
+
+
+# -------------------- the mode gate --------------------
+
+
+def test_icl_request_is_never_masked() -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits, modes=[False])
+ out = talker.compute_logits(torch.zeros((1, 4)), _metadata(0))
+ assert _banned(out[0]) == set(), "an ICL onset is real audio and must not be suppressed"
+
+
+def test_mixed_mode_batch_masks_only_the_x_vector_request() -> None:
+ logits = torch.zeros((2, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits, modes=[True, False])
+ out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0))
+ assert _banned(out[0]) == set(_SILENCE_IDS), "the x-vector-only request must be masked"
+ assert _banned(out[1]) == set(), "the ICL request in the same batch must not be"
+
+
+def test_unrecorded_mode_is_left_alone() -> None:
+ logits = torch.zeros((2, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits, modes=[True, None])
+ out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0))
+ assert _banned(out[0]) == set(_SILENCE_IDS)
+ assert _banned(out[1]) == set(), "a request that cannot be classified must not be masked"
+
+
+def test_stale_batch_req_ids_skip_the_ban() -> None:
+ logits = torch.zeros((2, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ # One req_id for a two-row batch: the recorded order no longer indexes the
+ # logits, so masking by position would hit the wrong request.
+ object.__setattr__(talker, "_batch_req_ids", ["req0"])
+ out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0, 0))
+ assert _banned(out[0]) == set(), "a req_id/logits length mismatch must disable the ban"
+ assert _banned(out[1]) == set()
+
+
+# -------------------- metadata plumbing --------------------
+
+
+def test_missing_sampling_metadata_does_not_crash() -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((1, 4)), None)
+ assert _banned(out[0]) == set()
+
+
+@pytest.mark.parametrize(
+ "metadata",
+ [SimpleNamespace(output_token_ids=None), SimpleNamespace(), SimpleNamespace(output_token_ids=[])],
+)
+def test_absent_output_token_ids_warns_without_masking(metadata) -> None:
+ logits = torch.zeros((1, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((1, 4)), metadata)
+ assert _banned(out[0]) == set(), "the ban must skip, not crash, when history is unavailable"
+
+
+def test_batch_length_mismatch_skips_the_ban() -> None:
+ logits = torch.zeros((2, _VOCAB))
+ talker = _make_talker(ban_frames=3, logits=logits)
+ out = talker.compute_logits(torch.zeros((2, 4)), _metadata(0))
+ assert _banned(out[0]) == set(), "a runner shape mismatch must disable the ban, not misindex it"
+ assert _banned(out[1]) == set()
+
+
+# -------------------- mode recording and eviction --------------------
+
+
+def _mode_talker(ban_frames: int = 3):
+ talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+ object.__setattr__(talker, "_silence_ban_frames", ban_frames)
+ object.__setattr__(talker, "_req_x_vector_only", {})
+ object.__setattr__(talker, "_batch_req_ids", [])
+ object.__setattr__(talker, "_prompt_builder", SimpleNamespace(preprocess_batch=lambda **_kwargs: None))
+ return talker
+
+
+def _info(*, task_type: str = "Base", xvec: bool = True, icl: bool | None = None) -> dict:
+ info: dict = {"task_type": [task_type], "x_vector_only_mode": [xvec]}
+ if icl is not None:
+ info["voice_clone_prompt"] = [{"icl_mode": icl}]
+ return info
+
+
+def test_preprocess_batch_records_the_generation_mode() -> None:
+ talker = _mode_talker()
+ buffer = {
+ "req0": _info(xvec=True),
+ "req1": _info(xvec=False),
+ "req2": _info(task_type="CustomVoice"),
+ }
+ talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu"))
+ assert talker._req_x_vector_only == {"req0": True, "req1": False, "req2": None}
+
+
+def test_recorded_mode_matches_the_prompt_builder_resolution() -> None:
+ """The recording path must agree with the resolution the prefill path uses;
+ a divergence here would mask requests the builder treats as in-context."""
+ infos = [
+ _info(xvec=True),
+ _info(xvec=False),
+ _info(xvec=True, icl=True),
+ _info(xvec=False, icl=False),
+ _info(task_type="VoiceDesign"),
+ ]
+ talker = _mode_talker()
+ buffer = {f"req{i}": info for i, info in enumerate(infos)}
+ talker.preprocess_batch(req_ids=list(buffer), model_intermediate_buffer=buffer, device=torch.device("cpu"))
+ for req_id, info in buffer.items():
+ assert talker._req_x_vector_only[req_id] == resolve_x_vector_only(info)
+
+
+def test_preprocess_batch_records_nothing_when_the_ban_is_off() -> None:
+ talker = _mode_talker(ban_frames=0)
+ buffer = {"req0": _info(xvec=True)}
+ talker.preprocess_batch(req_ids=["req0"], model_intermediate_buffer=buffer, device=torch.device("cpu"))
+ assert talker._req_x_vector_only == {}
+
+
+def test_set_batch_req_ids_evicts_departed_requests() -> None:
+ talker = _mode_talker()
+ object.__setattr__(talker, "_req_x_vector_only", {"req0": True, "req1": False, "req2": True})
+ talker.set_batch_req_ids(["req2", "req0"])
+ assert talker._batch_req_ids == ["req2", "req0"], "batch order must be preserved, not sorted"
+ assert talker._req_x_vector_only == {"req0": True, "req2": True}, "req1 has left the batch"
+
+
+# -------------------- _init_silence_mask sanity checks --------------------
+#
+# The silence vocabulary is derived from the checkpoint at load time rather than
+# hardcoded. A derivation that comes back empty, oversized, or out of range means
+# the encode did not produce silence codes, so the ban disables itself instead of
+# masking a guess.
+
+_TALKER_VOCAB = 2048
+
+
+def _make_init_talker(*, codes, ban_frames: int = 3, vocab: int = _TALKER_VOCAB):
+ """Talker stub exposing only what ``_init_silence_mask`` touches.
+
+ ``codes`` is either the list ``_encode_ref_audio_batch`` should return or an
+ exception instance it should raise.
+ """
+ talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+
+ def _encode(_wavs, _sr, *, device):
+ if isinstance(codes, Exception):
+ raise codes
+ return codes
+
+ def _parameters():
+ return iter([torch.zeros(1)])
+
+ object.__setattr__(talker, "_silence_ban_frames", ban_frames)
+ object.__setattr__(talker, "_silence_mask", torch.zeros((vocab,), dtype=torch.bool))
+ object.__setattr__(talker, "talker_config", SimpleNamespace(vocab_size=vocab))
+ object.__setattr__(talker, "_encode_ref_audio_batch", _encode)
+ object.__setattr__(talker, "parameters", _parameters)
+ return talker
+
+
+def _codes_from(ids) -> list[torch.Tensor]:
+ """One [T, Q] codec tensor whose codebook-0 column holds ``ids``."""
+ cb0 = torch.tensor(list(ids), dtype=torch.long).unsqueeze(1)
+ return [torch.cat([cb0, torch.zeros_like(cb0)], dim=1)]
+
+
+def test_plausible_derivation_populates_the_mask() -> None:
+ talker = _make_init_talker(codes=_codes_from([4, 4, 9, 17]))
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 3
+ assert set(torch.nonzero(talker._silence_mask).flatten().tolist()) == {4, 9, 17}
+
+
+def test_empty_derivation_disables_the_ban() -> None:
+ talker = _make_init_talker(codes=[])
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 0, "an all-False mask must not stay enabled"
+ assert not talker._silence_mask.any()
+
+
+def test_oversized_derivation_disables_the_ban() -> None:
+ limit = max(64, _TALKER_VOCAB // 20)
+ talker = _make_init_talker(codes=_codes_from(range(limit + 1)))
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 0, "a speech-sized vocabulary must not be masked"
+ assert not talker._silence_mask.any()
+
+
+def test_out_of_range_derivation_disables_the_ban() -> None:
+ talker = _make_init_talker(codes=_codes_from([4, 9, _TALKER_VOCAB + 1]))
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 0, "a vocabulary-size disagreement must not be half-applied"
+ assert not talker._silence_mask.any()
+
+
+def test_encoder_failure_disables_the_ban() -> None:
+ talker = _make_init_talker(codes=RuntimeError("encoder unavailable"))
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 0
+ assert not talker._silence_mask.any()
+
+
+def test_disabled_ban_skips_derivation_entirely() -> None:
+ talker = _make_init_talker(codes=RuntimeError("must not be called"), ban_frames=0)
+ talker._init_silence_mask()
+ assert talker._silence_ban_frames == 0
diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py
index 050a5b724..53a90452e 100644
--- a/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py
+++ b/tests/model_executor/models/step_audio2/test_step_audio2_token2wav_async_chunk.py
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
+from unittest.mock import MagicMock
import pytest
import torch
@@ -76,8 +77,8 @@ def test_step_audio2_token2wav_async_chunk_batch_guard():
model._forward_async_chunk(
input_ids=torch.tensor([1, 2, 3]),
runtime_additional_information=[
- {"left_context_size": 0},
- {"left_context_size": 1},
+ {"meta": {"left_context_size": 0}},
+ {"meta": {"left_context_size": 1}},
],
)
@@ -116,7 +117,7 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch):
out = model.forward(
input_ids=torch.tensor([10, 11, 12]),
positions=torch.tensor([0, 1, 2]),
- runtime_additional_information=[{"left_context_size": 1}],
+ runtime_additional_information=[{"meta": {"left_context_size": 1}}],
)
audio_list = out.multimodal_outputs["model_outputs"]
@@ -126,12 +127,24 @@ def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch):
assert calls == {"setup": 1, "stream": 1, "reset": 1}
+def test_step_audio2_token2wav_routes_non_final_chunk_to_async_forward():
+ model = _build_model()
+ model._forward_async_chunk = MagicMock()
+
+ model.forward(
+ input_ids=torch.tensor([10, 11, 12]),
+ positions=torch.tensor([0, 1, 2]),
+ runtime_additional_information=[{"meta": {"left_context_size": 0}}],
+ )
+ model._forward_async_chunk.assert_called_once()
+
+
def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk():
model = _build_model()
out = model._forward_async_chunk(
input_ids=torch.tensor([], dtype=torch.int64),
- runtime_additional_information=[{"left_context_size": 1}],
+ runtime_additional_information=[{"meta": {"left_context_size": 1}}],
)
audio_list = out.multimodal_outputs["model_outputs"]
diff --git a/tests/model_executor/models/test_encoder_quant_config.py b/tests/model_executor/models/test_encoder_quant_config.py
index 28a1efa71..694e92a02 100644
--- a/tests/model_executor/models/test_encoder_quant_config.py
+++ b/tests/model_executor/models/test_encoder_quant_config.py
@@ -1,12 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Regression test for #2686: pre-quantized methods must not apply
quant config to vision / audio encoders.
-For modelopt FP8/FP4/MXFP8 checkpoints the Thinker LM is the only
-quantized component. Vision and audio encoder weights are BF16 with no
-FP8 scale tensors — passing quant_config to them causes FP8 kernels to
-run on BF16 weights, producing garbage embeddings.
+For serialized ModelOpt and SVDQuant checkpoints the target model component is
+the only quantized component. Vision and audio encoder weights are BF16 with no
+serialized scale or correction tensors, so the same config must not be applied
+to them.
**Import note (pytest + ``tests/model_executor`` only):** :mod:`vllm_omni.quantization`
:file:`__init__` pulls ``factory``/``inc`` and reproduces vLLM CustomOp double
@@ -58,8 +58,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize("method", sorted(PRE_QUANTIZED_METHODS))
def test_pre_quantized_returns_none(method: str) -> None:
"""visual_quant_config and audio_quant_config must be None for
- pre-quantized methods (modelopt, modelopt_fp4, modelopt_mxfp8,
- modelopt_mixed)."""
+ pre-quantized methods (ModelOpt variants and SVDQuant)."""
mock_config = MagicMock()
mock_config.get_name.return_value = method
@@ -101,5 +100,5 @@ def test_component_config_passed_through() -> None:
def test_pre_quantized_methods_contains_expected() -> None:
"""Guard against accidental removal of a known pre-quantized method."""
- expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed"}
+ expected = {"modelopt", "modelopt_fp4", "modelopt_mxfp8", "modelopt_mixed", "svdquant"}
assert PRE_QUANTIZED_METHODS == expected
diff --git a/tests/model_executor/models/test_gepard_window.py b/tests/model_executor/models/test_gepard_window.py
new file mode 100644
index 000000000..67c61c7ab
--- /dev/null
+++ b/tests/model_executor/models/test_gepard_window.py
@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU coverage for Gepard's prompt layout, streaming-decode window and config.
+
+The e2e test needs a GPU, NeMo and the checkpoint, so it only runs nightly.
+These pin the parts that are pure logic and break silently: which frames each
+decode covers, which samples survive the lookback trim, the prompt layout, and
+that the packaged deploy config sits where the example looks for it.
+
+The codec is replaced by a shift-invariant stand-in whose output depends only
+on the codes, which makes "chunked emissions concatenate to one whole-history
+decode" an exact assertion — the contract lookback exists to provide.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+import vllm_omni
+from vllm_omni.config.stage_config import load_deploy_config
+from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+from vllm_omni.model_executor.models.gepard.gepard_talker import (
+ CHUNK_FRAMES,
+ FIRST_CHUNK_FRAMES,
+ GepardTalkerForConditionalGeneration,
+ _GepardState,
+)
+from vllm_omni.model_executor.models.gepard.nanocodec import (
+ LOOKBACK_FRAMES,
+ NanoCodec,
+ apply_end_of_speech_tail,
+)
+from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SAMPLES_PER_FRAME = 1024
+NUM_HEADS = 32
+SAMPLE_RATE = 22050
+# Text-token count at or above which the layout is left unrepeated, read from
+# the checkpoint defaults so the boundary cases follow the config.
+APPLY_BELOW = GepardConfig().text_repetition_apply_below
+
+
+class _ShiftInvariantCodec:
+ """Frame tagged ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits."""
+
+ def __init__(self) -> None:
+ self.windows: list[int] = [] # frame count of every decode call
+
+ def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
+ assert codes.shape[0] == 1, "streaming decode is per-request"
+ assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
+ tags = codes[0, 0, :].to(torch.int64) # (T,)
+ self.windows.append(int(tags.shape[0]))
+ offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+ audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
+ return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME
+
+
+def _make_codec() -> tuple[NanoCodec, _ShiftInvariantCodec]:
+ codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
+ fake = _ShiftInvariantCodec()
+ codec._codec = fake # bypasses load(); NeMo is never imported
+ assert codec.is_loaded
+ return codec, fake
+
+
+def _make_talker(codec: NanoCodec) -> GepardTalkerForConditionalGeneration:
+ """A real talker with only the audio-path attributes populated.
+
+ ``__init__`` needs a VllmConfig and builds the backbone; the methods under
+ test touch only these four attributes.
+ """
+ talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
+ nn.Module.__init__(talker) # sets up _parameters/_buffers/_modules only
+ talker._codec = codec
+ talker._audio_queue = []
+ talker._active_states = {}
+ talker._deferred_cleanup_ids = set()
+ return talker
+
+
+def _frame(tag: int) -> torch.Tensor:
+ """One committed frame: 32 per-dimension codes, all carrying the tag."""
+ return torch.full((NUM_HEADS,), tag, dtype=torch.long)
+
+
+def _expected_samples(num_frames: int) -> torch.Tensor:
+ offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+ tags = torch.arange(num_frames, dtype=torch.int64).unsqueeze(1)
+ return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)
+
+
+def _drain(talker) -> torch.Tensor:
+ """Concatenate and clear the queued deltas, as make_omni_output does."""
+ if not talker._audio_queue:
+ return torch.empty(0)
+ out = torch.cat([a.reshape(-1) for _, a in talker._audio_queue])
+ talker._audio_queue.clear()
+ return out
+
+
+def _generate(talker, state, num_frames: int, *, stop_at_end: bool) -> None:
+ """Drive the commit loop the way forward() does, one frame per step."""
+ for tag in range(num_frames):
+ state.frames.append(_frame(tag))
+ state.frame_count += 1
+ state.past_first_step = True
+ talker._emit_audio(state, is_final=stop_at_end and tag == num_frames - 1)
+
+
+# Counts straddling both cadence thresholds, where an off-by-one hides.
+@pytest.mark.parametrize(
+ "num_frames",
+ [
+ 1,
+ 2,
+ FIRST_CHUNK_FRAMES - 1,
+ FIRST_CHUNK_FRAMES,
+ FIRST_CHUNK_FRAMES + 1,
+ CHUNK_FRAMES - 1,
+ CHUNK_FRAMES,
+ CHUNK_FRAMES + 1,
+ FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
+ 79,
+ 200,
+ ],
+)
+def test_chunked_emissions_equal_a_whole_history_decode(num_frames: int) -> None:
+ codec, _fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+
+ _generate(talker, state, num_frames, stop_at_end=True)
+
+ torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
+ assert state.emitted_frames == num_frames
+
+
+def test_every_frame_is_emitted_exactly_once() -> None:
+ """The trim must not drop or duplicate samples at a chunk boundary."""
+ codec, _fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+ num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + 3
+
+ _generate(talker, state, num_frames, stop_at_end=True)
+ audio = _drain(talker)
+
+ assert audio.numel() == num_frames * SAMPLES_PER_FRAME
+ # Tags are globally unique, so a duplicated or dropped frame shows up here.
+ assert len(torch.unique(audio)) == audio.numel()
+
+
+def test_decode_windows_carry_lookback_context() -> None:
+ """Every decode after the first must re-decode LOOKBACK_FRAMES of context."""
+ codec, fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+ num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES
+
+ _generate(talker, state, num_frames, stop_at_end=True)
+
+ assert len(fake.windows) >= 3, "expected a first chunk plus steady-state chunks"
+ assert fake.windows[0] == FIRST_CHUNK_FRAMES, "the first chunk has no left context"
+ for window in fake.windows[1:]:
+ assert window == CHUNK_FRAMES + LOOKBACK_FRAMES
+
+
+def test_codec_runs_once_per_chunk_not_once_per_frame() -> None:
+ codec, fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+ num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES
+
+ _generate(talker, state, num_frames, stop_at_end=True)
+
+ assert len(fake.windows) == 4
+ assert sum(fake.windows) < num_frames * 2 # far below one decode per frame
+
+
+def test_final_emit_drains_a_partial_chunk() -> None:
+ """A request can end mid-chunk (max_tokens, STOP); the frames committed
+ since the last cadence boundary must still reach the queue.
+
+ Arithmetic only: this drives ``_emit_audio`` directly, so it says nothing
+ about *whether* the engine reaches it on the request's last step. That
+ wiring lives in ``test_gepard_wiring.py`` — a bug there leaves this green.
+ """
+ codec, _fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+ num_frames = FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5
+
+ _generate(talker, state, num_frames, stop_at_end=False)
+ assert state.emitted_frames < num_frames, "a partial chunk should still be pending"
+
+ talker._emit_audio(state, is_final=True)
+
+ torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
+ assert state.emitted_frames == num_frames
+
+
+def test_emitting_with_nothing_pending_is_a_no_op() -> None:
+ codec, fake = _make_codec()
+ talker = _make_talker(codec)
+ state = _GepardState(request_id="req-0")
+
+ _generate(talker, state, 4, stop_at_end=True)
+ calls_after_flush = len(fake.windows)
+ talker._audio_queue.clear()
+
+ talker._emit_audio(state, is_final=True)
+
+ assert len(fake.windows) == calls_after_flush, "codec ran with no new frames"
+ assert talker._audio_queue == []
+
+
+def test_concurrent_requests_keep_separate_windows() -> None:
+ codec, _fake = _make_codec()
+ talker = _make_talker(codec)
+ states = [_GepardState(request_id=f"req-{i}") for i in range(2)]
+
+ for state in states:
+ _generate(talker, state, FIRST_CHUNK_FRAMES, stop_at_end=True)
+
+ by_req: dict[str, list[torch.Tensor]] = {}
+ for req_id, audio in talker._audio_queue:
+ by_req.setdefault(req_id, []).append(audio.reshape(-1))
+ assert set(by_req) == {"req-0", "req-1"}
+ for chunks in by_req.values():
+ torch.testing.assert_close(torch.cat(chunks), _expected_samples(FIRST_CHUNK_FRAMES))
+
+
+def test_end_of_speech_tail_fades_and_pads() -> None:
+ audio = torch.ones(SAMPLE_RATE) # 1 s of full-scale signal
+ out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=10.0, silence_ms=20.0)
+
+ fade = int(SAMPLE_RATE * 10.0 / 1000.0)
+ pad = int(SAMPLE_RATE * 20.0 / 1000.0)
+ assert out.numel() == SAMPLE_RATE + pad
+ assert torch.all(out[-pad:] == 0.0), "trailing pad must be silent"
+ assert out[SAMPLE_RATE - fade - 1] == 1.0, "the fade must not reach back past fade_ms"
+ # Monotonically ramped to zero — the pad alone would leave the click.
+ ramp = out[SAMPLE_RATE - fade : SAMPLE_RATE]
+ assert torch.all(ramp[:-1] >= ramp[1:])
+ assert ramp[-1] == 0.0
+
+
+def test_end_of_speech_tail_is_a_no_op_by_default() -> None:
+ audio = torch.ones(128)
+ out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=0.0, silence_ms=0.0)
+ assert out is audio
+
+
+def _repeat_count(ids: list[int], cfg: GepardConfig) -> int:
+ """How many ``[SOT ... EOT]`` copies the layout carries."""
+ return ids.count(cfg.start_of_text)
+
+
+def test_gepard_prompt_layout() -> None:
+ """The prompt carries the speaker slots and exactly one SOS, at the end."""
+ cfg = GepardConfig()
+ ids = build_gepard_prompt_ids([1, 2, 3], config=cfg)
+
+ slots = [cfg.speaker_token_base + i for i in range(cfg.num_speaker_prefix)]
+ assert ids[: cfg.num_speaker_prefix] == slots
+ assert ids.count(cfg.start_of_speech) == 1, "only the canonical copy may carry SOS"
+ assert ids[-1] == cfg.start_of_speech, "the first frame is sampled from the SOS position"
+
+
+def test_short_text_repeats_to_reach_the_target_token_mass() -> None:
+ """A short text is repeated until its region holds ~target_text_tokens."""
+ cfg = GepardConfig()
+ text = [1, 2, 3, 4] # target 16 / 4 -> 4 copies
+ ids = build_gepard_prompt_ids(text, config=cfg)
+
+ assert _repeat_count(ids, cfg) == 4
+ assert ids.count(cfg.start_of_speech) == 1, "only the last copy opens the speech region"
+ assert ids[-1] == cfg.start_of_speech
+
+
+@pytest.mark.parametrize(
+ ("num_text_tokens", "repeated"),
+ [
+ (APPLY_BELOW - 1, True),
+ (APPLY_BELOW, False),
+ (APPLY_BELOW + 1, False),
+ ],
+)
+def test_apply_below_is_the_repetition_boundary(num_text_tokens: int, repeated: bool) -> None:
+ """Both sides of the threshold: at apply_below the text is left alone.
+
+ The layout has to match training here or WER collapses, and the boundary
+ is where an off-by-one would be invisible in a single happy-path clip.
+ """
+ cfg = GepardConfig()
+ ids = build_gepard_prompt_ids(list(range(1, num_text_tokens + 1)), config=cfg)
+
+ assert (_repeat_count(ids, cfg) > 1) is repeated
+
+
+def test_max_repeats_caps_a_one_token_text() -> None:
+ """target/1 would ask for 16 copies; the cap keeps the prefill bounded."""
+ cfg = GepardConfig()
+ ids = build_gepard_prompt_ids([1], config=cfg)
+
+ assert _repeat_count(ids, cfg) == cfg.text_repetition_max_repeats
+
+
+def test_repetition_thresholds_are_read_from_the_config() -> None:
+ """The numbers come from the checkpoint sidecar, not from literals.
+
+ A checkpoint whose text_repetition block differs must change the layout;
+ duplicating the values in the builder would silently ignore it.
+ """
+ cfg = GepardConfig(text_repetition={"enabled": True, "target_text_tokens": 9, "apply_below": 8, "max_repeats": 2})
+ assert _repeat_count(build_gepard_prompt_ids([1, 2, 3], config=cfg), cfg) == 2 # ceil(9/3)=3, capped at 2
+ assert _repeat_count(build_gepard_prompt_ids(list(range(8)), config=cfg), cfg) == 1 # at apply_below
+
+
+def test_repetition_can_be_disabled() -> None:
+ cfg = GepardConfig(text_repetition={"enabled": False})
+ ids = build_gepard_prompt_ids([1], config=cfg)
+
+ assert _repeat_count(ids, cfg) == 1
+ assert ids[-1] == cfg.start_of_speech
+
+
+def test_empty_text_is_rejected() -> None:
+ """An empty text would still assemble a valid layout and be voiced.
+
+ ``[slots, SOT, EOT, SOS]`` parses, prefills and generates: the failure is
+ audible, not raised, which is the worst shape for a caller to debug.
+ """
+ with pytest.raises(ValueError, match="at least one text token"):
+ build_gepard_prompt_ids([], config=GepardConfig())
+
+
+def test_packaged_deploy_config_is_where_the_example_looks_for_it() -> None:
+ """The example resolves its default off ``vllm_omni.__file__``; tests reach
+ the same YAML through ``tests.helpers.stage_config``, so nothing covered the
+ packaged location until this.
+
+ ``pipeline: gepard`` is the load-bearing key: the checkpoint self-identifies
+ as qwen3_5_text, so without the pin the architectures fallback routes it to
+ the diffusion registry.
+
+ This asserts the location, not the example's constant — importing the
+ example would pull ``Omni`` and NeMo into a CPU test. Running the quick
+ start verbatim stays a release check.
+ """
+ path = Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml"
+ assert path.is_file(), f"packaged deploy config missing at {path}"
+
+ deploy = load_deploy_config(path)
+ assert deploy.pipeline == "gepard"
+ assert deploy.stages, "deploy config declares no stages"
diff --git a/tests/model_executor/models/test_gepard_wiring.py b/tests/model_executor/models/test_gepard_wiring.py
new file mode 100644
index 000000000..9a08ddf0a
--- /dev/null
+++ b/tests/model_executor/models/test_gepard_wiring.py
@@ -0,0 +1,732 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU coverage for the seam between Gepard's talker and the AR runner.
+
+``test_gepard_window.py`` pins the decode arithmetic *below* this layer, and
+the e2e test exercises the layer *above* it but only along the path where the
+stop head fires. Neither covers the wiring itself: whether anything calls the
+final emit on a request's last step, and whether the payload it produces
+survives the runner's sparse-audio routing.
+
+These drive the real ``preprocess`` -> ``forward`` -> ``make_omni_output``
+sequence through a stand-in for ``GPUARModelRunner.execute_model``, and use the
+shared ``sparse_audio.resolve_sparse_mm_routing`` helper used by the runner to
+decide what actually reaches the caller. Only the LM backbone and the frame
+sampler are stubbed out; every method under test is the real one, so a test
+here goes red when the engine
+path breaks — which a test that calls the flush by hand cannot do.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+from vllm_omni.model_executor.models.gepard.gepard_talker import (
+ CHUNK_FRAMES,
+ FIRST_CHUNK_FRAMES,
+ GepardTalkerForConditionalGeneration,
+ _ReqInfo,
+)
+from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec
+from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+from vllm_omni.worker import sparse_audio
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SAMPLES_PER_FRAME = 1024
+NUM_HEADS = 32
+SAMPLE_RATE = 22050
+LOGGER_NAME = "vllm_omni.model_executor.models.gepard.gepard_talker"
+
+# Each request tags its frames from a distinct base, so a sample that reaches
+# the wrong request is visible in the value, not just the length.
+_TAG_STRIDE = 100_000
+
+
+class _TagCodec:
+ """Frame carrying tag ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits.
+
+ Shift-invariant, so "what the caller received" can be compared against one
+ whole-history decode exactly. The tag rides channel 1: channel 0 is head0,
+ which the talker keeps inside its real (small) vocab.
+ """
+
+ def __init__(self) -> None:
+ self.decode_calls = 0
+
+ def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
+ assert codes.shape[0] == 1, "streaming decode is per-request"
+ assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
+ self.decode_calls += 1
+ tags = codes[0, 1, :].to(torch.int64) # (T,)
+ offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+ audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
+ return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME
+
+
+class _StubBackbone(nn.Module):
+ """Stands in for the Qwen3.5 backbone; only shapes matter to this seam."""
+
+ def __init__(self, hidden: int) -> None:
+ super().__init__()
+ self.hidden = hidden
+
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return torch.zeros(int(input_ids.shape[0]), self.hidden)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ positions: torch.Tensor | None = None,
+ intermediate_tensors=None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ src = inputs_embeds if inputs_embeds is not None else input_ids
+ return torch.zeros(int(src.shape[0]), self.hidden)
+
+
+def _expected_samples(tag_base: int, num_frames: int) -> torch.Tensor:
+ offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+ tags = (torch.arange(num_frames, dtype=torch.int64) + tag_base).unsqueeze(1)
+ return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)
+
+
+class _MiniEngine:
+ """Drives the talker the way ``GPUARModelRunner.execute_model`` does.
+
+ Reproduces the three orderings the dropped-tail bug lived in:
+
+ * ``on_requests_finished`` fires at the START of the step AFTER the one
+ that finished the request (it reads ``scheduler_output.finished_req_ids``);
+ * a finished request is not scheduled again, so it is absent from
+ ``req_ids_output_copy`` and from that step's routing;
+ * with nothing left to schedule there is no step at all, so no further
+ ``forward`` runs;
+ * a stop-token finish gets one more scheduled step, a budget finish does
+ not -- under async scheduling only the latter is guarded against it.
+ """
+
+ def __init__(self, talker: GepardTalkerForConditionalGeneration, *, max_model_len: int) -> None:
+ self.talker = talker
+ self.max_model_len = max_model_len
+ self.reqs: dict[str, dict] = {}
+ self.order: list[str] = []
+ self._finished_last_step: list[str] = []
+ self.delivered: dict[str, list[torch.Tensor]] = {}
+ self.committed_frames: dict[str, int] = {}
+ self.steps = 0
+ self._tags_handed_out = 0
+
+ def add_request(
+ self,
+ req_id: str,
+ *,
+ max_tokens: int,
+ stop_after: int | None = None,
+ text_len: int = 3,
+ seed: int | None = None,
+ ) -> None:
+ """Submit a request. An id already seen is a resubmit of that id.
+
+ Tags come from a counter, so a resubmitted id gets frames the previous
+ request could not have produced.
+ """
+ prompt_ids = build_gepard_prompt_ids(list(range(10, 10 + text_len)), config=self.talker.config)
+ self._tags_handed_out += 1
+ self.reqs[req_id] = {
+ "prompt_ids": torch.tensor(prompt_ids, dtype=torch.long),
+ "prompt_len": len(prompt_ids),
+ "max_tokens": max_tokens,
+ "stop_after": stop_after,
+ "stop_fired": False,
+ "seed": seed,
+ "computed": 0,
+ "generated": 0,
+ "finished": False,
+ "ghost_step_pending": False,
+ "next_token": None,
+ "tag_base": _TAG_STRIDE * self._tags_handed_out,
+ }
+ if req_id not in self.order:
+ self.order.append(req_id)
+ self.delivered[req_id] = []
+ self.committed_frames[req_id] = 0
+
+ def abort(self, req_id: str) -> None:
+ """Kill a request the way an aborted client does: no STOP, no budget."""
+ self.reqs[req_id]["finished"] = True
+ self._finished_last_step.append(req_id)
+
+ def _live(self) -> list[str]:
+ return [r for r in self.order if not self.reqs[r]["finished"]]
+
+ def step(self) -> bool:
+ """Run one engine step. False when there was nothing to schedule."""
+ live = self._live()
+ if not live:
+ return False
+ talker = self.talker
+ self.steps += 1
+
+ # --- runner: notify the model of requests finished since last step ---
+ if self._finished_last_step:
+ talker.on_requests_finished(set(self._finished_last_step))
+ self._finished_last_step = []
+
+ # --- runner: per-request preprocess, in input_batch.req_ids order ---
+ embeds_parts, ids_parts = [], []
+ for req_id in live:
+ r = self.reqs[req_id]
+ is_prefill = r["computed"] < r["prompt_len"]
+ if is_prefill:
+ ids = r["prompt_ids"]
+ else:
+ ids = torch.tensor([r["next_token"] or 0], dtype=torch.long)
+ span = int(ids.shape[0])
+ out_ids, out_embeds, _ = talker.preprocess(
+ input_ids=ids,
+ input_embeds=None,
+ request_id=req_id,
+ _omni_prompt_len=r["prompt_len"],
+ _omni_num_computed_tokens=r["computed"],
+ _omni_is_prefill=is_prefill,
+ _omni_max_tokens=r["max_tokens"],
+ _omni_seed=r["seed"],
+ )
+ ids_parts.append(out_ids)
+ embeds_parts.append(out_embeds)
+ r["span"] = span
+
+ inputs_embeds = torch.cat(embeds_parts, dim=0)
+ input_ids = torch.cat(ids_parts, dim=0)
+ positions = torch.arange(int(input_ids.shape[0]), dtype=torch.long)
+
+ # --- runner: forward, then build the multimodal payload ---
+ hidden = talker.forward(input_ids, positions, inputs_embeds=inputs_embeds)
+ results = list(talker._results_queue)
+ omni_out = talker.make_omni_output(hidden)
+ mm = omni_out.multimodal_outputs
+ # compute_logits drains this; it does not touch the audio path.
+ talker._results_queue.clear()
+
+ for req_id in live:
+ state = talker._active_states.get(req_id)
+ if state is not None:
+ self.committed_frames[req_id] = state.frame_count
+
+ # --- runner: sparse-audio routing decides what reaches the caller ---
+ downstream, sparse_index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=list(live),
+ downstream_req_ids=list(live),
+ multimodal_outputs=mm,
+ )
+ assert is_sparse, "gepard always marks meta.sparse_audio"
+ for rid in downstream:
+ self.delivered[rid].append(mm["model_outputs"][sparse_index[rid]].reshape(-1))
+
+ # --- runner/scheduler bookkeeping: advance and detect finishes ---
+ stopped = {rid for rid, _head0, do_stop in results if do_stop}
+ for req_id in live:
+ r = self.reqs[req_id]
+ r["computed"] += r["span"]
+ r["generated"] = r["computed"] - r["prompt_len"] + 1
+ r["next_token"] = 0
+ length_cap = min(r["max_tokens"], self.max_model_len - r["prompt_len"])
+ if r["ghost_step_pending"]:
+ # This step WAS the extra one. The engine finished the request
+ # on the stop token last step and discards whatever came back
+ # here, whether or not the model chose to stop again.
+ r["finished"] = True
+ self._finished_last_step.append(req_id)
+ elif req_id in stopped:
+ # Ends on the stop token: one more step is already scheduled.
+ r["ghost_step_pending"] = True
+ elif r["generated"] >= length_cap:
+ r["finished"] = True
+ self._finished_last_step.append(req_id)
+ return True
+
+ def run(self, max_steps: int = 4000) -> None:
+ for _ in range(max_steps):
+ if not self.step():
+ return
+ raise AssertionError("engine did not drain")
+
+ def audio_for(self, req_id: str) -> torch.Tensor:
+ chunks = self.delivered[req_id]
+ return torch.cat(chunks) if chunks else torch.empty(0)
+
+
+def _make_talker(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec]:
+ """A real talker with only the audio-path collaborators populated.
+
+ ``__init__`` needs a VllmConfig and builds the Qwen3.5 backbone; the seam
+ under test needs none of that, so the backbone and the frame sampler are
+ stubbed and everything else is the class's own code.
+ """
+ cfg = GepardConfig()
+ # Fixed, not cfg.get_text_config().hidden_size: nothing in this seam reads
+ # the real width, and a small one keeps the test off the backbone config.
+ hidden = 64
+ codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
+ fake = _TagCodec()
+ codec._codec = fake # bypasses load(); NeMo is never imported
+ assert codec.is_loaded
+
+ talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
+ nn.Module.__init__(talker)
+ talker.config = cfg
+ talker.model = _StubBackbone(hidden)
+ # requires_grad=False: preprocess writes the prefix into inputs_embeds
+ # in place, and production runs that under inference_mode.
+ talker.null_prefix = nn.Parameter(torch.zeros(cfg.num_speaker_prefix, hidden), requires_grad=False)
+ talker._codec = codec
+ talker._sample_rate = SAMPLE_RATE
+ talker._max_model_len = max_model_len
+ talker._active_states = {}
+ talker._pending_requests = []
+ talker._results_queue = []
+ talker._audio_queue = []
+ talker._deferred_cleanup_ids = set()
+ return talker, fake
+
+
+def _install_scripted_sampling(talker: GepardTalkerForConditionalGeneration, engine: _MiniEngine) -> None:
+ """Replace the 32-head sampler with a scripted one.
+
+ Frames carry a per-request, per-step tag so the caller's audio is
+ self-describing. STOP fires once, on the request's configured frame: on the
+ extra step the engine then runs, re-firing is likely but not guaranteed,
+ and the model may not assume it.
+ """
+ head0_vocab = talker.config.head0_vocab_size
+
+ # Signature mirrors the real ``_sample_frame`` minus ``self``; the scripted
+ # sampler ignores ``generators`` because seeding is decided in
+ # ``_build_with_real_sampler``, which keeps the model's own sampler.
+ def _sample_frame(rows: torch.Tensor, generators: list[torch.Generator | None] | None = None):
+ n = int(rows.shape[0])
+ head0 = torch.zeros(n, dtype=torch.long)
+ heads = torch.zeros(n, NUM_HEADS - 1, dtype=torch.long)
+ stop = torch.zeros(n, dtype=torch.bool)
+ for i, (req_id, _span, _samples_frame, *_rest) in enumerate(talker._pending_requests):
+ state = talker._active_states.get(req_id)
+ frame_index = state.frame_count if state is not None else 0
+ tag = engine.reqs[req_id]["tag_base"] + frame_index
+ head0[i] = tag % head0_vocab
+ heads[i] = tag
+ r = engine.reqs[req_id]
+ if r["stop_after"] is not None and frame_index >= r["stop_after"] and not r["stop_fired"]:
+ r["stop_fired"] = True
+ stop[i] = True
+ return head0, heads, stop
+
+ talker._sample_frame = _sample_frame
+ talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden)
+
+
+def _build(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _TagCodec, _MiniEngine]:
+ talker, codec = _make_talker(max_model_len=max_model_len)
+ engine = _MiniEngine(talker, max_model_len=max_model_len)
+ _install_scripted_sampling(talker, engine)
+ return talker, codec, engine
+
+
+def _build_with_real_sampler(*, max_model_len: int = 4096) -> tuple[GepardTalkerForConditionalGeneration, _MiniEngine]:
+ """Same seam, but ``_sample_frame`` is the model's own 32-head sampler.
+
+ The stub backbone hands every row the same hidden state, so the sampled
+ codes — and therefore the audio the codec tags them into — vary with
+ nothing but the Gumbel draws. That makes this the one place where "did the
+ caller's seed reach the audio" is a decidable question.
+ """
+ talker, _codec = _make_talker(max_model_len=max_model_len)
+ engine = _MiniEngine(talker, max_model_len=max_model_len)
+ cfg = talker.config
+ talker.vocab_sizes = list(cfg.audio_head_levels)
+ talker.num_heads = cfg.num_audio_heads
+ talker.head0_vocab = cfg.head0_vocab_size
+ talker.stop_token = cfg.stop_token
+ talker.stop_threshold = cfg.stop_threshold
+ talker.temperature = cfg.temperature
+ talker._greedy = False
+ gather_idx, mask = GepardTalkerForConditionalGeneration._build_gather_mask(talker.vocab_sizes)
+ talker._cb_gather_idx = gather_idx
+ talker._cb_mask = mask
+ # Fixed weights: two engine runs must differ only by their RNG.
+ torch.manual_seed(0)
+ talker.fused_codebook_head = nn.Linear(talker.model.hidden, sum(talker.vocab_sizes))
+ talker.stop_head = nn.Linear(talker.model.hidden, 1)
+ # Never stop: these requests end on their token budget, so the frame count
+ # is fixed and only the sampled codes can differ.
+ torch.nn.init.constant_(talker.stop_head.bias, -20.0)
+ talker._audio_frame_embed = lambda h0, h: torch.zeros(int(h0.shape[0]), talker.model.hidden)
+ return talker, engine
+
+
+# --------------------------------------------------------------------------
+# The tail must reach the caller, whichever way the request ends.
+# --------------------------------------------------------------------------
+
+
+def test_last_in_flight_request_delivers_every_committed_frame() -> None:
+ """The only request runs out of budget: there is no later step to flush in.
+
+ This is the case a deferred flush cannot serve at all — nothing runs after
+ the final forward.
+ """
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5)
+ engine.run()
+
+ frames = engine.committed_frames["req-0"]
+ assert frames > 0
+ torch.testing.assert_close(
+ engine.audio_for("req-0"),
+ _expected_samples(engine.reqs["req-0"]["tag_base"], frames),
+ )
+
+
+@pytest.mark.parametrize(
+ "max_tokens",
+ [
+ 1,
+ FIRST_CHUNK_FRAMES,
+ FIRST_CHUNK_FRAMES + 1,
+ 30,
+ FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
+ 100,
+ 150,
+ ],
+)
+def test_budget_truncated_request_delivers_every_committed_frame(max_tokens: int) -> None:
+ """Lowering max_tokens must shorten the clip, never silently cut it back
+ to the last cadence boundary."""
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=max_tokens)
+ engine.run()
+
+ frames = engine.committed_frames["req-0"]
+ assert frames == max_tokens, "every sampled step commits one frame when STOP never fires"
+ torch.testing.assert_close(
+ engine.audio_for("req-0"),
+ _expected_samples(engine.reqs["req-0"]["tag_base"], frames),
+ )
+
+
+def test_stop_head_path_still_delivers_every_committed_frame() -> None:
+ """Regression guard: the path that already worked must keep working."""
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=4000, stop_after=FIRST_CHUNK_FRAMES + 7)
+ engine.run()
+
+ frames = engine.committed_frames["req-0"]
+ assert frames == FIRST_CHUNK_FRAMES + 7, "the STOP frame itself is not committed"
+ torch.testing.assert_close(
+ engine.audio_for("req-0"),
+ _expected_samples(engine.reqs["req-0"]["tag_base"], frames),
+ )
+
+
+def test_the_step_after_a_stop_token_commits_nothing(caplog: pytest.LogCaptureFixture) -> None:
+ """Async scheduling runs one more step on a request that already stopped.
+
+ Committing its frame would strand it — nothing ships after the stop step —
+ and would make the truncation warning fire on a request that delivered
+ everything it produced, which is the one thing that warning must not do.
+ """
+ talker, _codec, engine = _build()
+ stop_after = FIRST_CHUNK_FRAMES + 6
+ engine.add_request("req-0", max_tokens=4000, stop_after=stop_after)
+
+ target = logging.getLogger(LOGGER_NAME)
+ target.addHandler(caplog.handler)
+ prev = target.level
+ target.setLevel(logging.WARNING)
+ try:
+ engine.run()
+ # The cleanup that would warn runs on the step after the request
+ # leaves, and here there is none. Drive it the way the runner would.
+ talker.on_requests_finished({"req-0"})
+ talker._flush_deferred_cleanup()
+ finally:
+ target.removeHandler(caplog.handler)
+ target.setLevel(prev)
+
+ assert engine.reqs["req-0"]["ghost_step_pending"], "the harness did not run the extra step"
+ frames = engine.committed_frames["req-0"]
+ assert frames == stop_after, f"the extra step committed a frame: {frames} > {stop_after}"
+ torch.testing.assert_close(
+ engine.audio_for("req-0"),
+ _expected_samples(engine.reqs["req-0"]["tag_base"], frames),
+ )
+ assert not [m for m in caplog.messages if "undelivered frame" in m], (
+ f"a fully delivered request was reported as truncated; saw {caplog.messages}"
+ )
+
+
+def test_a_resubmitted_request_id_keeps_its_new_state(caplog: pytest.LogCaptureFixture) -> None:
+ """An aborted id can come back before its cleanup has run.
+
+ The runner treats the two as distinct requests, so the pending free must
+ not take the second one's state: without that, the state is popped at the
+ end of the very forward that created it and every later decode falls back
+ to the zero-embedding path — the request keeps producing frames, so only
+ the audio is wrong.
+ """
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=4000)
+ for _ in range(FIRST_CHUNK_FRAMES + 2):
+ engine.step()
+ engine.abort("req-0")
+
+ # Same id, resubmitted before the model has been told about the abort.
+ engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5)
+
+ target = logging.getLogger(LOGGER_NAME)
+ target.addHandler(caplog.handler)
+ prev = target.level
+ target.setLevel(logging.WARNING)
+ try:
+ engine.run()
+ finally:
+ target.removeHandler(caplog.handler)
+ target.setLevel(prev)
+
+ frames = engine.committed_frames["req-0"]
+ assert frames == engine.reqs["req-0"]["max_tokens"], "the second request lost frames"
+ torch.testing.assert_close(
+ engine.audio_for("req-0"),
+ _expected_samples(engine.reqs["req-0"]["tag_base"], frames),
+ msg="the second request did not receive its own audio",
+ )
+ # The abort still gets reported: the cleanup that would have said so is the
+ # one this cancelled.
+ assert any("resubmitted" in m for m in caplog.messages), (
+ f"the aborted request's truncation went unreported; saw {caplog.messages}"
+ )
+
+
+def test_a_preempted_request_is_reported_rather_than_resumed(caplog: pytest.LogCaptureFixture) -> None:
+ """Recompute preemption re-prefills a request that was mid-generation.
+
+ This model cannot survive that, and the point of the test is that it says
+ so. The recomputed ids are the prompt plus the head0 codes sampled so far,
+ which go through the text embedding table, and each frame's other 31 codes
+ are never in the token stream at all — there is nothing to stitch the two
+ halves back together with, so generation restarts from frame 0.
+
+ The harness re-prefills with the bare prompt, which understates the damage.
+ What is under test is the report and the reset, not the corruption, which
+ needs a real backbone to observe.
+ """
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=4000)
+ for _ in range(FIRST_CHUNK_FRAMES + 2):
+ engine.step()
+ assert engine.committed_frames["req-0"] > 0
+
+ # Preemption the way the scheduler does it: computed tokens back to 0, so
+ # the next step re-prefills, and no id reaches on_requests_finished.
+ engine.reqs["req-0"]["computed"] = 0
+
+ target = logging.getLogger(LOGGER_NAME)
+ target.addHandler(caplog.handler)
+ prev = target.level
+ target.setLevel(logging.WARNING)
+ try:
+ engine.step()
+ finally:
+ target.removeHandler(caplog.handler)
+ target.setLevel(prev)
+
+ assert any("preempted" in m for m in caplog.messages), (
+ f"a preempted request restarted its audio silently; saw {caplog.messages}"
+ )
+ assert not any("resubmitted" in m for m in caplog.messages), (
+ "a preemption was reported as an id resubmit; the two need different words"
+ )
+ assert engine.committed_frames["req-0"] == 1, "the re-prefill must start the frame history over"
+
+
+def test_concurrent_requests_each_receive_their_own_full_audio() -> None:
+ """One request finishes while another stays live.
+
+ The earlier finisher is the case where a deferred flush *does* run but the
+ payload is filtered out again, because its id has left the output batch.
+ """
+ _talker, _codec, engine = _build()
+ engine.add_request("req-short", max_tokens=FIRST_CHUNK_FRAMES + 3)
+ engine.add_request("req-long", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 11)
+ engine.run()
+
+ for req_id in ("req-short", "req-long"):
+ frames = engine.committed_frames[req_id]
+ assert frames == engine.reqs[req_id]["max_tokens"]
+ torch.testing.assert_close(
+ engine.audio_for(req_id),
+ _expected_samples(engine.reqs[req_id]["tag_base"], frames),
+ msg=f"{req_id} lost or gained samples",
+ )
+
+
+def test_codec_still_runs_once_per_chunk_not_once_per_frame() -> None:
+ """Emitting on the final step must not degrade into per-frame decoding."""
+ _talker, codec, engine = _build()
+ num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES
+ engine.add_request("req-0", max_tokens=num_frames)
+ engine.run()
+
+ assert codec.decode_calls == 4, f"expected one decode per chunk, got {codec.decode_calls}"
+
+
+# --------------------------------------------------------------------------
+# Deferred cleanup is cleanup: it must not try to produce output.
+# --------------------------------------------------------------------------
+
+
+def test_deferred_cleanup_never_queues_audio() -> None:
+ """Anything queued from the cleanup hook is unroutable by construction.
+
+ By the time it runs the id has left ``req_ids_output_copy``, so a payload
+ under it is dropped; queueing there only risks leaking one request's audio
+ into the next step's payload for a different request.
+ """
+ talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=FIRST_CHUNK_FRAMES + 4)
+ engine.add_request("req-1", max_tokens=FIRST_CHUNK_FRAMES + CHUNK_FRAMES)
+ while engine.step():
+ if engine.reqs["req-0"]["finished"] and not engine.reqs["req-1"]["finished"]:
+ break
+
+ talker._audio_queue.clear()
+ talker.on_requests_finished({"req-0"})
+ talker._flush_deferred_cleanup()
+
+ assert talker._audio_queue == [], "cleanup emitted audio that routing cannot deliver"
+ assert "req-0" not in talker._active_states, "cleanup must still free the state"
+
+
+def test_cleanup_warns_when_frames_are_left_undelivered(caplog: pytest.LogCaptureFixture) -> None:
+ """An abort is the one case with no last step to emit on. Say so."""
+ _talker, _codec, engine = _build()
+ engine.add_request("req-0", max_tokens=4000)
+ engine.add_request("req-1", max_tokens=4000)
+ for _ in range(FIRST_CHUNK_FRAMES + 3):
+ engine.step()
+ engine.abort("req-0")
+
+ target = logging.getLogger(LOGGER_NAME)
+ target.addHandler(caplog.handler)
+ prev = target.level
+ target.setLevel(logging.WARNING)
+ try:
+ engine.step() # cleanup for req-0 runs at the start of this step
+ finally:
+ target.removeHandler(caplog.handler)
+ target.setLevel(prev)
+
+ assert any("undelivered frame" in m for m in caplog.messages), (
+ f"aborted request dropped audio without a warning; saw {caplog.messages}"
+ )
+
+
+# --------------------------------------------------------------------------
+# A request's seed must reach the in-model sampler.
+# --------------------------------------------------------------------------
+
+_SEEDED_FRAMES = FIRST_CHUNK_FRAMES + 5
+
+
+def _run_seeded(seed: int | None, *, neighbour: bool = False, neighbour_seed: int | None = None) -> torch.Tensor:
+ """Audio delivered to ``req-a``, optionally sharing the batch with another."""
+ _talker, engine = _build_with_real_sampler()
+ engine.add_request("req-a", max_tokens=_SEEDED_FRAMES, seed=seed)
+ if neighbour:
+ engine.add_request("req-b", max_tokens=_SEEDED_FRAMES, seed=neighbour_seed)
+ engine.run()
+ return engine.audio_for("req-a")
+
+
+def test_a_seeded_request_is_reproducible_within_one_process() -> None:
+ """The bug this pins: the 32 heads are sampled in ``forward``, so a seed
+ that only reaches vLLM's sampler leaves the audio at the mercy of the
+ global RNG — two identical requests in one process then diverge."""
+ torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234))
+
+
+def test_different_seeds_produce_different_audio() -> None:
+ """Guards the opposite failure: a seed that is threaded but never used."""
+ assert not torch.equal(_run_seeded(1234), _run_seeded(4321))
+
+
+def test_a_seeded_request_is_unaffected_by_its_batch() -> None:
+ """Each seeded request draws its own noise, so the noise it sees cannot
+ depend on who else was scheduled alongside it.
+
+ The stub backbone is what makes that checkable. A real one is a bf16
+ batched matmul whose rows are not reproducible across batch shapes, so
+ end-to-end batch invariance is not attainable there; this pins the half
+ the model actually controls.
+ """
+ torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True, neighbour_seed=999))
+
+
+def test_an_unseeded_neighbour_does_not_disturb_a_seeded_request() -> None:
+ """The mixed case: one row draws from the global RNG, the other must not."""
+ torch.testing.assert_close(_run_seeded(1234), _run_seeded(1234, neighbour=True))
+
+
+# --------------------------------------------------------------------------
+# Budget arithmetic, straight through the helper.
+# --------------------------------------------------------------------------
+
+
+def _info(**kwargs) -> _ReqInfo:
+ base = {
+ "request_id": "req-0",
+ "is_prefill": False,
+ "prompt_len": 20,
+ "num_computed_tokens": 20,
+ "max_tokens": 10,
+ "seed": None,
+ }
+ base.update(kwargs)
+ return _ReqInfo(**base)
+
+
+def test_prefill_step_is_the_first_output_token() -> None:
+ talker, _codec, _engine = _build()
+ first = _info(is_prefill=True, num_computed_tokens=0, max_tokens=1)
+ assert talker._is_last_output_token(first, span=20, samples_frame=True) is True
+ more = _info(is_prefill=True, num_computed_tokens=0, max_tokens=2)
+ assert talker._is_last_output_token(more, span=20, samples_frame=True) is False
+
+
+def test_partial_prefill_chunk_does_not_consume_budget() -> None:
+ """vLLM discards that row's token, so it must not count against max_tokens."""
+ talker, _codec, _engine = _build()
+ info = _info(is_prefill=True, num_computed_tokens=0, prompt_len=40, max_tokens=1)
+ assert talker._is_last_output_token(info, span=10, samples_frame=False) is False
+
+
+def test_budget_is_capped_by_max_model_len() -> None:
+ """A max_tokens larger than the context window never gets reached."""
+ talker, _codec, _engine = _build(max_model_len=30)
+ # prompt 20 + 10 generated fills the window, whatever max_tokens says.
+ at_cap = _info(num_computed_tokens=28, max_tokens=1000)
+ assert talker._is_last_output_token(at_cap, span=1, samples_frame=True) is True
+ below_cap = _info(num_computed_tokens=27, max_tokens=1000)
+ assert talker._is_last_output_token(below_cap, span=1, samples_frame=True) is False
+
+
+def test_missing_budget_is_not_treated_as_exhausted() -> None:
+ """A runner that does not thread max_tokens must not trigger a fade."""
+ talker, _codec, _engine = _build(max_model_len=0)
+ assert talker._is_last_output_token(_info(max_tokens=None), span=1, samples_frame=True) is False
diff --git a/tests/model_executor/models/test_minimax_h3_token_tags.py b/tests/model_executor/models/test_minimax_h3_token_tags.py
new file mode 100644
index 000000000..69ee7cb84
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_h3_token_tags.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
+
+from vllm_omni.model_executor.models.minimax_h3.text_encoder import MiniMaxH3TextEncoder
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _model_without_weights() -> MiniMaxH3TextEncoder:
+ model = MiniMaxH3TextEncoder.__new__(MiniMaxH3TextEncoder)
+ model.config = SimpleNamespace(
+ vision_start_token_id=1,
+ vision_end_token_id=2,
+ image_token_id=3,
+ video_token_id=4,
+ )
+ model._token_tags = None
+ return model
+
+
+def test_embed_failure_clears_token_tags():
+ model = _model_without_weights()
+ model._token_tags = torch.tensor([0])
+
+ with (
+ patch.object(Qwen3VLForConditionalGeneration, "embed_input_ids", side_effect=RuntimeError("failed")),
+ pytest.raises(RuntimeError, match="failed"),
+ ):
+ model.embed_input_ids(torch.tensor([5, 3]))
+
+ assert model._token_tags is None
+
+
+def test_invalid_output_length_consumes_token_tags():
+ model = _model_without_weights()
+ model._token_tags = torch.tensor([1, 0])
+
+ with pytest.raises(RuntimeError, match="token tags exceed"):
+ model.make_omni_output(torch.zeros(1, 8))
+
+ assert model._token_tags is None
+
+
+def test_make_omni_output_uses_shared_stage_payload_fields():
+ model = _model_without_weights()
+ model._token_tags = torch.tensor([1, 0])
+ hidden = torch.zeros(2, 5120)
+
+ output = model.make_omni_output(hidden)
+
+ assert output.multimodal_outputs is not None
+ assert torch.equal(output.multimodal_outputs["hidden_states"]["output"], hidden)
+ assert torch.equal(output.multimodal_outputs["meta"]["token_role_ids"], torch.tensor([[1], [0]]))
+ assert "encoder_hidden_states" not in output.multimodal_outputs
+ assert "token_tags" not in output.multimodal_outputs
diff --git a/tests/model_executor/models/test_minimax_music3_batch_payload.py b/tests/model_executor/models/test_minimax_music3_batch_payload.py
new file mode 100644
index 000000000..531500cd4
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_music3_batch_payload.py
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 batched payload partitioning.
+
+The stage-0 payload is split per request by indexing list values: a list is
+indexed by batch row, anything else is handed to every request unchanged. So
+per-request metadata has to be carried as lists. Merging it into one dict made
+every request in a batch receive the last request's seed, window offset and
+terminal flag, which crosses seeds between songs, misplaces acoustic windows,
+and lets one request's terminal flag end another request's stream early.
+
+These drive ``make_omni_output`` directly with a hand-built batch, so no GPU
+and no engine are needed.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+ MiniMaxMusic3TalkerForConditionalGeneration as Talker,
+)
+from vllm_omni.utils.mm_outputs import to_payload_element
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_WIDTH = 4
+
+
+def _info(*, seed: int, start: int, frames: int, last: bool) -> dict:
+ """One batch row's intermediate buffer, as ``postprocess`` leaves it."""
+ return {
+ "latent": torch.full((frames, _WIDTH), float(seed)),
+ "codes": {"audio": torch.full((frames, 8), seed)},
+ "meta": {
+ "audio_seed": seed,
+ "num_processed_tokens": start,
+ "codec_chunk_frames": frames,
+ "last_chunk": last,
+ },
+ }
+
+
+def _payload_for(output, row: int) -> dict:
+ """Partition the batch payload down to one request, as the runner does."""
+ flat = {}
+ for key, value in output.multimodal_outputs.items():
+ if isinstance(value, dict):
+ for qual, val in value.items():
+ flat[f"{key}.{qual}"] = val
+ else:
+ flat[key] = value
+ return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()}
+
+
+class TestBatchedMetadataPartitioning:
+ """Two requests in one batch must keep their own metadata."""
+
+ @staticmethod
+ def _two_request_batch():
+ infos = [
+ _info(seed=111, start=100, frames=3, last=False),
+ _info(seed=222, start=200, frames=5, last=True),
+ ]
+ talker = Talker.__new__(Talker)
+ return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+ def test_seeds_are_not_crossed(self) -> None:
+ """A song decoded with another request's seed is a different song."""
+ output = self._two_request_batch()
+ assert _payload_for(output, 0)["meta.audio_seed"] == 111
+ assert _payload_for(output, 1)["meta.audio_seed"] == 222
+
+ def test_window_offsets_stay_with_their_request(self) -> None:
+ """The start frame decides where a window lands in the song."""
+ output = self._two_request_batch()
+ assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100
+ assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200
+
+ def test_terminal_flag_does_not_leak_across_requests(self) -> None:
+ """One request finishing must not end another request's stream."""
+ output = self._two_request_batch()
+ assert _payload_for(output, 0)["meta.last_chunk"] is False
+ assert _payload_for(output, 1)["meta.last_chunk"] is True
+
+ def test_span_length_matches_its_own_latent(self) -> None:
+ output = self._two_request_batch()
+ for row, frames in ((0, 3), (1, 5)):
+ entry = _payload_for(output, row)
+ assert entry["meta.codec_chunk_frames"] == frames
+ assert entry["latent"].shape[0] == frames
+
+ def test_latent_and_codes_follow_the_same_row(self) -> None:
+ """The tensors were already per-request; pin that they stay aligned."""
+ output = self._two_request_batch()
+ for row, seed in ((0, 111), (1, 222)):
+ entry = _payload_for(output, row)
+ assert entry["latent"][0, 0].item() == pytest.approx(float(seed))
+ assert int(entry["codes.audio"][0, 0]) == seed
+
+
+class TestPartialBatches:
+ """Only some rows carry a span on any given step."""
+
+ def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None:
+ infos = [{}, _info(seed=777, start=300, frames=2, last=True)]
+ talker = Talker.__new__(Talker)
+ output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+ idle = _payload_for(output, 0)
+ assert idle["latent"].numel() == 0
+ assert idle["meta.audio_seed"] is None
+ assert idle["meta.last_chunk"] is None
+
+ active = _payload_for(output, 1)
+ assert active["meta.audio_seed"] == 777
+ assert active["meta.last_chunk"] is True
+
+ def test_batch_with_no_spans_emits_no_payload(self) -> None:
+ talker = Talker.__new__(Talker)
+ output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}])
+ assert output.multimodal_outputs is None
+
+ def test_consumed_entries_are_cleared_from_the_buffer(self) -> None:
+ """The runner never clears this buffer, so a span left behind would be
+ re-emitted on every later step."""
+ infos = [_info(seed=1, start=0, frames=2, last=False)]
+ talker = Talker.__new__(Talker)
+ talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos)
+ assert "latent" not in infos[0]
+ assert "meta" not in infos[0]
diff --git a/tests/model_executor/models/test_minimax_music3_frame_state.py b/tests/model_executor/models/test_minimax_music3_frame_state.py
new file mode 100644
index 000000000..73dc9215b
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_music3_frame_state.py
@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping.
+
+The talker emits overlapping windows while dropping frames it can no longer
+need, and flushes whatever is left as one span when the song ends. Getting the
+eviction boundary wrong by one frame silently corrupts the conditioning handed
+to the acoustic stage from the second window onward, which is audible but hard
+to attribute, so the lifecycle is simulated here frame by frame.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+ AR_CHUNK_FRAMES,
+ AR_CHUNK_HOP_FRAMES,
+ NUM_CODEBOOKS,
+)
+from vllm_omni.model_executor.models.minimax_music3.frame_state import (
+ FrameBuffer,
+ RequestARState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# One scalar per frame instead of 32768 dims: the bookkeeping is what is under
+# test, and a frame's value doubles as its identity.
+_WIDTH = 4
+
+
+def _frame(index: int) -> torch.Tensor:
+ return torch.full((1, _WIDTH), float(index))
+
+
+def _identities(chunk: torch.Tensor) -> list[int]:
+ return [int(row[0].item()) for row in chunk]
+
+
+class TestFrameBuffer:
+ def test_slice_returns_the_requested_frames(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(10):
+ buffer.append(_frame(index))
+ assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6]
+
+ def test_slice_is_relative_to_the_live_base(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(10):
+ buffer.append(_frame(index))
+ buffer.discard_before(4)
+ assert buffer.base_frame == 4
+ assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7]
+
+ def test_slicing_a_discarded_range_raises(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(10):
+ buffer.append(_frame(index))
+ buffer.discard_before(5)
+ with pytest.raises(RuntimeError, match="unavailable"):
+ buffer.slice(3, 6)
+
+ def test_slicing_past_the_end_raises(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(4):
+ buffer.append(_frame(index))
+ with pytest.raises(RuntimeError, match="unavailable"):
+ buffer.slice(2, 9)
+
+ def test_discard_is_clamped_at_both_ends(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(5):
+ buffer.append(_frame(index))
+ buffer.discard_before(-3)
+ assert buffer.base_frame == 0
+ buffer.discard_before(99)
+ assert buffer.base_frame == 5
+ assert buffer.end_frame == 5
+
+ def test_peak_tracks_the_high_water_mark(self) -> None:
+ buffer = FrameBuffer()
+ for index in range(10):
+ buffer.append(_frame(index))
+ buffer.discard_before(8)
+ assert buffer.peak_frames == 10
+
+
+class TestWindowLifecycle:
+ """Drive the state exactly as the talker does, one frame per step."""
+
+ @staticmethod
+ def _run(total_frames: int) -> tuple[list[list[int]], int]:
+ """Generate ``total_frames`` frames, returning the emitted spans.
+
+ Returns:
+ The frame identities of every emitted span, and the peak buffer
+ occupancy reached along the way.
+ """
+ state = RequestARState(seed=1, max_frames=total_frames)
+ spans: list[list[int]] = []
+ for index in range(total_frames):
+ state.frames.append(_frame(index))
+ state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+ state.generated_frames += 1
+ window = state.ready_window()
+ if window is not None:
+ spans.append(_identities(state.take_window(window, is_final=False)))
+
+ # Terminal flush: everything not yet emitted, as one contiguous span.
+ start = state.emitted_windows * AR_CHUNK_HOP_FRAMES
+ end = state.frames.end_frame
+ if end > start:
+ spans.append(_identities(state.frames.slice(start, end)))
+ return spans, state.frames.peak_frames
+
+ def test_short_song_is_one_span(self) -> None:
+ spans, _ = self._run(50)
+ assert spans == [list(range(50))]
+
+ def test_every_frame_is_delivered_exactly_once_in_order(self) -> None:
+ """Spans overlap by design, but their union must be the whole song.
+
+ Reconstruct by taking each span's non-overlapping remainder, which is
+ what the acoustic stage does after cropping.
+ """
+ for total in (250, 437, 600, 901):
+ spans, _ = self._run(total)
+ rebuilt: list[int] = []
+ for span in spans:
+ rebuilt.extend(frame for frame in span if frame >= len(rebuilt))
+ assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames"
+
+ def test_windows_advance_by_the_hop(self) -> None:
+ spans, _ = self._run(600)
+ for index, span in enumerate(spans[:-1]):
+ assert span[0] == index * AR_CHUNK_HOP_FRAMES
+ assert len(span) == AR_CHUNK_FRAMES
+
+ def test_mid_stream_windows_are_never_the_final_frame(self) -> None:
+ """A window is held back until a further frame proves it is not last."""
+ total = 600
+ spans, _ = self._run(total)
+ for span in spans[:-1]:
+ assert span[-1] < total - 1
+
+ def test_terminal_span_covers_the_tail(self) -> None:
+ total = 437
+ spans, _ = self._run(total)
+ assert spans[-1][-1] == total - 1
+
+ def test_terminal_span_is_bounded(self) -> None:
+ """The tail is at most one full window plus one hop."""
+ for total in (250, 301, 437, 600, 901, 1000):
+ spans, _ = self._run(total)
+ assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+ def test_buffer_stays_bounded_over_a_long_song(self) -> None:
+ """Peak occupancy must not grow with song length; that is the point
+ of evicting retired frames."""
+ _, peak_short = self._run(600)
+ _, peak_long = self._run(3000)
+ assert peak_long == peak_short
+ assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+ def test_emitted_window_count_matches_the_window_plan(self) -> None:
+ """What the stream emitted must reconcile with chunk_windows()."""
+ for total in (250, 437, 600):
+ spans, _ = self._run(total)
+ planned = chunk_windows(total)
+ # Every window is accounted for: the pre-terminal spans are one
+ # window each, and the terminal span carries whatever remains.
+ assert len(spans) <= len(planned)
+ assert spans[0][0] == planned[0].start
+
+
+class TestSamplingPositions:
+ def test_each_frame_reserves_one_slot_per_codebook(self) -> None:
+ """Codebook streams must never collide, or codes repeat across frames."""
+ state = RequestARState(seed=1, max_frames=10)
+ positions = [state.advance_position() for _ in range(5)]
+ assert positions == [0, 8, 16, 24, 32]
+
+ def test_draw_slots_are_globally_unique(self) -> None:
+ state = RequestARState(seed=1, max_frames=64)
+ slots: set[int] = set()
+ for _ in range(64):
+ base = state.advance_position()
+ for codebook in range(NUM_CODEBOOKS):
+ slot = base + codebook
+ assert slot not in slots
+ slots.add(slot)
+ assert len(slots) == 64 * NUM_CODEBOOKS
+
+
+class TestCodeRetention:
+ def test_codes_span_matches_the_frame_span(self) -> None:
+ state = RequestARState(seed=1, max_frames=10)
+ for index in range(10):
+ state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+ codes = state.take_codes(2, 6)
+ assert codes is not None
+ assert codes.shape == (4, NUM_CODEBOOKS)
+ assert [int(row[0]) for row in codes] == [2, 3, 4, 5]
+
+ def test_out_of_range_span_is_dropped_not_raised(self) -> None:
+ """Codes are diagnostic; the audio path must not fail without them."""
+ state = RequestARState(seed=1, max_frames=10)
+ for index in range(4):
+ state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+ assert state.take_codes(2, 99) is None
+ assert state.take_codes(0, 0) is None
diff --git a/tests/model_executor/models/test_minimax_music3_prompt.py b/tests/model_executor/models/test_minimax_music3_prompt.py
new file mode 100644
index 000000000..8cd9e41b0
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_music3_prompt.py
@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.
+
+The prompt transformations are not cosmetic: they decide token positions, and
+token positions decide the audio. These tests pin the documented behaviours
+that a well-meaning cleanup would otherwise "fix".
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import (
+ chunk_windows,
+ crop_sample_bounds,
+ overlap_mel_length,
+)
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+ AR_CHUNK_FRAMES,
+ AR_CHUNK_HOP_FRAMES,
+ BLEND_MEL_FRAMES,
+ DAV_HOP_SAMPLES,
+ HOP_MEL_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.prompt import (
+ SPECIAL_TOKEN_IDS,
+ build_cfg_null_token_ids,
+ build_prompt,
+ clean_caption,
+ normalize_lyrics,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestNormalizeLyrics:
+ def test_prepends_start_tag(self) -> None:
+ assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")
+
+ def test_lowercases_structure_tags(self) -> None:
+ assert "[verse]" in normalize_lyrics("[Verse]\nhello")
+ assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")
+
+ def test_text_after_a_leading_tag_is_dropped(self) -> None:
+ """Documented contract: a tag must sit on its own line.
+
+ Lyrics written next to a tag are silently lost, which is why the
+ serving docs tell callers to put tags on their own line.
+ """
+ out = normalize_lyrics("[Verse] Walking down the street")
+ assert "Walking" not in out
+ assert "[verse]" in out
+
+ def test_text_on_its_own_line_survives(self) -> None:
+ out = normalize_lyrics("[Verse]\nWalking down the street")
+ assert "Walking down the street" in out
+
+ def test_bracket_space_becomes_newline(self) -> None:
+ assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"
+
+
+class TestCleanCaption:
+ def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
+ """A caption must not be able to inject a real special token."""
+ assert clean_caption("<|bpm 92|>") == "bpm is 92"
+ assert "<|" not in clean_caption("<|bpm 92|>")
+
+ def test_lone_tag_keeps_its_inner_text(self) -> None:
+ assert clean_caption("<|instrumental|>") == "instrumental"
+
+ def test_markdown_is_stripped(self) -> None:
+ assert clean_caption("## Heading") == "Heading"
+ assert clean_caption("**bold**") == "bold"
+ assert clean_caption("- item") == "item"
+
+ def test_blank_lines_collapse_at_two_not_three(self) -> None:
+ """The source collapses runs of two or more, not three or more.
+
+ This shifts blank-line token positions and therefore the audio, so it
+ is pinned deliberately.
+ """
+ assert clean_caption("a\n\nb") == "a\nb"
+ assert clean_caption("a\n\n\n\nb") == "a\nb"
+
+
+class TestBuildPrompt:
+ def test_framing_tokens_are_present_and_ordered(self) -> None:
+ prompt = build_prompt("a bright pop song", "[Verse]\nhello")
+ assert prompt.startswith("<|im_start|><|caption_start|>")
+ assert prompt.endswith("<|im_end|><|audio_start|>")
+ assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")
+
+
+class TestBuildCfgNullTokenIds:
+ @staticmethod
+ def _prompt_ids(body_len: int) -> list[int]:
+ return (
+ [SPECIAL_TOKEN_IDS["<|im_start|>"]]
+ + list(range(1000, 1000 + body_len))
+ + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
+ )
+
+ def test_length_is_preserved(self) -> None:
+ """Equal length is what lets the two rows decode in lockstep."""
+ ids = self._prompt_ids(12)
+ assert len(build_cfg_null_token_ids(ids)) == len(ids)
+
+ def test_body_is_replaced_and_framing_survives(self) -> None:
+ ids = self._prompt_ids(12)
+ null = build_cfg_null_token_ids(ids)
+ assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
+ assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
+ assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
+ assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}
+
+ def test_input_is_not_mutated(self) -> None:
+ ids = self._prompt_ids(6)
+ before = list(ids)
+ build_cfg_null_token_ids(ids)
+ assert ids == before
+
+ def test_minimal_prompt_with_empty_body(self) -> None:
+ ids = [
+ SPECIAL_TOKEN_IDS["<|im_start|>"],
+ SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
+ SPECIAL_TOKEN_IDS["<|im_end|>"],
+ SPECIAL_TOKEN_IDS["<|audio_start|>"],
+ ]
+ assert build_cfg_null_token_ids(ids) == ids
+
+ def test_too_short_is_rejected(self) -> None:
+ with pytest.raises(ValueError, match="too short"):
+ build_cfg_null_token_ids([1, 2, 3])
+
+ def test_wrong_framing_is_rejected(self) -> None:
+ with pytest.raises(ValueError, match="framing"):
+ build_cfg_null_token_ids([9, 9, 9, 9, 9])
+
+
+class TestChunkWindows:
+ def test_no_frames_yields_nothing(self) -> None:
+ assert chunk_windows(0) == []
+
+ def test_short_song_is_one_final_window(self) -> None:
+ windows = chunk_windows(50)
+ assert len(windows) == 1
+ assert (windows[0].start, windows[0].end) == (0, 50)
+ assert windows[0].is_first and windows[0].is_last
+
+ def test_exactly_one_window(self) -> None:
+ windows = chunk_windows(AR_CHUNK_FRAMES)
+ assert len(windows) == 1
+ assert windows[0].is_last
+
+ def test_windows_advance_by_the_hop(self) -> None:
+ windows = chunk_windows(450)
+ starts = [w.start for w in windows]
+ assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
+ assert windows[-1].is_last
+ assert not any(w.is_last for w in windows[:-1])
+
+ def test_windows_cover_every_frame(self) -> None:
+ frames = 437
+ windows = chunk_windows(frames)
+ covered: set[int] = set()
+ for window in windows:
+ covered.update(range(window.start, window.end))
+ assert covered == set(range(frames))
+
+ def test_final_window_is_clamped_to_the_song(self) -> None:
+ windows = chunk_windows(437)
+ assert windows[-1].end == 437
+
+ def test_only_the_first_window_is_first(self) -> None:
+ windows = chunk_windows(600)
+ assert windows[0].is_first
+ assert not any(w.is_first for w in windows[1:])
+
+
+class TestCropBounds:
+ def test_single_window_song_is_not_cropped(self) -> None:
+ window = chunk_windows(50)[0]
+ assert crop_sample_bounds(window) == (0, 0)
+
+ def test_first_window_keeps_its_head(self) -> None:
+ left, right = crop_sample_bounds(chunk_windows(600)[0])
+ assert left == 0
+ assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
+
+ def test_last_window_keeps_its_tail(self) -> None:
+ left, right = crop_sample_bounds(chunk_windows(600)[-1])
+ assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
+ assert right == 0
+
+ def test_middle_window_is_cropped_on_both_sides(self) -> None:
+ windows = chunk_windows(900)
+ middle = windows[len(windows) // 2]
+ left, right = crop_sample_bounds(middle)
+ assert left > 0 and right > 0
+
+ def test_overlap_is_half_a_window(self) -> None:
+ assert overlap_mel_length() == HOP_MEL_FRAMES // 2
diff --git a/tests/model_executor/models/test_minimax_music3_repo_root.py b/tests/model_executor/models/test_minimax_music3_repo_root.py
new file mode 100644
index 000000000..182635ab7
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_music3_repo_root.py
@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Repository-root resolution for MiniMax Music 3's multi-component checkpoint.
+
+The acoustic stage declares no ``model_subdir`` because it reads the repo root
+itself, so on a hub-id deployment its model path is ``MiniMaxAI/MiniMax-Music3``
+rather than a directory. Resolving that as a relative path silently points at
+the server's working directory (issue #6638).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.weights import (
+ _COMPONENT_WEIGHT_NAME,
+ _REQUIRED_COMPONENTS,
+ _ROOT_MARKERS,
+ resolve_repo_root,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_root(path):
+ # A real snapshot carries weights for every component a stage loads, not
+ # just the folders that identify the root. Building only ``_ROOT_MARKERS``
+ # here would let a snapshot missing ``condition_encoder`` pass as complete,
+ # which is the gap this suite regresses.
+ for component in _REQUIRED_COMPONENTS:
+ folder = path / component
+ folder.mkdir(parents=True, exist_ok=True)
+ (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+ return path
+
+
+def test_resolve_repo_root_accepts_the_root_itself(tmp_path):
+ root = _make_root(tmp_path / "snapshot")
+ assert resolve_repo_root(str(root)) == root
+
+
+def test_resolve_repo_root_walks_up_from_a_model_subdir(tmp_path):
+ root = _make_root(tmp_path / "snapshot")
+ (root / "language_model").mkdir()
+ assert resolve_repo_root(str(root / "language_model")) == root
+
+
+def test_resolve_repo_root_resolves_a_hub_id_to_a_local_snapshot(monkeypatch, tmp_path):
+ """A hub id must resolve through the cache, not against the working directory."""
+ from huggingface_hub import HfApi
+
+ root = _make_root(tmp_path / "snapshots" / "deadbeef")
+ seen = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ seen.append(kwargs)
+ return str(root)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root
+ # Cache-first, and only the component folders.
+ assert seen[0]["local_files_only"] is True
+ assert "transformer/*" in seen[0]["allow_patterns"]
+ assert "qwen_7B/*" not in seen[0]["allow_patterns"]
+
+
+def test_resolve_repo_root_falls_back_to_the_hub_when_the_cache_is_incomplete(monkeypatch, tmp_path):
+ from huggingface_hub import HfApi
+
+ root = _make_root(tmp_path / "snapshots" / "deadbeef")
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if kwargs.get("local_files_only"):
+ raise OSError("incomplete snapshot")
+ return str(root)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == root
+ assert len(calls) == 2
+
+
+def test_resolve_repo_root_retries_online_when_the_local_lookup_returns_a_partial_root(monkeypatch, tmp_path):
+ """Hub clients differ on partial caches: some raise, some return the root.
+
+ A returned partial root must fall through to the online call all the same,
+ or startup fails without ever consulting the Hub.
+ """
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+ (snapshot / "language_model").mkdir(parents=True)
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if not kwargs.get("local_files_only"):
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot
+ assert [bool(call.get("local_files_only")) for call in calls] == [True, False]
+
+
+def test_resolve_repo_root_redownloads_weightless_marker_dirs(monkeypatch, tmp_path):
+ """Marker folders without their component weights are an interrupted download."""
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+ for marker in _ROOT_MARKERS:
+ (snapshot / marker).mkdir(parents=True)
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ if not kwargs.get("local_files_only"):
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root("MiniMaxAI/MiniMax-Music3") == snapshot
+ assert any(not call.get("local_files_only") for call in calls)
+
+
+def test_resolve_repo_root_reports_the_original_reference_when_unresolvable(monkeypatch):
+ from huggingface_hub import HfApi
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ raise OSError("offline")
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ with pytest.raises(FileNotFoundError, match="MiniMaxAI/MiniMax-Music3"):
+ resolve_repo_root("MiniMaxAI/MiniMax-Music3")
+
+
+def test_resolve_repo_root_downloads_components_for_a_cold_cache_model_subdir(monkeypatch, tmp_path):
+ """Cold-cache stage 0: the talker gets an EXISTING ``language_model`` dir.
+
+ Stage init pre-downloads only ``language_model/`` and ``tokenizer/``, so
+ the marker walk fails although the repo id is recoverable from the cache
+ layout and one snapshot_download away from working.
+ """
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+ (snapshot / "language_model").mkdir(parents=True)
+ (snapshot / "tokenizer").mkdir()
+ seen = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ seen.append(repo_id)
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root(str(snapshot / "language_model")) == snapshot
+ assert seen == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_does_not_guess_a_repo_id_for_plain_directories(monkeypatch, tmp_path):
+ """A directory outside the HF cache has no recoverable repo id."""
+ from huggingface_hub import HfApi
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ raise AssertionError("must not reach the Hub for a plain directory")
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ plain = tmp_path / "language_model"
+ plain.mkdir()
+ with pytest.raises(FileNotFoundError, match="expected sibling folders"):
+ resolve_repo_root(str(plain))
+
+
+def test_resolve_repo_root_rejects_a_snapshot_missing_condition_encoder(monkeypatch, tmp_path):
+ """``condition_encoder`` is loaded by the acoustic stage but is not a root marker.
+
+ A cache holding only the marker folders satisfies the name probe while the
+ acoustic stage would still fail, so completeness must cover every required
+ component rather than the identification markers alone.
+ """
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+ for marker in _ROOT_MARKERS:
+ folder = snapshot / marker
+ folder.mkdir(parents=True)
+ (folder / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+ repaired = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ if not kwargs.get("local_files_only"):
+ repaired.append(repo_id)
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root(str(snapshot)) == snapshot
+ assert repaired == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_rejects_a_partial_shard_run(monkeypatch, tmp_path):
+ """One shard of a multi-shard component satisfies a bare glob but cannot load."""
+ from huggingface_hub import HfApi
+
+ stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "deadbeef"
+ for component in _REQUIRED_COMPONENTS:
+ (snapshot / component).mkdir(parents=True)
+ (snapshot / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+ # Replace one component with an interrupted 1-of-3 shard run.
+ partial = snapshot / _ROOT_MARKERS[0]
+ (partial / _COMPONENT_WEIGHT_NAME).unlink()
+ (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x")
+ repaired = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ if not kwargs.get("local_files_only"):
+ repaired.append(repo_id)
+ for index in (2, 3):
+ (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x")
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root(str(snapshot)) == snapshot
+ assert repaired == ["MiniMaxAI/MiniMax-Music3"]
+
+
+def test_resolve_repo_root_threads_revision_and_download_dir(monkeypatch, tmp_path):
+ """A pinned revision must not be dropped when fetching the components.
+
+ Resolving components against the default branch while the AR backbone is
+ pinned elsewhere would combine weights from two different commits.
+ """
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "snapshot"
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root("MiniMaxAI/MiniMax-Music3", revision="abc123", download_dir="/tmp/hub-cache") == snapshot
+ assert calls
+ assert all(call.get("revision") == "abc123" for call in calls)
+ assert all(call.get("cache_dir") == "/tmp/hub-cache" for call in calls)
+
+
+def test_resolve_repo_root_reuses_the_snapshot_revision_from_the_cache_path(monkeypatch, tmp_path):
+ """The commit is recoverable from ``snapshots/`` and must be reused.
+
+ Re-resolving the repo id alone would fetch the default branch rather than
+ the snapshot the caller is already sitting in.
+ """
+ from huggingface_hub import HfApi
+
+ snapshot = tmp_path / "models--MiniMaxAI--MiniMax-Music3" / "snapshots" / "cafebabe"
+ (snapshot / "language_model").mkdir(parents=True)
+ calls = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ calls.append(kwargs)
+ _make_root(snapshot)
+ return str(snapshot)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root(str(snapshot / "language_model")) == snapshot
+ assert calls
+ assert all(call.get("revision") == "cafebabe" for call in calls)
+
+
+def test_resolve_repo_root_accepts_shard_names_the_loader_can_read(monkeypatch, tmp_path):
+ """The completeness predicate must not be stricter than ``load_component_state``.
+
+ The loader globs ``-*.safetensors`` and reads whatever it finds, so a
+ component whose weights do not use the ``--of-`` naming is
+ loadable. Judging it incomplete would force a needless re-download and, with
+ no network, fail on weights that are present and readable.
+ """
+ from huggingface_hub import HfApi
+
+ stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+ root = tmp_path / "snapshot"
+ for component in _REQUIRED_COMPONENTS:
+ (root / component).mkdir(parents=True)
+ (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+ unsharded = root / _ROOT_MARKERS[0]
+ (unsharded / _COMPONENT_WEIGHT_NAME).unlink()
+ (unsharded / f"{stem}-fp16.safetensors").write_bytes(b"x")
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ raise AssertionError("must not reach the Hub for readable weights")
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ assert resolve_repo_root(str(root)) == root
+
+
+def test_resolve_repo_root_still_rejects_a_partial_run_beside_an_extra_file(monkeypatch, tmp_path):
+ """An unparsable extra must not mask a genuinely short ``-of-`` run."""
+ from huggingface_hub import HfApi
+
+ stem = _COMPONENT_WEIGHT_NAME.removesuffix(".safetensors")
+ root = tmp_path / "snapshot"
+ for component in _REQUIRED_COMPONENTS:
+ (root / component).mkdir(parents=True)
+ (root / component / _COMPONENT_WEIGHT_NAME).write_bytes(b"x")
+ partial = root / _ROOT_MARKERS[0]
+ (partial / _COMPONENT_WEIGHT_NAME).unlink()
+ (partial / f"{stem}-00001-of-00003.safetensors").write_bytes(b"x")
+ (partial / f"{stem}-extra.safetensors").write_bytes(b"x")
+ repaired = []
+
+ def fake_snapshot_download(self, repo_id, **kwargs):
+ repaired.append(repo_id)
+ for index in (2, 3):
+ (partial / f"{stem}-0000{index}-of-00003.safetensors").write_bytes(b"x")
+ return str(root)
+
+ monkeypatch.setattr(HfApi, "snapshot_download", fake_snapshot_download)
+
+ with pytest.raises(FileNotFoundError, match="components are incomplete"):
+ resolve_repo_root(str(root))
+ assert repaired == []
diff --git a/tests/model_executor/models/test_minimax_music3_row_binding.py b/tests/model_executor/models/test_minimax_music3_row_binding.py
new file mode 100644
index 000000000..94ff30199
--- /dev/null
+++ b/tests/model_executor/models/test_minimax_music3_row_binding.py
@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 row-to-request binding.
+
+vLLM compacts and reorders the persistent batch whenever a request finishes or
+is preempted, so a batch row number is only meaningful within the step that
+produced it. Every failure pinned here was a real defect: state keyed by row
+survived into a step where the row belonged to someone else, or state keyed by
+one spelling of the request id was released under another and leaked.
+
+These exercise the pure helpers and the binding arithmetic directly, so they
+need no GPU and no engine.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+ DEFAULT_MAX_AUDIO_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+ _CFG_UNCOND_SUFFIX,
+ _first_int,
+ _pair_id_from_request,
+ _request_key,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestRequestKey:
+ def test_engine_list_form_is_unwrapped(self) -> None:
+ """The engine injects global_request_id as a single-element list.
+
+ Keying state on the list's ``str()`` would produce "['req-1']", which
+ never matches the bare id that ``on_requests_finished`` delivers, so
+ every finished song would leak its frame buffer.
+ """
+ assert _request_key(["req-1"]) == "req-1"
+
+ def test_bare_string_is_unchanged(self) -> None:
+ assert _request_key("req-1") == "req-1"
+
+ def test_list_and_bare_forms_agree(self) -> None:
+ """Allocation and release must land on the same key."""
+ assert _request_key(["req-1"]) == _request_key("req-1")
+
+ def test_empty_and_missing_are_stable(self) -> None:
+ assert _request_key(None) == ""
+ assert _request_key([]) == ""
+
+ def test_tuple_form(self) -> None:
+ assert _request_key(("req-9",)) == "req-9"
+
+
+class TestPairIdDerivation:
+ def test_conditioned_request_is_its_own_pair(self) -> None:
+ assert _pair_id_from_request("req-1") == "req-1"
+
+ def test_companion_maps_to_its_partner(self) -> None:
+ assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1"
+
+ def test_pair_members_agree_on_the_pair_id(self) -> None:
+ cond = "req-42"
+ uncond = f"req-42{_CFG_UNCOND_SUFFIX}"
+ assert _pair_id_from_request(cond) == _pair_id_from_request(uncond)
+
+ def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None:
+ assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond"
+
+
+class TestFirstInt:
+ def test_serving_layer_list_wrapping(self) -> None:
+ """tts_params values arrive wrapped in a single-element list."""
+ assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+ def test_bare_value(self) -> None:
+ assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+ def test_missing_falls_back(self) -> None:
+ assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+ assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+ def test_nonsense_falls_back_rather_than_raising(self) -> None:
+ assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+ assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+ assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+
+def _feedback_plan(
+ request_ids: list[str],
+ computed: list[int],
+ scheduled: list[int],
+ prompt_tokens: dict[str, int],
+ has_codes: set[str],
+) -> list[int]:
+ """Mirror of the talker's feedback row selection.
+
+ Returns the flat token offsets whose embedding gets replaced by an audio
+ frame. Kept as a standalone function so the arithmetic can be checked
+ without constructing a model.
+ """
+ offsets: list[int] = []
+ running = 0
+ for count in scheduled:
+ offsets.append(running)
+ running += count
+
+ rows: list[int] = []
+ for row, request_id in enumerate(request_ids):
+ if request_id not in has_codes:
+ continue
+ if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0):
+ continue
+ rows.append(offsets[row])
+ return rows
+
+
+class TestFeedbackRowSelection:
+ """The guard that decides which token embeddings become audio frames."""
+
+ def test_decode_only_step_replaces_every_row(self) -> None:
+ ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"]
+ rows = _feedback_plan(
+ ids,
+ computed=[120, 120, 90, 90],
+ scheduled=[1, 1, 1, 1],
+ prompt_tokens={i: 50 for i in ids},
+ has_codes=set(ids),
+ )
+ assert rows == [0, 1, 2, 3]
+
+ def test_prefill_row_is_never_overwritten(self) -> None:
+ """A prompt token must keep its real embedding.
+
+ The earlier guard compared the flat token count against a row budget,
+ so a short prompt slipped through and had its first embeddings
+ replaced with a previous request's audio frame.
+ """
+ rows = _feedback_plan(
+ ["a", "b"],
+ computed=[0, 300],
+ scheduled=[8, 1],
+ prompt_tokens={"a": 8, "b": 40},
+ has_codes={"b"},
+ )
+ # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single
+ # decode token sits at flat offset 8.
+ assert rows == [8]
+
+ def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None:
+ """The exact shape that defeated the old token-count guard."""
+ rows = _feedback_plan(
+ ["a"],
+ computed=[0],
+ scheduled=[10],
+ prompt_tokens={"a": 10},
+ has_codes={"a"},
+ )
+ assert rows == []
+
+ def test_chunked_prefill_midway_is_not_fed_back(self) -> None:
+ rows = _feedback_plan(
+ ["a"],
+ computed=[2048],
+ scheduled=[2048],
+ prompt_tokens={"a": 5000},
+ has_codes={"a"},
+ )
+ assert rows == []
+
+ def test_final_prefill_chunk_is_not_fed_back(self) -> None:
+ """The step that completes a prompt still carries prompt tokens."""
+ rows = _feedback_plan(
+ ["a"],
+ computed=[4096],
+ scheduled=[904],
+ prompt_tokens={"a": 5000},
+ has_codes={"a"},
+ )
+ assert rows == []
+
+ def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None:
+ """Tokens are a flat concatenation, so offsets are a running sum."""
+ ids = ["p", "d1", "d2"]
+ rows = _feedback_plan(
+ ids,
+ computed=[0, 500, 700],
+ scheduled=[64, 1, 1],
+ prompt_tokens={"p": 64, "d1": 40, "d2": 40},
+ has_codes={"d1", "d2"},
+ )
+ assert rows == [64, 65]
+
+ def test_request_without_codes_yet_is_skipped(self) -> None:
+ """The first decode step after prefill has no previous frame."""
+ rows = _feedback_plan(
+ ["a"],
+ computed=[50],
+ scheduled=[1],
+ prompt_tokens={"a": 50},
+ has_codes=set(),
+ )
+ assert rows == []
+
+ def test_reordered_batch_binds_by_request_not_position(self) -> None:
+ """The same two requests in swapped rows get swapped offsets.
+
+ This is the compaction case: holding codes in a row-indexed buffer
+ would have handed each request the other's frame.
+ """
+ prompt_tokens = {"a": 50, "b": 50}
+ first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"})
+ swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"})
+ assert first == [0]
+ assert swapped == [1]
diff --git a/tests/model_executor/models/test_nemotron_voicechat_registration.py b/tests/model_executor/models/test_nemotron_voicechat_registration.py
index 09d85b271..935ffe96d 100644
--- a/tests/model_executor/models/test_nemotron_voicechat_registration.py
+++ b/tests/model_executor/models/test_nemotron_voicechat_registration.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""L1 registration + config-shim tests for NemotronVoiceChat (CPU, no weights)."""
import subprocess
@@ -59,7 +59,6 @@ def test_pipeline_registered_with_three_stages() -> None:
assert code2wav.execution_type == StageExecutionType.LLM_GENERATION
assert code2wav.final_output and code2wav.final_output_type == "audio"
assert code2wav.input_sources == (1,)
- assert pipeline.validate() == []
def test_architectures_registered() -> None:
@@ -67,15 +66,6 @@ def test_architectures_registered() -> None:
assert arch in _OMNI_MODELS, f"{arch} missing from _OMNI_MODELS"
-def test_full_payload_whitelists_cover_code2wav_only() -> None:
- from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
-
- assert ("NemotronVoiceChatCode2Wav", "code2wav") in _FULL_PAYLOAD_INPUT_STAGES
- # The thinker->talker hop is token-path only: the talker must NOT wait on a
- # connector payload or stage 1 hangs.
- assert ("NemotronVoiceChatTalkerForConditionalGeneration", "talker") not in _FULL_PAYLOAD_INPUT_STAGES
-
-
def test_config_parses_nemo_layout_without_hub_access(monkeypatch) -> None:
# Provide the thinker text config explicitly so the test never touches HF.
cfg = NemotronVoiceChatConfig(
diff --git a/tests/model_executor/models/test_nemotron_voicechat_shapes.py b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
index 8bd279bb8..9b3c01559 100644
--- a/tests/model_executor/models/test_nemotron_voicechat_shapes.py
+++ b/tests/model_executor/models/test_nemotron_voicechat_shapes.py
@@ -225,3 +225,39 @@ def test_torchaudio_mel_filterbank_matches_librosa() -> None:
reference = librosa.filters.mel(sr=16000, n_fft=512, n_mels=128, fmin=0.0, fmax=8000.0, norm="slaney")
assert ta.shape == reference.shape
assert abs(ta - reference).max() < 1e-6
+
+
+def test_code2wav_reuses_and_clears_per_request_streaming_cache() -> None:
+ import torch
+ from torch import nn
+
+ from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_code2wav import (
+ NemotronVoiceChatCode2Wav,
+ )
+
+ class Codec(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.anchor = nn.Parameter(torch.zeros(()))
+ self.caches = []
+
+ def decode(self, codes, lengths, *, cache):
+ self.caches.append(cache)
+ wav = torch.zeros((1, int(codes.shape[1]) * 1764), device=codes.device)
+ return wav, lengths * 1764
+
+ model = NemotronVoiceChatCode2Wav.__new__(NemotronVoiceChatCode2Wav)
+ nn.Module.__init__(model)
+ model._sample_rate, model._num_quantizers, model._codebook_size, model._wav_per_frame = 22050, 31, 1024, 1764
+ model._duplex_codec_caches = {}
+ model.audio_codec = Codec()
+
+ def info(codes):
+ return [{"meta": {"codec_streaming": True, "request_id": "req"}, "codes": {"audio": codes}}]
+
+ model.forward(runtime_additional_information=info(torch.zeros((1, 31), dtype=torch.long)))
+ model.forward(runtime_additional_information=info(torch.zeros((2, 31), dtype=torch.long)))
+ assert model.audio_codec.caches[0] is model.audio_codec.caches[1]
+ assert "req" in model._duplex_codec_caches
+ model.on_requests_finished({"req"})
+ assert "req" not in model._duplex_codec_caches
diff --git a/tests/model_executor/models/voxcpm2/test_talker_output_marker.py b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py
new file mode 100644
index 000000000..1a76525fe
--- /dev/null
+++ b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py
@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""VoxCPM2 dense-mode audio outputs must carry the sparse-alignment marker.
+
+Dense mode (`_uses_sparse_audio_outputs()` False) still yields a strict
+SUBSET of the batch whenever some requests emit no audio in a step (e.g.
+prefill phase). The runner routes per-request `model_outputs` lists by batch
+position unless `meta.sparse_audio` declares `meta.req_id` alignment — so a
+subset without the marker misroutes audio across requests (RFC #5450 C3
+review follow-up).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+torch = pytest.importorskip("torch")
+
+
+@functools.lru_cache(maxsize=1)
+def _talker_cls():
+ """Defer talker import (pulls vLLM model_executor) until first use."""
+ from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+ VoxCPM2TalkerForConditionalGeneration,
+ )
+
+ return VoxCPM2TalkerForConditionalGeneration
+
+
+def _make_dense_talker():
+ cls = _talker_cls()
+ talker = cls.__new__(cls)
+ # Dense mode: emit every step, no delayed copy, no batched VAE decode.
+ talker._audio_emit_every = 1
+ talker._vae_decode_every = 1
+ talker._enable_delayed_audio_copy = False
+ talker._coalesce_audio_d2h = False
+ talker._sample_rate = 16000
+ talker._audio_queue = []
+ return talker
+
+
+def test_dense_subset_audio_carries_sparse_alignment_marker():
+ talker = _make_dense_talker()
+ # Batch had [r0, r1, r2]; only r1/r2 produced audio this step.
+ talker._audio_queue = [("r1", torch.zeros(4)), ("r2", torch.ones(4))]
+
+ out = talker.make_omni_output(torch.zeros(1))
+
+ mm = out.multimodal_outputs
+ assert mm["meta"] == {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}
+ assert len(mm["model_outputs"]) == 2
+ assert len(mm["sr"]) == 2
+ # (The coalesce-D2H sibling branch carries the same marker; it needs
+ # device-resident chunks, so its marker contract is pinned here via the
+ # shared dense-subset shape rather than a GPU-only test.)
+
+
+def test_dense_full_batch_audio_still_carries_marker_and_order():
+ # Full coverage is the degenerate subset: the marker stays truthful and
+ # `meta.req_id` order matches the emitted list order.
+ talker = _make_dense_talker()
+ talker._audio_queue = [("r0", torch.zeros(2)), ("r1", torch.ones(2))]
+
+ out = talker.make_omni_output(torch.zeros(1))
+
+ mm = out.multimodal_outputs
+ assert mm["meta"]["req_id"] == ["r0", "r1"]
+ assert mm["meta"]["sparse_audio"] == ["1"]
diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
index 8ff67b5ff..c7686ca31 100644
--- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
+++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Tests for CUDA Graph acoustic transformer wrapper numerical equivalence.
@@ -18,11 +18,12 @@ import pytest
import torch
import torch.nn as nn
+from tests.helpers.mark import hardware_marks
+
pytestmark = [
pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
pytest.mark.core_model,
- pytest.mark.cuda,
- pytest.mark.L4,
+ *hardware_marks(res={"cuda": "L4"}),
]
DEVICE = torch.device("cuda:0")
diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
index d5971efe7..08090d954 100644
--- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
+++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections import defaultdict
from types import SimpleNamespace
@@ -7,6 +7,7 @@ from types import SimpleNamespace
import pytest
import torch
+from vllm_omni.data_entry_keys import serialize_payload
from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
talker2code2wav_async_chunk,
text2flow_full_payload,
@@ -153,6 +154,36 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
assert payload.embed.embedding is not None
+def test_talker2code2wav_async_chunk_decodes_bfloat16_conditioning():
+ speech_feat = torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.bfloat16)
+ additional_information = serialize_payload(
+ {
+ "embed": {
+ "speech_token": torch.tensor([[11]], dtype=torch.long),
+ "speech_feat": speech_feat,
+ "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.bfloat16),
+ }
+ }
+ )
+ request = SimpleNamespace(
+ external_req_id="rid-bfloat16",
+ output_token_ids=[1],
+ additional_information=additional_information,
+ is_finished=lambda: True,
+ )
+
+ payload = talker2code2wav_async_chunk(
+ transfer_manager=_transfer_manager(),
+ multimodal_output=None,
+ request=request,
+ is_finished=True,
+ )
+
+ assert payload is not None
+ assert payload.embed.speech_feat.dtype == torch.bfloat16
+ assert torch.equal(payload.embed.speech_feat, speech_feat)
+
+
def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
transfer_manager = _transfer_manager(chunk_frames=25)
request = SimpleNamespace(
diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py
index f0da872b9..afa7b46c4 100644
--- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py
+++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections import defaultdict
from types import SimpleNamespace
@@ -72,6 +72,16 @@ def _codes(payload) -> list[int]:
return payload.codes.audio.tolist()
+def test_empty_full_payload_releases_consumer_wait_gate() -> None:
+ payload = tts2code2wav_full_payload(_manager(), None, _request("req"))
+
+ assert _codes(payload) == []
+ assert payload.meta.code_flat_numel == 0
+ assert payload.meta.left_context_size == 0
+ assert payload.meta.last_chunk is True
+ assert payload.meta.finished.item() is True
+
+
@pytest.mark.parametrize(("count", "emitted"), [(24, False), (25, True), (26, True)])
def test_first_chunk_threshold_is_25_generated_codes(count: int, emitted: bool) -> None:
manager = _manager()
@@ -158,10 +168,14 @@ def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None:
def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
manager = _manager()
request = _request("req")
- request.additional_information = {
+ request.model_intermediate_buffer = {
"codes": {"ref": [0.1, -0.1]},
"meta": {"ref_audio_sr": 16000},
}
+ request.additional_information = {
+ "codes": {"ref": [0.9]},
+ "meta": {"ref_audio_sr": 8000},
+ }
payload = tts2code2wav_async_chunk(
manager,
@@ -183,10 +197,33 @@ def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
assert payload.meta.turn_end is True
+def test_first_chunk_falls_back_to_legacy_reference_fields() -> None:
+ manager = _manager()
+ request = _request("req")
+ request.model_intermediate_buffer = {
+ "meta": {"ref_audio_sr": 16000},
+ }
+ request.additional_information = {
+ "codes": {"ref": [0.9]},
+ "meta": {"ref_audio_sr": 8000},
+ }
+
+ payload = tts2code2wav_async_chunk(
+ manager,
+ _duplex_delta(*range(7), turn_end=True),
+ request,
+ True,
+ )
+
+ assert payload is not None
+ assert payload.codes.ref.tolist() == pytest.approx([0.9])
+ assert payload.meta.ref_audio_sr == 16000
+
+
def test_full_payload_forwards_all_codes_and_request_metadata() -> None:
manager = _manager()
request = _request("req")
- request.additional_information = {
+ request.model_intermediate_buffer = {
"codes": {"ref": [0.1, -0.1]},
"meta": {
"ref_audio_sr": 16000,
@@ -423,6 +460,12 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
request,
True,
)
+ duplicate_boundary = tts2code2wav_async_chunk(
+ manager,
+ _duplex_delta(turn_id=7, turn_end=True),
+ request,
+ True,
+ )
next_turn = tts2code2wav_async_chunk(
manager,
_duplex_delta(20, turn_id=8),
@@ -435,6 +478,10 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
assert turn_end.meta.last_chunk is True
assert turn_end.meta.turn_end is True
assert turn_end.meta.is_segment_finished.item() is True
+ assert turn_end.meta.replace_runtime_additional_information is True
+ assert duplicate_boundary is not None
+ assert duplicate_boundary.meta.is_segment_finished.item() is True
+ assert duplicate_boundary.meta.replace_runtime_additional_information is True
assert next_turn.meta.cache_epoch == turn_end.meta.cache_epoch + 1
assert next_turn.meta.chunk_seq == 0
assert next_turn.meta.last_chunk is False
diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py
index 652215a08..795047f1c 100644
--- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py
+++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from types import SimpleNamespace
import pytest
@@ -131,6 +134,7 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None:
assert converted["prompt_token_ids"] == [0, 0, 0]
assert info["meta"]["replace_streaming_prompt"] is True
assert info["meta"]["next_stage_prompt_len"] == 3
+ assert info["meta"]["next_stage_generation_tokens"] == 26
assert info["meta"]["turn_start"] is True
assert info["meta"]["segment_end"] is True
assert info["duplex"]["epoch"] == 3
@@ -203,6 +207,7 @@ def test_native_duplex_continuation_appends_only_new_talker_condition() -> None:
assert "replace_streaming_prompt" not in second_input["model_intermediate_buffer"]["meta"]
assert third_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True
assert second_input["model_intermediate_buffer"]["meta"]["next_stage_prompt_len"] == 3
+ assert second_input["model_intermediate_buffer"]["meta"]["next_stage_generation_tokens"] == 26
assert second_input["prompt_token_ids"] == [0, 0, 0]
diff --git a/tests/model_executor/stage_input_processors/test_minimax_h3.py b/tests/model_executor/stage_input_processors/test_minimax_h3.py
new file mode 100644
index 000000000..599be81f2
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_minimax_h3.py
@@ -0,0 +1,387 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
+
+import os
+import shutil
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+ _load_audio,
+ resolve_minimax_h3_diffusion_model_path,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
+ resolve_minimax_h3_model_root,
+ resolve_minimax_h3_partition,
+)
+from vllm_omni.model_executor.models.minimax_h3.conditioning import (
+ MINIMAX_H3_CONDITION_LABELS_KEY,
+ MINIMAX_H3_PRESENTATION_TASK_KEY,
+)
+from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
+ minimax_h3_ref2va_presentation,
+ minimax_h3_ref2va_video_presentation,
+)
+from vllm_omni.model_executor.models.minimax_h3.reference_video import (
+ MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
+ deserialize_prepared_reference_videos,
+)
+from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
+ MiniMaxH3MultiModalProcessor,
+ _build_minimax_h3_presentation,
+)
+from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
+ _audio_items,
+ _diffusion_sampling_params,
+ prepare_text_encoder_prompt,
+ text_encoder2diffusion,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _SegmentTokenizer:
+ _special_ids = {
+ "<|vision_start|>": 1,
+ "<|vision_end|>": 2,
+ "<|image_pad|>": 3,
+ "<|video_pad|>": 4,
+ }
+
+ def __call__(self, text, *, add_special_tokens):
+ assert add_special_tokens is False
+ return {"input_ids": [100 + len(text), 1000 + sum(text.encode())]}
+
+ def convert_tokens_to_ids(self, token):
+ return self._special_ids[token]
+
+
+def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker):
+ processor = object.__new__(MiniMaxH3MultiModalProcessor)
+ sentinel = ([1, 2, 3], object(), True)
+ apply_processor = mocker.patch.object(
+ MiniMaxH3MultiModalProcessor,
+ "_apply_hf_processor",
+ return_value=sentinel,
+ )
+ inputs = object()
+ timing_ctx = object()
+
+ result = processor._cached_apply_hf_processor(inputs, timing_ctx)
+
+ assert result is sentinel
+ apply_processor.assert_called_once_with(inputs, timing_ctx)
+
+
+@pytest.mark.parametrize(
+ ("value", "expected_count"),
+ [
+ ((torch.zeros(16), 16_000), 1),
+ ([np.zeros(16), 16_000], 1),
+ ([(torch.zeros(16), 16_000), (torch.ones(16), 24_000)], 2),
+ (["first.wav", "second.wav"], 2),
+ ],
+)
+def test_audio_items_preserves_waveform_pairs(value, expected_count):
+ assert len(_audio_items(value)) == expected_count
+
+
+def test_h3_selects_the_single_explicit_diffusion_stage_params() -> None:
+ stage_zero = SimpleNamespace(extra_args={"task": "t2va"})
+ diffusion = OmniDiffusionSamplingParams(extra_args={"task": "ref2va"})
+
+ assert _diffusion_sampling_params([stage_zero, diffusion]) is diffusion
+
+
+def test_h3_rejects_missing_or_ambiguous_diffusion_stage_params() -> None:
+ with pytest.raises(RuntimeError, match="exactly one OmniDiffusionSamplingParams"):
+ _diffusion_sampling_params([SimpleNamespace(extra_args={"task": "t2va"})])
+ with pytest.raises(RuntimeError, match="got 2"):
+ _diffusion_sampling_params([OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()])
+
+
+def test_fused_audio_loader_accepts_list_waveform_pair():
+ waveform, sample_rate = _load_audio([[0.0, 0.5, -0.5], 16_000])
+ assert sample_rate == 16_000
+ torch.testing.assert_close(waveform, torch.tensor([0.0, 0.5, -0.5]))
+
+
+def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
+ prompt = {
+ "prompt": "hello",
+ "multi_modal_data": {
+ "image": Image.new("RGB", (256, 256)),
+ "audio": [np.zeros(16), 16_000],
+ },
+ }
+ sampling = OmniDiffusionSamplingParams(
+ height=256,
+ width=448,
+ extra_args={"task": "ref2va"},
+ )
+
+ transformed = prepare_text_encoder_prompt(prompt, [sampling])
+
+ assert transformed["prompt"] == "hello"
+ assert len(transformed["multi_modal_data"]["image"]) == 1
+ assert "audio" not in transformed["multi_modal_data"]
+ assert transformed["mm_processor_kwargs"][MINIMAX_H3_PRESENTATION_TASK_KEY] == "ref2va"
+ assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
+ ("image", 1),
+ ("audio", 1),
+ ]
+
+
+def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor():
+ prompt = {
+ "prompt": "hello",
+ "additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}},
+ }
+ sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"})
+
+ transformed = prepare_text_encoder_prompt(prompt, [sampling])
+
+ assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"]
+
+
+def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch):
+ prepare_videos = Mock()
+
+ def prepare(_videos, *, target_frame_count, workdir, start_time_seconds):
+ prepare_videos(_videos, target_frame_count, workdir, start_time_seconds)
+ prepared_path = os.path.join(workdir, "prepared.mp4")
+ open(prepared_path, "wb").close()
+ return [
+ {
+ "original_path": "/tmp/input.mp4",
+ "prepared_path": prepared_path,
+ "input_has_audio": True,
+ "width": 448,
+ "height": 256,
+ "start_time_seconds": 0.0,
+ "duration_seconds": 5.2,
+ "audio_duration_seconds": 5.2,
+ }
+ ]
+
+ sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]})
+ monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare)
+ monkeypatch.setattr(
+ "vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames",
+ sample_frames,
+ )
+ prompt = {
+ "prompt": "hello",
+ "multi_modal_data": {"video": "/tmp/input.mp4"},
+ }
+ sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"})
+
+ transformed = prepare_text_encoder_prompt(prompt, [sampling])
+
+ prepare_videos.assert_called_once()
+ _, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args
+ assert target_frame_count == 124
+ assert start_time_seconds is None
+ descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY]
+ described_dir, described_videos = deserialize_prepared_reference_videos(descriptor)
+ assert described_dir == artifact_dir
+ assert os.path.isfile(described_videos[0]["prepared_path"])
+ sample_frames.assert_called_once_with(described_videos[0]["prepared_path"])
+ assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"}
+ assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
+ ("audio", 1),
+ ("video", 1),
+ ]
+ shutil.rmtree(artifact_dir)
+
+
+def test_stage_wire_rejects_multiple_text_encoder_sources():
+ with pytest.raises(RuntimeError, match="exactly one text-encoder source"):
+ text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"})
+
+
+def test_stage_wire_rejects_request_id_mismatch():
+ source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()])
+ prompt = {
+ "prompt": "hello",
+ "additional_information": {"global_request_id": ["req-1"]},
+ }
+
+ with pytest.raises(RuntimeError, match="request ID does not match"):
+ text_encoder2diffusion([source], prompt=prompt)
+
+
+def test_stage_wire_rejects_multiple_completions():
+ source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()])
+
+ with pytest.raises(RuntimeError, match="exactly one completion"):
+ text_encoder2diffusion([source], prompt={"prompt": "hello"})
+
+
+def _source_output(payload: dict) -> SimpleNamespace:
+ completion = SimpleNamespace(multimodal_output=payload)
+ return SimpleNamespace(request_id="request-1", outputs=[completion])
+
+
+def _diffusion_prompt() -> dict:
+ return {
+ "request_id": "request-1",
+ "prompt": "test prompt",
+ "additional_information": {},
+ }
+
+
+def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None:
+ hidden = torch.randn(4, 5120)
+ token_role_ids = torch.tensor([[1], [1], [0], [0]])
+
+ result = text_encoder2diffusion(
+ [
+ _source_output(
+ {
+ "hidden_states": {"output": hidden},
+ "meta": {"token_role_ids": token_role_ids},
+ }
+ )
+ ],
+ _diffusion_prompt(),
+ )
+
+ conditioning = result["additional_information"]["text_encoder_output"]
+ torch.testing.assert_close(conditioning["hidden_states"], hidden)
+ torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1))
+
+
+@pytest.mark.parametrize(
+ ("payload", "message"),
+ [
+ ({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"),
+ ({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"),
+ ({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"),
+ ({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"),
+ (
+ {
+ "hidden_states": {"output": torch.randn(4, 5120)},
+ "meta": {"token_role_ids": torch.ones(4)},
+ },
+ r"must have shape \[tokens, 1\]",
+ ),
+ ],
+)
+def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None:
+ with pytest.raises(RuntimeError, match=message):
+ text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
+
+
+def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
+ tokenizer = _SegmentTokenizer()
+ labels = [("image", 1), ("audio", 1)]
+ image_grid = torch.tensor([[1, 4, 4]])
+
+ actual = _build_minimax_h3_presentation(
+ tokenizer,
+ prompt="hello",
+ task="ref2va",
+ condition_labels=labels,
+ image_grid_thw=image_grid,
+ video_grid_thw=None,
+ video_timestamps=None,
+ merge_size=2,
+ )
+ expected = minimax_h3_ref2va_presentation(
+ tokenizer,
+ prompt="hello",
+ condition_labels=labels,
+ image_token_count=[4],
+ )
+
+ torch.testing.assert_close(actual[0], expected[0])
+ torch.testing.assert_close(actual[1], expected[1])
+
+
+def test_ref2va_video_tokens_and_tags_match_fused_without_outer_markers():
+ tokenizer = _SegmentTokenizer()
+ labels = [("audio", 1), ("video", 1)]
+ video_grid = torch.tensor([[2, 4, 4]])
+ timestamps = [[0.2, 0.4]]
+
+ actual = _build_minimax_h3_presentation(
+ tokenizer,
+ prompt="hello",
+ task="ref2va",
+ condition_labels=labels,
+ image_grid_thw=None,
+ video_grid_thw=video_grid,
+ video_timestamps=timestamps,
+ merge_size=2,
+ )
+ expected = minimax_h3_ref2va_video_presentation(
+ tokenizer,
+ prompt="hello",
+ condition_labels=labels,
+ image_token_count=None,
+ video_block_token_counts=[[4, 4]],
+ video_block_timestamps=timestamps,
+ )
+
+ torch.testing.assert_close(actual[0], expected[0])
+ torch.testing.assert_close(actual[1], expected[1])
+ assert int((actual[0] == tokenizer._special_ids["<|vision_start|>"]).sum()) == 2
+ assert int((actual[0] == tokenizer._special_ids["<|vision_end|>"]).sum()) == 2
+
+
+def test_checkpoint_resolver_selects_local_partition(tmp_path):
+ root = tmp_path / "MiniMax-H3"
+ (root / "FL2VA" / "text_encoder").mkdir(parents=True)
+ (root / "Ref2VA" / "text_encoder").mkdir(parents=True)
+
+ assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
+ assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
+ assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder")
+ assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
+
+
+def test_checkpoint_resolver_rejects_unknown_task(tmp_path):
+ with pytest.raises(ValueError, match="task_type must be one of"):
+ resolve_minimax_h3_model_root(str(tmp_path), None, "unknown")
+
+
+def test_partition_resolver_preserves_consumer_auto_default(tmp_path):
+ root = tmp_path / "MiniMax-H3"
+ ref2va = root / "Ref2VA"
+ ref2va.mkdir(parents=True)
+
+ assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="fl2va") == "fl2va"
+ assert resolve_minimax_h3_partition(str(root), "auto", auto_partition="combined") == "combined"
+ assert resolve_minimax_h3_partition(str(ref2va), "auto", auto_partition="combined") == "ref2va"
+
+
+def test_diffusion_resolver_selects_startup_partition(tmp_path):
+ root = tmp_path / "MiniMax-H3"
+ fl2va = root / "FL2VA"
+ ref2va = root / "Ref2VA"
+ fl2va.mkdir(parents=True)
+ ref2va.mkdir()
+ (fl2va / "model_index.json").write_text("{}")
+ (ref2va / "model_index.json").write_text("{}")
+
+ assert resolve_minimax_h3_diffusion_model_path(str(root), None, "fl2va") == str(fl2va)
+ assert resolve_minimax_h3_diffusion_model_path(str(root), None, "ref2va") == str(ref2va)
+ assert resolve_minimax_h3_diffusion_model_path(str(root), None, None) == str(fl2va)
+ assert resolve_minimax_h3_diffusion_model_path(str(root), None, "combined") == str(root)
+ assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, None) == str(ref2va)
+
+
+def test_diffusion_resolver_normalizes_partial_partition_directory(tmp_path):
+ root = tmp_path / "MiniMax-H3"
+ ref2va = root / "Ref2VA"
+ (ref2va / "text_encoder").mkdir(parents=True)
+
+ assert resolve_minimax_h3_diffusion_model_path(str(ref2va), None, "ref2va") == str(ref2va)
diff --git a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
index dcdff5bbb..172faf9fb 100644
--- a/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
+++ b/tests/model_executor/stage_input_processors/test_nemotron_voicechat.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
@@ -285,9 +285,46 @@ def test_talker2code2wav_async_chunk_requires_prompt_len() -> None:
talker2code2wav_async_chunk(tm, {"codes": {"audio": torch.zeros(5, 31, dtype=torch.long)}}, req)
-# =========================
-# Talker drain-per-wake (async scheduling review P1/P2)
-# =========================
+def test_talker2code2wav_duplex_keeps_cumulative_history_across_wakes() -> None:
+ from vllm_omni.model_executor.stage_input_processors.nemotron_voicechat import (
+ talker2code2wav_async_chunk,
+ )
+
+ tm = _fake_transfer_manager(codec_chunk_frames=2)
+ prompt_len = 3
+ req = _streaming_request(
+ prompt_len=prompt_len,
+ generated=[],
+ info={"meta": {"codec_streaming": torch.tensor([True])}, "nvc_logical_prompt_len": prompt_len},
+ )
+ prompt_rows = torch.zeros((prompt_len - 1, 31), dtype=torch.long)
+ first = torch.full((1, 31), 101, dtype=torch.long)
+ second = torch.full((1, 31), 202, dtype=torch.long)
+
+ for rows, finished in ((prompt_rows, False), (torch.cat([prompt_rows, first]), False)):
+ payload = talker2code2wav_async_chunk(
+ tm,
+ {"codes": {"audio": rows}, "meta": {"codec_streaming": torch.tensor([True])}},
+ req,
+ is_finished=finished,
+ )
+ assert payload is not None and not bool(payload.meta.finished)
+
+ payload = talker2code2wav_async_chunk(
+ tm,
+ {
+ "codes": {"audio": torch.cat([prompt_rows, first, second])},
+ "meta": {"codec_streaming": torch.tensor([True])},
+ },
+ req,
+ is_finished=True,
+ )
+
+ assert payload is not None
+ assert not bool(payload.meta.finished)
+ assert not bool(payload.meta.is_segment_finished)
+ assert torch.equal(payload.codes.audio, torch.cat([first, second]))
+ assert tm.request_payload["req-0"]["nvc_frames_seen"] == 2
class _FakeTTSModelOutput:
@@ -342,7 +379,14 @@ def _bare_talker():
return talker
-def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id: str = "req-0") -> dict:
+def _async_info(
+ timeline: list[int],
+ prompt_len: int,
+ finished: bool,
+ request_id: str = "req-0",
+ *,
+ codec_streaming: bool = False,
+) -> dict:
return {
"request_id": request_id,
"additional_information": {
@@ -351,6 +395,7 @@ def _async_info(timeline: list[int], prompt_len: int, finished: bool, request_id
"num_processed_tokens": prompt_len,
"finished": torch.tensor(finished),
"request_id": request_id,
+ "codec_streaming": codec_streaming,
},
},
}
@@ -369,11 +414,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
session = talker._sessions["req-0"]
assert session["step"] == 1 and not session["sync_mode"]
- # First decode wake: drains ALL prompt-region steps (t=1..P-1) at once.
- _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+ # First resumed wake: drains ALL prompt-region steps (t=1..P-1) at once.
+ # vLLM 0.28 reports each resumed one-token segment as prefill; the model
+ # session must still advance instead of restarting its EAR-TTS warmup.
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=True, **info)
assert session["step"] == prompt_len # t advanced 1 -> P in one wake
assert update["codes"]["audio"].shape == (prompt_len - 1, 31)
- assert float(embeds[0, 0]) == 0.0 # not finished
+ assert float(embeds[0, 0]) == 0.0 # offline stream waits for upstream terminal
# Coalesced chunk: TWO new frame tokens arrive in ONE wake (delayed save
# thread merged two thinker steps). Both must be consumed by this wake.
@@ -383,14 +430,13 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert update["codes"]["audio"].shape == (prompt_len + 1, 31)
assert float(embeds[0, 0]) == 0.0
- # Zero-progress wake (no new positions, not finished): plain CONTINUE, no
- # fabricated frames, no codes update.
+ # Zero-progress wake: never fabricate frames or end the offline stream
+ # before the thinker sends its terminal marker.
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info2)
assert session["step"] == prompt_len + 2
assert update == {}
assert float(embeds[0, 0]) == 0.0
- # Final chunk: one more token + finished marker -> drain + stop flag.
info3 = _async_info([pad] * prompt_len + [7, 8, 9], prompt_len, finished=True)
_, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info3)
assert session["step"] == prompt_len + 3
@@ -398,6 +444,18 @@ def test_talker_drains_all_received_positions_per_wake() -> None:
assert float(embeds[0, 0]) == 1.0 # stop
+def test_talker_native_codec_stream_ends_each_received_segment() -> None:
+ talker = _bare_talker()
+ ids = torch.zeros(1, dtype=torch.long)
+ info = _async_info([12, 12, 7], 2, finished=False, codec_streaming=True)
+ talker.preprocess(ids, None, _omni_is_prefill=True, **info)
+
+ _, embeds, update = talker.preprocess(ids, None, _omni_is_prefill=False, **info)
+
+ assert update["codes"]["audio"].shape == (2, 31)
+ assert float(embeds[0, 0]) == 1.0
+
+
def test_talker_sync_mode_steps_once_per_wake() -> None:
"""The verified sync parity path keeps its one-TTS-step-per-engine-step cadence."""
talker = _bare_talker()
diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
index 013c0b99e..996486d49 100644
--- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
+++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
@@ -8,9 +8,12 @@ import pytest
import torch
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
+ AdaptiveChunkController,
+ compute_adaptive_emit,
compute_dynamic_initial_chunk_size,
compute_ramp_emit,
max_ic_for_chunk_size,
+ parse_adaptive_config,
parse_chunk_ramp,
ramp_chunk_size,
ramp_cumulative,
@@ -41,7 +44,17 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None):
)
-def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frames=0, chunk_ramp=None):
+def _tm(
+ *,
+ chunk_frames=25,
+ left_context=25,
+ max_num_seqs=1,
+ initial_chunk_frames=0,
+ chunk_ramp=None,
+ adaptive=False,
+ adaptive_min=2,
+ adaptive_margin=200.0,
+):
extra = {
"codec_chunk_frames": chunk_frames,
"codec_left_context_frames": left_context,
@@ -49,11 +62,16 @@ def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1, initial_chunk_frame
}
if chunk_ramp is not None:
extra["codec_chunk_ramp"] = chunk_ramp
+ if adaptive:
+ extra["codec_chunk_adaptive"] = True
+ extra["codec_chunk_min_frames"] = adaptive_min
+ extra["codec_chunk_safety_margin_ms"] = adaptive_margin
return SimpleNamespace(
code_prompt_token_ids=defaultdict(list),
scheduler_max_num_seqs=max_num_seqs,
put_req_chunk=defaultdict(int),
ramp_chunk_count=defaultdict(int),
+ _adaptive_states={},
request_payload={},
connector=SimpleNamespace(config={"extra": extra}),
)
@@ -605,18 +623,32 @@ class TestRampHelpers:
def test_parse_ramp_warns_on_tail_mismatch(self, caplog):
import logging
- with caplog.at_level(logging.WARNING):
+ target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
+ target_logger.addHandler(caplog.handler)
+ prev_level = target_logger.level
+ target_logger.setLevel(logging.WARNING)
+ try:
result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25)
+ finally:
+ target_logger.removeHandler(caplog.handler)
+ target_logger.setLevel(prev_level)
assert result == [4, 4, 8]
- assert any("reintroduces" in r.message for r in caplog.records)
+ assert "reintroduces" in caplog.text
def test_parse_ramp_no_warn_on_tail_match(self, caplog):
import logging
- with caplog.at_level(logging.WARNING):
+ target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
+ target_logger.addHandler(caplog.handler)
+ prev_level = target_logger.level
+ target_logger.setLevel(logging.WARNING)
+ try:
result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25)
+ finally:
+ target_logger.removeHandler(caplog.handler)
+ target_logger.setLevel(prev_level)
assert result == [4, 4, 8, 16, 25]
- assert not any("reintroduces" in r.message for r in caplog.records)
+ assert "reintroduces" not in caplog.text
@pytest.mark.parametrize(
"index,ramp,steady,expected",
@@ -941,3 +973,488 @@ class TestChunkRampEmission:
p_seg2_0 = self._emit(tm, rid, 4)
assert p_seg2_0 is not None
assert p_seg2_0.meta.left_context_size == 0
+
+
+class TestAdaptiveController:
+ def test_compute_next_chunk_size_basic(self):
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=20.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=10,
+ )
+ now = 0.5
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=100.0)
+ assert 2 <= size <= 25
+
+ def test_compute_next_chunk_size_clamp_min(self):
+ """min_frames floor holds when ramp_floor (last_target + delta) is
+ below min_frames. Uses a small EWMA so delta=2 and last_target=0,
+ i.e. ramp_floor=2 < min_frames=4 -> greedy=min_frames wins."""
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=10.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=1,
+ )
+ # buffer = 80 - 70 = 10ms (positive but < margin=200)
+ # greedy: available=0, max_n=0, greedy=min_frames=4
+ # ramp_floor: last_target=0 + delta=max(2,int(10/15))=2 -> 2 < 4
+ now = 0.07
+ size = ctrl.compute_next_chunk_size(now, min_frames=4, max_frames=25, safety_margin_ms=200.0)
+ assert size == 4 # min_frames floor dominates ramp_floor=2
+
+ def test_compute_next_chunk_size_clamp_max(self):
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=10.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=100,
+ )
+ now = 1.0
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert size == 25
+
+ def test_hysteresis_locks_steady(self):
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=10.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=50,
+ )
+ now = 0.1
+ size1 = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert ctrl.steady_state_reached
+ assert size1 == 25
+ size2 = ctrl.compute_next_chunk_size(now + 10.0, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert size2 == 25
+
+ def test_ramp_floor_overrides_greedy_when_buffer_low(self):
+ """When buffer is negative (overloaded), ramp_floor still forces
+ gradual growth to avoid stagnating at min_frames (which would cause
+ too many Code2Wav decode passes and RTF explosion)."""
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=50.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=2,
+ last_target=4,
+ )
+ # buffer = 160 - 200 = -40ms (negative)
+ # greedy: available=0, max_n=0, greedy=2
+ # ramp_floor: last_target=4 + delta=max(2,int(50/15))=3 -> 7
+ now = 0.2
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert not ctrl.steady_state_reached
+ assert size == 7 # ramp_floor (unconditional) dominates greedy=2
+
+ def test_ramp_floor_dominates_when_last_target_set(self):
+ """Once last_target is set (post-emit), the unconditional ramp_floor
+ (last_target + ewma-scaled delta) dominates greedy whenever greedy is
+ smaller. This is intended: chunk size grows monotonically toward max
+ to minimize Code2Wav decode passes (pre-W5; see W5 coupling)."""
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=50.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=4,
+ last_target=4,
+ )
+ # buffer = 320 - 100 = 220ms (positive)
+ # greedy: available=170, max_n=int(170/50)=3, greedy=3
+ # ramp_floor: 4 + delta=max(2,int(50/15))=3 -> 7
+ now = 0.1
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert size == 7 # ramp_floor (7) dominates greedy (3)
+
+ def test_greedy_dominates_when_ramp_floor_low(self):
+ """When last_target=0 (fresh, no prior emit) and buffer is positive
+ but below the hysteresis threshold, greedy selects an intermediate
+ size larger than ramp_floor and greedy dominates. This is the only
+ regime where greedy (not ramp_floor or hysteresis) sets the size."""
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=50.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=8,
+ last_target=0,
+ )
+ # buffer = 640 - 100 = 540ms (positive, < hysteresis 1300ms)
+ # greedy: available=490, max_n=int(490/50)=9, greedy=9
+ # ramp_floor: 0 + delta=max(2,int(50/15))=3 -> 3
+ now = 0.1
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert not ctrl.steady_state_reached
+ assert size == 9 # greedy (9) dominates ramp_floor (3)
+
+ def test_record_emit_ewma_update(self):
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=20.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=4,
+ )
+ ctrl.record_emit(now=0.5, n_frames=8, accumulated=12, target_size=8)
+ assert ctrl.frames_emitted_this_segment == 12
+ assert ctrl.accumulated_at_last_emit == 12
+ assert ctrl.last_emit_time == 0.5
+ assert ctrl.last_target == 8 # D6: stores target, not actual
+ measured = (0.5 - 0.0) * 1000.0 / 8
+ expected = 0.3 * measured + 0.7 * 20.0
+ assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01
+
+ def test_record_emit_first_then_second(self):
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=0.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=0,
+ )
+ ctrl.record_emit(now=0.4, n_frames=4, accumulated=4, target_size=4)
+ assert ctrl.frame_time_ewma_ms == 0.0
+ assert ctrl.frames_emitted_this_segment == 4
+ assert ctrl.last_target == 4
+
+ ctrl.record_emit(now=0.8, n_frames=4, accumulated=8, target_size=4)
+ expected = (0.8 - 0.4) * 1000.0 / 4
+ assert abs(ctrl.frame_time_ewma_ms - expected) < 0.01
+
+ def test_record_emit_collects_underrun(self):
+ """record_emit mirrors stream_client.py's playback-cursor underrun
+ model: per_chunk_gap_ms records each chunk's lateness, total_gap_s
+ accumulates gaps > 5ms. Deterministic given explicit `now`."""
+ ctrl = AdaptiveChunkController(
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=0,
+ )
+ # chunk0: now=0.0 (= first_emit_time), n=2 -> 0.16s audio
+ ctrl.record_emit(now=0.0, n_frames=2, accumulated=2, target_size=2)
+ assert ctrl.chunk_sizes == [2]
+ assert ctrl.per_chunk_gap_ms == [0.0]
+ assert ctrl.total_gap_s == 0.0
+ assert abs(ctrl.playback_cursor_s - 0.16) < 1e-9
+ # chunk1: now=0.2 (200ms wall) -> arrival 0.2, late by 0.2-0.16=0.04s
+ ctrl.record_emit(now=0.2, n_frames=2, accumulated=4, target_size=2)
+ assert ctrl.chunk_sizes == [2, 2]
+ assert abs(ctrl.per_chunk_gap_ms[1] - 40.0) < 1e-6
+ assert abs(ctrl.total_gap_s - 0.04) < 1e-9
+ assert abs(ctrl.playback_cursor_s - 0.36) < 1e-9
+
+ def test_ewma_zero_skips_hysteresis(self):
+ """When EWMA=0, hysteresis must not trigger (steady_gen_ms=0
+ would make any positive buffer satisfy the condition)."""
+ ctrl = AdaptiveChunkController(
+ frame_time_ewma_ms=0.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=10,
+ )
+ now = 0.001
+ size = ctrl.compute_next_chunk_size(now, min_frames=2, max_frames=25, safety_margin_ms=50.0)
+ assert not ctrl.steady_state_reached
+ assert size == 2 # min_frames (EWMA=0 → max_n=min_frames)
+
+ def test_ramp_divisor_configurable(self):
+ """ramp_divisor/ramp_delta_min are controller fields wired from config.
+ Defaults (15/2) preserve the original formula; changing them changes
+ delta. Lets a deploy tune the ramp via yaml without code changes, while
+ the unit tests (which use defaults) stay green."""
+ # default divisor=15: ewma=50 -> delta=max(2,int(50/15))=3 -> floor=4+3=7
+ c1 = AdaptiveChunkController(
+ frame_time_ewma_ms=50.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=2,
+ last_target=4,
+ )
+ assert c1.compute_next_chunk_size(0.2, 2, 25, 50.0) == 7
+ # divisor=10: delta=max(2,int(50/10))=5 -> floor=4+5=9
+ c2 = AdaptiveChunkController(
+ frame_time_ewma_ms=50.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=2,
+ last_target=4,
+ ramp_divisor=10,
+ )
+ assert c2.compute_next_chunk_size(0.2, 2, 25, 50.0) == 9
+ # delta_min=4 + ewma=10 (int(10/15)=0): delta=max(4,0)=4 -> floor=4+4=8
+ c3 = AdaptiveChunkController(
+ frame_time_ewma_ms=10.0,
+ first_emit_time=0.0,
+ last_emit_time=0.0,
+ frames_emitted_this_segment=2,
+ last_target=4,
+ ramp_delta_min=4,
+ )
+ assert c3.compute_next_chunk_size(0.2, 2, 25, 50.0) == 8
+
+
+class TestAdaptiveEmitHelper:
+ @pytest.mark.parametrize(
+ "length,accumulated,target,finished,expected_emit,expected_ctx",
+ [
+ (4, 0, 4, False, True, 4),
+ (3, 0, 4, False, False, 0),
+ (12, 4, 8, False, True, 8),
+ (10, 4, 8, False, False, 0),
+ # D3: finished flushes ALL remaining frames, not just target_size
+ (6, 4, 8, True, True, 2),
+ (4, 4, 8, True, True, 0),
+ (20, 4, 8, True, True, 16), # new_frames=16 > target=8 → flush 16
+ (30, 0, 25, True, True, 30), # finished, new_frames=30 > target=25 → flush 30
+ # B1: non-finished path also ships ALL accumulated frames (new_frames),
+ # not just target_size. When new_frames > target_size (e.g. scheduler
+ # stall dropped target below queued frames), returning target_size
+ # would cause the [-context_length:] slice to skip surplus frames.
+ (10, 0, 4, False, True, 10), # new_frames=10 > target=4 → ship 10
+ (8, 2, 4, False, True, 6), # new_frames=6 > target=4 → ship 6
+ (100, 0, 25, False, True, 100), # new_frames=100 > target=25 → ship 100
+ ],
+ )
+ def test_compute_adaptive_emit(self, length, accumulated, target, finished, expected_emit, expected_ctx):
+ emit, ctx = compute_adaptive_emit(length, accumulated, target, finished)
+ assert emit == expected_emit
+ assert ctx == expected_ctx
+
+
+class TestAdaptiveConfigParsing:
+ def test_disabled_by_default(self):
+ enabled, min_f, margin, divisor, delta_min = parse_adaptive_config({})
+ assert enabled is False
+ assert min_f == 2
+ assert margin == 50.0
+ assert divisor == 15
+ assert delta_min == 2
+
+ def test_enabled(self):
+ cfg = {"codec_chunk_adaptive": True}
+ enabled, *_ = parse_adaptive_config(cfg)
+ assert enabled is True
+
+ def test_custom_values(self):
+ cfg = {
+ "codec_chunk_adaptive": True,
+ "codec_chunk_min_frames": 4,
+ "codec_chunk_safety_margin_ms": 300.0,
+ "codec_chunk_ramp_divisor": 10,
+ "codec_chunk_ramp_delta_min": 3,
+ }
+ enabled, min_f, margin, divisor, delta_min = parse_adaptive_config(cfg)
+ assert enabled is True
+ assert min_f == 4
+ assert margin == 300.0
+ assert divisor == 10
+ assert delta_min == 3
+
+ def test_min_frames_clamped(self):
+ cfg = {"codec_chunk_adaptive": True, "codec_chunk_min_frames": 0}
+ _, min_f, _, _, _ = parse_adaptive_config(cfg)
+ assert min_f == 1
+
+ def test_ramp_divisor_clamped(self):
+ cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_divisor": 0}
+ _, _, _, divisor, _ = parse_adaptive_config(cfg)
+ assert divisor == 1
+
+ def test_ramp_delta_min_clamped(self):
+ cfg = {"codec_chunk_adaptive": True, "codec_chunk_ramp_delta_min": 0}
+ _, _, _, _, delta_min = parse_adaptive_config(cfg)
+ assert delta_min == 1
+
+
+class TestAdaptiveEmission:
+ def _emit(self, tm, rid, n_frames, finished=False):
+ tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
+ return talker2code2wav_async_chunk(
+ transfer_manager=tm,
+ multimodal_output={"codes": {"audio": torch.zeros((0,))}},
+ request=_req(rid, finished=finished),
+ is_finished=finished,
+ )
+
+ def test_adaptive_chunk0_uses_dynamic_ic(self):
+ tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-ic"
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ assert len(p0.codes.audio) == _Q * 2
+
+ def test_adaptive_chunk0_holds_until_ic(self):
+ tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-hold"
+ p = self._emit(tm, rid, 1)
+ assert p is None
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ assert len(p0.codes.audio) == _Q * 2
+
+ def test_adaptive_finished_flush(self):
+ tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-flush"
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ tm.ramp_chunk_count[rid] = 1
+
+ p_fin = self._emit(tm, rid, 5, finished=True)
+ assert p_fin is not None
+ assert p_fin.meta.finished.item() is True
+
+ def test_adaptive_finished_no_new_frames(self):
+ tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-no-new"
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ tm.ramp_chunk_count[rid] = 1
+
+ p_fin = self._emit(tm, rid, 2, finished=True)
+ assert p_fin is not None
+ assert p_fin.meta.finished.item() is True
+ assert p_fin.codes.audio.numel() == 0
+
+ def test_adaptive_precedence_over_fixed_ramp(self):
+ tm = _tm(chunk_ramp=[4, 4, 8, 16, 25], adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-precedence"
+
+ p = self._emit(tm, rid, 1)
+ assert p is None
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ assert len(p0.codes.audio) == _Q * 2
+
+ def test_adaptive_backward_compat(self):
+ tm = _tm(initial_chunk_frames=1)
+ rid = "no-adaptive"
+
+ p0 = self._emit(tm, rid, 1)
+ assert p0 is not None
+ assert len(p0.codes.audio) == _Q * 1
+
+ def test_adaptive_segment_reset(self):
+ tm = _tm(adaptive=True, adaptive_min=2, max_num_seqs=8)
+ rid = "adaptive-segment"
+
+ p0 = self._emit(tm, rid, 2)
+ assert p0 is not None
+ tm.ramp_chunk_count[rid] = 1
+ tm.put_req_chunk[rid] = 1
+
+ tm.code_prompt_token_ids[rid] = []
+ tm.ramp_chunk_count.pop(rid, None)
+ tm._adaptive_states.pop(rid, None)
+
+ p_seg2 = self._emit(tm, rid, 2)
+ assert p_seg2 is not None
+ assert p_seg2.meta.left_context_size == 0
+
+ def test_adaptive_with_ref_code_first_chunk(self):
+ tm = _tm(adaptive=True, adaptive_min=2, left_context=25, max_num_seqs=8)
+ rid = "adaptive-ref"
+ tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(2)]
+ ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+
+ payload = talker2code2wav_async_chunk(
+ transfer_manager=tm,
+ multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+ request=_req(rid, finished=False),
+ is_finished=False,
+ )
+
+ assert payload is not None
+ assert payload.meta.ref_context_size == 2
+ assert payload.meta.ref_context_included is True
+ assert payload.meta.left_context_size == 2
+
+
+class TestAdaptiveAccumulationPath:
+ """Regression test for the frame-skipping bug (PR #6001 review B1).
+
+ Appends one frame per call with a controlled clock, asserting that
+ the concatenation of every emitted chunk equals the accumulated
+ frame sequence, in order, with no gaps. Each frame has a distinct
+ value so the assertion can detect a skipped frame rather than just
+ a short total.
+ """
+
+ def test_no_frames_skipped_when_target_drops(self, monkeypatch):
+ """When greedy target_size is high (large buffer) and then drops
+ (clock jumps), new_frames > target_size at the emit point.
+ compute_adaptive_emit now returns new_frames (not target_size),
+ so all accumulated frames are shipped and the [-context_length:]
+ slice stays contiguous with the previous emit. last_target is set
+ to target_size (not n_frames) to prevent the ramp floor from
+ ratcheting up on transient overshoot.
+
+ Clock schedule (max_num_seqs=4 gives IC=2 via dynamic IC):
+ t=0.0 chunk 0 (IC=2), emit 2. ewma=0.
+ t=0.05 emit target=4 (ramp_floor). ewma=12.5 (direct assign).
+ t=0.05 greedy dominates (target=18), 17 frames accumulate (holds).
+ t=100.0 clock jumps -> greedy collapses, target drops to
+ ramp_floor=6. new_frames=18 > 6 -> emit 18 (all new frames).
+ last_target=6 (not 18, preventing floor ratchet).
+ """
+ tm = _tm(
+ adaptive=True,
+ adaptive_min=2,
+ max_num_seqs=4,
+ )
+ rid = "adaptive-accum"
+ tm.code_prompt_token_ids[rid] = []
+
+ clock = [0.0]
+ monkeypatch.setattr(
+ "vllm_omni.model_executor.stage_input_processors.qwen3_tts.time.monotonic",
+ lambda: clock[0],
+ )
+
+ emitted_indices: list[int] = []
+ total_frames = 40
+
+ def append_and_emit(frame_idx, finished=False):
+ tm.code_prompt_token_ids[rid].append([frame_idx] * _Q)
+ p = talker2code2wav_async_chunk(
+ transfer_manager=tm,
+ multimodal_output={"codes": {"audio": torch.zeros((0,))}},
+ request=_req(rid, finished=finished),
+ is_finished=finished,
+ )
+ if p is not None:
+ if p.codes.audio.numel() > 0:
+ n = p.codes.audio.numel() // _Q
+ emitted_indices.extend(p.codes.audio[:n].tolist())
+ tm.ramp_chunk_count[rid] += 1
+ return p
+
+ # Phase 1 (t=0.0): chunk 0, IC=2
+ for i in range(2):
+ append_and_emit(i)
+
+ # Phase 2 (t=0.05): build EWMA, emit target=4 at length=6
+ clock[0] = 0.05
+ for i in range(2, 6):
+ append_and_emit(i)
+
+ # Phase 3 (t=0.05): greedy target=18 (large buffer), accumulate
+ # 17 frames while holding (new_frames 1..17 < 18)
+ for i in range(6, 23):
+ append_and_emit(i)
+
+ # Phase 4 (t=100.0): clock jumps, greedy collapses, target drops
+ # to ramp_floor=6. new_frames=18 > 6 -> bug trigger point.
+ clock[0] = 100.0
+ for i in range(23, total_frames):
+ append_and_emit(i, finished=(i == total_frames - 1))
+
+ # Assert: all frames 0..39 emitted, in order, no gaps.
+ expected = list(range(total_frames))
+ assert emitted_indices == expected, (
+ f"Frame loss detected! "
+ f"Missing: {sorted(set(expected) - set(emitted_indices))}, "
+ f"Extra: {sorted(set(emitted_indices) - set(expected))}"
+ )
diff --git a/tests/model_executor/test_cfg_pairing.py b/tests/model_executor/test_cfg_pairing.py
new file mode 100644
index 000000000..c00313a2a
--- /dev/null
+++ b/tests/model_executor/test_cfg_pairing.py
@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the model-neutral CFG request-pairing scheduler patches.
+
+The patches keep a cond/uncond request pair step-locked inside one engine:
+admitted together, kept at equal ``num_computed_tokens``, finished together.
+They are installed onto a lightweight fake scheduler here; vLLM's real
+``Scheduler`` needs a full engine (and a GPU) to construct.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm_omni.model_executor.models.common import cfg_pairing
+from vllm_omni.model_executor.models.common.cfg_pairing import (
+ _equalize_cfg_pair_progress,
+ _patch_scheduler,
+ cfg_uncond_suffixes,
+ register_cfg_uncond_suffix,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ───────── fakes ─────────
+
+
+class _FakeRequest:
+ """Minimal stand-in for ``vllm.v1.request.Request`` (identity equality)."""
+
+ def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None:
+ self.request_id = request_id
+ self.sampling_params = SimpleNamespace(extra_args=extra_args)
+ self.num_computed_tokens = 0
+ self.step_tokens = step_tokens
+ self.skip_reading_prefix_cache = False
+
+
+class _FakeWaitingQueue(deque):
+ """FCFS waiting queue: vLLM's is also a ``deque`` subclass."""
+
+ def remove_requests(self, requests: list[_FakeRequest]) -> None:
+ keep = [req for req in self if not any(req is dropped for dropped in requests)]
+ self.clear()
+ self.extend(keep)
+
+ def prepend_request(self, request: _FakeRequest) -> None:
+ self.appendleft(request)
+
+
+@dataclass
+class _FakeSchedulerOutput:
+ num_scheduled_tokens: dict[str, int] = field(default_factory=dict)
+ total_num_scheduled_tokens: int = 0
+
+
+def _make_scheduler_cls() -> type:
+ """Fresh class per test so the in-place patches never stack."""
+
+ class FakeScheduler:
+ def __init__(self) -> None:
+ self.waiting = _FakeWaitingQueue()
+ self.running: list[_FakeRequest] = []
+ self.requests: dict[str, _FakeRequest] = {}
+ self.finished: list[set[str]] = []
+ self.observed_waiting: list[list[str]] = []
+ self.max_num_scheduled_tokens = 512
+ self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096)
+
+ def add_request(self, request: _FakeRequest) -> None:
+ self.requests[request.request_id] = request
+ self.waiting.append(request)
+
+ def finish_requests(self, request_ids: Any, finished_status: Any) -> None:
+ if request_ids is None:
+ ids = set(self.requests)
+ elif isinstance(request_ids, str):
+ ids = {request_ids}
+ else:
+ ids = set(request_ids)
+ self.finished.append(ids)
+ self.running = [req for req in self.running if req.request_id not in ids]
+ for req_id in ids:
+ self.requests.pop(req_id, None)
+
+ def schedule(self) -> _FakeSchedulerOutput:
+ self.observed_waiting.append([req.request_id for req in self.waiting])
+ while self.waiting:
+ self.running.append(self.waiting.popleft())
+ output = _FakeSchedulerOutput()
+ for req in self.running:
+ req.num_computed_tokens += req.step_tokens
+ output.num_scheduled_tokens[req.request_id] = req.step_tokens
+ output.total_num_scheduled_tokens += req.step_tokens
+ return output
+
+ return FakeScheduler
+
+
+@pytest.fixture
+def scheduler() -> Any:
+ cls = _make_scheduler_cls()
+ _patch_scheduler(cls)
+ return cls()
+
+
+@pytest.fixture(autouse=True)
+def _restore_suffix_registry():
+ original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES)
+ yield
+ cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear()
+ cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original)
+
+
+def _add_pair(
+ scheduler: Any,
+ pair_id: str | None = None,
+ base_id: str = "req-1",
+ suffix: str = "__cfg_uncond",
+) -> tuple[_FakeRequest, _FakeRequest]:
+ cond_extra: dict[str, Any] = {"cfg_role": "cond"}
+ uncond_extra: dict[str, Any] = {"cfg_role": "uncond"}
+ if pair_id is not None:
+ cond_extra["cfg_pair_id"] = pair_id
+ uncond_extra["cfg_pair_id"] = pair_id
+ cond = _FakeRequest(base_id, cond_extra)
+ uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra)
+ scheduler.add_request(cond)
+ scheduler.add_request(uncond)
+ return cond, uncond
+
+
+# ───────── pair-id derivation ─────────
+
+
+def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None:
+ cond, uncond = _add_pair(scheduler)
+
+ assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}
+ assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"}
+
+
+def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None:
+ register_cfg_uncond_suffix("::negative")
+ assert "::negative" in cfg_uncond_suffixes()
+
+ cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative")
+
+ assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}}
+
+
+def test_register_rejects_empty_suffix() -> None:
+ with pytest.raises(ValueError):
+ register_cfg_uncond_suffix("")
+
+
+def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None:
+ # Degenerate input: the companion never joins the cond pair, so the pair
+ # stays incomplete and the hold budget eventually releases it.
+ scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"}))
+
+ assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}}
+
+
+def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None:
+ cond, uncond = _add_pair(scheduler, pair_id="pair-A")
+
+ assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}}
+ assert cond.request_id not in scheduler._cfg_pairs
+ assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A"
+
+
+def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None:
+ cond, uncond = _add_pair(scheduler)
+
+ assert cond.skip_reading_prefix_cache is True
+ assert uncond.skip_reading_prefix_cache is True
+
+
+# ───────── admission holding ─────────
+
+
+def test_lone_member_is_held_then_released(scheduler: Any) -> None:
+ scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+ scheduler.schedule()
+
+ assert scheduler.observed_waiting == [[]] # the base scheduler saw nothing
+ assert [req.request_id for req in scheduler.waiting] == ["req-1"]
+ assert scheduler.running == []
+
+
+def test_complete_pair_is_admitted_together(scheduler: Any) -> None:
+ scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+ scheduler.schedule()
+ scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+ scheduler.schedule()
+
+ assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"]
+ assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"]
+ assert list(scheduler.waiting) == []
+
+
+def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2)
+ scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+
+ scheduler.schedule()
+ scheduler.schedule()
+ assert scheduler.running == []
+ assert scheduler._cfg_pairs # still hoping for the partner
+
+ scheduler.schedule()
+
+ assert scheduler._cfg_pairs == {}
+ assert scheduler._cfg_req_to_pair == {}
+ assert [req.request_id for req in scheduler.running] == ["req-1"]
+
+
+def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None:
+ scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"}))
+ scheduler.add_request(_FakeRequest("other", None))
+ scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}))
+
+ scheduler.schedule()
+
+ assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"]
+
+
+# ───────── progress equalization ─────────
+
+
+def _equalize_fixture(
+ cond_computed: int,
+ uncond_computed: int,
+ cond_sched: int,
+ uncond_sched: int,
+) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]:
+ cond = _FakeRequest("req-1", {"cfg_role": "cond"})
+ uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})
+ cond.num_computed_tokens = cond_computed
+ uncond.num_computed_tokens = uncond_computed
+ fake = SimpleNamespace(
+ _cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}},
+ requests={cond.request_id: cond, uncond.request_id: uncond},
+ )
+ output = _FakeSchedulerOutput(
+ num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched},
+ total_num_scheduled_tokens=cond_sched + uncond_sched,
+ )
+ return fake, output, cond, uncond
+
+
+def test_equalize_pulls_the_ahead_member_back() -> None:
+ fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4)
+
+ _equalize_cfg_pair_progress(fake, output)
+
+ assert cond.num_computed_tokens == 6
+ assert uncond.num_computed_tokens == 6
+ assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4}
+ assert output.total_num_scheduled_tokens == 5
+
+
+def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None:
+ fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4)
+
+ _equalize_cfg_pair_progress(fake, output)
+
+ assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8)
+ assert output.total_num_scheduled_tokens == 8
+
+
+def test_equalize_skips_an_infeasible_rollback() -> None:
+ # Clawing back 4 tokens would leave the cond row with 0 scheduled tokens,
+ # which the scheduler output cannot express; leave the pair to _drop_split_pairs.
+ fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4)
+
+ _equalize_cfg_pair_progress(fake, output)
+
+ assert cond.num_computed_tokens == 10
+ assert output.total_num_scheduled_tokens == 8
+
+
+# ───────── finish ─────────
+
+
+def test_finish_requests_finishes_the_partner(scheduler: Any) -> None:
+ cond, uncond = _add_pair(scheduler)
+ scheduler.schedule()
+
+ scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED")
+
+ assert scheduler.finished == [{cond.request_id, uncond.request_id}]
+ assert scheduler._cfg_pairs == {}
+ assert scheduler._cfg_req_to_pair == {}
+ assert scheduler.running == []
+
+
+def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None:
+ _add_pair(scheduler)
+ scheduler.schedule()
+
+ scheduler.finish_requests(None, "FINISHED_ABORTED")
+
+ assert scheduler._cfg_pairs == {}
+ assert scheduler._cfg_req_to_pair == {}
+
+
+# ───────── non-CFG traffic ─────────
+
+
+def test_plain_requests_are_untouched(scheduler: Any) -> None:
+ scheduler.add_request(_FakeRequest("plain-1", None))
+ scheduler.add_request(_FakeRequest("plain-2", {"seed": 3}))
+
+ output = scheduler.schedule()
+
+ assert scheduler._cfg_pairs == {}
+ assert scheduler.observed_waiting == [["plain-1", "plain-2"]]
+ assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1}
+ assert output.total_num_scheduled_tokens == 2
+ # No pair ever existed, so the hold bookkeeping is never allocated.
+ assert not hasattr(scheduler, "_cfg_hold_counts")
+
+
+def test_default_suffix_matches_the_audex_stage_processor() -> None:
+ from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX
+
+ assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes()
+
+
+def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None:
+ from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied
+
+ assert apply_cfg_patches is cfg_pairing.apply_cfg_patches
+ assert cfg_patches_applied is cfg_pairing.cfg_patches_applied
+
+
+# ───────── install gate ─────────
+
+
+def test_gate_triggers_on_a_cfg_logits_processor() -> None:
+ from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+ class AudexCFGLogitsProcessor:
+ pass
+
+ assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor]))
+ # vLLM keeps unresolved entries as dotted strings.
+ assert _stage_declares_cfg_pairs(
+ SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"])
+ )
+
+
+def test_gate_triggers_on_a_declared_cfg_role_default() -> None:
+ from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+ assert _stage_declares_cfg_pairs(
+ SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale"))
+ )
+
+
+def test_gate_stays_off_for_non_cfg_stages() -> None:
+ from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs
+
+ assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=()))
+ assert not _stage_declares_cfg_pairs(
+ SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",))
+ )
+ assert not _stage_declares_cfg_pairs(SimpleNamespace())
+
+
+def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None:
+ from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys
+
+ assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == (
+ "cfg_role",
+ "cfg_scale",
+ )
+ assert _sampling_extra_args_keys({"temperature": 0.1}) == ()
+ assert _sampling_extra_args_keys(None) == ()
diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py
index d57671adc..463bedfc3 100644
--- a/tests/model_extras/test_model_extras.py
+++ b/tests/model_extras/test_model_extras.py
@@ -20,6 +20,7 @@ from vllm_omni.model_extras import (
get_output_tensor_range,
get_x_to_text_model_family,
should_init_extra_args_for_non_diffusion_stages,
+ should_preserve_reference_image_size,
)
@@ -150,6 +151,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
"max_sequence_length",
"use_resolution_template",
"use_duration_template",
+ "negative_metadata_mode",
"use_system_prompt",
"system_prompt",
"negative_prompt",
@@ -192,23 +194,6 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_magi_human_extra_registry_declares_request_and_response_params() -> None:
- assert get_extra_body_params("MagiHumanPipeline") == frozenset(
- {
- "seconds",
- "audio_path",
- "image_path",
- "sr_height",
- "sr_width",
- "sr_num_inference_steps",
- }
- )
- assert get_extra_output_params("MagiHumanPipeline") == frozenset()
- assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
-
-
@pytest.mark.core_model
@pytest.mark.cpu
def test_ltx_extra_registry_declares_official_guidance_params() -> None:
@@ -230,13 +215,60 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
"audio_rescale_scale",
"video_stg_blocks",
"audio_stg_blocks",
+ "sigmas",
+ "stage_1_sigmas",
+ "stage_2_sigmas",
+ "image_crf",
}
)
- assert get_extra_body_params("LTX2Pipeline") == expected
- assert get_extra_output_params("LTX2Pipeline") == frozenset()
- assert get_extra_body_params("LTX2DistilledPipeline") == expected
- assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
+ for pipeline_name in (
+ "LTX2Pipeline",
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledOneStagePipeline",
+ "LTX2DistilledTwoStagePipeline",
+ "LTX2DistilledPipeline",
+ ):
+ assert get_extra_body_params(pipeline_name) == expected
+ assert get_extra_output_params(pipeline_name) == frozenset()
+
+
+@pytest.mark.parametrize(
+ ("model_version", "expected"),
+ [("2", False), ("2.3", False), ("2.5", True)],
+)
+def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: bool) -> None:
+ (tmp_path / "model_index.json").write_text(
+ '{"_class_name": "LTX2Pipeline", "model_version": "' + model_version + '"}'
+ )
+
+ assert (
+ should_preserve_reference_image_size(
+ None,
+ model=str(tmp_path),
+ )
+ is expected
+ )
+
+
+def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None:
+ captured = {}
+
+ def fake_policy(*, model, revision=None):
+ captured.update(model=model, revision=revision)
+ return True
+
+ monkeypatch.setattr(
+ "vllm_omni.diffusion.models.ltx2.ltx2_components.preserves_reference_image_size",
+ fake_policy,
+ )
+
+ assert should_preserve_reference_image_size(
+ "LTX2Pipeline",
+ model="org/model",
+ revision="pinned-revision",
+ )
+ assert captured == {"model": "org/model", "revision": "pinned-revision"}
@pytest.mark.core_model
@@ -337,47 +369,6 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
assert "negative_prompt" not in result
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_audiox_extra_registry_declares_request_and_response_params() -> None:
- assert get_extra_body_params("AudioXPipeline") == frozenset(
- {
- "audiox_task",
- "seconds_start",
- "seconds_total",
- "sigma_min",
- "sigma_max",
- "cfg_rescale",
- "video_path",
- "audio_path",
- }
- )
- assert get_extra_output_params("AudioXPipeline") == frozenset({"audiox_task"})
- assert should_init_extra_args_for_non_diffusion_stages("AudioXPipeline") is False
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_audiox_declared_extra_args_route_into_sampling_params() -> None:
- params = OmniDiffusionSamplingParams()
- declared = get_extra_body_params("AudioXPipeline")
- apply_declared_extra_args(
- params,
- declared,
- {
- "audiox_task": "t2a",
- "seconds_total": 10.0,
- "sigma_min": 0.03,
- "unknown": "ignored",
- },
- )
- assert params.extra_args == {
- "audiox_task": "t2a",
- "seconds_total": 10.0,
- "sigma_min": 0.03,
- }
-
-
@pytest.mark.core_model
@pytest.mark.cpu
def test_helios_extra_registry_declares_request_and_response_params() -> None:
@@ -421,6 +412,13 @@ def test_vace_extra_registry_has_no_pipeline_params() -> None:
def test_unknown_pipeline_has_empty_extra_registry() -> None:
assert get_extra_body_params("UnknownPipeline") == frozenset()
assert get_extra_output_params("UnknownPipeline") == frozenset()
+ assert (
+ should_preserve_reference_image_size(
+ "UnknownPipeline",
+ model="org/model",
+ )
+ is False
+ )
assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False
diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py
index ab71a02cd..7760fa25d 100644
--- a/tests/model_tests/conftest.py
+++ b/tests/model_tests/conftest.py
@@ -1,12 +1,14 @@
+import json
from shutil import rmtree
import pytest
+from filelock import ReadWriteLock
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
@pytest.fixture(scope="session")
-def tiny_model_paths(request, run_level):
+def tiny_model_paths(request, run_level, tmp_path_factory):
"""Build or download the tiny models for the selected tests.
At core_model level, builds tiny models via the builder function.
@@ -15,27 +17,43 @@ def tiny_model_paths(request, run_level):
NOTE: this is session scoped to avoid churn in tiny model creation,
but will ensure all the tiny models you need are created for the selected tests
before it starts to execute them."""
+ base_path = tmp_path_factory.getbasetemp().parent / "tiny_models"
+ use_lock = ReadWriteLock(base_path.with_suffix(".use.lock"))
+ build_lock = ReadWriteLock(base_path.with_suffix(".build.lock"))
+ data_file = base_path.with_suffix(".json")
+
model_paths = {}
built_paths = []
- print("Initializing tiny models...")
- for item in request.session.items:
- if not hasattr(item, "callspec"):
- continue
- model_name = item.callspec.params.get("model_name")
- if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
- continue
- if model_name not in model_paths:
- settings = DIFFUSION_TEST_SETTINGS[model_name]
- if run_level == "core_model":
- print(f"Calling tiny model builder for: {model_name}")
- path = settings.builder()
- built_paths.append(path)
- else:
- print(f"Run level is {run_level}; {model_name} will use the full model")
- path = settings.model
- model_paths[model_name] = path
- yield model_paths
- for path in built_paths:
- print(f"Removing tiny model: {path}")
- rmtree(path, ignore_errors=True)
+ with use_lock.read_lock():
+ with build_lock.write_lock():
+ if data_file.is_file():
+ model_paths = json.loads(data_file.read_text())
+ else:
+ print("Initializing tiny models...")
+ for item in request.session.items:
+ if not hasattr(item, "callspec"):
+ continue
+ model_name = item.callspec.params.get("model_name")
+ if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
+ continue
+ if model_name not in model_paths:
+ settings = DIFFUSION_TEST_SETTINGS[model_name]
+ if run_level == "core_model":
+ print(f"Calling tiny model builder for: {model_name}")
+ path = settings.builder()
+ built_paths.append(path)
+ else:
+ print(f"Run level is {run_level}; {model_name} will use the full model")
+ path = settings.model
+ model_paths[model_name] = path
+ data_file.write_text(json.dumps(model_paths))
+
+ yield model_paths
+
+ with use_lock.write_lock():
+ if built_paths:
+ data_file.unlink(missing_ok=True)
+ for path in built_paths:
+ print(f"Removing tiny model: {path}")
+ rmtree(path, ignore_errors=True)
diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py
index e47ea1895..7c1c9df67 100644
--- a/tests/model_tests/diffusion/diff_model_builders.py
+++ b/tests/model_tests/diffusion/diff_model_builders.py
@@ -1,7 +1,16 @@
from functools import partial
from pathlib import Path
-from tests.helpers.tiny_model import build_tiny_from_configs
+import torch
+from diffusers import (
+ AutoencoderKLWan,
+ DPMSolverMultistepScheduler,
+ SanaVideoPipeline,
+ SanaVideoTransformer3DModel,
+)
+from transformers import Gemma2Config, Gemma2Model, GemmaTokenizerFast
+
+from tests.helpers.tiny_model import _get_tiny_model_path, build_tiny_from_configs
TINY_CONFIGS_DIR = Path(__file__).parent / "tiny_configs"
@@ -40,6 +49,72 @@ def tiny_ltx2_builder() -> str:
return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
+def tiny_sana_video_builder() -> str:
+ """Build a tiny 480p SANA-Video model without downloading model weights.
+
+ The tokenizer is the only component loaded from the upstream repository.
+ All modules with weights are initialized locally from intentionally small
+ configs, and the scheduler is constructed from its weight-free config.
+ """
+ model_id = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
+ model_dir = _get_tiny_model_path("SanaVideoPipeline")
+
+ tokenizer = GemmaTokenizerFast.from_pretrained(model_id, subfolder="tokenizer")
+ scheduler = DPMSolverMultistepScheduler(
+ algorithm_type="dpmsolver++",
+ flow_shift=8.0,
+ prediction_type="flow_prediction",
+ use_flow_sigmas=True,
+ )
+ text_encoder = Gemma2Model(
+ Gemma2Config(
+ vocab_size=256000,
+ hidden_size=32,
+ intermediate_size=64,
+ num_hidden_layers=1,
+ num_attention_heads=4,
+ num_key_value_heads=2,
+ head_dim=8,
+ max_position_embeddings=512,
+ sliding_window=256,
+ layer_types=["sliding_attention"],
+ )
+ )
+ transformer = SanaVideoTransformer3DModel(
+ in_channels=16,
+ out_channels=16,
+ num_attention_heads=2,
+ attention_head_dim=16,
+ num_layers=1,
+ num_cross_attention_heads=2,
+ cross_attention_head_dim=16,
+ cross_attention_dim=32,
+ caption_channels=32,
+ mlp_ratio=2.0,
+ sample_size=30,
+ patch_size=(1, 2, 2),
+ )
+ vae = AutoencoderKLWan(
+ base_dim=8,
+ decoder_base_dim=8,
+ z_dim=16,
+ dim_mult=[1, 1, 1, 1],
+ num_res_blocks=1,
+ temperal_downsample=[False, True, True],
+ latents_mean=[0.0] * 16,
+ latents_std=[1.0] * 16,
+ )
+ pipeline = SanaVideoPipeline(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer,
+ scheduler=scheduler,
+ )
+ pipeline.to(dtype=torch.bfloat16).save_pretrained(model_dir)
+ return model_dir
+
+
def _shrink_flux_clip_text_encoder(config: dict) -> dict:
config["num_hidden_layers"] = 2
return config
@@ -54,14 +129,18 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
- text_config = config["text_config"]
+ """Shrink a Qwen2.5-VL text encoder (nested Qwen-Image or flat LongCat layout)."""
+ # Nested: text fields live under text_config, with top-level mirrors.
+ # Flat: text fields live at the top level.
+ text_config = config["text_config"] if "text_config" in config else config
old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
text_config["num_hidden_layers"] = 2
text_config["intermediate_size"] = 64
text_config["hidden_size"] = hidden_size
text_config["num_attention_heads"] = 2
text_config["num_key_value_heads"] = 2
- text_config["layer_types"] = text_config["layer_types"][:2]
+ if "layer_types" in text_config:
+ text_config["layer_types"] = text_config["layer_types"][:2]
factor = old_head_dim / (hidden_size / 2)
mrope_section = text_config["rope_scaling"]["mrope_section"]
text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
@@ -105,6 +184,38 @@ def tiny_qwen_image_edit_plus_builder() -> str:
)
+def tiny_longcat_image_builder() -> str:
+ return build_tiny_from_configs(
+ "LongCatImagePipeline",
+ "meituan-longcat/LongCat-Image",
+ transform={
+ "text_encoder": _shrink_qwen_text_encoder_config,
+ "transformer": partial(
+ _shrink_dit_rope_config,
+ num_single_layers=2,
+ default_axes_dims_rope=[16, 56, 56],
+ joint_attention_dim=64,
+ ),
+ },
+ )
+
+
+def tiny_longcat_image_edit_builder() -> str:
+ return build_tiny_from_configs(
+ "LongCatImageEditPipeline",
+ "meituan-longcat/LongCat-Image-Edit",
+ transform={
+ "text_encoder": _shrink_qwen_text_encoder_config,
+ "transformer": partial(
+ _shrink_dit_rope_config,
+ num_single_layers=2,
+ default_axes_dims_rope=[16, 56, 56],
+ joint_attention_dim=64,
+ ),
+ },
+ )
+
+
def tiny_flux_builder() -> str:
return build_tiny_from_configs(
"FluxPipeline",
@@ -131,19 +242,18 @@ def tiny_flux_kontext_builder() -> str:
def tiny_flux2_builder() -> str:
def shrink_text_encoder(config: dict) -> dict:
- # The real checkpoint ships language_model.lm_head.weight as its own
- # tensor even though tie_word_embeddings defaults to True; keep it
- # untied here too so transformers doesn't drop it on save.
config["tie_word_embeddings"] = False
- config["text_config"]["num_hidden_layers"] = 31
+ config["text_config"]["num_hidden_layers"] = 3
+ config["text_config"]["hidden_size"] = 32
config["text_config"]["intermediate_size"] = 64
- config["text_config"]["num_attention_heads"] = 4
- config["text_config"]["head_dim"] = 32
+ config["text_config"]["num_attention_heads"] = 2
+ config["text_config"]["head_dim"] = 16
config["text_config"]["num_key_value_heads"] = 2
+ config["text_config"]["text_encoder_out_layers"] = [0, 1, 2]
config["vision_config"]["num_hidden_layers"] = 2
config["vision_config"]["intermediate_size"] = 64
- config["vision_config"]["num_attention_heads"] = 4
- config["vision_config"]["head_dim"] = 32
+ config["vision_config"]["num_attention_heads"] = 2
+ config["vision_config"]["head_dim"] = 16
return config
return build_tiny_from_configs(
@@ -151,6 +261,6 @@ def tiny_flux2_builder() -> str:
"black-forest-labs/FLUX.2-dev",
transform={
"text_encoder": shrink_text_encoder,
- "transformer": partial(_shrink_dit_rope_config, num_single_layers=2),
+ "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, joint_attention_dim=96),
},
)
diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py
index cbb37729a..f24564f62 100644
--- a/tests/model_tests/diffusion/model_settings.py
+++ b/tests/model_tests/diffusion/model_settings.py
@@ -49,11 +49,37 @@ DIFFUSION_TEST_SETTINGS = {
[DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
],
),
+ "SanaVideoPipeline": DiffusionModelTestOpts(
+ model="Efficient-Large-Model/SANA-Video_2B_480p_diffusers",
+ builder=diff_model_builders.tiny_sana_video_builder,
+ supported_tasks=[DiffusionTasks.TEXT_TO_VIDEO],
+ check_multi_output=False,
+ ),
"QwenImagePipeline": DiffusionModelTestOpts(
model="Qwen/Qwen-Image",
builder=diff_model_builders.tiny_qwen_image_builder,
supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
),
+ "LongCatImagePipeline": DiffusionModelTestOpts(
+ model="meituan-longcat/LongCat-Image",
+ builder=diff_model_builders.tiny_longcat_image_builder,
+ supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
+ extra_test_groups=[
+ [DiffusionAccs.TEA_CACHE],
+ [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+ [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+ ],
+ ),
+ "LongCatImageEditPipeline": DiffusionModelTestOpts(
+ model="meituan-longcat/LongCat-Image-Edit",
+ builder=diff_model_builders.tiny_longcat_image_edit_builder,
+ supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+ extra_test_groups=[
+ [DiffusionAccs.TEA_CACHE],
+ [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+ [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+ ],
+ ),
"FluxPipeline": DiffusionModelTestOpts(
model="black-forest-labs/FLUX.1-schnell",
builder=diff_model_builders.tiny_flux_builder,
diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py
index 2be2338c1..4fb31fb5c 100644
--- a/tests/model_tests/diffusion/task_runners.py
+++ b/tests/model_tests/diffusion/task_runners.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Definitions running and validating individual tasks, e.g., text to image,
image to image, and so on. These are called by the core test runner.
@@ -6,12 +9,13 @@ image to image, and so on. These are called by the core test runner.
import base64
import io
from dataclasses import replace
+from typing import Any
import numpy as np
import pytest
from PIL import Image
-from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.runtime import DiffusionResponse, OmniServer, OnlineOmniClient, dummy_messages_from_mix_data
from tests.model_tests.diffusion.config_types import DiffusionTasks
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -48,7 +52,7 @@ IMAGE_GEN_EXTRA_BODY = {
}
# Online form_data for video generation requests (multipart /v1/videos API)
-VIDEO_GEN_FORM_DATA = {
+VIDEO_GEN_FORM_DATA: dict[str, Any] = {
"height": HEIGHT,
"width": WIDTH,
"num_inference_steps": 4,
@@ -171,7 +175,7 @@ def _build_online_image_data_url() -> str:
def _run_online_t2i(
- server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+ server: OmniServer, client: OnlineOmniClient, extra_body: dict | None = None
) -> list[DiffusionResponse]:
"""Run a text to image request through the server."""
messages = dummy_messages_from_mix_data(content_text=PROMPT)
@@ -184,7 +188,7 @@ def _run_online_t2i(
def _run_online_i2i(
- server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+ server: OmniServer, client: OnlineOmniClient, extra_body: dict | None = None
) -> list[DiffusionResponse]:
"""Run an image to image request through the server."""
image_data_url = _build_online_image_data_url()
@@ -278,7 +282,7 @@ def run_and_validate_multi_output(omni: Omni, task_type: DiffusionTasks):
def _run_online_t2v(
- server: OmniServer, client: OpenAIClientHandler, form_data: dict | None = None
+ server: OmniServer, client: OnlineOmniClient, form_data: dict | None = None
) -> list[DiffusionResponse]:
"""Run a text to video request through the server's /v1/videos API."""
data = dict(form_data or VIDEO_GEN_FORM_DATA)
@@ -287,7 +291,7 @@ def _run_online_t2v(
return client.send_video_diffusion_request({"form_data": data})
-def _run_online_i2v(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+def _run_online_i2v(server: OmniServer, client: OnlineOmniClient) -> list[DiffusionResponse]:
"""Run an image to video request through the server's /v1/videos API."""
data = dict(VIDEO_GEN_FORM_DATA)
data.setdefault("prompt", PROMPT)
@@ -306,17 +310,17 @@ def _get_online_videos(responses: list[DiffusionResponse]) -> list:
### Online task runners
-def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+def run_and_validate_online_text_to_image_request(server: OmniServer, client: OnlineOmniClient):
"""Run and validate a text to image request through the server."""
_validate_images(_get_online_images(_run_online_t2i(server, client)))
-def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+def run_and_validate_online_image_to_image_request(server: OmniServer, client: OnlineOmniClient):
"""Run and validate an image to image request through the server."""
_validate_images(_get_online_images(_run_online_i2i(server, client)))
-def run_and_validate_online_determinism(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks):
+def run_and_validate_online_determinism(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks):
"""Checks for determinism through the server, dispatching by task type.
Video is skipped online only: send_video_diffusion_request returns encoded
@@ -342,7 +346,7 @@ def run_and_validate_online_determinism(server: OmniServer, client: OpenAIClient
raise ValueError(f"Task type {task_type} is not yet supported for determinism checks")
-def run_and_validate_online_multi_output(server: OmniServer, client: OpenAIClientHandler, task_type: DiffusionTasks):
+def run_and_validate_online_multi_output(server: OmniServer, client: OnlineOmniClient, task_type: DiffusionTasks):
"""Checks for multi-output through the server, dispatching by task type.
Video is skipped online only: send_video_diffusion_request explicitly
@@ -368,11 +372,11 @@ def run_and_validate_online_multi_output(server: OmniServer, client: OpenAIClien
raise ValueError(f"Task type {task_type} is not yet supported for multi-output checks")
-def run_and_validate_online_text_to_video_request(server: OmniServer, client: OpenAIClientHandler):
+def run_and_validate_online_text_to_video_request(server: OmniServer, client: OnlineOmniClient):
"""Run and validate a text to video request through the server."""
_get_online_videos(_run_online_t2v(server, client))
-def run_and_validate_online_image_to_video_request(server: OmniServer, client: OpenAIClientHandler):
+def run_and_validate_online_image_to_video_request(server: OmniServer, client: OnlineOmniClient):
"""Run and validate an image to video request through the server."""
_get_online_videos(_run_online_i2v(server, client))
diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py
index be064f603..f9a5072f9 100644
--- a/tests/model_tests/diffusion/test_alignment.py
+++ b/tests/model_tests/diffusion/test_alignment.py
@@ -26,7 +26,11 @@ EXCLUDED_MODELS = [
"ZImagePipeline",
"OvisImagePipeline",
"WanPipeline",
+ "WanDMDPipeline",
"WanVACEPipeline",
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledOneStagePipeline",
+ "LTX2DistilledTwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -37,14 +41,12 @@ EXCLUDED_MODELS = [
"WanS2VPipeline",
"WanT2VDMD2Pipeline",
"WanI2VDMD2Pipeline",
- "LongCatImagePipeline",
"LongCatVideoAvatarPipeline",
"BagelPipeline",
"BooguImagePipeline",
"LancePipeline",
"MingImagePipeline",
"InternVLAA1Pipeline",
- "LongCatImageEditPipeline",
"StableDiffusion3Pipeline",
"HunyuanImage3ForCausalMM",
"ErnieImagePipeline",
@@ -55,24 +57,31 @@ EXCLUDED_MODELS = [
"OmniGen2Pipeline",
"HeliosPipeline",
"HeliosPyramidPipeline",
- "DreamIDOmniPipeline",
"SenseNovaU1Pipeline",
- "AudioXPipeline",
"HunyuanVideo15Pipeline",
"HunyuanVideo15ImageToVideoPipeline",
"LingBotVideoPipeline",
- "MagiHumanPipeline",
+ # The common harness does not yet define an image-to-video task or supply
+ # an input image. SANA I2V is covered by
+ # tests/diffusion/models/sana_video/test_pipeline_sana_video.py and
+ # tests/e2e/online_serving/test_sana_video_expansion.py. SANA T2V is exercised through
+ # the common tiny-model setting in model_settings.py.
+ "SanaImageToVideoPipeline",
+ # Requires camera/action assets and persistent AR state; covered by its
+ # dedicated contract tests and real-checkpoint E2E smoke instead.
+ "LingBotWorldCausalDMDPipeline",
"OmniVoicePipeline",
"OmniVoice",
"Cosmos3OmniDiffusersPipeline",
"Cosmos3OmniPipeline",
"DiffusersAdapterPipeline",
"HiDreamImagePipeline",
+ "HiDreamO1ImagePipeline",
"DreamZeroPipeline",
"StableDiffusionXLPipeline",
"Gr00tN1d7Pipeline",
- "SoulXSingerPipeline",
- "SoulXSingerSVCPipeline",
+ "Pi0Pipeline",
+ "SanaWmPipeline",
]
diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py
index b0eb26c9b..884710f2e 100644
--- a/tests/model_tests/diffusion/test_common_offline.py
+++ b/tests/model_tests/diffusion/test_common_offline.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import pytest
from tests.model_tests.diffusion.case_filtering import get_parametrized_options
@@ -17,7 +20,7 @@ from tests.model_tests.diffusion.task_runners import (
)
# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
@pytest.mark.parametrize(
@@ -50,7 +53,7 @@ def test_pipeline_on_supported_tasks(
)
try:
for task_type in supported_tasks:
- with subtests.test(msg=task_type):
+ with subtests.test(msg=task_type.value):
if task_type == DiffusionTasks.TEXT_TO_IMAGE:
run_and_validate_text_to_image_request(omni)
elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py
index 464adc13a..9c7ef8a0a 100644
--- a/tests/model_tests/diffusion/test_common_online.py
+++ b/tests/model_tests/diffusion/test_common_online.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""
Analogous to test_common_offline, but for server tests. Validates the full
online serving stack (CLI arg parsing, subprocess, API routing, response
@@ -6,7 +9,7 @@ encoding) using tiny models.
import pytest
-from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+from tests.helpers.runtime import OmniServer, OnlineOmniClient
from tests.model_tests.diffusion.case_filtering import get_parametrized_options
from tests.model_tests.diffusion.config_types import (
DiffusionAccs,
@@ -24,7 +27,7 @@ from tests.model_tests.diffusion.task_runners import (
)
# NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
@pytest.mark.parametrize(
@@ -50,7 +53,7 @@ def test_online_on_supported_tasks(
# TODO: We may want to revisit run_level validation here,
# because checks for things like image size etc should not
# depend on whether or not the weights are real or random
- client = OpenAIClientHandler(
+ client = OnlineOmniClient(
host=server.host,
port=server.port,
api_key="EMPTY",
@@ -58,7 +61,7 @@ def test_online_on_supported_tasks(
log_stats=server.log_stats,
)
for task_type in supported_tasks:
- with subtests.test(msg=task_type):
+ with subtests.test(msg=task_type.value):
if task_type == DiffusionTasks.TEXT_TO_IMAGE:
run_and_validate_online_text_to_image_request(server, client)
elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py
index 99dbb4650..08cbb2adc 100644
--- a/tests/platforms/npu/test_310p_patches.py
+++ b/tests/platforms/npu/test_310p_patches.py
@@ -414,6 +414,108 @@ def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling(
assert sample_calls == [{0: generator}]
+def test_310p_attention_forward_runs_fused_qkv_runtime(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Exercise the mocked 310P fused-QKV path at runtime on CPU.
+
+ The reviewer explicitly asked for more than source inspection: this test now
+ instantiates the patched attention, runs ``prepare_qkv_weights()``, verifies
+ the packed weight/bias cache, and drives the mocked 310P flash-attention op.
+ """
+ module, _ = _load_qwen3_tts_patch(monkeypatch)
+
+ trans_calls = []
+ rotary_calls = []
+ flash_calls = []
+
+ def maybe_trans_nz(weight: torch.Tensor) -> torch.Tensor:
+ trans_calls.append(weight.detach().clone())
+ return weight + 1.0
+
+ def npu_rotary_mul(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+ rotary_calls.append((hidden_states.detach().clone(), cos.detach().clone(), sin.detach().clone()))
+ return hidden_states + cos - sin
+
+ def npu_flash_attention(**kwargs) -> None:
+ flash_calls.append(
+ {key: value.detach().clone() if torch.is_tensor(value) else value for key, value in kwargs.items()}
+ )
+ kwargs["out"].copy_((kwargs["query"] + kwargs["key"] + kwargs["value"]).to(kwargs["out"].dtype))
+
+ monkeypatch.setattr(module, "maybe_trans_nz", maybe_trans_nz)
+ monkeypatch.setattr(module, "aligned_16", lambda tensor: tensor)
+ monkeypatch.setattr(module.torch_npu, "npu_rotary_mul", npu_rotary_mul, raising=False)
+ monkeypatch.setattr(module.torch_npu, "_npu_flash_attention", npu_flash_attention, raising=False)
+
+ config = SimpleNamespace(
+ num_attention_heads=2,
+ num_key_value_heads=2,
+ head_dim=2,
+ hidden_size=4,
+ num_code_groups=2,
+ attention_bias=True,
+ rms_norm_eps=1e-6,
+ )
+ attention = module._Qwen3CodePredictorAttention310P(config)
+ attention.num_heads = 2
+ attention.num_kv_heads = 2
+ attention.head_dim = 2
+ attention.scaling = 0.5
+ attention._q_size = attention.num_heads * attention.head_dim
+ attention._kv_size = attention.num_kv_heads * attention.head_dim
+ attention.q_norm = torch.nn.Identity()
+ attention.k_norm = torch.nn.Identity()
+ attention.o_proj = torch.nn.Identity()
+ attention.qkv_proj = torch.nn.Linear(4, attention._q_size + 2 * attention._kv_size, bias=True)
+ with torch.no_grad():
+ attention.qkv_proj.weight.copy_(torch.arange(48, dtype=torch.float32).reshape(12, 4) / 10)
+ attention.qkv_proj.bias.copy_(torch.arange(12, dtype=torch.float32) / 10)
+ attention._split_qkv = lambda qkv: qkv.split([attention._q_size, attention._kv_size, attention._kv_size], dim=-1)
+
+ attention.prepare_qkv_weights()
+
+ assert len(trans_calls) == 1
+ assert trans_calls[0].is_contiguous()
+ expected_fused_weight = attention.qkv_proj.weight.detach().clone().contiguous() + 1.0
+ torch.testing.assert_close(attention._fused_qkv_weight, expected_fused_weight)
+ assert attention._fused_qkv_bias is attention.qkv_proj.bias
+
+ hidden_states = (torch.arange(12, dtype=torch.float32).reshape(1, 3, 4) - 3) / 10
+ cos = torch.full((1, 3, 2), 0.25, dtype=torch.float32)
+ sin = torch.full((1, 3, 2), -0.5, dtype=torch.float32)
+ attention_mask = torch.arange(9, dtype=torch.float32).reshape(3, 3)
+
+ output = attention(hidden_states, (cos, sin), attention_mask=attention_mask)
+
+ expected_qkv = torch.nn.functional.linear(hidden_states, expected_fused_weight, attention.qkv_proj.bias)
+ expected_q_raw, expected_k_raw, expected_v_raw = expected_qkv.split(
+ [attention._q_size, attention._kv_size, attention._kv_size],
+ dim=-1,
+ )
+ expected_q = expected_q_raw.view(1, 3, attention.num_heads, attention.head_dim).transpose(1, 2)
+ expected_k = expected_k_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2)
+ expected_v = expected_v_raw.view(1, 3, attention.num_kv_heads, attention.head_dim).transpose(1, 2)
+ expected_cos = cos.unsqueeze(1)
+ expected_sin = sin.unsqueeze(1)
+ expected_q = expected_q + expected_cos - expected_sin
+ expected_k = expected_k + expected_cos - expected_sin
+ expected_q_f = expected_q.transpose(1, 2).reshape(3, attention.num_heads, attention.head_dim)
+ expected_k_f = expected_k.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim)
+ expected_v_f = expected_v.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim)
+ expected_flash = (expected_q_f + expected_k_f + expected_v_f).to(torch.float16).to(torch.float32)
+
+ assert len(rotary_calls) == 2
+ assert len(flash_calls) == 1
+ torch.testing.assert_close(flash_calls[0]["query"], expected_q_f)
+ torch.testing.assert_close(flash_calls[0]["key"], expected_k_f)
+ torch.testing.assert_close(flash_calls[0]["value"], expected_v_f)
+ torch.testing.assert_close(flash_calls[0]["mask"], attention_mask)
+ torch.testing.assert_close(flash_calls[0]["seq_len"], torch.tensor([3], dtype=torch.int32))
+ assert flash_calls[0]["scale_value"] == pytest.approx(0.5)
+ assert flash_calls[0]["num_heads"] == 2
+ assert flash_calls[0]["num_kv_heads"] == 2
+ torch.testing.assert_close(output, expected_flash.reshape(1, 3, 4))
+
+
def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
module, _ = _load_qwen3_tts_patch(monkeypatch)
talker = module._Qwen3TTSTalker310P(vllm_config=object())
diff --git a/tests/platforms/npu/test_diffusion_attn_backend_selector.py b/tests/platforms/npu/test_diffusion_attn_backend_selector.py
new file mode 100644
index 000000000..1136ef2e3
--- /dev/null
+++ b/tests/platforms/npu/test_diffusion_attn_backend_selector.py
@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the NPU diffusion attention backend selector.
+
+Regression coverage for the eager ``mindiesd`` import in
+``NPUOmniPlatform.get_diffusion_attn_backend_cls``: the import must run only
+for backends that reach mindiesd kernels (FLASH_ATTN, and RAINFUSION_ATTN via
+its dense FlashAttention fallback), so that a broken optional mindiesd
+install cannot block backends that never use it (e.g. TORCH_SDPA).
+
+The platform module is loaded from source with fake vllm/vllm-ascend
+dependencies, so the tests run on CPU without NPU or mindiesd installed.
+"""
+
+from __future__ import annotations
+
+import builtins
+import importlib.util
+import sys
+import types
+from enum import Enum
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_BACKEND_NAMES = (
+ "FLASH_ATTN",
+ "FLASH_ATTN_HUB",
+ "FLASH_ATTN_3_HUB",
+ "TORCH_SDPA",
+ "RAINFUSION_ATTN",
+)
+
+
+def _repo_root() -> Path:
+ marker = Path("vllm_omni") / "platforms" / "npu" / "platform.py"
+ for parent in Path(__file__).resolve().parents:
+ if (parent / marker).is_file():
+ return parent
+ raise FileNotFoundError(f"could not locate repo root containing {marker}")
+
+
+def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
+ module = types.ModuleType(name)
+ for key, value in attrs.items():
+ setattr(module, key, value)
+ monkeypatch.setitem(sys.modules, name, module)
+ return module
+
+
+def _make_fake_backend_enum():
+ members = {
+ name: SimpleNamespace(get_path=lambda n=name: f"fake.backends.{n.lower()}.Backend") for name in _BACKEND_NAMES
+ }
+
+ class _FakeEnumMeta(type):
+ def __getitem__(cls, name):
+ return members[name]
+
+ def __getattr__(cls, name):
+ try:
+ return members[name]
+ except KeyError:
+ raise AttributeError(name) from None
+
+ class FakeDiffusionAttentionBackendEnum(metaclass=_FakeEnumMeta):
+ pass
+
+ return FakeDiffusionAttentionBackendEnum
+
+
+class _FakeOmniPlatformEnum(Enum):
+ CUDA = "cuda"
+ NPU = "npu"
+
+
+class _FakeNPUPlatform:
+ pass
+
+
+class _FakeOmniPlatform:
+ _omni_enum = _FakeOmniPlatformEnum.NPU
+
+ @classmethod
+ def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
+ pass
+
+
+def _load_platform_module(monkeypatch: pytest.MonkeyPatch):
+ """Load vllm_omni/platforms/npu/platform.py with fake dependencies."""
+ _install_fake_module(monkeypatch, "vllm")
+ _install_fake_module(
+ monkeypatch,
+ "vllm.logger",
+ init_logger=lambda name: SimpleNamespace(
+ debug=lambda *a, **k: None,
+ info=lambda *a, **k: None,
+ warning=lambda *a, **k: None,
+ error=lambda *a, **k: None,
+ ),
+ )
+ _install_fake_module(monkeypatch, "vllm_ascend")
+ _install_fake_module(monkeypatch, "vllm_ascend.platform", NPUPlatform=_FakeNPUPlatform)
+ for pkg in (
+ "vllm_omni",
+ "vllm_omni.diffusion",
+ "vllm_omni.diffusion.attention",
+ "vllm_omni.diffusion.attention.backends",
+ "vllm_omni.platforms",
+ ):
+ _install_fake_module(monkeypatch, pkg)
+ _install_fake_module(
+ monkeypatch,
+ "vllm_omni.diffusion.attention.backends.registry",
+ DiffusionAttentionBackendEnum=_make_fake_backend_enum(),
+ )
+ _install_fake_module(
+ monkeypatch,
+ "vllm_omni.platforms.interface",
+ OmniPlatform=_FakeOmniPlatform,
+ OmniPlatformEnum=_FakeOmniPlatformEnum,
+ )
+
+ path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "platform.py"
+ spec = importlib.util.spec_from_file_location("npu_platform_under_test", path)
+ if spec is None or spec.loader is None:
+ raise RuntimeError(f"cannot load import spec for {path}")
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module.NPUOmniPlatform
+
+
+@pytest.fixture
+def mindiesd_tracker(monkeypatch: pytest.MonkeyPatch):
+ """Make mindiesd importable and record eager import attempts."""
+ monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+ attempts = []
+ real_import = builtins.__import__
+
+ def fake_import(name, *args, **kwargs):
+ if name == "mindiesd":
+ attempts.append(name)
+ return types.ModuleType("mindiesd")
+ return real_import(name, *args, **kwargs)
+
+ monkeypatch.setattr(builtins, "__import__", fake_import)
+ monkeypatch.setattr(
+ "importlib.util.find_spec",
+ lambda name, *a, **k: object() if name == "mindiesd" else None,
+ )
+ return attempts
+
+
+@pytest.mark.parametrize(
+ ("selected_backend", "expect_mindiesd_import"),
+ [
+ ("FLASH_ATTN", True),
+ # HF-hub variants fall back to local FLASH_ATTN on NPU.
+ ("FLASH_ATTN_HUB", True),
+ ("FLASH_ATTN_3_HUB", True),
+ ("RAINFUSION_ATTN", True),
+ # Backends that never reach mindiesd must not trigger the import.
+ ("TORCH_SDPA", False),
+ ],
+)
+def test_eager_mindiesd_import_scoped_to_mindiesd_backends(
+ monkeypatch: pytest.MonkeyPatch,
+ mindiesd_tracker: list[str],
+ selected_backend: str,
+ expect_mindiesd_import: bool,
+):
+ platform_cls = _load_platform_module(monkeypatch)
+
+ path = platform_cls.get_diffusion_attn_backend_cls(selected_backend, head_size=64)
+
+ assert (len(mindiesd_tracker) > 0) is expect_mindiesd_import
+ expected = "FLASH_ATTN" if selected_backend in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB") else selected_backend
+ assert path == f"fake.backends.{expected.lower()}.Backend"
+
+
+def test_explicit_flash_attn_without_mindiesd_still_resolves(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ platform_cls = _load_platform_module(monkeypatch)
+ monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+ monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
+
+ path = platform_cls.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+ assert path == "fake.backends.flash_attn.Backend"
+
+
+def test_default_backend_prefers_flash_attn_when_mindiesd_available(
+ monkeypatch: pytest.MonkeyPatch,
+ mindiesd_tracker: list[str],
+):
+ platform_cls = _load_platform_module(monkeypatch)
+
+ path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
+
+ assert path == "fake.backends.flash_attn.Backend"
+ assert len(mindiesd_tracker) > 0
+
+
+def test_default_backend_falls_back_to_sdpa_without_mindiesd(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ platform_cls = _load_platform_module(monkeypatch)
+ monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+ monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
+
+ path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
+
+ assert path == "fake.backends.torch_sdpa.Backend"
diff --git a/tests/platforms/npu/test_diffusion_platform.py b/tests/platforms/npu/test_diffusion_platform.py
new file mode 100644
index 000000000..10e43e034
--- /dev/null
+++ b/tests/platforms/npu/test_diffusion_platform.py
@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("available", [False, True])
+def test_dense_flash_capability_tracks_mindiesd(monkeypatch, available):
+ pytest.importorskip("vllm_ascend")
+ from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+ monkeypatch.setattr(
+ importlib.util,
+ "find_spec",
+ lambda name: object() if name == "mindiesd" and available else None,
+ )
+
+ assert NPUOmniPlatform.supports_diffusion_dense_flash_attention() is available
+
+
+def test_paged_config_uses_ascend_kernel_block_size() -> None:
+ pytest.importorskip("vllm_ascend")
+ from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
+
+ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+ from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+ vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
+ NPUOmniPlatform.configure_diffusion_vllm_config(
+ vllm_config,
+ SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER),
+ )
+
+ assert vllm_config.cache_config.block_size == AscendAttentionBackend.get_supported_kernel_block_sizes()[0]
+
+
+def test_strict_ulysses_paged_backend_bypasses_ascend_pcp() -> None:
+ pytest.importorskip("vllm_ascend")
+ from vllm_ascend.attention.attention_v1 import (
+ AscendAttentionBackend,
+ AscendAttentionBackendImpl,
+ AscendAttentionMetadataBuilder,
+ )
+
+ from vllm_omni.platforms.npu.platform import NPUOmniPlatform
+
+ assert (
+ NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+ AscendAttentionBackend,
+ ulysses_degree=1,
+ )
+ is AscendAttentionBackend
+ )
+ backend = NPUOmniPlatform.get_diffusion_paged_kv_attn_backend(
+ AscendAttentionBackend,
+ ulysses_degree=2,
+ )
+ assert backend.get_impl_cls() is AscendAttentionBackendImpl
+ assert backend.get_builder_cls() is AscendAttentionMetadataBuilder
diff --git a/tests/platforms/npu/test_graph_tools.py b/tests/platforms/npu/test_graph_tools.py
new file mode 100644
index 000000000..703b8f630
--- /dev/null
+++ b/tests/platforms/npu/test_graph_tools.py
@@ -0,0 +1,375 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.platforms.npu.graph_tools import (
+ CapturedDeviceGraph,
+ NPUExactGraphRunner,
+)
+from vllm_omni.platforms.npu.models import minicpmo_4_5_code2wav as code2wav_patch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_captured_graph_replay_clones_persistent_outputs():
+ static_input = torch.zeros(1)
+ static_output = torch.zeros(1)
+
+ class _Graph:
+ def replay(self):
+ static_output.copy_(static_input * 2)
+
+ graph = CapturedDeviceGraph(
+ graph=_Graph(),
+ static_inputs=(static_input,),
+ static_outputs=(static_output,),
+ )
+
+ first = graph.replay((torch.tensor([2.0]),))[0]
+ second = graph.replay((torch.tensor([3.0]),))[0]
+
+ torch.testing.assert_close(first, torch.tensor([4.0]))
+ torch.testing.assert_close(second, torch.tensor([6.0]))
+ assert first.data_ptr() != second.data_ptr()
+
+
+def test_capture_uses_vllm_global_graph_pool(monkeypatch):
+ runner = NPUExactGraphRunner()
+ global_pool = object()
+ seen_pools = []
+ synchronizations = 0
+
+ class _Graph:
+ def replay(self):
+ pass
+
+ @contextmanager
+ def graph_context(graph, *, pool):
+ del graph
+ seen_pools.append(pool)
+ yield
+
+ def synchronize():
+ nonlocal synchronizations
+ synchronizations += 1
+
+ fake_npu = SimpleNamespace(
+ NPUGraph=_Graph,
+ graph=graph_context,
+ synchronize=synchronize,
+ )
+ monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+ import vllm.platforms
+
+ monkeypatch.setattr(
+ vllm.platforms,
+ "current_platform",
+ SimpleNamespace(get_global_graph_pool=lambda: global_pool),
+ )
+
+ captured = runner.capture(
+ (torch.tensor([2.0]),),
+ lambda value: (value + 1,),
+ )
+
+ assert seen_pools == [global_pool]
+ assert synchronizations == 2
+ torch.testing.assert_close(captured.static_outputs[0], torch.tensor([3.0]))
+
+
+def test_exact_signature_dispatch_captures_then_replays(monkeypatch):
+ runner = NPUExactGraphRunner(max_graphs=2)
+ replay_count = 0
+
+ class _FunctionalGraph:
+ def __init__(self, compute):
+ self.compute = compute
+
+ def replay(self, inputs):
+ nonlocal replay_count
+ replay_count += 1
+ return tuple(output.detach().clone() for output in self.compute(*inputs))
+
+ monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+ monkeypatch.setattr(
+ runner,
+ "capture",
+ lambda inputs, compute: _FunctionalGraph(compute),
+ )
+
+ def compute(value):
+ return (value * 2,)
+
+ first = runner.run("unit", (torch.tensor([2.0]),), (False,), compute)
+ second = runner.run("unit", (torch.tensor([3.0]),), (False,), compute)
+
+ torch.testing.assert_close(first[0], torch.tensor([4.0]))
+ torch.testing.assert_close(second[0], torch.tensor([6.0]))
+ assert runner.stats == {"captures": 1, "failed": 0, "hits": 1}
+ assert replay_count == 1
+
+
+def test_capture_failure_is_fatal_for_stage_process(monkeypatch):
+ runner = NPUExactGraphRunner(
+ component_name="test component",
+ disable_config_hint="disable the test graph",
+ )
+ captures = 0
+ computes = 0
+
+ monkeypatch.setattr(runner, "_eligible", lambda inputs: True)
+
+ def fail_capture(inputs, compute):
+ nonlocal captures
+ captures += 1
+ raise RuntimeError("unsupported op")
+
+ monkeypatch.setattr(runner, "capture", fail_capture)
+
+ def compute(value):
+ nonlocal computes
+ computes += 1
+ return (value + 1,)
+
+ with pytest.raises(RuntimeError, match="restart the stage process"):
+ runner.run("unit", (torch.tensor([1.0]),), (False,), compute)
+ with pytest.raises(RuntimeError, match="cannot continue"):
+ runner.run(
+ "different-shape",
+ (torch.tensor([1.0, 2.0]),),
+ (True,),
+ compute,
+ )
+
+ assert captures == 1
+ assert computes == 1
+ assert runner.stats == {"captures": 0, "failed": 1, "hits": 0}
+
+
+def test_code2wav_runtime_disables_internal_format_and_jit(monkeypatch):
+ monkeypatch.delenv("ASCEND_LAUNCH_BLOCKING", raising=False)
+ compile_modes = []
+ fake_npu = SimpleNamespace(
+ config=SimpleNamespace(allow_internal_format=True),
+ set_compile_mode=lambda **kwargs: compile_modes.append(kwargs),
+ )
+ monkeypatch.setattr(torch, "npu", fake_npu, raising=False)
+
+ code2wav_patch.prepare_code2wav_graph_runtime()
+
+ assert fake_npu.config.allow_internal_format is False
+ assert compile_modes == [{"jit_compile": False}]
+
+
+def test_code2wav_runtime_rejects_launch_blocking(monkeypatch):
+ monkeypatch.setenv("ASCEND_LAUNCH_BLOCKING", "1")
+
+ with pytest.raises(RuntimeError, match="ASCEND_LAUNCH_BLOCKING=1"):
+ code2wav_patch.prepare_code2wav_graph_runtime()
+
+
+@pytest.mark.parametrize(
+ ("enabled", "max_graphs", "expected_prepared"),
+ [
+ (True, 7, 1),
+ (False, 7, 0),
+ (True, 0, 0),
+ ],
+)
+def test_code2wav_patch_reads_stage_additional_config(
+ monkeypatch,
+ enabled,
+ max_graphs,
+ expected_prepared,
+):
+ prepared = 0
+
+ class _Backend:
+ speech_window = torch.zeros(1)
+ flow = SimpleNamespace(training=False)
+
+ model = SimpleNamespace(
+ backend=None,
+ _extra_config=lambda: {},
+ vllm_config=SimpleNamespace(
+ additional_config={
+ "code2wav_enable_npu_graph": enabled,
+ "code2wav_max_npu_graphs": max_graphs,
+ }
+ ),
+ )
+
+ def build_backend(instance):
+ instance.backend = _Backend()
+
+ def prepare_runtime():
+ nonlocal prepared
+ prepared += 1
+
+ monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+ monkeypatch.setattr(code2wav_patch, "prepare_code2wav_graph_runtime", prepare_runtime)
+ code2wav_patch._patched_build_backend(model)
+
+ assert prepared == expected_prepared
+ if expected_prepared:
+ graph_runner = code2wav_patch._backend_graph_runners[model.backend]
+ assert isinstance(graph_runner, NPUExactGraphRunner)
+ assert graph_runner.max_graphs == max_graphs
+ else:
+ assert model.backend not in code2wav_patch._backend_graph_runners
+
+
+@pytest.mark.parametrize("value", [True, "true", "false", "0", "off"])
+def test_code2wav_patch_rejects_bfloat16_attention_cache(monkeypatch, value):
+ built = False
+
+ model = SimpleNamespace(
+ backend=None,
+ _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value},
+ )
+
+ def build_backend(_instance):
+ nonlocal built
+ built = True
+
+ monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+
+ with pytest.raises(ValueError, match="currently supported only on CUDA"):
+ code2wav_patch._patched_build_backend(model)
+
+ assert not built
+
+
+@pytest.mark.parametrize("value", [False, None, 0, ""])
+def test_code2wav_patch_allows_disabled_bfloat16_attention_cache(monkeypatch, value):
+ built = False
+
+ model = SimpleNamespace(
+ backend=None,
+ _extra_config=lambda: {"code2wav_bfloat16_attention_cache": value},
+ )
+
+ def build_backend(instance):
+ nonlocal built
+ built = True
+ instance.backend = object()
+
+ monkeypatch.setattr(code2wav_patch, "_original_build_backend", build_backend)
+ code2wav_patch._patched_build_backend(model)
+
+ assert built
+
+
+@pytest.mark.parametrize("with_cache", [False, True])
+def test_code2wav_estimator_graph_dispatch_is_platform_owned(monkeypatch, with_cache):
+ calls = []
+
+ class _Backend:
+ _trt_stepper = None
+ _cfm_graph_wrapper = None
+
+ class _Runner:
+ def run(self, operation, inputs, constants, compute):
+ calls.append((operation, inputs, constants))
+ return compute(*inputs)
+
+ backend = _Backend()
+ runner = _Runner()
+ code2wav_patch._backend_graph_runners[backend] = runner
+
+ def graphable_estimator_step(
+ instance,
+ estimator,
+ *,
+ x,
+ mu,
+ time_embedding,
+ speakers,
+ cond,
+ cnn_cache,
+ att_cache,
+ ):
+ del instance, estimator, time_embedding, speakers, cond
+ cache_marker = x.new_tensor(0 if cnn_cache is None else 1)
+ return x + mu, cache_marker, cache_marker.clone()
+
+ monkeypatch.setattr(code2wav_patch, "_original_estimator_step", lambda *args, **kwargs: None)
+ monkeypatch.setattr(code2wav_patch, "_graphable_estimator_step", graphable_estimator_step)
+ value = torch.tensor([2.0])
+ cache = torch.tensor([1.0]) if with_cache else None
+ estimator = SimpleNamespace(t_embedder=lambda time: time + 1)
+ outputs = code2wav_patch._patched_estimator_step(
+ backend,
+ estimator,
+ x=value,
+ mu=value,
+ time=value,
+ speakers=value,
+ cond=value,
+ cnn_cache=cache,
+ att_cache=cache,
+ )
+
+ assert calls[0][0] == "cfm_estimator"
+ assert calls[0][2] == (with_cache,)
+ assert len(calls[0][1]) == (7 if with_cache else 5)
+ torch.testing.assert_close(calls[0][1][2], torch.tensor([[3.0]]))
+ torch.testing.assert_close(outputs[0], torch.tensor([4.0]))
+ torch.testing.assert_close(outputs[1], torch.tensor(1.0 if with_cache else 0.0))
+
+
+def test_code2wav_platform_wraps_flow_execution_context(monkeypatch):
+ entered = []
+
+ class _Backend:
+ pass
+
+ class _Context:
+ def __init__(self, device, require_math):
+ self.device = device
+ self.require_math = require_math
+
+ def __enter__(self):
+ entered.append((self.device, self.require_math))
+
+ def __exit__(self, exc_type, exc, traceback):
+ return False
+
+ backend = _Backend()
+ code2wav_patch._backend_graph_runners[backend] = NPUExactGraphRunner()
+ monkeypatch.setattr(
+ code2wav_patch,
+ "_flow_execution_context",
+ lambda device, *, require_math: _Context(device, require_math),
+ )
+ monkeypatch.setattr(
+ code2wav_patch,
+ "_original_setup_batch",
+ lambda instance, features, batch_size: ("setup", batch_size),
+ )
+ monkeypatch.setattr(
+ code2wav_patch,
+ "_original_decode_batch",
+ lambda instance, tokens, features, states, *, last_chunk, flush_encoder: (
+ "decode",
+ last_chunk,
+ flush_encoder,
+ ),
+ )
+ features = SimpleNamespace(speech_tokens=torch.zeros(1))
+
+ assert code2wav_patch._patched_setup_batch(backend, features, 3) == ("setup", 3)
+ assert code2wav_patch._patched_decode_batch(
+ backend,
+ torch.zeros(1),
+ features,
+ [],
+ last_chunk=True,
+ ) == ("decode", True, False)
+ assert entered == [(torch.device("cpu"), True), (torch.device("cpu"), True)]
diff --git a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
new file mode 100644
index 000000000..4ae544380
--- /dev/null
+++ b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Ascend hardware validation for MiniCPM-o Code2Wav NPUGraph replay."""
+
+from math import prod
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytest.importorskip("vllm_ascend")
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
+ BatchedToken2Wav,
+)
+from vllm_omni.platforms.npu.graph_tools import NPUExactGraphRunner
+from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
+ _graphable_estimator_step,
+ prepare_code2wav_graph_runtime,
+)
+from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
+ npu_token2wav_sdpa_context,
+)
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+
+pytestmark = [
+ pytest.mark.core_model,
+ pytest.mark.omni,
+ *hardware_marks(res={"npu": "A3"}, num_cards=1),
+]
+
+
+class _Flow(nn.Module):
+ def __init__(self, estimator: nn.Module):
+ super().__init__()
+ self.encoder = nn.Identity()
+ self.decoder = SimpleNamespace(estimator=estimator)
+
+
+class _Token2Wav:
+ def __init__(self, estimator: nn.Module):
+ self.flow = _Flow(estimator).eval()
+ self.hift = nn.Identity().eval()
+ self.float16 = False
+ self.n_timesteps = 2
+ self.mel_cache_len = 1
+ self.source_cache_len = 2
+ self.speech_window = torch.ones(4, device="npu")
+
+
+def test_npu_runner_inherits_runtime_snapshot_replacement():
+ assert issubclass(OmniNPUModelRunner, OmniGPUModelRunner)
+ assert OmniNPUModelRunner._replace_intermediate_buffer is OmniGPUModelRunner._replace_intermediate_buffer
+ runner = object.__new__(OmniNPUModelRunner)
+ request = SimpleNamespace(additional_information_cpu=None)
+ runner.requests = {"req": request}
+ runner.model_intermediate_buffer = {
+ "req": {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"last_chunk": True},
+ }
+ }
+
+ runner._replace_intermediate_buffer("req", {"meta": {"is_segment_finished": True}})
+
+ assert runner.model_intermediate_buffer["req"] == {"meta": {"is_segment_finished": True}}
+ assert request.additional_information_cpu == runner.model_intermediate_buffer["req"]
+
+
+def test_minicpmo_cfm_estimator_npugraph_matches_eager():
+ decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
+ prepare_code2wav_graph_runtime()
+
+ estimator = (
+ decoder_dit.DiT(
+ in_channels=6,
+ out_channels=2,
+ depth=2,
+ num_heads=2,
+ head_dim=8,
+ hidden_size=16,
+ mlp_ratio=2.0,
+ )
+ .npu()
+ .eval()
+ )
+ graph_runner = NPUExactGraphRunner(max_graphs=4)
+ adapter = BatchedToken2Wav(_Token2Wav(estimator))
+
+ def inputs(seed: int):
+ def make(shape: tuple[int, ...], offset: float):
+ values = torch.arange(
+ prod(shape),
+ device="npu",
+ dtype=torch.float32,
+ )
+ return torch.sin(values * 0.17 + seed + offset).reshape(shape)
+
+ return (
+ make((2, 2, 6), 0.0),
+ make((2, 2, 6), 0.5),
+ make((2, 1, 16), 1.0),
+ make((2, 1), 1.5),
+ make((2, 1, 6), 2.0),
+ )
+
+ def compute(*values, caches=None):
+ cnn_cache, att_cache = (None, None) if caches is None else caches
+ return _graphable_estimator_step(
+ adapter,
+ estimator,
+ x=values[0],
+ mu=values[1],
+ time_embedding=values[2],
+ speakers=values[3],
+ cond=values[4],
+ cnn_cache=cnn_cache,
+ att_cache=att_cache,
+ )
+
+ with torch.inference_mode(), npu_token2wav_sdpa_context(require_math=True):
+ warm_inputs = inputs(11)
+ warm_outputs = graph_runner.run(
+ "cfm_estimator",
+ warm_inputs,
+ (False,),
+ lambda *values: compute(*values),
+ )
+ replay_inputs = inputs(13)
+ expected = compute(*replay_inputs)
+ actual = graph_runner.run(
+ "cfm_estimator",
+ replay_inputs,
+ (False,),
+ lambda *values: compute(*values),
+ )
+ for eager, replayed in zip(expected, actual, strict=True):
+ torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+
+ cached_inputs = (*inputs(17), warm_outputs[1], warm_outputs[2])
+ graph_runner.run(
+ "cfm_estimator",
+ cached_inputs,
+ (True,),
+ lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+ )
+ replay_cached_inputs = (
+ *inputs(19),
+ warm_outputs[1] + 0.1,
+ warm_outputs[2] + 0.1,
+ )
+ expected_cached = compute(
+ *replay_cached_inputs[:5],
+ caches=(replay_cached_inputs[5], replay_cached_inputs[6]),
+ )
+ actual_cached = graph_runner.run(
+ "cfm_estimator",
+ replay_cached_inputs,
+ (True,),
+ lambda *values: compute(*values[:5], caches=(values[5], values[6])),
+ )
+ for eager, replayed in zip(expected_cached, actual_cached, strict=True):
+ torch.testing.assert_close(replayed, eager, rtol=1e-3, atol=1e-3)
+ assert graph_runner.stats == {"captures": 2, "failed": 0, "hits": 2}
diff --git a/tests/platforms/npu/test_step_audio2_token2wav.py b/tests/platforms/npu/test_step_audio2_token2wav.py
index ae86f92d9..23baa2c3e 100644
--- a/tests/platforms/npu/test_step_audio2_token2wav.py
+++ b/tests/platforms/npu/test_step_audio2_token2wav.py
@@ -7,6 +7,7 @@ from __future__ import annotations
import importlib.util
import sys
import types
+from contextlib import contextmanager
from pathlib import Path
from types import SimpleNamespace
@@ -194,3 +195,71 @@ def test_hift_patch_reports_incompatible_layout(monkeypatch: pytest.MonkeyPatch)
with pytest.raises(TypeError, match=r"m_source\.l_sin_gen\._f02sine"):
module.patch_step_audio2_hift_for_npu(SimpleNamespace())
+
+
+def test_sdpa_context_preserves_exceptions_from_token2wav_body(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ module = _load_patch_module(monkeypatch)
+ fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+ fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+ @contextmanager
+ def fake_math_context(*, require_available=False):
+ del require_available
+ yield
+
+ fake_attention_patch.npu_math_sdpa_context = fake_math_context
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+ fake_attention_patch,
+ )
+
+ with pytest.raises(RuntimeError, match="original graph failure"):
+ with module.npu_token2wav_sdpa_context():
+ raise RuntimeError("original graph failure")
+
+
+def test_sdpa_context_requires_math_for_npu_graph(monkeypatch: pytest.MonkeyPatch) -> None:
+ module = _load_patch_module(monkeypatch)
+ required = []
+ fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+ fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = lambda: None
+
+ @contextmanager
+ def fake_math_context(*, require_available=False):
+ required.append(require_available)
+ yield
+
+ fake_attention_patch.npu_math_sdpa_context = fake_math_context
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+ fake_attention_patch,
+ )
+
+ with module.npu_token2wav_sdpa_context(require_math=True):
+ pass
+
+ assert required == [True]
+
+
+def test_sdpa_context_does_not_hide_patch_setup_failures(monkeypatch: pytest.MonkeyPatch) -> None:
+ module = _load_patch_module(monkeypatch)
+ fake_attention_patch = types.ModuleType("vllm_omni.platforms.npu.models.cosyvoice2_dit_attn")
+
+ def fail_patch():
+ raise RuntimeError("attention patch failed")
+
+ fake_attention_patch.apply_cosyvoice2_dit_attn_npu_patch = fail_patch
+ fake_attention_patch.npu_math_sdpa_context = contextmanager(lambda: iter((None,)))
+ monkeypatch.setitem(
+ sys.modules,
+ "vllm_omni.platforms.npu.models.cosyvoice2_dit_attn",
+ fake_attention_patch,
+ )
+
+ with pytest.raises(RuntimeError, match="attention patch failed"):
+ with module.npu_token2wav_sdpa_context():
+ pass
diff --git a/tests/platforms/test_record_device_event.py b/tests/platforms/test_record_device_event.py
index a979ca940..9b81f9cd4 100644
--- a/tests/platforms/test_record_device_event.py
+++ b/tests/platforms/test_record_device_event.py
@@ -1,9 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for OmniPlatform.record_device_event implementations."""
import pytest
+from tests.helpers.mark import hardware_test
+
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -38,6 +40,71 @@ class TestCudaOmniPlatformRecordDeviceEvent:
assert result is None
+class TestXPUOmniPlatformRecordDeviceEvent:
+ """XPU must record a device-agnostic ``torch.Event``, not ``torch.xpu.Event``.
+
+ The async diffusion output thread gates its D2H copy with
+ ``torch.Stream.wait_event`` on a generic ``torch.Stream``, and that binding
+ silently drops the dependency for a ``torch.xpu.Event`` — the copy then races
+ the still-running forward and the host reads a partially written tensor
+ (garbage rows at the bottom of the image).
+ """
+
+ def test_records_device_agnostic_torch_event(self, mocker):
+ # vLLM's XPUPlatform imports vllm_xpu_kernels at module scope, and that
+ # package is absent on non-XPU images.
+ pytest.importorskip("vllm_xpu_kernels")
+
+ from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+ mock_event = mocker.MagicMock()
+ mocker.patch("torch.Event", return_value=mock_event)
+ xpu_event = mocker.patch("torch.xpu.Event")
+
+ result = XPUOmniPlatform.record_device_event()
+
+ assert result is mock_event
+ mock_event.record.assert_called_once()
+ xpu_event.assert_not_called()
+
+ @hardware_test(res={"xpu": "B60"})
+ def test_recorded_event_orders_side_stream_d2h(self):
+ """On-device: the event must make a side-stream D2H wait for compute."""
+ import torch
+
+ if not torch.xpu.is_available():
+ pytest.skip("XPU not available")
+
+ from vllm_omni.platforms.xpu.platform import XPUOmniPlatform
+
+ device = torch.device("xpu", 0)
+ # Normalized so the matmul chain stays finite; the ``* 0.0`` term keeps
+ # every element exactly 7.0 while still depending on the whole chain.
+ base = torch.randn(4096, 4096, device=device) / 64.0
+ chained = base.clone()
+ for _ in range(40):
+ chained = chained @ base
+ tensor = torch.full_like(chained, 7.0) + chained * 0.0
+
+ event = XPUOmniPlatform.record_device_event()
+ assert event is not None
+
+ # Mirrors WorkerProc._async_output_loop.
+ d2h_stream = torch.Stream(device=device)
+ d2h_stream.wait_event(event)
+ previous = torch.accelerator.current_stream()
+ torch.accelerator.set_stream(d2h_stream)
+ try:
+ host = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True)
+ host.copy_(tensor, non_blocking=True)
+ finally:
+ torch.accelerator.set_stream(previous)
+ d2h_stream.synchronize()
+
+ copied_early = int((host != 7.0).sum())
+ assert copied_early == 0, f"{copied_early} elements copied before compute finished"
+
+
class TestNPUOmniPlatformRecordDeviceEvent:
"""Test NPUOmniPlatform.record_device_event with mocked torch.npu."""
@@ -55,8 +122,8 @@ class TestNPUOmniPlatformRecordDeviceEvent:
result = NPUOmniPlatform.record_device_event()
assert result is None
- def test_synchronizes_stream_then_records_event(self, mocker):
- """NPU should sync stream first, then create and record an Event."""
+ def test_synchronizes_stream_then_records_generic_event(self, mocker):
+ """NPU should return an event consumable by the generic torch.Stream."""
try:
from vllm_omni.platforms.npu.platform import NPUOmniPlatform
except ModuleNotFoundError:
@@ -66,14 +133,16 @@ class TestNPUOmniPlatformRecordDeviceEvent:
mock_stream = mocker.MagicMock()
mock_torch = mocker.MagicMock()
mock_torch.npu.current_stream.return_value = mock_stream
- mock_torch.npu.Event.return_value = mock_event
+ mock_torch.Event.return_value = mock_event
mocker.patch("vllm_omni.platforms.npu.platform.torch", mock_torch)
result = NPUOmniPlatform.record_device_event()
# Stream should be synced first (HCCL ordering)
mock_stream.synchronize.assert_called_once()
- # Then event should be created and recorded
- mock_torch.npu.Event.assert_called_once()
+ # Then a backend-neutral event should be created and recorded. The
+ # worker waits on it through the public torch.Stream wrapper.
+ mock_torch.Event.assert_called_once()
+ mock_torch.npu.Event.assert_not_called()
mock_event.record.assert_called_once()
assert result is mock_event
diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py
index 486a629d5..47c6ac89f 100644
--- a/tests/tools/test_check_tts_adapter.py
+++ b/tests/tools/test_check_tts_adapter.py
@@ -17,9 +17,7 @@ import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit"))
from check_tts_adapter import ( # noqa: E402
- MAX_LEGACY_DETECTORS,
MAX_MODEL_TYPE_BRANCHES,
- _legacy_detectors,
_model_type_branches,
main,
)
@@ -60,12 +58,6 @@ def test_dict_dispatch_is_a_known_gap():
assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0
-def test_legacy_detector_counting():
- src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))'
- found = _legacy_detectors(ast.parse(src))
- assert [name for _, name in found] == ["a", "b"]
-
-
def test_gate_passes_on_current_tree():
"""The recorded budgets must match reality, or the ratchet is decorative."""
assert main([]) == 0
@@ -74,6 +66,4 @@ def test_gate_passes_on_current_tree():
def test_budgets_are_the_real_counts():
root = Path(__file__).resolve().parents[2]
serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py"
- adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py"
assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES
- assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS
diff --git a/tests/utils/test_speaker_cache.py b/tests/utils/test_speaker_cache.py
index 373baacb4..ea50c5881 100644
--- a/tests/utils/test_speaker_cache.py
+++ b/tests/utils/test_speaker_cache.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import threading
import pytest
@@ -144,7 +147,7 @@ class TestSpeakerEmbeddingCacheBehavior:
class TestSingleton:
- def test_singleton_identity(self, fresh_speaker_cache):
+ def test_singleton_identity(self, clean_speaker_cache):
a = get_speaker_cache()
b = get_speaker_cache()
assert a is b
diff --git a/tests/utils/test_speaker_cache_integration.py b/tests/utils/test_speaker_cache_integration.py
index 8a24659f0..01b886725 100644
--- a/tests/utils/test_speaker_cache_integration.py
+++ b/tests/utils/test_speaker_cache_integration.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Integration tests for the process-wide speaker cache across serving + models."""
import pytest
@@ -26,7 +29,7 @@ class TestSpeakerCacheIntegration:
assert cache.get(fish_key) is None
assert cache.get(bob_key) is not None
- def test_singleton_shared_across_call_sites(self, fresh_speaker_cache):
+ def test_singleton_shared_across_call_sites(self, clean_speaker_cache):
cache_a = get_speaker_cache()
cache_b = get_speaker_cache()
assert cache_a is cache_b
diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py
index f00da2797..443ec0138 100644
--- a/tests/utils/test_tracking_parser.py
+++ b/tests/utils/test_tracking_parser.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Tests for TrackingArgumentParser and related utilities."""
import argparse
@@ -35,7 +38,7 @@ _TEST_PIPELINE = PipelineConfig(
stages=(
StagePipelineConfig(stage_id=0, model_stage="stage0"),
StagePipelineConfig(stage_id=1, model_stage="stage1"),
- StagePipelineConfig(stage_id=2, model_stage="stage2"),
+ StagePipelineConfig(stage_id=2, model_stage="stage2", final_output=True),
),
)
diff --git a/tests/worker/test_capture_replay_contract.py b/tests/worker/test_capture_replay_contract.py
index 504aaa7b0..a0ca095f3 100644
--- a/tests/worker/test_capture_replay_contract.py
+++ b/tests/worker/test_capture_replay_contract.py
@@ -14,9 +14,9 @@ cannot be masked by an adjacent ``else``/``elif`` that happens to read it. It
matches by attribute chain (``self.input_ids.gpu``), so local variable renames
do not break it.
-Scope: the BASE runner (``OmniGPUModelRunner``) only. The generation runner is
-KNOWN to lack the ``has_preprocess`` branch; extend this
-check to the generation file when it adds the branch.
+Scope: the BASE runner (``OmniGPUModelRunner``) and the generation runner's
+``_dummy_run`` override (which gained the branch in RFC #5450 C6; its
+``_preprocess`` is inherited from the base).
"""
import ast
@@ -25,6 +25,7 @@ import textwrap
import pytest
+from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -69,11 +70,24 @@ def _has_preprocess_branch_reads_both_buffers(method) -> bool:
return any(_REQUIRED_BUFFERS <= _attr_chains_in(node.body) for node in guarded)
-@pytest.mark.parametrize("method_name", ["_dummy_run", "_preprocess"])
-def test_has_preprocess_branch_reads_shared_buffers(method_name):
- method = getattr(OmniGPUModelRunner, method_name)
+@pytest.mark.parametrize(
+ ("runner_cls", "method_name"),
+ [
+ (OmniGPUModelRunner, "_dummy_run"),
+ (OmniGPUModelRunner, "_preprocess"),
+ (GPUGenerationModelRunner, "_dummy_run"),
+ ],
+)
+def test_has_preprocess_branch_reads_shared_buffers(runner_cls, method_name):
+ method = getattr(runner_cls, method_name)
assert _has_preprocess_branch_reads_both_buffers(method), (
- f"{OmniGPUModelRunner.__name__}.{method_name}: the has_preprocess branch no "
+ f"{runner_cls.__name__}.{method_name}: the has_preprocess branch no "
"longer reads BOTH self.input_ids.gpu and self.inputs_embeds.gpu — capture "
"and replay would diverge."
)
+
+
+def test_generation_dummy_run_overrides_with_branch():
+ # Guard against the override being deleted (which would silently fall back
+ # to the base implementation and vacuously pass the parametrized check).
+ assert "_dummy_run" in GPUGenerationModelRunner.__dict__
diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py
index 116b5ee4d..5e3cfe865 100644
--- a/tests/worker/test_gpu_ar_model_runner.py
+++ b/tests/worker/test_gpu_ar_model_runner.py
@@ -1,7 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
import copy
+import re
from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
from types import SimpleNamespace
+from typing import Any
+from unittest.mock import patch
import numpy as np
import pytest
@@ -14,12 +22,15 @@ from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechReques
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker import sparse_audio
from vllm_omni.worker.gpu_ar_model_runner import (
ExecuteModelState,
GPUARModelRunner,
OmniAsyncGPUModelRunnerOutput,
)
+from vllm_omni.worker.output import payload_build
from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
+from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -200,6 +211,15 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
async def build(self, request, sampling_params_list, has_inline_ref_audio):
return PreparedRequest(prompt={"prompt": request.input}, tts_params={}, model_type="higgs_audio_v3")
+ def apply_sampling_overrides(
+ self,
+ sampling_params_list: list,
+ request: "OpenAICreateSpeechRequest",
+ prompt: dict[str, Any] | None = None,
+ request_id: str | None = None,
+ ) -> list:
+ return sampling_params_list
+
engine = SimpleNamespace(
errored=False,
default_sampling_params_list=[SamplingParams(temperature=1.0, top_p=0.95, top_k=50)],
@@ -291,11 +311,13 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
def test_sparse_mm_req_ids_requires_sparse_audio_marker():
- assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
- assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
- assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
- assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}) == ["r1"]
+ assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"sparse_audio": ["1"]}, "meta.req_id": ["r1"]}) == ["r1"]
def test_runner_assisted_full_attention_metadata_request_is_opt_in():
@@ -875,6 +897,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
lambda *args, **kwargs: (
0,
None,
+ None,
[],
None,
["r1", "r2"],
@@ -1188,3 +1211,643 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow
assert torch.equal(output.inter_stage_outputs[0]["codes.ref"], ref_codes[0])
assert torch.equal(output.inter_stage_outputs[2]["codes.ref"], ref_codes[2])
+
+
+class TestPromptIdsClampPaddingConvention:
+ """Pin `sampling_utils.clamp_prompt_ids_to_penalty_padding` (the clamp
+ `sample_tokens` applies before penalties).
+
+ The audit (RFC #5450 C2) originally flagged the clamp as an off-by-one
+ (OOB index). It is not: upstream penalties allocate `vocab_size + 1` bins
+ and drop the last column, so `vocab_size` is the designed padding value
+ (vllm/model_executor/layers/utils.py::get_token_bin_counts_and_mask).
+ These tests call the runner's own helper, so a future "fix" to the clamp
+ fails here.
+ """
+
+ LOGITS_VOCAB = 8
+ BATCH_VOCAB = 10 # wider input_batch vocab; its pad value is BATCH_VOCAB
+
+ def _padded_prompt(self):
+ # Two real tokens + two padding slots padded with the batch vocab size.
+ return torch.tensor([[1, 2, self.BATCH_VOCAB, self.BATCH_VOCAB]])
+
+ def test_clamp_to_logits_vocab_keeps_padding_inert(self):
+ from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+ # The runner's own helper, not a re-implementation.
+ clamped = clamp_prompt_ids_to_penalty_padding(self._padded_prompt(), self.LOGITS_VOCAB)
+
+ _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1)
+
+ expected = torch.zeros(1, self.LOGITS_VOCAB, dtype=torch.bool)
+ expected[0, 1] = True
+ expected[0, 2] = True
+ # Pad ids landed in the dropped (vocab_size-th) bin: no penalty effect.
+ assert torch.equal(mask, expected)
+
+ def test_clamp_one_lower_would_corrupt_last_token_penalties(self):
+ from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+ # The "fix" the audit first proposed - shown here to be the actual bug:
+ # padding would count as occurrences of the last real vocab token.
+ clamped = self._padded_prompt().clamp(max=self.LOGITS_VOCAB - 1)
+
+ _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1)
+
+ assert bool(mask[0, self.LOGITS_VOCAB - 1])
+
+ def test_unclamped_padding_overflows_penalty_bins(self):
+ from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+
+ # Why the clamp exists: the batch-level pad value (BATCH_VOCAB) is out
+ # of range even for the vocab_size+1 penalty bins of a narrower head.
+ with pytest.raises(RuntimeError):
+ get_token_bin_counts_and_mask(self._padded_prompt(), self.LOGITS_VOCAB, 1)
+
+
+class TestSparseAudioMarkerRobustness:
+ """RFC #5450 C3: `sparse_audio.is_sparse_audio_marker` must not crash on tensor/array
+ markers, and the combined prefix-cache payload builder must never ship one
+ request's data as another's."""
+
+ def test_marker_list_and_str_forms(self):
+ marker = sparse_audio.is_sparse_audio_marker
+ assert marker(["1"]) is True
+ assert marker(["0"]) is False
+ assert marker("true") is True
+ assert marker("off") is False
+
+ def test_marker_enforced_canonical_forms_no_error(self):
+ # The designed carrier (list-of-str; bare str accepted) and marker
+ # absence (None) are the ONLY silent forms.
+ marker = sparse_audio.is_sparse_audio_marker
+ with patch.object(sparse_audio.logger, "error") as err:
+ assert marker(["1"]) is True
+ assert marker(["0"]) is False
+ assert marker(["0", "1"]) is True
+ assert marker("true") is True
+ assert marker("off") is False
+ assert marker(None) is False
+ err.assert_not_called()
+
+ def test_marker_illegal_carriers_raise_one_error_per_type(self):
+ # Enforcement, not normalization: any carrier outside the design
+ # contract is a producer bug - it raises InvalidSparseDeclarationError
+ # (caught inside the module by resolve_sparse_mm_routing, which
+ # fails closed)
+ # and logs at error level once per offending type. Interpreting
+ # deviants would need a second truth predicate, and predicate
+ # divergence is exactly how the np.array(["0"])-routes-as-sparse bug
+ # happened.
+ marker = sparse_audio.is_sparse_audio_marker
+ for value in (torch.tensor(1), torch.tensor([1, 1]), np.array(["1"]), 1, 1.0, True, ["1", 1]):
+ with (
+ patch.object(sparse_audio, "_logged_offenders", set()),
+ patch.object(sparse_audio.logger, "error") as err,
+ ):
+ with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+ marker(value)
+ assert err.call_count == 1, f"expected one error for {value!r}"
+ with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+ marker(value)
+ assert err.call_count == 1 # deduped per type, not per step
+
+ def test_marker_device_tensor_rejected_without_read(self):
+ # Deviant carriers are never read element-wise, so a device tensor
+ # can never introduce a D2H sync here. Meta tensors hold no data -
+ # any read (`.item()`/`bool()`) raises a RuntimeError, so seeing OUR
+ # exception type (not RuntimeError) proves the data is never touched.
+ marker = sparse_audio.is_sparse_audio_marker
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+ marker(torch.ones((), device="meta"))
+ with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+ marker(torch.ones(3, device="meta"))
+
+ def test_invalid_marker_fails_closed_at_routing(self):
+ # A present-but-invalid sparse declaration most likely belongs to a
+ # producer that INTENDED sparse output; falling back to dense would
+ # batch-index-assign its payloads to the WRONG requests. Routing
+ # must fail closed: no payload routed for the step (empty sparse
+ # set), never dense misrouting.
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0", "r1"],
+ downstream_req_ids=["r0", "r1"],
+ multimodal_outputs={
+ "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+ "model_outputs": [torch.zeros(1)],
+ },
+ )
+ assert downstream == []
+ assert index == {}
+ assert is_sparse is True
+
+ def test_fail_closed_routing_error_remains_observable(self):
+ def route_invalid_marker():
+ return sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0"],
+ downstream_req_ids=["r0"],
+ multimodal_outputs={
+ "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+ "model_outputs": [torch.zeros(1)],
+ },
+ )
+
+ with (
+ patch.object(sparse_audio, "_logged_offenders", set()),
+ patch.object(sparse_audio, "_fail_closed_last_log_at", None),
+ patch.object(sparse_audio, "_fail_closed_suppressed", 0),
+ patch.object(sparse_audio.time, "monotonic", side_effect=[100.0, 120.0, 161.0]),
+ patch.object(sparse_audio.logger, "error") as err,
+ ):
+ for _ in range(3):
+ assert route_invalid_marker() == ([], {}, True)
+
+ routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]]
+ assert len(routing_errors) == 2
+ assert [call.args[1] for call in routing_errors] == [0, 1]
+
+ def test_fail_closed_logging_is_thread_safe(self):
+ workers = 8
+ ready = Barrier(workers)
+
+ def route_invalid_marker():
+ return sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0"],
+ downstream_req_ids=["r0"],
+ multimodal_outputs={
+ "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])},
+ },
+ )
+
+ def route_concurrently(_):
+ ready.wait(timeout=10)
+ return route_invalid_marker()
+
+ with (
+ patch.object(sparse_audio, "_logged_offenders", set()),
+ patch.object(sparse_audio, "_fail_closed_last_log_at", None),
+ patch.object(sparse_audio, "_fail_closed_suppressed", 0),
+ patch.object(sparse_audio.time, "monotonic", return_value=100.0) as clock,
+ patch.object(sparse_audio.logger, "error") as err,
+ ):
+ with ThreadPoolExecutor(max_workers=workers) as executor:
+ results = list(executor.map(route_concurrently, range(workers)))
+ assert results == [([], {}, True)] * workers
+
+ clock.return_value = 161.0
+ assert route_invalid_marker() == ([], {}, True)
+
+ routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]]
+ assert len(routing_errors) == 2
+ assert [call.args[1] for call in routing_errors] == [0, workers - 1]
+ classifier_errors = [call for call in err.call_args_list if "Illegal sparse-audio marker" in call.args[0]]
+ assert len(classifier_errors) == 1
+
+ def test_duplicate_sparse_metadata_must_be_consistent(self):
+ equal = {
+ "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+ "meta.req_id": ["r1"],
+ "meta.sparse_audio": ["true"],
+ }
+ assert sparse_audio.resolve_sparse_mm_req_ids(equal) == ["r1"]
+
+ invalid_declarations = (
+ {
+ "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+ "meta.req_id": ["r2"],
+ },
+ {
+ "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+ "meta.req_id": "r1",
+ },
+ {
+ "meta": {"req_id": ["r1"], "sparse_audio": ["0"]},
+ "meta.sparse_audio": ["1"],
+ },
+ {"meta": {"req_id": ["r1", "r1"], "sparse_audio": ["1"]}},
+ )
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ for declaration in invalid_declarations:
+ with pytest.raises(sparse_audio.InvalidSparseDeclarationError):
+ sparse_audio.resolve_sparse_mm_req_ids(declaration)
+
+ def test_conflicting_sparse_metadata_fails_closed_at_public_boundary(self):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r1", "r2"],
+ downstream_req_ids=["r1", "r2"],
+ multimodal_outputs={
+ "meta": {"req_id": ["r1"], "sparse_audio": ["1"]},
+ "meta.req_id": ["r2"],
+ },
+ )
+ assert downstream == []
+ assert index == {}
+ assert is_sparse is True
+
+ def test_flat_marker_with_nested_req_ids_is_not_erased(self):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0", "r1"],
+ downstream_req_ids=["r0", "r1"],
+ multimodal_outputs={"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]},
+ )
+ assert downstream == ["r1"]
+ assert index == {"r1": 0}
+ assert is_sparse is True
+
+ def test_sparse_marker_with_unusable_req_ids_fails_closed(self):
+ # The marker says "sparse" but the alignment list is unusable - the
+ # declaration as a whole is out of contract; same fail-closed shape.
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0"],
+ downstream_req_ids=["r0"],
+ multimodal_outputs={"meta": {"req_id": "r0", "sparse_audio": ["1"]}},
+ )
+ assert downstream == []
+ assert is_sparse is True
+
+ def test_nested_marker_without_req_id_fails_closed(self):
+ # {"meta": {"sparse_audio": ["1"]}} - marker set, req_id missing.
+ # The flattened-encoding fallback must not ERASE the nested marker
+ # (which used to silently demote the declaration to dense); this is
+ # an unusable-req_id invalid declaration -> fail closed.
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0"],
+ downstream_req_ids=["r0"],
+ multimodal_outputs={"meta": {"sparse_audio": ["1"]}},
+ )
+ assert downstream == []
+ assert is_sparse is True
+
+ def test_invalid_marker_on_non_audio_pipeline_keeps_dense_routing(self):
+ # Non-audio pipelines never consult sparse routing; the carrier
+ # error is logged by the classifier but routing is unaffected.
+ with patch.object(sparse_audio, "_logged_offenders", set()):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="latent",
+ req_ids_output_copy=["r0", "r1"],
+ downstream_req_ids=["r0", "r1"],
+ multimodal_outputs={"meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}},
+ )
+ assert downstream == ["r0", "r1"]
+ assert is_sparse is False
+
+ def test_legal_marker_still_routes_sparse(self):
+ downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing(
+ engine_output_type="audio",
+ req_ids_output_copy=["r0", "r1", "r2"],
+ downstream_req_ids=["r0", "r1", "r2"],
+ multimodal_outputs={"meta": {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}},
+ )
+ assert downstream == ["r1", "r2"]
+ assert index == {"r1": 0, "r2": 1}
+ assert is_sparse is True
+
+ def test_combined_payload_unwraps_aligned_list(self):
+ combined = {"codes.audio": {"req-1": [torch.zeros(1), torch.ones(1)]}}
+ payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=1)
+ assert torch.equal(payload["codes.audio"], torch.ones(1))
+
+ def test_combined_payload_broadcasts_qwen_request_invariant_singletons(self):
+ req_ids = ("r0", "r1")
+ tts_embeddings = {
+ "embed.tts_bos": torch.full((1, 1, 2), 1.0),
+ "embed.tts_eos": torch.full((1, 1, 2), 2.0),
+ "embed.tts_pad": torch.full((1, 1, 2), 3.0),
+ }
+ combined = {key: {rid: [value.clone()] for rid in req_ids} for key, value in tts_embeddings.items()}
+
+ for idx, rid in enumerate(req_ids):
+ payload = payload_build.build_omni_mm_payload(
+ combined_multimodal_outputs=combined,
+ mm_cpu=None,
+ rid=rid,
+ idx=idx,
+ start=idx,
+ end=idx + 1,
+ audio_sparse_output=False,
+ sparse_mm_index={},
+ hidden_seq_len=len(req_ids),
+ scheduled_seq_len=len(req_ids),
+ )
+ assert payload.keys() == tts_embeddings.keys()
+ for key, expected in tts_embeddings.items():
+ assert torch.equal(payload[key], expected)
+
+ def test_combined_payload_drops_sparse_protocol_metadata(self):
+ combined = {
+ "meta.sparse_audio": {"req-1": ["1"]},
+ "meta.req_id": {"req-1": ["req-1"]},
+ "sr": {"req-1": 22050},
+ }
+ payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=2)
+ assert payload == {"sr": 22050}
+
+ def test_combined_payload_misalignment_drops_key_not_request_zero(self):
+ # A short multi-element list used to silently return `v[0]` - request
+ # 0's payload shipped under request `rid`. Never substitute another
+ # request's data; the key is dropped with an error log (same policy
+ # as the non-combined sparse mismatch) rather than raised - v1
+ # runners don't catch per-request exceptions, so a raise would turn a
+ # per-request protocol break into a whole-stage outage.
+ combined = {
+ "codes.audio": {"req-3": [torch.zeros(1), torch.ones(1)]},
+ "sr": {"req-3": 22050},
+ }
+ payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-3", list_idx=2)
+ assert "codes.audio" not in payload
+ assert payload["sr"] == 22050
+
+ def test_direct_path_sparse_mismatch_drops_key_and_keeps_payload_shape(self):
+ # Mirror of the combined-path policy on the non-prefix-cache path:
+ # a sparse index beyond the per-request list drops that key (error
+ # log) and the rest of the request's payload survives.
+ payload = payload_build.build_omni_mm_payload(
+ combined_multimodal_outputs=None,
+ mm_cpu={"codes.audio": [torch.zeros(1)]},
+ rid="r2",
+ idx=1,
+ start=0,
+ end=1,
+ audio_sparse_output=True,
+ sparse_mm_index={"r2": 1},
+ hidden_seq_len=1,
+ scheduled_seq_len=1,
+ )
+ assert payload == {}
+
+ def test_combined_payload_uses_sparse_index_under_sparse_routing(self):
+ # Batch [r0, r1, r2], sparse outputs only for [r1, r2]: producers
+ # align per-request lists to the SPARSE order, so r1 (batch idx 1,
+ # sparse idx 0) must get element 0 and r2 (batch idx 2, sparse idx 1)
+ # element 1. Indexing by batch idx shipped r2's payload to r1 and
+ # went out of range for r2. The sparse->list_idx resolution lives in
+ # `_build_omni_mm_payload` (the combined helper itself is
+ # sparse-agnostic), so this exercises the caller.
+ sparse_mm_index = {"r1": 0, "r2": 1}
+ combined = {
+ "codes.audio": {
+ "r1": [torch.zeros(1), torch.ones(1)],
+ "r2": [torch.zeros(1), torch.ones(1)],
+ }
+ }
+ payloads = {}
+ for rid, idx in (("r1", 1), ("r2", 2)):
+ payloads[rid] = payload_build.build_omni_mm_payload(
+ combined_multimodal_outputs=combined,
+ mm_cpu=None,
+ rid=rid,
+ idx=idx,
+ start=0,
+ end=1,
+ audio_sparse_output=True,
+ sparse_mm_index=sparse_mm_index,
+ hidden_seq_len=1,
+ scheduled_seq_len=1,
+ )
+ assert torch.equal(payloads["r1"]["codes.audio"], torch.zeros(1))
+ assert torch.equal(payloads["r2"]["codes.audio"], torch.ones(1))
+
+
+class TestMergeModelKvTransferMetadata:
+ """RFC #5450 C4: model KV-transfer metadata must merge into a COPY of the
+ scheduler's finished-request data - `scheduler_output` can be shared with
+ the engine-core process and must never be mutated by the runner."""
+
+ def _runner_with_model(self, model):
+ runner = object.__new__(GPUARModelRunner)
+ runner.model = model
+ return runner
+
+ def test_merge_does_not_mutate_scheduler_data(self):
+ class _Model:
+ def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+ return {"talker_codes": [7], "seen_tokens": num_computed_tokens}
+
+ original_meta = {"a": 1}
+ original = {"r1": {"seq_len": 3, "block_ids": [4], "custom_metadata": original_meta}}
+ runner = self._runner_with_model(_Model())
+
+ merged = runner._merge_model_kv_transfer_metadata(original)
+
+ # The manager sees the merged view...
+ assert merged["r1"]["custom_metadata"] == {"a": 1, "talker_codes": [7], "seen_tokens": 3}
+ assert merged["r1"]["seq_len"] == 3
+ # ...while the engine-shared originals are untouched.
+ assert original["r1"]["custom_metadata"] is original_meta
+ assert original_meta == {"a": 1}
+ assert "talker_codes" not in original["r1"]["custom_metadata"]
+
+ def test_merge_without_model_meta_passes_entry_through(self):
+ class _Model:
+ def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+ return None
+
+ data = {"seq_len": 2, "block_ids": [1]}
+ runner = self._runner_with_model(_Model())
+
+ merged = runner._merge_model_kv_transfer_metadata({"r1": data})
+
+ assert merged["r1"] is data
+
+ def test_merge_failure_raises_instead_of_shipping_partial_transfer(self):
+ # RFC #5450 C7: the bare `except Exception -> logger.warning` used to
+ # drop the model's custom metadata and let the KV transfer proceed
+ # incomplete; the failure then surfaced much later on the consumer
+ # stage. It must fail at the defect instead.
+ class _Model:
+ def get_kv_transfer_metadata(self, req_id, num_computed_tokens):
+ raise RuntimeError("boom")
+
+ data = {"seq_len": 2, "block_ids": [1], "custom_metadata": {"a": 1}}
+ runner = self._runner_with_model(_Model())
+
+ with pytest.raises(RuntimeError, match="boom"):
+ runner._merge_model_kv_transfer_metadata({"r1": data})
+ assert data["custom_metadata"] == {"a": 1}
+
+
+class TestDownstreamPayloadMemoization:
+ """RFC #5450 C7: `_request_needs_downstream_stage_payload` must not
+ memoize before the deciding marker exists - the final-stage id arrives via
+ `model_intermediate_buffer`, which can be unpopulated on the first call."""
+
+ def _runner(self, stages):
+ runner = object.__new__(GPUARModelRunner)
+ runner._downstream_payload_cache = {}
+ runner._request_final_stage_id = lambda rid: stages[rid]
+ return runner
+
+ def test_missing_marker_defaults_true_without_memoizing(self):
+ stages = {"r1": None}
+ runner = self._runner(stages)
+
+ assert runner._request_needs_downstream_stage_payload("r1") is True
+ # Not cached: the answer must refresh once the marker arrives.
+ assert "r1" not in runner._downstream_payload_cache
+
+ stages["r1"] = 0 # marker lands: this stage IS the final stage
+ assert runner._request_needs_downstream_stage_payload("r1") is False
+ assert runner._downstream_payload_cache["r1"] is False
+
+ def test_memoizes_once_marker_known(self):
+ stages = {"r1": 2}
+ runner = self._runner(stages)
+
+ assert runner._request_needs_downstream_stage_payload("r1") is True
+ assert runner._downstream_payload_cache["r1"] is True
+ # Cached value now authoritative for the request's lifetime.
+ stages["r1"] = 0
+ assert runner._request_needs_downstream_stage_payload("r1") is True
+
+
+# Matches a real ASSIGNMENT of the capability (optionally `self.`-qualified
+# and/or annotated), not a bare mention: a comment, a docstring, or a read
+# such as `getattr(self.model, "prefer_model_sampler", False)` must not count
+# as declaring it. `=(?!=)` keeps `==` comparisons out.
+_PREFER_MODEL_SAMPLER_ASSIGNMENT = re.compile(
+ r"^[ \t]*(?:self\.)?prefer_model_sampler[ \t]*(?::[^=\n]*)?=(?!=)[ \t]*(?P[^\n]*)",
+ re.MULTILINE,
+)
+
+
+def _declares_prefer_model_sampler(source: str) -> bool:
+ """Whether `source` opts INTO the model-sampler contract.
+
+ An explicit `prefer_model_sampler = False` is an opt-OUT and does not
+ count; any other right-hand side does, including runtime-conditional
+ forms (minicpmo_4_5 assigns `self.model_stage in {...}`), which can
+ evaluate True.
+ """
+ for match in _PREFER_MODEL_SAMPLER_ASSIGNMENT.finditer(source):
+ if match.group("rhs").split("#")[0].strip() != "False":
+ return True
+ return False
+
+
+class TestPreferModelSamplerNoneFallback:
+ """RFC #5450 C7: a `prefer_model_sampler` model returning None from
+ `sample()` DECLINES the step and the runner falls back to the default
+ sampler. The fallback is load-bearing (cosyvoice3/glm_tts decline for
+ empty-logits / inapplicable-stage steps), so it must keep working - now
+ as a documented contract with a one-time visibility log."""
+
+ def _runner(self, model_sample_result, default_result):
+ runner = object.__new__(GPUARModelRunner)
+ runner.input_batch = SimpleNamespace(
+ sampling_metadata="smd",
+ update_async_output_token_ids=lambda: None,
+ )
+ runner.model = SimpleNamespace(
+ prefer_model_sampler=True,
+ sample=lambda logits, metadata: model_sample_result,
+ )
+ runner._sampling_metadata_for_model_sampler = lambda smd: smd
+ runner._resolve_duplex_sampling_hook = lambda: None
+ runner.sampler = lambda logits, sampling_metadata: default_result
+ return runner
+
+ def test_none_falls_back_to_default_sampler_and_warns_once(self):
+ import logging
+
+ from vllm_omni.worker import gpu_ar_model_runner as ar_module
+
+ class _UniqueDecliningSampler:
+ prefer_model_sampler = True
+
+ @staticmethod
+ def sample(logits, metadata):
+ return None
+
+ sentinel = object()
+ runner = self._runner(model_sample_result=None, default_result=sentinel)
+ runner.model = _UniqueDecliningSampler()
+
+ records: list[logging.LogRecord] = []
+
+ class _Capture(logging.Handler):
+ def emit(self, record):
+ records.append(record)
+
+ handler = _Capture(level=logging.WARNING)
+ ar_module.logger.addHandler(handler)
+ try:
+ for _ in range(3):
+ out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None)
+ assert out is sentinel
+ finally:
+ ar_module.logger.removeHandler(handler)
+
+ fallback_records = [r for r in records if "falling back to the default sampler" in r.getMessage()]
+ assert len(fallback_records) == 1
+ assert "_UniqueDecliningSampler" in fallback_records[0].getMessage()
+
+ def test_model_sampler_output_short_circuits_default(self):
+ model_out = object()
+ runner = self._runner(model_sample_result=model_out, default_result=object())
+
+ out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None)
+
+ assert out is model_out
+
+ def test_declaration_matcher_ignores_mentions_and_opt_outs(self):
+ # Guards the guard: the inventory below is only meaningful if
+ # "declares" means an actual opt-in assignment.
+ assert _declares_prefer_model_sampler(" prefer_model_sampler = True")
+ assert _declares_prefer_model_sampler(" prefer_model_sampler: bool = True")
+ assert _declares_prefer_model_sampler(' self.prefer_model_sampler = stage in {"llm"}')
+ assert not _declares_prefer_model_sampler(" prefer_model_sampler = False")
+ assert not _declares_prefer_model_sampler(" prefer_model_sampler = False # opted out")
+ assert not _declares_prefer_model_sampler(" # prefer_model_sampler = True (not anymore)")
+ assert not _declares_prefer_model_sampler(' """Does not use prefer_model_sampler."""')
+ assert not _declares_prefer_model_sampler(' x = getattr(model, "prefer_model_sampler", False)')
+ assert not _declares_prefer_model_sampler(" if prefer_model_sampler == True:")
+
+ def test_in_tree_declarer_inventory_is_exact(self):
+ # The contract's audience: an exact inventory, so ADDING a declarer
+ # fails this test and the new model owner consciously signs up for the
+ # documented None-fallback semantics (update the set when they do).
+ import pathlib
+
+ import vllm_omni
+
+ models_dir = pathlib.Path(vllm_omni.__file__).parent / "model_executor" / "models"
+ declarers = {
+ p.parent.name
+ for p in models_dir.rglob("*.py")
+ if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore"))
+ }
+ expected = {
+ "cosyvoice3",
+ "glm_tts",
+ "higgs_audio_v2",
+ "higgs_audio_v3",
+ "hunyuan_image3",
+ "minicpmo_4_5",
+ "minimax_music3",
+ "nemotron_voicechat",
+ }
+ assert declarers == expected, (
+ "The set of models declaring `prefer_model_sampler` changed:\n"
+ f" added: {sorted(declarers - expected) or 'none'}\n"
+ f" removed: {sorted(expected - declarers) or 'none'}\n\n"
+ "This is a deliberate consent gate, not a broken test. Declaring "
+ "`prefer_model_sampler` opts a model into the runner's documented "
+ "None-fallback contract: `model.sample()` may return None to decline "
+ "a step, and the runner then falls back to the DEFAULT sampler "
+ "(see `GPUARModelRunner._sample`; RFC #5450 C7.3).\n\n"
+ "If you ADDED a declarer: confirm your model relies on -- or at "
+ "least tolerates -- that fallback, then add its directory name to "
+ "`expected` above. If you REMOVED one, drop its name."
+ )
diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py
index a6df4732f..b97168161 100644
--- a/tests/worker/test_gpu_generation_model_runner.py
+++ b/tests/worker/test_gpu_generation_model_runner.py
@@ -16,14 +16,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _DummyInputBatch:
- def __init__(self):
- self.req_ids = ["req-1"]
- self.req_id_to_index = {"req-1": 0}
- self.num_reqs = 1
+ def __init__(self, num_reqs: int = 1):
+ self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]
+ self.req_id_to_index = {rid: i for i, rid in enumerate(self.req_ids)}
+ self.num_reqs = num_reqs
self.vocab_size = 10
-def _make_runner(multimodal_outputs):
+def _make_runner(multimodal_outputs, num_reqs: int = 1):
runner = object.__new__(GPUGenerationModelRunner)
runner.execute_model_state = ExecuteModelState(
None,
@@ -40,7 +40,7 @@ def _make_runner(multimodal_outputs):
None,
)
runner.kv_connector_output = None
- runner.input_batch = _DummyInputBatch()
+ runner.input_batch = _DummyInputBatch(num_reqs)
runner.use_async_scheduling = False
runner.device = torch.device("cpu")
runner.supports_mm_inputs = False
@@ -111,6 +111,12 @@ def _make_guard_runner():
runner.cache_config = CacheConfig()
runner.input_batch = _DummyInputBatch()
runner.model = object()
+ runner.parallel_config = SimpleNamespace(
+ distributed_executor_backend=None,
+ data_parallel_size=1,
+ )
+ runner._dummy_run_calls = []
+ runner._dummy_run = lambda n: runner._dummy_run_calls.append(n)
runner._update_states = lambda scheduler_output: None
runner.synchronize_input_prep = contextlib.nullcontext
runner.attach_omni_connector_output = lambda result: result
@@ -121,9 +127,112 @@ def _make_guard_runner():
def test_execute_model_skips_non_positive_scheduled_span(monkeypatch, total):
"""#5196: a negative span is truthy, so it used to reach `_prepare_inputs`,
whose `assert total_num_scheduled_tokens > 0` killed the stage EngineCore."""
+ monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
runner = _make_guard_runner()
output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total))
assert output is EMPTY_MODEL_RUNNER_OUTPUT
+ assert runner._dummy_run_calls == []
+
+
+@pytest.mark.parametrize("total", [-1, 0])
+def test_execute_model_zero_tokens_runs_dp_dummy_run(monkeypatch, total):
+ """A 0-token step under external_launcher + DP>1 must `_dummy_run(1)` so
+ `coordinate_batch_across_dp` stays in sync (it used to be shadowed by an
+ unconditional early return and never ran -> DP hang risk)."""
+ monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
+ monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
+ runner = _make_guard_runner()
+ runner.parallel_config = SimpleNamespace(
+ distributed_executor_backend="external_launcher",
+ data_parallel_size=2,
+ )
+
+ output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total))
+
+ assert output is EMPTY_MODEL_RUNNER_OUTPUT
+ assert runner._dummy_run_calls == [1]
+
+
+def test_execute_model_zero_tokens_kv_connector_no_forward(monkeypatch):
+ """With a KV-transfer group, a 0-token step must go through
+ `kv_connector_no_forward` (also shadowed by the old early return)."""
+ monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
+ monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: True)
+ runner = _make_guard_runner()
+ sentinel = object()
+ calls = []
+ runner.kv_connector_no_forward = lambda scheduler_output, vllm_config: (
+ calls.append(scheduler_output),
+ sentinel,
+ )[1]
+ runner.vllm_config = object()
+
+ output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))
+
+ assert output is sentinel
+ assert len(calls) == 1
+
+
+def test_execute_model_zero_tokens_ec_producer_takes_encoder_path(monkeypatch):
+ """AR lock-step: the EC-producer block sits ABOVE the zero-token block
+ (matching GPUARModelRunner), so a 0-token step on an EC producer returns
+ the empty ENCODER output. Before the C1 reorder, the unconditional early
+ return won and produced EMPTY_MODEL_RUNNER_OUTPUT instead — this is the
+ one cell whose behaviour the reorder changed."""
+ monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: True)
+ monkeypatch.setattr(gen_runner_module, "get_ec_transfer", lambda: SimpleNamespace(is_consumer=False))
+ monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
+ sentinel = object()
+ monkeypatch.setattr(gen_runner_module, "make_empty_encoder_model_runner_output", lambda so: sentinel)
+ runner = _make_guard_runner()
+ runner.encoder_cache = {}
+ runner.maybe_get_ec_connector_output = lambda scheduler_output, encoder_cache: contextlib.nullcontext()
+ encoder_calls = []
+ runner._execute_mm_encoder = lambda so: encoder_calls.append(so)
+
+ output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))
+
+ assert output is sentinel
+ assert len(encoder_calls) == 1
+ assert runner._dummy_run_calls == []
+
+
+class TestSampleTokensBatched:
+ """RFC #5450 C5: `sample_tokens` must slice per request for num_reqs > 1.
+ The tensor branch's asserts used to force num_reqs == 1 and the list
+ branch built a length-1 payload misaligned with `req_ids`."""
+
+ def test_tensor_output_sliced_per_request(self):
+ multimodal_outputs = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3)
+ runner = _make_runner(multimodal_outputs, num_reqs=2)
+
+ output = GPUGenerationModelRunner.sample_tokens(runner)
+
+ assert len(output.multimodal_outputs) == 2
+ assert torch.equal(output.multimodal_outputs[0]["model_outputs"], multimodal_outputs[0])
+ assert torch.equal(output.multimodal_outputs[1]["model_outputs"], multimodal_outputs[1])
+
+ def test_list_output_one_entry_per_request(self):
+ entries = [torch.zeros(2, 1), torch.ones(2, 1)]
+ runner = _make_runner(list(entries), num_reqs=2)
+
+ output = GPUGenerationModelRunner.sample_tokens(runner)
+
+ assert len(output.multimodal_outputs) == 2
+ assert torch.equal(output.multimodal_outputs[0]["model_outputs"], entries[0])
+ assert torch.equal(output.multimodal_outputs[1]["model_outputs"], entries[1])
+
+ def test_tensor_output_batch_mismatch_raises(self):
+ runner = _make_runner(torch.randn(1, 2, 3), num_reqs=2)
+
+ with pytest.raises(ValueError, match="leading dim 1 .* 2 requests"):
+ GPUGenerationModelRunner.sample_tokens(runner)
+
+ def test_list_output_batch_mismatch_raises(self):
+ runner = _make_runner([torch.randn(2, 1)], num_reqs=2)
+
+ with pytest.raises(ValueError, match="length 1 .* 2 requests"):
+ GPUGenerationModelRunner.sample_tokens(runner)
diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py
index 383379202..a30dcdb58 100644
--- a/tests/worker/test_native_duplex_hooks.py
+++ b/tests/worker/test_native_duplex_hooks.py
@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
+import ast
import inspect
from types import SimpleNamespace
@@ -848,6 +852,161 @@ def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
assert result["prompt_suffix_len"] == 0
+def _stage0_vision_runtime():
+ import torch
+
+ from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
+ MiniCPMO45Stage0DuplexRuntime,
+ )
+
+ class _StageModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.embed = torch.nn.Embedding(256, 2)
+
+ def get_input_embeddings(self):
+ return self.embed
+
+ def get_audio_hidden_states(self, _data):
+ return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]
+
+ runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
+ runtime.stage_model = _StageModel()
+ runtime.thinker = runtime.stage_model
+ runtime.tokenizer = SimpleNamespace(
+ unk_token_id=0,
+ convert_tokens_to_ids=lambda token: {
+ "": 1,
+ "": 2,
+ "<|listen|>": 3,
+ "<|speak|>": 4,
+ "<|tts_bos|>": 5,
+ "<|tts_eos|>": 6,
+ "<|tts_pad|>": 7,
+ "<|chunk_eos|>": 8,
+ "<|chunk_tts_eos|>": 9,
+ "<|turn_eos|>": 10,
+ "<|audio|>": 11,
+ "": 12,
+ "": 13,
+ "": 14,
+ "": 15,
+ }.get(token, 0),
+ encode=lambda text, add_special_tokens=False: [201, 5],
+ )
+ runtime.processor = SimpleNamespace(get_streaming_chunk_size=lambda: 4)
+ runtime.device = "cpu"
+ runtime._init_token_ids()
+ # One 2-row block per frame keeps the expected token_ids readable; the real
+ # resampler emits VISION_EMBEDS_PER_FRAME rows per slice.
+ runtime._stage_vision_embeddings = lambda frames: [
+ [torch.full((2, 2), float(index), dtype=torch.float32)] for index in range(len(frames))
+ ]
+ return runtime
+
+
+def test_minicpmo_stage0_puts_every_frame_of_an_append_in_one_unit():
+ """Official streaming_prefill feeds the whole frame_list into one unit.
+
+ Stacked pairs used to be spread one frame per unit, so a two-frame append
+ that closed a single unit left a frame over, failed the prefill, and still
+ advanced the session state.
+ """
+ from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
+ _MiniCPMO45Stage0SessionState,
+ )
+
+ runtime = _stage0_vision_runtime()
+ state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-stacked-frames")
+
+ result = runtime._stage_prefill_embeddings_only(
+ state,
+ np.zeros(4, dtype=np.float32),
+ video_frames=["frame-prev", "frame-cur"],
+ seq=1,
+ )
+
+ assert result["success"] is True
+ # , then each frame as + embeds + , then the audio.
+ assert result["input_token_ids"] == [1, 12, 0, 0, 13, 12, 0, 0, 13, 11]
+ assert state.audio_chunk_idx == 1
+
+
+def test_minicpmo_stage0_wraps_hd_slices_after_the_source_image():
+ """Official stacked pair is max_slice_nums=[2, 1]: source + slices, then audio."""
+ import torch
+
+ from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
+ _MiniCPMO45Stage0SessionState,
+ )
+
+ runtime = _stage0_vision_runtime()
+ state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-hd-slices")
+ runtime._stage_vision_embeddings = lambda frames: [
+ [torch.ones((2, 2)), torch.full((2, 2), 2.0), torch.full((2, 2), 3.0)],
+ [torch.full((2, 2), 4.0)],
+ ]
+
+ result = runtime._stage_prefill_embeddings_only(
+ state,
+ np.zeros(4, dtype=np.float32),
+ video_frames=["base", "stack"],
+ seq=1,
+ )
+
+ assert result["success"] is True
+ # src p0 p1
+ # stack audio
+ assert result["input_token_ids"] == [
+ 1,
+ 12,
+ 0,
+ 0,
+ 13,
+ 14,
+ 0,
+ 0,
+ 15,
+ 14,
+ 0,
+ 0,
+ 15,
+ 12,
+ 0,
+ 0,
+ 13,
+ 11,
+ ]
+
+
+def test_minicpmo_stage0_reports_frames_that_close_no_unit_instead_of_dropping_them():
+ """A frame only binds to the unit its append closes.
+
+ Frames are not carried across appends, so one arriving early used to vanish
+ under a generic "audio not enough" buffering result.
+ """
+ from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
+ _MiniCPMO45Stage0SessionState,
+ )
+
+ runtime = _stage0_vision_runtime()
+ state = _MiniCPMO45Stage0SessionState(session_id="sid-stage0-early-frame")
+
+ assert runtime._stage_prefill_embeddings_only(state, np.zeros(4, dtype=np.float32), seq=1)["success"] is True
+
+ early = runtime._stage_prefill_embeddings_only(
+ state,
+ np.zeros(2, dtype=np.float32),
+ video_frames=["frame-cur"],
+ seq=2,
+ )
+
+ assert early["success"] is False
+ assert "closes no model unit" in early["reason"]
+ # The half-filled append must not advance the unit counter.
+ assert state.audio_chunk_idx == 1
+
+
def test_minicpmo_stage0_speech_append_sets_pending_context_once():
import torch
@@ -2466,3 +2625,128 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio():
assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5]
assert len(state.context_embeds) == 6
+
+
+# ---------------------------------------------------------------------------
+# Cross-platform wiring: every AR model runner must reach the duplex hook.
+#
+# GPUARModelRunner and NPUARModelRunner are siblings -- neither inherits the
+# other's _sample -- so the NPU runner once copied GPU _sample without the
+# prepare_duplex_sampling call. Parse source instead of importing NPUARModelRunner:
+# that class needs vllm_ascend, which CPU test images do not have.
+# ---------------------------------------------------------------------------
+
+_REQUIRED_DUPLEX_CALLS = {
+ "__init__": "_init_duplex_sampling_state",
+ "load_model": "_resolve_duplex_sampling_hook",
+ "_update_states": "_update_duplex_sampling_states",
+ "_sample": "_apply_duplex_sampling",
+}
+
+_MODEL_SAMPLER_CALL = "model_sample"
+
+
+def _repo_root():
+ from pathlib import Path
+
+ return Path(__file__).parents[2]
+
+
+def _ar_runner_sources():
+ root = _repo_root()
+ paths = [root / "vllm_omni/worker/gpu_ar_model_runner.py"]
+ paths.extend(sorted(root.glob("vllm_omni/platforms/*/worker/*_ar_model_runner.py")))
+ return paths
+
+
+def _ar_runner_classdefs():
+ root = _repo_root()
+ found = {}
+ for path in _ar_runner_sources():
+ module = ast.parse(path.read_text())
+ for node in module.body:
+ if isinstance(node, ast.ClassDef) and node.name.endswith("ARModelRunner"):
+ found[node.name] = (node, str(path.relative_to(root)))
+ assert found, "no AR model runners discovered; the glob in _ar_runner_sources() is stale"
+ return found
+
+
+def _method_calls(method):
+ calls: dict[str, int] = {}
+ for node in ast.walk(method):
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
+ calls.setdefault(node.func.attr, node.lineno)
+ return calls
+
+
+def _wires_duplex_sampling_directly(classdef) -> bool:
+ bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+ if "DuplexSamplingRunnerMixin" not in bases:
+ return False
+ methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
+ return all(
+ method_name in methods and required_call in _method_calls(methods[method_name])
+ for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items()
+ )
+
+
+def test_every_ar_runner_reaches_the_duplex_sampling_hook():
+ classdefs = _ar_runner_classdefs()
+ wired = {name for name, (classdef, _path) in classdefs.items() if _wires_duplex_sampling_directly(classdef)}
+
+ for class_name, (classdef, path) in sorted(classdefs.items()):
+ if class_name in wired:
+ continue
+ base_names = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+ assert base_names & wired, (
+ f"{path}::{class_name} neither wires the duplex sampling hook nor derives from a "
+ f"runner that does (wired runners: {sorted(wired)}). Inherit DuplexSamplingRunnerMixin "
+ f"and call {sorted(_REQUIRED_DUPLEX_CALLS.values())} from "
+ f"{sorted(_REQUIRED_DUPLEX_CALLS)} respectively."
+ )
+
+
+@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
+def test_ar_runner_calls_every_duplex_sampling_hook_site(class_name: str):
+ classdefs = _ar_runner_classdefs()
+ assert class_name in classdefs, f"{class_name} not found by _ar_runner_sources()"
+ classdef, path = classdefs[class_name]
+
+ bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)}
+ assert "DuplexSamplingRunnerMixin" in bases, (
+ f"{path}::{class_name} must inherit DuplexSamplingRunnerMixin so the duplex "
+ f"sampling hook is wired the same way on every platform"
+ )
+
+ methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)}
+ for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items():
+ method = methods.get(method_name)
+ assert method is not None, f"{path}::{class_name}.{method_name} is missing; it must call {required_call}"
+ calls = _method_calls(method)
+ assert required_call in calls, f"{path}::{class_name}.{method_name} must call self.{required_call}"
+
+
+@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"])
+def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: str):
+ classdef, path = _ar_runner_classdefs()[class_name]
+ sample = next(node for node in classdef.body if isinstance(node, ast.FunctionDef) and node.name == "_sample")
+
+ apply_lineno = _method_calls(sample).get("_apply_duplex_sampling")
+ sampler_lineno = next(
+ (
+ node.lineno
+ for node in ast.walk(sample)
+ if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == _MODEL_SAMPLER_CALL
+ ),
+ None,
+ )
+ assert apply_lineno is not None, f"{path}::{class_name}._sample must call self._apply_duplex_sampling"
+ assert sampler_lineno is not None, (
+ f"{path}::{class_name}._sample no longer calls {_MODEL_SAMPLER_CALL}(); update _MODEL_SAMPLER_CALL "
+ f"so the ordering check keeps testing something"
+ )
+ assert apply_lineno < sampler_lineno, (
+ f"{path}::{class_name}._sample calls self._apply_duplex_sampling at line {apply_lineno}, after "
+ f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}. The hook masks the logits and publishes the "
+ f"row -> session map the sampler reads, so running it afterwards is a silent no-op."
+ )
diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py
index ca282994a..81d2a6b8a 100644
--- a/tests/worker/test_omni_gpu_model_runner.py
+++ b/tests/worker/test_omni_gpu_model_runner.py
@@ -1,9 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from contextlib import contextmanager
from types import SimpleNamespace
import pytest
import torch
from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
from vllm_omni.worker.gpu_model_runner import (
@@ -107,6 +111,23 @@ class DummyReqState:
pass
+def test_model_forward_passes_request_ids_to_decode_metadata(monkeypatch):
+ received = {}
+ model = SimpleNamespace(
+ supports_omni_decode_step_metadata=True,
+ update_decode_step_metadata=lambda **kwargs: received.update(kwargs),
+ )
+ runner = object.__new__(OmniGPUModelRunner)
+ runner.model = model
+ runner.input_batch = DummyInputBatch(["request-a", "request-b"])
+ runner._build_model_kwargs_extra = lambda: {}
+ monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: torch.zeros(1))
+
+ OmniGPUModelRunner._model_forward(runner, input_ids=torch.ones(2, dtype=torch.long))
+
+ assert received["req_ids"] == ["request-a", "request-b"]
+
+
class MiMoAudioForConditionalGeneration(torch.nn.Module):
"""Dummy model whose class name must exactly match the production check."""
@@ -543,6 +564,80 @@ def test_update_additional_information_deserializes_new_request_payload():
)
+def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot():
+ from vllm_omni.engine.serialization import serialize_additional_information
+
+ runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+ runner.model.replace_runtime_additional_information = True
+ terminal = {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+ }
+ peer = {
+ "codes": {"audio": torch.tensor([9])},
+ "meta": {"cache_epoch": 3, "chunk_seq": 1, "last_chunk": False},
+ }
+ runner.model_intermediate_buffer.update(r1=terminal, r2=peer)
+ marker = {
+ "meta": {
+ "finished": False,
+ "is_segment_finished": True,
+ "request_finished": False,
+ "replace_runtime_additional_information": True,
+ }
+ }
+ new_req = SimpleNamespace(
+ req_id="r1",
+ model_intermediate_buffer=marker,
+ additional_information=serialize_additional_information(terminal),
+ )
+
+ OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1")
+ OmniGPUModelRunner._update_additional_information(
+ runner,
+ SimpleNamespace(
+ scheduled_new_reqs=[new_req],
+ scheduled_cached_reqs=SimpleNamespace(),
+ ),
+ )
+
+ info = runner.model_intermediate_buffer["r1"]
+ assert "codes" not in info
+ assert info["meta"] == {
+ **marker["meta"],
+ "num_processed_tokens": 0,
+ "resumable": True,
+ }
+ assert runner.requests["r1"].additional_information_cpu == info
+ assert runner.model_intermediate_buffer["r2"] == peer
+
+
+def test_cached_empty_marker_replaces_terminal_chunk_snapshot():
+ runner = _make_runner(req_ids=("r1",), hidden_size=4)
+ runner.model.replace_runtime_additional_information = True
+ runner.model_intermediate_buffer["r1"] = {
+ "codes": {"audio": torch.tensor([1, 2])},
+ "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+ }
+ marker = {
+ "meta": {
+ "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+ "replace_runtime_additional_information": True,
+ }
+ }
+
+ OmniGPUModelRunner._update_additional_information(
+ runner,
+ SimpleNamespace(
+ scheduled_new_reqs=[],
+ scheduled_cached_reqs=SimpleNamespace(additional_information={"r1": marker}),
+ ),
+ )
+
+ assert runner.model_intermediate_buffer["r1"] == marker
+ assert runner.requests["r1"].additional_information_cpu == marker
+
+
def test_update_intermediate_buffer_skips_empty_update():
"""Validate that an empty update dict is a no-op."""
runner = _make_runner(req_ids=("r1",), hidden_size=4)
diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py
index 96e1ea5d3..305b0778f 100644
--- a/tests/worker/test_worker_base.py
+++ b/tests/worker/test_worker_base.py
@@ -199,7 +199,14 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
worker = object.__new__(OmniGPUWorkerBase)
worker.rank = 0
worker.device = "cuda:0"
- monkeypatch.setattr(base, "current_omni_platform", _fake_platform())
+ order: list[str] = []
+ platform = SimpleNamespace(
+ get_current_memory_usage=lambda device: 0,
+ empty_cache=lambda: order.append("empty_cache"),
+ synchronize=lambda: order.append("sync"),
+ collect=lambda: None,
+ )
+ monkeypatch.setattr(base, "current_omni_platform", platform)
calls = {}
@@ -209,12 +216,17 @@ def test_sleep_offload_tags_by_level(monkeypatch, level, expected_tags):
return cls()
def sleep(self, offload_tags):
+ order.append("allocator_sleep")
calls["offload_tags"] = offload_tags
monkeypatch.setattr(cumem_mod, "CuMemAllocator", FakeAllocator)
assert OmniGPUWorkerBase.sleep(worker, level=level) is True
assert calls["offload_tags"] == expected_tags
+ assert order[0] == "sync"
+ assert "allocator_sleep" in order
+ assert order.index("sync") < order.index("allocator_sleep")
+ assert "empty_cache" not in order
def test_wake_up_forwards_tags_to_allocator(monkeypatch):
diff --git a/tools/baselines/omni_hash_baseline.py b/tools/baselines/omni_hash_baseline.py
index 1bb93099b..a1644aeb3 100644
--- a/tools/baselines/omni_hash_baseline.py
+++ b/tools/baselines/omni_hash_baseline.py
@@ -147,7 +147,7 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
# (see tests/e2e/accuracy/qwen3_omni). Confirm the exact knob at bench time.
omni = Omni(
model,
- stage_configs_path=yaml_path,
+ deploy_config=yaml_path,
stage_init_timeout=1800,
init_timeout=3600,
async_chunk=async_mode,
@@ -159,11 +159,10 @@ def _run_once(model: str, yaml_path: str, *, async_mode: bool) -> list[dict]:
wav: bytes = b""
# Batch size 1: submit a single prompt so scheduling is deterministic.
for step in omni.generate([_text_prompt(question)], _sampling_params(), py_generator=False):
- out = step.request_output
if step.final_output_type == "text":
- text = out.outputs[0].text # exact text, no strip -> truly bit-identical
+ text = step.outputs[0].text # exact text, no strip -> truly bit-identical
elif step.final_output_type == "audio":
- wav = _audio_bytes(out.outputs[0].multimodal_output["audio"])
+ wav = _audio_bytes(step.outputs[0].multimodal_output["audio"])
if text is None:
raise RuntimeError(f"prompt_id={prompt_id} async={async_mode}: produced no text output")
if not wav:
diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py
index 740f896f5..06f09c791 100644
--- a/tools/nightly/buildkite_testcase_statistics.py
+++ b/tools/nightly/buildkite_testcase_statistics.py
@@ -1,8 +1,7 @@
#!/usr/bin/env python3
"""
-Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml,
-and test-weekly.yml; collect test cases (including parametrized) via pytest --collect-only -q
-and produce an HTML report.
+Parse pytest commands from Buildkite .buildkite/cuda/test-{ready,merge,nightly,weekly}.yml;
+collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report.
Leaf steps may live under ``group:`` blocks (nested ``steps``); those are flattened with the
group title carried into the report.
@@ -30,7 +29,7 @@ import yaml
# Repo root (parent of the directory containing this script)
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
-BUILDKITE_DIR = REPO_ROOT / ".buildkite"
+BUILDKITE_DIR = REPO_ROOT / ".buildkite" / "cuda"
PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml", "test-weekly.yml"]
@@ -256,13 +255,34 @@ def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_l
rel_dir = Path(target.rstrip("/"))
dir_path = (repo_root / rel_dir).resolve()
if dir_path.exists() and dir_path.is_dir():
- return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 120
+ return [str(dir_path)], target.replace("\\", "/").rstrip("/"), 300
if not extra:
raise RuntimeError(f"Failed to parse pytest sidecar args from line: {raw_line!r}")
if not (repo_root / "tests").exists():
raise FileNotFoundError("tests/ directory not found under repo root")
- return ["tests/"], "tests/", 120
+ # Broad ``tests/`` collection loads many conftests; allow more time on busy CI agents.
+ return ["tests/"], "tests/", 300
+
+
+def _format_pytest_collect_failure(
+ path_args: list[str],
+ extra: list[str],
+ result: subprocess.CompletedProcess[str],
+) -> str:
+ """pytest often puts collection errors on stdout; include both streams + exit code."""
+ stdout = (result.stdout or "").strip()
+ stderr = (result.stderr or "").strip()
+ parts = [
+ f"pytest --collect-only failed for {path_args} with args {extra} (exit={result.returncode})",
+ ]
+ if stderr:
+ parts.append(f"stderr:\n{stderr}")
+ if stdout:
+ parts.append(f"stdout:\n{stdout}")
+ if not stderr and not stdout:
+ parts.append("(no stdout/stderr captured)")
+ return "\n".join(parts)
def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]:
@@ -273,6 +293,8 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
flags ``--ignore``, ``--test-config-file``, ``-m``, ``-k``, ``--run-level`` are taken from the
full ``raw_line`` and appended to the same invocation so collection matches CI (perf scripts
parametrize from ``--test-config-file`` at import time).
+
+ Pytest exit code 5 (no tests collected) is treated as an empty result, not a hard failure.
"""
extra = _pytest_collect_sidecar_args(raw_line)
path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line)
@@ -286,8 +308,15 @@ def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]
text=True,
timeout=timeout_quiet,
)
+ # 5 == no tests collected under the given -m/-k filters.
+ if result.returncode == 5:
+ print(
+ f"pytest --collect-only: no tests matched for {path_args} with args {extra}",
+ file=sys.stderr,
+ )
+ return []
if result.returncode != 0:
- raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}")
+ raise RuntimeError(_format_pytest_collect_failure(path_args, extra, result))
print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}")
return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "")
@@ -962,7 +991,7 @@ def main() -> None:
"--buildkite-dir",
type=Path,
default=BUILDKITE_DIR,
- help="Path to .buildkite directory",
+ help="Path to Buildkite CUDA pipeline directory (default: .buildkite/cuda)",
)
default_out = REPO_ROOT / "buildkite_testcase_statistics.html"
parser.add_argument(
@@ -1023,7 +1052,12 @@ def main() -> None:
names = []
skipped_ids: set[str] = set()
for target, raw in targets:
- collected = collect_test_names(target, REPO_ROOT, raw)
+ try:
+ collected = collect_test_names(target, REPO_ROOT, raw)
+ except (RuntimeError, FileNotFoundError, subprocess.TimeoutExpired) as exc:
+ # Keep generating the report for other steps; surface the root cause.
+ print(f"WARNING: collect failed for {module_key} target={target!r}: {exc}", file=sys.stderr)
+ collected = []
skipped_ids.update(get_skip_status(target, raw, REPO_ROOT))
names.extend(collected)
count = len(names) - len(skipped_ids)
diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh
index a95eb8876..b5b6b92ce 100644
--- a/tools/nightly/run_nightly_jobs.sh
+++ b/tools/nightly/run_nightly_jobs.sh
@@ -20,6 +20,8 @@
# acc — label contains "Accuracy Test" (incl. GEBench / GEdit-Bench style)
# function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.)
# stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below)
+# Uses --pytest-mark / PYTEST_MARK for pytest -m (hardware via JSON marks).
+# When unset, -m defaults to H800 (CUDA stability fleet; use --pytest-mark A3 on NPU).
# local — pytest -sv -m " and local_model" from repo root (no YAML step extract)
# When LABEL_SUBSTR is set, also runs matching tests/**/test_*.py and perf JSON configs under
# tests/dfx/perf/tests/*.json via run_benchmark.py / run_diffusion_benchmark.py.
@@ -42,8 +44,10 @@
# whose filename contains the substring (benchmark runner chosen by JSON family).
#
# stability (when included in TEST_TYPE):
-# From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py
-# model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six
+# From repo root: pytest -s -v --run-level full_model -m "" tests/dfx/stability/scripts/...
+# -m comes from --pytest-mark / PYTEST_MARK when set (e.g. H800 or A3).
+# Otherwise -m defaults to H800. PYTEST_MARK is rejected unless stability is enabled.
+# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22 + hunyuan_image; all → all five
# LABEL_SUBSTR: if set, script path / job key / filename must contain it
#
# Requirements: bash, python3, PyYAML (pip install pyyaml)
@@ -73,6 +77,7 @@
# MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above
# LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename;
# local: substring of test_*.py basename or tests/dfx/perf/tests/*.json filename
+# PYTEST_MARK - pytest -m expression for stability jobs only (also --pytest-mark)
# DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest
# RUN_JOB_TIMEOUT_KILL_AFTER - seconds after SIGTERM before SIGKILL on inline pytest
# timeout (default: 60; baked into generated job scripts)
@@ -82,6 +87,7 @@ set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LABEL_SUBSTR="${LABEL_SUBSTR:-}"
+PYTEST_MARK="${PYTEST_MARK:-}"
TEST_TYPE_ENV="${TEST_TYPE:-all}"
MODEL_TYPE_ENV="${MODEL_TYPE:-all}"
TEST_TYPE_CLI_PARTS=()
@@ -204,6 +210,10 @@ while [[ $# -gt 0 ]]; do
LABEL_SUBSTR="$2"
shift 2
;;
+ --pytest-mark)
+ PYTEST_MARK="$2"
+ shift 2
+ ;;
--model-type)
MODEL_TYPE_FROM_CLI=1
_split_append_csv_array MODEL_TYPE_CLI_PARTS "$2"
@@ -339,7 +349,7 @@ fi
mkdir -p "${LOG_DIR}/jobs"
LOG_DIR="$(cd "${LOG_DIR}" && pwd)"
-export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN
+export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN PYTEST_MARK
# Drop stale wrapper scripts, old tee logs, and perf manifest so a filtered run does not
# mix with the previous outputs.
@@ -588,7 +598,6 @@ def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -
STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [
("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)),
("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)),
- ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)),
("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)),
("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)),
("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)),
@@ -601,6 +610,26 @@ def _stability_model_matches(model_types: list[str], families: tuple[str, ...])
return bool(set(model_types) & set(families))
+STABILITY_DEFAULT_PYTEST_MARK = "H800"
+
+
+def _shell_quote_pytest_mark(marker_expr: str) -> str:
+ """Shell-safe pytest ``-m`` value for generated bash job scripts."""
+ expr = (marker_expr or "").strip()
+ if not expr:
+ raise ValueError("pytest -m expression must not be empty")
+ return shlex.quote(expr)
+
+
+def _stability_pytest_mark(model_types: list[str]) -> str:
+ """Return explicit --pytest-mark or the default CUDA stability hardware marker."""
+ del model_types # MODEL_TYPE filters scripts; default -m is hardware-specific.
+ override = (os.environ.get("PYTEST_MARK") or "").strip()
+ if override:
+ return override
+ return STABILITY_DEFAULT_PYTEST_MARK
+
+
def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None:
body = "\n".join(script_lines) + "\n"
if DRY_RUN:
@@ -713,7 +742,7 @@ def run_local_mode(
if rel_paths:
paths_arg = " ".join(shlex.quote(p) for p in rel_paths)
- pytest_line = f'pytest -sv -m "{marker_expr}" {paths_arg}'
+ pytest_line = f"pytest -sv -m {_shell_quote_pytest_mark(marker_expr)} {paths_arg}"
slug = _local_job_slug(LABEL_SUBSTR)
key = f"local_pytest_{slug}"
header = (
@@ -733,7 +762,7 @@ def run_local_mode(
for json_rel, runner_rel in perf_pairs:
json_name = Path(json_rel).name
pytest_line = (
- f'pytest -s -v -m "{marker_expr}" '
+ f"pytest -s -v -m {_shell_quote_pytest_mark(marker_expr)} "
f"{shlex.quote(runner_rel)} "
f"--test-config-file {shlex.quote(json_rel)}"
)
@@ -764,7 +793,7 @@ def run_local_mode(
)
return matched
- pytest_line = f'pytest -sv -m "{marker_expr}"'
+ pytest_line = f"pytest -sv -m {_shell_quote_pytest_mark(marker_expr)}"
key = "local_pytest"
header = f"# Local marker tests — MODEL_TYPE={','.join(model_types)}"
script_lines = [
@@ -780,6 +809,8 @@ def run_local_mode(
def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int:
matched = 0
+ marker_expr = _stability_pytest_mark(model_types)
+ mark_from_cli = bool((os.environ.get("PYTEST_MARK") or "").strip())
for key, rel_posix, families in STABILITY_CASES:
if not _stability_model_matches(model_types, families):
continue
@@ -795,11 +826,17 @@ def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int:
if not rel_path.is_file():
print(f"# skip (missing file): {rel_path}", file=sys.stderr)
continue
- pytest_line = f"pytest -s -v --run-level full_model {rel_posix}"
+ pytest_line = (
+ f"pytest -s -v --run-level full_model "
+ f"-m {_shell_quote_pytest_mark(marker_expr)} "
+ f"{shlex.quote(rel_posix)}"
+ )
matched += 1
script_lines = [
"#!/usr/bin/env bash",
f"# Stability: {key} — {rel_posix}",
+ f"# pytest -m {marker_expr!r}"
+ + (" (--pytest-mark)" if mark_from_cli else f" (default {STABILITY_DEFAULT_PYTEST_MARK})"),
"set -euo pipefail",
f'cd "{REPO_ROOT}"',
pytest_line,
@@ -817,6 +854,14 @@ def main() -> None:
yaml_test_types = [t for t in test_types if t not in ("stability", "local")]
needs_yaml = bool(yaml_test_types)
+ pytest_mark = (os.environ.get("PYTEST_MARK") or "").strip()
+ if pytest_mark and not want_stability:
+ print(
+ "PYTEST_MARK / --pytest-mark is only supported when --test-type includes stability",
+ file=sys.stderr,
+ )
+ sys.exit(2)
+
jobs_dir = LOG_DIR / "jobs"
if not DRY_RUN:
jobs_dir.mkdir(parents=True, exist_ok=True)
diff --git a/tools/pre_commit/check_buildkite.py b/tools/pre_commit/check_buildkite.py
new file mode 100644
index 000000000..362253494
--- /dev/null
+++ b/tools/pre_commit/check_buildkite.py
@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Validate .buildkite YAML against the official Buildkite schema.
+
+Test pipelines (test-ready.yml, etc.) use Omni-only keys such as
+``mirror_hardwares`` and ``source_file_dependencies``. Those are expanded or
+stripped the same way ``upload_pipeline.py`` does before upload, then the
+rendered YAML is checked with check-jsonschema.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+ROOT = Path(__file__).resolve().parents[2]
+SCRIPTS = ROOT / ".buildkite" / "common" / "scripts"
+sys.path.insert(0, str(SCRIPTS))
+
+from upload_pipeline import _render_test_pipeline # noqa: E402
+
+SKIP_FILES = {
+ ".buildkite/common/ci_mirror_hardwares.yml",
+ ".buildkite/npu/pipeline-npu-a3.yml",
+}
+
+# AMD pipelines are uploaded by a different path and use list-valued
+# mirror_hardwares plus agent_pool/grade, which CUDA's expander rejects.
+AMD_STRIP_KEYS = frozenset({"mirror_hardwares", "agent_pool", "grade", "source_file_dependencies"})
+
+
+def _normalize(path: str) -> str:
+ return path.replace("\\", "/")
+
+
+def _strip_keys(obj: Any, keys: frozenset[str]) -> Any:
+ if isinstance(obj, dict):
+ return {k: _strip_keys(v, keys) for k, v in obj.items() if k not in keys}
+ if isinstance(obj, list):
+ return [_strip_keys(item, keys) for item in obj]
+ return obj
+
+
+def _prepare_pipeline(path: Path) -> dict[str, Any]:
+ rel = _normalize(str(path.relative_to(ROOT) if path.is_absolute() else path))
+ doc = yaml.safe_load(path.read_text(encoding="utf-8"))
+ if not isinstance(doc, dict):
+ raise ValueError(f"{rel}: expected a mapping (pipeline YAML), got {type(doc).__name__}")
+
+ if "/amd/" in f"/{rel}":
+ return _strip_keys(doc, AMD_STRIP_KEYS)
+
+ return _render_test_pipeline(doc, changed_files=None)
+
+
+def _validate(path: Path, doc: dict[str, Any]) -> int:
+ with tempfile.NamedTemporaryFile(
+ mode="w",
+ suffix=".yml",
+ encoding="utf-8",
+ delete=False,
+ ) as handle:
+ yaml.safe_dump(doc, handle, sort_keys=False)
+ tmp = Path(handle.name)
+
+ try:
+ result = subprocess.run(
+ [
+ sys.executable,
+ "-m",
+ "check_jsonschema",
+ "--builtin-schema",
+ "vendor.buildkite",
+ str(tmp),
+ ],
+ capture_output=True,
+ text=True,
+ )
+ finally:
+ tmp.unlink(missing_ok=True)
+
+ if result.returncode == 0:
+ return 0
+
+ rel = _normalize(str(path))
+ sys.stderr.write(f"{rel}: Buildkite schema validation failed\n")
+ sys.stderr.write(result.stdout or "")
+ sys.stderr.write(result.stderr or "")
+ return 1
+
+
+def main(argv: list[str]) -> int:
+ rc = 0
+ for raw in argv:
+ path = Path(raw)
+ rel = _normalize(str(path))
+ if rel in SKIP_FILES:
+ continue
+ if not path.is_file():
+ path = ROOT / path
+ try:
+ doc = _prepare_pipeline(path)
+ except Exception as exc: # noqa: BLE001 - surface expander errors as hook failures
+ sys.stderr.write(f"{rel}: {exc}\n")
+ rc = 1
+ continue
+ rc |= _validate(path, doc)
+ return rc
+
+
+if __name__ == "__main__":
+ raise SystemExit(main(sys.argv[1:]))
diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py
new file mode 100644
index 000000000..926e102fa
--- /dev/null
+++ b/tools/pre_commit/check_forbidden_imports.py
@@ -0,0 +1,367 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ban dangerous / policy-violating imports unless a file is allowlisted.
+
+Ported from vllm/tools/pre_commit/check_forbidden_imports.py. Existing
+vllm_omni/ usages are grandfathered; add a path to the matching
+allowed_files set only after a security/policy review.
+"""
+
+import sys
+from dataclasses import dataclass, field
+
+import regex as re
+
+# Hub entry points that must go through the vLLM-tagged repo_utils helpers.
+_HF_NAMES = (
+ r"HfApi|HfFileSystem|hf_hub_download|snapshot_download"
+ r"|list_repo_files|file_exists|try_to_load_from_cache"
+ r"|list_repo_refs|repo_exists"
+)
+
+# Non-library trees: examples, tests, and tooling may keep stdlib imports.
+_NON_LIBRARY_DIRS = {
+ "examples/",
+ "tests/",
+ "benchmarks/",
+ "apps/",
+ "docs/",
+ "tools/",
+ "scripts/",
+ ".buildkite/",
+ ".github/",
+}
+
+
+@dataclass
+class ForbiddenImport:
+ pattern: str
+ tip: str
+ allowed_pattern: re.Pattern = re.compile(r"^$") # matches nothing by default
+ allowed_files: set[str] = field(default_factory=set)
+ allowed_dirs: set[str] = field(default_factory=set)
+
+
+CHECK_IMPORTS = {
+ "pickle/cloudpickle": ForbiddenImport(
+ pattern=(
+ r"^\s*(import\s+(?:[\w.]+\s*(?:as\s+\w+)?\s*,\s*)*(pickle|cloudpickle)\b"
+ r"|from\s+(pickle|cloudpickle)\s+import\b)"
+ ),
+ tip=("Avoid using pickle or cloudpickle or add this file to tools/pre_commit/check_forbidden_imports.py."),
+ allowed_files={
+ # STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST:
+ # pickle/cloudpickle are unsafe when deserializing untrusted data.
+ "tests/diffusion/attention/test_attention_sp.py",
+ "tests/helpers/process.py",
+ "vllm_omni/diffusion/distributed/group_coordinator.py",
+ },
+ ),
+ "base64": ForbiddenImport(
+ pattern=r"^\s*(?:import\s+base64(?:$|\s|,)|from\s+base64\s+import)",
+ tip=("Replace 'import base64' with 'import pybase64' or 'import pybase64 as base64'."),
+ allowed_pattern=re.compile(r"^\s*import\s+pybase64(\s*|\s+as\s+base64\s*)$"),
+ allowed_dirs=_NON_LIBRARY_DIRS,
+ allowed_files={
+ "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py",
+ "vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py",
+ "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py",
+ "vllm_omni/entrypoints/openai/api_server.py",
+ "vllm_omni/entrypoints/openai/audio_utils_mixin.py",
+ "vllm_omni/entrypoints/openai/image_api_utils.py",
+ "vllm_omni/entrypoints/openai/protocol/images.py",
+ "vllm_omni/entrypoints/openai/realtime_connection.py",
+ "vllm_omni/entrypoints/openai/serving_chat.py",
+ "vllm_omni/entrypoints/openai/serving_speech.py",
+ "vllm_omni/entrypoints/openai/serving_speech_stream.py",
+ "vllm_omni/entrypoints/openai/video_api_utils.py",
+ "vllm_omni/entrypoints/openai/video_stream_base.py",
+ "vllm_omni/entrypoints/openai/video_stream_session.py",
+ "vllm_omni/experimental/fullduplex/client.py",
+ "vllm_omni/experimental/fullduplex/engine/contracts.py",
+ "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py",
+ "vllm_omni/experimental/fullduplex/minicpmo45/adapter.py",
+ "vllm_omni/experimental/fullduplex/minicpmo45/input.py",
+ "vllm_omni/experimental/fullduplex/minicpmo45/runtime.py",
+ "vllm_omni/experimental/fullduplex/minicpmo45/stage0.py",
+ "vllm_omni/experimental/fullduplex/openai/audio.py",
+ "vllm_omni/experimental/fullduplex/openai/realtime_input.py",
+ "vllm_omni/experimental/fullduplex/openai/runtime_bridge.py",
+ "vllm_omni/experimental/fullduplex/openai/serving.py",
+ "vllm_omni/experimental/fullduplex/personaplex/input.py",
+ "vllm_omni/experimental/fullduplex/personaplex/runtime_extension.py",
+ "vllm_omni/experimental/fullduplex/personaplex/stage0.py",
+ "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py",
+ "vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py",
+ "vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_tokenizer.py",
+ "vllm_omni/tokenizers/mammoth_moda2_tokenizer.py",
+ "vllm_omni/transformers_utils/processors/audiox.py",
+ },
+ ),
+ "re": ForbiddenImport(
+ pattern=r"^\s*(?:import\s+re(?:$|\s|,)|from\s+re\s+import)",
+ tip="Replace 'import re' with 'import regex as re' or 'import regex'.",
+ allowed_pattern=re.compile(r"^\s*import\s+regex(\s*|\s+as\s+re\s*)$"),
+ allowed_dirs=_NON_LIBRARY_DIRS,
+ allowed_files={
+ "setup.py",
+ "vllm_omni/benchmarks/data_modules/daily_omni_eval.py",
+ "vllm_omni/config/stage_config.py",
+ "vllm_omni/diffusion/layers/mot/ops/mot_gemm.py",
+ "vllm_omni/diffusion/model_loader/diffusers_loader.py",
+ "vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py",
+ "vllm_omni/diffusion/models/dreamid_omni/fusion.py",
+ "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py",
+ "vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py",
+ "vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py",
+ "vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py",
+ "vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py",
+ "vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py",
+ "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py",
+ "vllm_omni/diffusion/models/minimax_h3/encoder.py",
+ "vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py",
+ "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py",
+ "vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py",
+ "vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py",
+ "vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py",
+ "vllm_omni/distributed/omni_connectors/utils/env.py",
+ "vllm_omni/entrypoints/cli/logo.py",
+ "vllm_omni/entrypoints/openai/serving_speech.py",
+ "vllm_omni/experimental/fullduplex/joyvl/memory/memory.py",
+ "vllm_omni/model_executor/models/common/whisper_vq.py",
+ "vllm_omni/model_executor/models/fish_speech/prompt_utils.py",
+ "vllm_omni/model_executor/models/glm_tts/text_frontend.py",
+ "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py",
+ "vllm_omni/model_executor/models/indextts2/indextts2_talker.py",
+ "vllm_omni/model_executor/models/indextts2/text_processing_v2_5.py",
+ "vllm_omni/model_executor/models/indextts2/utils/common.py",
+ "vllm_omni/model_executor/models/ming_flash_omni/text_processing.py",
+ "vllm_omni/model_executor/models/ming_flash_omni/vision_encoder.py",
+ "vllm_omni/model_executor/models/ming_tts/prompt_assembly.py",
+ "vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/ear_tts_commons.py",
+ "vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py",
+ "vllm_omni/outputs/output_modality.py",
+ "vllm_omni/quantization/mxfp4_config.py",
+ "vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py",
+ "vllm_omni/utils/custom_voice_io.py",
+ },
+ ),
+ "triton": ForbiddenImport(
+ pattern=r"^(from|import)\s+triton(\s|\.|$)",
+ tip="Use 'from vllm.triton_utils import triton' instead.",
+ allowed_pattern=re.compile(r"from vllm.triton_utils import (triton|tl|tl, triton)"),
+ ),
+ "tilelang": ForbiddenImport(
+ pattern=r"^(from|import)\s+tilelang(\s|\.|$)",
+ tip="Use 'from vllm.tilelang_utils import tilelang, T' instead.",
+ allowed_pattern=re.compile(
+ r"from\s+vllm\.tilelang_utils\s+import\s+"
+ r"(tilelang|T|T, tilelang|tilelang, T)\b"
+ ),
+ ),
+ "huggingface_hub repo API": ForbiddenImport(
+ # Catch `from huggingface_hub import `, including parenthesized,
+ # multi-line imports.
+ pattern=(
+ r"^\s*from\s+huggingface_hub\s+import\s*\([^)]*\b(?:" + _HF_NAMES + r")\b"
+ r"|"
+ r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b"
+ ),
+ tip=(
+ "Use the shared, vLLM-tagged helpers from "
+ "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), "
+ "list_repo_files, file_exists) instead of calling "
+ "huggingface_hub directly."
+ ),
+ allowed_dirs=_NON_LIBRARY_DIRS,
+ allowed_files={
+ "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py",
+ "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py",
+ "vllm_omni/benchmarks/data_modules/seed_tts_eval.py",
+ "vllm_omni/diffusion/lora/loader.py",
+ "vllm_omni/diffusion/model_loader/hub_prefetch.py",
+ "vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py",
+ "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py",
+ "vllm_omni/diffusion/models/helios/pipeline_helios.py",
+ "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py",
+ "vllm_omni/diffusion/models/ltx2/ltx2_components.py",
+ "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py",
+ "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py",
+ "vllm_omni/diffusion/models/pi0/pipeline_pi0.py",
+ "vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py",
+ "vllm_omni/diffusion/models/soulx_singer/utils.py",
+ "vllm_omni/diffusion/models/utils.py",
+ "vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py",
+ "vllm_omni/engine/arg_utils.py",
+ "vllm_omni/entrypoints/openai/api_server.py",
+ "vllm_omni/entrypoints/openai/serving_speech.py",
+ "vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py",
+ "vllm_omni/experimental/fullduplex/personaplex/runtime.py",
+ "vllm_omni/experimental/fullduplex/personaplex/serving/server.py",
+ "vllm_omni/model_executor/model_loader/weight_utils.py",
+ "vllm_omni/model_executor/models/audex/checkpoint.py",
+ "vllm_omni/model_executor/models/bagel/bagel.py",
+ "vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py",
+ "vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py",
+ "vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py",
+ "vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py",
+ "vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py",
+ "vllm_omni/model_executor/models/glm_tts/glm_tts.py",
+ "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py",
+ "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py",
+ "vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py",
+ "vllm_omni/model_executor/models/indextts2/preprocess_utils.py",
+ "vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py",
+ "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py",
+ "vllm_omni/model_executor/models/ming_tts/speaker_extractor.py",
+ "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py",
+ "vllm_omni/model_executor/models/omnivoice/omnivoice.py",
+ "vllm_omni/model_executor/models/personaplex/personaplex_mimi.py",
+ "vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py",
+ "vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py",
+ "vllm_omni/transformers_utils/configs/higgs_audio_v3.py",
+ },
+ ),
+}
+
+
+def _normalize_path(path: str) -> str:
+ return path.replace("\\", "/")
+
+
+def check_file(path: str) -> int:
+ path = _normalize_path(path)
+ with open(path, encoding="utf-8") as f:
+ content = f.read()
+ return_code = 0
+ # Check all patterns in the whole file
+ for import_name, forbidden_import in CHECK_IMPORTS.items():
+ # Skip files that are allowed for this import
+ if path in forbidden_import.allowed_files:
+ continue
+ # Skip directories that are allowed for this import
+ if any(path.startswith(prefix) for prefix in forbidden_import.allowed_dirs):
+ continue
+ # Search for forbidden imports
+ for match in re.finditer(forbidden_import.pattern, content, re.MULTILINE):
+ # Check if it's allowed
+ if forbidden_import.allowed_pattern.match(match.group()):
+ continue
+ # Skip matches inside a comment
+ line_start = content.rfind("\n", 0, match.start()) + 1
+ if "#" in content[line_start : match.start()]:
+ continue
+ # Calculate line number from match position
+ line_num = content[: match.start() + 1].count("\n") + 1
+ print(
+ f"{path}:{line_num}: "
+ "\033[91merror:\033[0m " # red color
+ f"Found forbidden import: {import_name}. {forbidden_import.tip}"
+ )
+ return_code = 1
+ return return_code
+
+
+def main():
+ returncode = 0
+ for path in sys.argv[1:]:
+ returncode |= check_file(path)
+ return returncode
+
+
+def test_regex():
+ def matches(rule: str, content: str) -> bool:
+ return bool(re.search(CHECK_IMPORTS[rule].pattern, content, re.MULTILINE))
+
+ pickle_cases = [
+ # Should match
+ ("import pickle", True),
+ ("import cloudpickle", True),
+ ("import pickle as pkl", True),
+ ("import cloudpickle as cpkl", True),
+ ("from pickle import *", True),
+ ("from cloudpickle import dumps", True),
+ ("from pickle import dumps, loads", True),
+ ("from cloudpickle import (dumps, loads)", True),
+ (" import pickle", True),
+ ("\timport cloudpickle", True),
+ ("from pickle import loads", True),
+ ("import pickle, os", True),
+ ("import os, pickle", True),
+ ("import os, pickle as pkl", True),
+ ("import pickle as pkl, os", True),
+ ("import os.path, cloudpickle", True),
+ ("import cloudpickle, json", True),
+ # Should not match
+ ("import somethingelse", False),
+ ("from somethingelse import pickle", False),
+ ("# import pickle", False),
+ ("print('import pickle')", False),
+ ("import pickleas as asdf", False),
+ ]
+ for i, (content, should_match) in enumerate(pickle_cases):
+ result = matches("pickle/cloudpickle", content)
+ assert result == should_match, f"pickle case {i} failed: {content!r} (expected {should_match}, got {result})"
+
+ tilelang_cases = [
+ # Should match
+ ("import tilelang", True),
+ ("import tilelang.language as T", True),
+ ("from tilelang.jit import JITImpl", True),
+ ("from tilelang.jit.kernel import JITKernel", True),
+ # Should not match: indented (local) imports are allowed, mirroring
+ # the "triton" rule, so mocked-module test imports are not flagged.
+ (" import tilelang", False),
+ (" from tilelang.jit import JITImpl", False),
+ ("from vllm.tilelang_utils import tilelang", False),
+ ("from vllm.tilelang_utils import T", False),
+ ("from vllm.tilelang_utils import T, tilelang", False),
+ ("from vllm.tilelang_utils import tilelang, T", False),
+ ("import tilelang_kernels", False),
+ ]
+ for i, (content, should_match) in enumerate(tilelang_cases):
+ result = matches("tilelang", content)
+ assert result == should_match, f"tilelang case {i} failed: {content!r} (expected {should_match}, got {result})"
+
+ hf_cases = [
+ # Should match
+ ("from huggingface_hub import snapshot_download", True),
+ ("from huggingface_hub import hf_hub_download", True),
+ ("from huggingface_hub import HfApi", True),
+ ("from huggingface_hub import HfFileSystem", True),
+ ("from huggingface_hub import list_repo_files", True),
+ ("from huggingface_hub import try_to_load_from_cache", True),
+ (" from huggingface_hub import snapshot_download", True),
+ ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True),
+ ("from huggingface_hub import (snapshot_download)", True),
+ # Parenthesized multi-line import must not bypass the hook
+ ("from huggingface_hub import (\n snapshot_download,\n)", True),
+ (
+ "from huggingface_hub import (\n PyTorchModelHubMixin,\n HfApi,\n)",
+ True,
+ ),
+ # Should not match
+ ("import huggingface_hub", False),
+ ("import huggingface_hub as hf", False),
+ ("from huggingface_hub import PyTorchModelHubMixin", False),
+ ("from huggingface_hub.constants import HF_HUB_CACHE", False),
+ ("from huggingface_hub.utils import EntryNotFoundError", False),
+ ("from vllm.transformers_utils.repo_utils import hf_api", False),
+ ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False),
+ ("# from huggingface_hub import snapshot_download", False),
+ ]
+ for i, (content, should_match) in enumerate(hf_cases):
+ result = matches("huggingface_hub repo API", content)
+ assert result == should_match, (
+ f"huggingface_hub case {i} failed: {content!r} (expected {should_match}, got {result})"
+ )
+
+ print("All regex tests passed.")
+
+
+if __name__ == "__main__":
+ if "--test-regex" in sys.argv:
+ test_regex()
+ else:
+ sys.exit(main())
diff --git a/tools/pre_commit/check_pickle_imports.py b/tools/pre_commit/check_pickle_imports.py
deleted file mode 100644
index 819a7c8c3..000000000
--- a/tools/pre_commit/check_pickle_imports.py
+++ /dev/null
@@ -1,83 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import sys
-
-import regex as re
-
-# List of files (relative to repo root) that are allowed to import pickle or
-# cloudpickle
-#
-# STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST:
-# The pickle and cloudpickle modules are known to be unsafe when deserializing
-# data from potentially untrusted parties. They have resulted in multiple CVEs
-# for vLLM and numerous vulnerabilities in the Python ecosystem more broadly.
-# Before adding new uses of pickle/cloudpickle, please consider safer
-# alternatives like msgpack or pydantic that are already in use in vLLM. Only
-# add to this list if absolutely necessary and after careful security review.
-ALLOWED_FILES = {
- "tests/helpers/process.py",
- "vllm_omni/diffusion/distributed/group_coordinator.py",
- "tests/diffusion/attention/test_attention_sp.py",
-}
-
-PICKLE_RE = re.compile(
- r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)"
- r"|from\s+(pickle|cloudpickle)\s+import\b)"
-)
-
-
-def scan_file(path: str) -> int:
- with open(path, encoding="utf-8") as f:
- for i, line in enumerate(f, 1):
- if PICKLE_RE.match(line):
- print(
- f"{path}:{i}: "
- "\033[91merror:\033[0m " # red color
- "Found pickle/cloudpickle import"
- )
- return 1
- return 0
-
-
-def main():
- returncode = 0
- for filename in sys.argv[1:]:
- if filename in ALLOWED_FILES:
- continue
- returncode |= scan_file(filename)
- return returncode
-
-
-def test_regex():
- test_cases = [
- # Should match
- ("import pickle", True),
- ("import cloudpickle", True),
- ("import pickle as pkl", True),
- ("import cloudpickle as cpkl", True),
- ("from pickle import *", True),
- ("from cloudpickle import dumps", True),
- ("from pickle import dumps, loads", True),
- ("from cloudpickle import (dumps, loads)", True),
- (" import pickle", True),
- ("\timport cloudpickle", True),
- ("from pickle import loads", True),
- # Should not match
- ("import somethingelse", False),
- ("from somethingelse import pickle", False),
- ("# import pickle", False),
- ("print('import pickle')", False),
- ("import pickleas as asdf", False),
- ]
- for i, (line, should_match) in enumerate(test_cases):
- result = bool(PICKLE_RE.match(line))
- assert result == should_match, f"Test case {i} failed: '{line}' (expected {should_match}, got {result})"
- print("All regex tests passed.")
-
-
-if __name__ == "__main__":
- if "--test-regex" in sys.argv:
- test_regex()
- else:
- sys.exit(main())
diff --git a/tools/pre_commit/check_spdx_header.py b/tools/pre_commit/check_spdx_header.py
new file mode 100644
index 000000000..5bfb26034
--- /dev/null
+++ b/tools/pre_commit/check_spdx_header.py
@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Require SPDX license headers on source files.
+
+Ported from vllm/tools/pre_commit/check_spdx_header.py. The copyright line
+uses the vLLM-Omni project name rather than upstream vLLM. Shell scripts
+(``.sh``) are included so wrappers such as ``shellcheck.sh`` cannot keep a
+stale upstream copyright that Python/Rust/proto files would have rewritten.
+``.pyi`` stubs use the same ``#`` header style as ``.py``.
+"""
+
+import sys
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+
+# Older files still carry the upstream vLLM copyright line. Treat it as stale
+# and rewrite it to the Omni line instead of inserting a duplicate header.
+LEGACY_COPYRIGHT_TEXT = "SPDX-FileCopyrightText: Copyright contributors to the vLLM project"
+
+
+@dataclass(frozen=True)
+class HeaderStyle:
+ """Comment syntax and preamble handling for an SPDX header."""
+
+ comment_prefix: str
+ preserve_shebang: bool = False
+
+
+class SPDXStatus(Enum):
+ """SPDX header status enumeration"""
+
+ EMPTY = "empty" # empty __init__.py
+ COMPLETE = "complete"
+ MISSING_LICENSE = "missing_license" # Only has copyright line
+ MISSING_COPYRIGHT = "missing_copyright" # Only has license line
+ MISSING_BOTH = "missing_both" # Completely missing
+
+
+LICENSE_TEXT = "SPDX-License-Identifier: Apache-2.0"
+COPYRIGHT_TEXT = "SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project"
+FILE_STYLES = {
+ ".py": HeaderStyle("#", preserve_shebang=True),
+ ".pyi": HeaderStyle("#"),
+ ".sh": HeaderStyle("#", preserve_shebang=True),
+ ".rs": HeaderStyle("//"),
+ ".proto": HeaderStyle("//"),
+}
+
+
+def file_style(file_path):
+ """Return the declared header style for a file."""
+ suffix = Path(file_path).suffix
+ try:
+ return FILE_STYLES[suffix]
+ except KeyError:
+ raise ValueError(f"Unsupported file type: {file_path}") from None
+
+
+def spdx_header(style):
+ """Return the SPDX header for a file style."""
+ license_line = f"{style.comment_prefix} {LICENSE_TEXT}"
+ copyright_line = f"{style.comment_prefix} {COPYRIGHT_TEXT}"
+ return license_line, copyright_line
+
+
+def legacy_copyright_line(style):
+ """Return the stale upstream vLLM copyright line for a file style."""
+ return f"{style.comment_prefix} {LEGACY_COPYRIGHT_TEXT}"
+
+
+def header_insertion_index(style, lines):
+ """Return the line index where a missing header should be inserted."""
+ if style.preserve_shebang and lines and lines[0].startswith("#!"):
+ return 1
+ return 0
+
+
+def check_spdx_header_status(file_path):
+ """Check SPDX header status of the file"""
+ style = file_style(file_path)
+ license_line, copyright_line = spdx_header(style)
+ with open(file_path, encoding="UTF-8") as file:
+ lines = file.readlines()
+ if not lines:
+ # Empty file
+ return SPDXStatus.EMPTY
+
+ has_license = False
+ has_copyright = False
+
+ # Check all lines for SPDX headers (not just the first two)
+ for raw_line in lines:
+ line = raw_line.strip()
+ if line == license_line:
+ has_license = True
+ elif line == copyright_line:
+ has_copyright = True
+
+ # Determine status based on what we found
+ if has_license and has_copyright:
+ return SPDXStatus.COMPLETE
+ elif has_license and not has_copyright:
+ # Only has license line (or a stale vLLM copyright line)
+ return SPDXStatus.MISSING_COPYRIGHT
+ elif not has_license and has_copyright:
+ # Only has copyright line
+ return SPDXStatus.MISSING_LICENSE
+ else:
+ # Completely missing both lines
+ return SPDXStatus.MISSING_BOTH
+
+
+def add_header(file_path, status):
+ """Add or supplement SPDX header based on status"""
+ style = file_style(file_path)
+ license_line, copyright_line = spdx_header(style)
+ stale_copyright_line = legacy_copyright_line(style)
+ full_spdx_header = f"{license_line}\n{copyright_line}"
+ with open(file_path, "r+", encoding="UTF-8") as file:
+ lines = file.readlines()
+ file.seek(0, 0)
+ file.truncate()
+
+ if status == SPDXStatus.MISSING_BOTH:
+ # Completely missing, add complete header
+ insertion_index = header_insertion_index(style, lines)
+ file.writelines(lines[:insertion_index])
+ file.write(full_spdx_header + "\n")
+ remaining_lines = lines[insertion_index:]
+ if remaining_lines and remaining_lines[0].strip():
+ file.write("\n")
+ file.writelines(remaining_lines)
+
+ elif status == SPDXStatus.MISSING_COPYRIGHT:
+ # Only has license line, need to add copyright line.
+ # Rewrite a stale upstream vLLM copyright line in place.
+ replaced = False
+ for i, line in enumerate(lines):
+ if line.strip() == stale_copyright_line:
+ lines[i] = f"{copyright_line}\n"
+ replaced = True
+ break
+ if not replaced:
+ for i, line in enumerate(lines):
+ if line.strip() == license_line:
+ lines.insert(i + 1, f"{copyright_line}\n")
+ break
+
+ file.writelines(lines)
+
+ elif status == SPDXStatus.MISSING_LICENSE:
+ # Only has copyright line, need to add license line
+ # Find the copyright line and add license line before it
+ for i, line in enumerate(lines):
+ if line.strip() == copyright_line:
+ # Insert license line before copyright line
+ lines.insert(i, f"{license_line}\n")
+ break
+ file.writelines(lines)
+
+
+def main():
+ """Main function"""
+ files_missing_both = []
+ files_missing_copyright = []
+ files_missing_license = []
+
+ for file_path in sys.argv[1:]:
+ if Path(file_path).suffix not in FILE_STYLES:
+ continue
+ status = check_spdx_header_status(file_path)
+
+ if status == SPDXStatus.MISSING_BOTH:
+ files_missing_both.append(file_path)
+ elif status == SPDXStatus.MISSING_COPYRIGHT:
+ files_missing_copyright.append(file_path)
+ elif status == SPDXStatus.MISSING_LICENSE:
+ files_missing_license.append(file_path)
+ else:
+ continue
+
+ # Collect all files that need fixing
+ all_files_to_fix = files_missing_both + files_missing_copyright + files_missing_license
+ if all_files_to_fix:
+ print("The following files are missing the SPDX header:")
+ if files_missing_both:
+ for file_path in files_missing_both:
+ print(f" {file_path}")
+ add_header(file_path, SPDXStatus.MISSING_BOTH)
+
+ if files_missing_copyright:
+ for file_path in files_missing_copyright:
+ print(f" {file_path}")
+ add_header(file_path, SPDXStatus.MISSING_COPYRIGHT)
+ if files_missing_license:
+ for file_path in files_missing_license:
+ print(f" {file_path}")
+ add_header(file_path, SPDXStatus.MISSING_LICENSE)
+
+ sys.exit(1 if all_files_to_fix else 0)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tools/pre_commit/check_test_marks.py b/tools/pre_commit/check_test_marks.py
index 2afef8b6e..6af0ae825 100644
--- a/tools/pre_commit/check_test_marks.py
+++ b/tools/pre_commit/check_test_marks.py
@@ -1,52 +1,45 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Run a pre-commit hook that fails if test files are modified or added
that (probably) never run in the CI. For now, this means that every tests file
needs to have a CI level marker (e.g., core_model, advanced_model, full_model,
local_model, slow, etc) and hardware mark / helper so that we ensure mutated
tests will actually be selected as long as there are pytest commands pointing
at the right paths.
+
+SKU markers (``H100``, ``L4``, … tagged ``[hardware-resource]`` in
+``pyproject.toml``) must be applied via ``hardware_test(`` / ``hardware_marks(``
+so ``cards_{n}`` is attached. Direct ``pytest.mark.H100`` is rejected.
+
+Platform names allowed as ``pytest.mark.cpu`` / ``cuda`` come from
+``get_supported_platforms()`` (``[hardware-platform]`` in ``pyproject.toml``).
+CI level names come from ``get_level_markers()`` (``[ci-level]``).
"""
+from __future__ import annotations
+
+import importlib.util
import os
import re
import sys
+from functools import lru_cache
+from pathlib import Path
+from types import ModuleType
-# CI level markers
-LEVEL_MARKERS = ("core_model", "advanced_model", "full_model", "local_model", "slow")
-
-# Hardware markers. These are platforms and resources that CI filters on.
-# NOTE: If new hardware marks are added to pyproject.toml etc, we need
-# to add them here as well.
-HARDWARE_MARKERS = (
- "cpu",
- "gpu",
- "cuda",
- "rocm",
- "xpu",
- "npu",
- "musa",
- "H100",
- "L4",
- "MI325",
- "B60",
- "S5000",
- "A2",
- "A3",
-)
-
-# Helpers from tests/helpers/mark.py that auto-apply hardware marks.
+# Helpers from tests/helpers/mark.py that auto-apply hardware + cards_* marks.
HARDWARE_HELPERS = ("hardware_test", "hardware_marks")
+# The helper implementation is the only file allowed to write pytest.mark..
+_ALLOWED_DIRECT_SKU_FILES = frozenset({"tests/helpers/mark.py"})
+
# Match mark.X since we could also do `from pytest import mark`.
# \b prevents matching prefixes (e.g., mark.slow vs mark.slow_test).
-LEVEL_RE = re.compile(r"mark\.(?:" + "|".join(LEVEL_MARKERS) + r")\b")
-HARDWARE_RE = re.compile(r"mark\.(?:" + "|".join(HARDWARE_MARKERS) + r")\b")
HELPER_RE = re.compile(r"(?:" + "|".join(HARDWARE_HELPERS) + r")\s*\(")
MISSING_LEVEL_MARKER = "Level"
MISSING_HARDWARE_MARKER = "Hardware"
+DIRECT_SKU_MARKER = "Direct SKU"
# Check if a file is located under tests/ and matches test_.py
# or _test.py, since pytest technically collects on both.
@@ -54,9 +47,68 @@ MISSING_HARDWARE_MARKER = "Hardware"
TEST_FILE_RE = re.compile(r"^tests/(?:.*/)?(?:test_[^/]*\.py$|[^/]*_test\.py$)")
+def _normalize_path(path: str) -> str:
+ return path.replace("\\", "/")
+
+
+def _repo_root() -> Path:
+ return Path(__file__).resolve().parents[2]
+
+
+@lru_cache(maxsize=1)
+def _mark_module() -> ModuleType:
+ """Load ``mark.py`` by file path (no pytest/vllm; skip helpers ``__init__``)."""
+ path = _repo_root() / "tests" / "helpers" / "mark.py"
+ spec = importlib.util.spec_from_file_location("_vllm_omni_mark", path)
+ if spec is None or spec.loader is None:
+ raise ImportError(f"Cannot load mark helpers from {path}")
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+def _mark_name_re(names: tuple[str, ...]) -> re.Pattern[str]:
+ if not names:
+ return re.compile(r"(?!)")
+ return re.compile(r"mark\.(?:" + "|".join(re.escape(n) for n in names) + r")\b")
+
+
+@lru_cache(maxsize=1)
+def level_markers() -> tuple[str, ...]:
+ """CI level names tagged ``[ci-level]`` (``core_model``, ``slow``, …)."""
+ return tuple(sorted(_mark_module().get_level_markers()))
+
+
+@lru_cache(maxsize=1)
+def platform_markers() -> tuple[str, ...]:
+ """Names tests may apply as ``pytest.mark.cpu`` / ``cuda`` (not SKUs)."""
+ return tuple(sorted(_mark_module().get_supported_platforms()))
+
+
+@lru_cache(maxsize=1)
+def sku_markers() -> tuple[str, ...]:
+ """SKU marker names tagged ``[hardware-resource]`` in ``pyproject.toml``."""
+ return tuple(sorted(_mark_module().get_hardware_mark_list()))
+
+
+@lru_cache(maxsize=1)
+def _level_re() -> re.Pattern[str]:
+ return _mark_name_re(level_markers())
+
+
+@lru_cache(maxsize=1)
+def _platform_re() -> re.Pattern[str]:
+ return _mark_name_re(platform_markers())
+
+
+@lru_cache(maxsize=1)
+def _sku_mark_re() -> re.Pattern[str]:
+ return _mark_name_re(sku_markers())
+
+
def is_test_file(path: str) -> bool:
"""Determine whether or not a path is pointing at a test file or not."""
- return bool(TEST_FILE_RE.match(path))
+ return bool(TEST_FILE_RE.match(_normalize_path(path)))
def read_test_file(path: str) -> str | None:
@@ -69,22 +121,31 @@ def read_test_file(path: str) -> str | None:
def has_level_marker(contents: str) -> bool:
"""Check if file contents contain at least one CI level marker."""
- return bool(LEVEL_RE.search(contents))
+ return bool(_level_re().search(contents))
def has_hardware_marker(contents: str) -> bool:
- """Check if file contents contain any hardware marker or helpers."""
- return bool(HARDWARE_RE.search(contents) or HELPER_RE.search(contents))
+ """Check if file contents contain a platform marker or hardware helper."""
+ return bool(_platform_re().search(contents) or HELPER_RE.search(contents))
+
+
+def has_direct_sku_marker(path: str, contents: str) -> bool:
+ """True when a test applies ``pytest.mark.`` instead of the helpers."""
+ if _normalize_path(path) in _ALLOWED_DIRECT_SKU_FILES:
+ return False
+ return bool(_sku_mark_re().search(contents))
def get_files_missing_markers(
staged_files: list[str],
) -> dict[str, list[str]]:
- """Return a dict mapping file path to list of missing marker types."""
+ """Return a dict mapping file path to list of missing / invalid marker types."""
results: dict[str, list[str]] = {}
for path in staged_files:
if is_test_file(path) and (contents := read_test_file(path)) is not None:
missing = []
+ if has_direct_sku_marker(path, contents):
+ missing.append(DIRECT_SKU_MARKER)
if not has_level_marker(contents):
missing.append(MISSING_LEVEL_MARKER)
if not has_hardware_marker(contents):
@@ -99,14 +160,17 @@ if __name__ == "__main__":
if missing:
file_lines = "\n".join(f" - {path} [{' and '.join(problems)}]" for path, problems in missing.items())
+ sku = ", ".join(sku_markers())
print(
"\033[91merror:\033[0m test files are missing pytest marks "
- "required for Buildkite CI collection.\n\n"
- f"Level marks, e.g.: {', '.join(LEVEL_MARKERS[:4])}\n"
- f"Hardware marks, e.g.: {', '.join(HARDWARE_MARKERS[:4])}, ...\n"
- f" or helpers: {', '.join(HARDWARE_HELPERS)}\n\n"
+ "required for Buildkite CI collection, or apply SKU marks directly.\n\n"
+ f"Level marks, e.g.: {', '.join(level_markers()[:4])}\n"
+ f"Hardware marks, e.g.: {', '.join(platform_markers()[:4])}, ...\n"
+ f" or helpers: {', '.join(HARDWARE_HELPERS)}\n"
+ f"Do not write pytest.mark. ({sku}). "
+ "Use hardware_test(...) / hardware_marks(...) so cards_* is attached.\n\n"
"The following files are missing marks:\n"
f"{file_lines}\n\n"
- "To skip: SKIP=check-test-ci-coverage git commit ..."
+ "To skip: SKIP=check-mark git commit ..."
)
sys.exit(1)
diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py
new file mode 100644
index 000000000..a298baae5
--- /dev/null
+++ b/tools/pre_commit/check_torch_cuda.py
@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Ban direct torch.cuda.* usage outside an allowlist.
+
+Modeled on vllm/tools/pre_commit/check_torch_cuda.py. Call sites should use
+current_omni_platform / OmniPlatform instead of CUDA-only torch.cuda APIs.
+"""
+
+import sys
+
+import regex as re
+
+# Match `torch.cuda.xxx` device/memory/seed helpers that have portable Omni
+# platform equivalents. CUDAGraph / is_available / stream capture are not in
+# this list (same split as upstream vLLM).
+_TORCH_CUDA_PATTERNS = [
+ r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|mem_get_info|set_device)\b",
+ # Match torch.cuda.device(...) including args. `device\()\b` misses
+ # `device()` (no word char after `)`) and `device(0)`.
+ r"\btorch\.cuda\.device\s*\(",
+ r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",
+ r"\bwith\storch\.cuda\.device\b",
+ # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally
+ r"\bcuda_device_count_stateless\(\)\b",
+]
+
+# Platform adapters must call native torch.cuda / torch.npu / torch.xpu APIs.
+ALLOWED_PREFIXES = {
+ "vllm_omni/platforms/",
+}
+
+# Grandfathered call sites. Remove a path after migrating it to
+# current_omni_platform / OmniPlatform. Do not add new files without review.
+ALLOWED_FILES = {
+ "benchmarks/kernels/mot_linear_benchmarks.py",
+ "examples/online_serving/qwen2_5_omni/gradio_demo.py",
+ "examples/online_serving/qwen3_omni/gradio_demo.py",
+ "tests/dfx/reliability/helpers.py",
+ "tests/diffusion/distributed/test_cfg_parallel.py",
+ "tests/diffusion/distributed/test_pipeline_parallel.py",
+ "tests/e2e/accuracy/wan22_i2v/run_wan22_i2v_diffusers_cp.py",
+ "tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py",
+ "tests/e2e/features/fullduplex/test_personaplex_elastic.py",
+ "tests/entrypoints/test_omni_sleep_mode.py",
+ "tools/configure_stage_memory.py",
+ "vllm_omni/diffusion/models/bagel/pipeline_bagel.py",
+ "vllm_omni/diffusion/models/lance/pipeline_lance.py",
+ "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py",
+ "vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py",
+ "vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py",
+ "vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py",
+ "vllm_omni/experimental/ar_diffusion/runner.py",
+ "vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py",
+ "vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py",
+ "vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py",
+ "vllm_omni/worker/gpu_ar_model_runner.py",
+}
+
+_GENERIC_TIP = (
+ "Found torch.cuda API call. Use current_omni_platform / OmniPlatform "
+ "(empty_cache, synchronize, get_device_count, current_device, "
+ "get_device_memory, set_device) instead."
+)
+_SEED_TIP = (
+ "Found CUDA manual-seed API call. Use "
+ "torch.Generator(device=current_omni_platform.device_type).manual_seed(...) "
+ "instead."
+)
+
+
+def _normalize_path(path: str) -> str:
+ return path.replace("\\", "/")
+
+
+def _is_allowed(path: str) -> bool:
+ if any(path.startswith(prefix) for prefix in ALLOWED_PREFIXES):
+ return True
+ return path in ALLOWED_FILES
+
+
+def scan_file(path: str) -> int:
+ with open(path, encoding="utf-8") as f:
+ content = f.read()
+ for pattern in _TORCH_CUDA_PATTERNS:
+ for match in re.finditer(pattern, content, re.MULTILINE):
+ line_start = content.rfind("\n", 0, match.start()) + 1
+ if "#" in content[line_start : match.start()]:
+ continue
+ line_num = content[: match.start() + 1].count("\n") + 1
+ matched_text = match.group(0)
+ tip = _SEED_TIP if "manual_seed" in matched_text else _GENERIC_TIP
+ print(f"{path}:{line_num}: \033[91merror:\033[0m {tip}")
+ return 1
+ return 0
+
+
+def main():
+ returncode = 0
+ for filename in sys.argv[1:]:
+ path = _normalize_path(filename)
+ if _is_allowed(path):
+ continue
+ returncode |= scan_file(path)
+ return returncode
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/tools/pre_commit/check_tts_adapter.py b/tools/pre_commit/check_tts_adapter.py
index 626323b4f..d46928662 100644
--- a/tools/pre_commit/check_tts_adapter.py
+++ b/tools/pre_commit/check_tts_adapter.py
@@ -6,10 +6,8 @@ Per-model behaviour is being moved out of ``serving_speech.py`` and into
nothing stopped a new model from adding one more ``self._tts_model_type == ...``
branch faster than the migration removed them.
-This gate makes the two debt counters monotonically non-increasing:
-
-1. ``self._tts_model_type`` comparisons in ``serving_speech.py``
-2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter)
+This gate makes the number of ``self._tts_model_type`` comparisons in
+``serving_speech.py`` monotonically non-increasing.
Lowering a budget after removing branches is expected and encouraged — the check
insists on it, so the recorded numbers cannot silently drift upward. Raising one
@@ -28,11 +26,9 @@ from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
-ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py")
# Budgets. These may only go DOWN. See the module docstring.
-MAX_MODEL_TYPE_BRANCHES = 27
-MAX_LEGACY_DETECTORS = 0
+MAX_MODEL_TYPE_BRANCHES = 9
_GUIDANCE = """
Move the per-model behaviour into its adapter under
@@ -109,20 +105,6 @@ def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]:
return found
-def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
- """Every ``LegacyDetector(...)`` construction."""
- found: list[tuple[int, str]] = []
- for node in ast.walk(tree):
- if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector":
- name = "?"
- for kw in node.keywords:
- if kw.arg == "name" and isinstance(kw.value, ast.Constant):
- name = str(kw.value.value)
- found.append((node.lineno, name))
- found.sort()
- return found
-
-
def _check(
label: str,
path: Path,
@@ -151,21 +133,18 @@ def _check(
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
- # pre-commit passes the staged filenames; we always audit the same two files,
+ # pre-commit passes the staged filenames; we always audit the same file,
# so they are accepted and ignored.
parser.add_argument("filenames", nargs="*")
parser.parse_args(argv)
serving_path = REPO_ROOT / SERVING_SPEECH
- adapters_path = REPO_ROOT / ADAPTERS_INIT
- for path in (serving_path, adapters_path):
- if not path.is_file():
- print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
- return 1
+ if not serving_path.is_file():
+ print(f"check_tts_adapter: expected file not found: {serving_path}", file=sys.stderr)
+ return 1
notices: list[str] = []
branches = _model_type_branches(ast.parse(serving_path.read_text()))
- detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
errors = _check(
"`_tts_model_type` comparisons",
@@ -176,15 +155,6 @@ def main(argv: list[str] | None = None) -> int:
"MAX_MODEL_TYPE_BRANCHES",
notices,
)
- errors += _check(
- "unmigrated model types in LEGACY_TTS_DETECTORS",
- ADAPTERS_INIT,
- len(detectors),
- MAX_LEGACY_DETECTORS,
- [f"line {lineno}: {name}" for lineno, name in detectors],
- "MAX_LEGACY_DETECTORS",
- notices,
- )
for notice in notices:
print(f"check_tts_adapter: {notice}", file=sys.stderr)
diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py
new file mode 100644
index 000000000..e24f62444
--- /dev/null
+++ b/tools/pre_commit/mypy.py
@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Run mypy on changed files.
+
+Ported from vllm/tools/pre_commit/mypy.py. Group files so tests can use
+``follow_imports=skip`` without pulling that setting into library code.
+
+Usage:
+ python tools/pre_commit/mypy.py
+
+Args:
+ python_version: Python version to use (e.g. "3.10") or "local" to use
+ the interpreter running this script.
+ changed_files: List of changed files to check.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import regex as re
+
+# tests/ uses follow_imports=skip so third-party and library errors do not
+# leak into test type-checking. After a directory is clean under skip, it
+# can join the default group (follow_imports from pyproject.toml).
+SEPARATE_GROUPS = [
+ "tests",
+]
+
+# Model trees are large third-party ports. Drop a prefix from this list when
+# that directory is ready to be type-checked (same ratchet as upstream vLLM).
+EXCLUDE = [
+ r"vllm_omni/model_executor/models/",
+ r"vllm_omni/diffusion/models/",
+]
+
+
+def _normalize_path(path: str) -> str:
+ return path.replace("\\", "/")
+
+
+def group_files(changed_files: list[str]) -> dict[str, list[str]]:
+ """Group changed files into different mypy calls."""
+ exclude_pattern = re.compile(f"^{'|'.join(EXCLUDE)}.*")
+ file_groups: dict[str, list[str]] = {"": []}
+ file_groups.update({k: [] for k in SEPARATE_GROUPS})
+ # Longest path first so a sub-directory is not shadowed by its parent.
+ separate_groups = sorted(SEPARATE_GROUPS, key=len, reverse=True)
+ for changed_file in changed_files:
+ changed_file = _normalize_path(changed_file)
+ if exclude_pattern.match(changed_file):
+ continue
+ for directory in separate_groups:
+ if re.match(f"^{directory}.*", changed_file):
+ file_groups[directory].append(changed_file)
+ break
+ else:
+ if changed_file.startswith(("vllm_omni/", "tests/")):
+ file_groups[""].append(changed_file)
+ return file_groups
+
+
+def mypy(
+ targets: list[str],
+ python_version: str | None,
+ follow_imports: str | None,
+ file_group: str,
+) -> int:
+ """Run mypy on the given targets."""
+ args = ["mypy"]
+ if python_version is not None:
+ args += ["--python-version", python_version]
+ if follow_imports is not None:
+ args += ["--follow-imports", follow_imports]
+ print(f"$ {' '.join(args)} {file_group}")
+ return subprocess.run(args + targets, check=False).returncode
+
+
+def main() -> int:
+ python_version = sys.argv[1]
+ file_groups = group_files(sys.argv[2:])
+
+ if python_version == "local":
+ python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+ returncode = 0
+ for file_group, changed_files in file_groups.items():
+ follow_imports = None if file_group == "" else "skip"
+ if changed_files:
+ returncode |= mypy(changed_files, python_version, follow_imports, file_group)
+ return returncode
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/tools/pre_commit/shellcheck.sh b/tools/pre_commit/shellcheck.sh
new file mode 100755
index 000000000..7794ffe51
--- /dev/null
+++ b/tools/pre_commit/shellcheck.sh
@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+#
+# Lint bash scripts for undefined vars, quoting, and similar bugs.
+# Do not download a binary: use a distro/package-manager install so the
+# artifact is signed (apt/dnf/brew) or at least an explicit local install.
+set -euo pipefail
+
+is_windows_exe() {
+ [[ "$1" == *.exe ]]
+}
+
+is_git_bash() {
+ case "$(uname -s)" in
+ MINGW*|MSYS*|CYGWIN*) return 0 ;;
+ *) return 1 ;;
+ esac
+}
+
+find_native_shellcheck() {
+ local cand
+ cand="$(command -v shellcheck 2>/dev/null || true)"
+ if [ -n "$cand" ] && ! is_windows_exe "$cand"; then
+ echo "$cand"
+ return 0
+ fi
+ return 1
+}
+
+install_hint() {
+ echo "Please install shellcheck with your package manager, then re-run:"
+ echo " Debian/Ubuntu/WSL: sudo apt-get install shellcheck"
+ echo " Fedora: sudo dnf install ShellCheck"
+ echo " macOS: brew install shellcheck"
+ echo " Git Bash: scoop install shellcheck"
+ echo " https://github.com/koalaman/shellcheck?tab=readme-ov-file#installing"
+}
+
+SHELLCHECK_BIN=""
+if SHELLCHECK_BIN="$(find_native_shellcheck)"; then
+ :
+elif is_git_bash && SHELLCHECK_BIN="$(command -v shellcheck.exe 2>/dev/null || true)" && [ -n "$SHELLCHECK_BIN" ]; then
+ :
+else
+ install_hint
+ exit 1
+fi
+
+should_lint() {
+ local f="${1//\\//}"
+ f="${f#./}"
+ case "$f" in
+ *.sh) ;;
+ *) return 1 ;;
+ esac
+ git check-ignore -q "$f" && return 1
+ return 0
+}
+
+run_shellcheck() {
+ local f
+ for f in "$@"; do
+ if should_lint "$f"; then
+ "$SHELLCHECK_BIN" -s bash "$f"
+ fi
+ done
+}
+
+if [ "$#" -gt 0 ]; then
+ run_shellcheck "$@"
+ exit 0
+fi
+
+# Direct invocation with no args: lint every tracked *.sh.
+while IFS= read -r -d '' f || [ -n "$f" ]; do
+ git check-ignore -q "$f" && continue
+ "$SHELLCHECK_BIN" -s bash "$f"
+done < <(find . -path ./.git -prune -o -name "*.sh" -print0)
diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py
index 512e86538..abb02b1d5 100644
--- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py
+++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
"""Daily-Omni Dataset loader for benchmark.
Daily-Omni is an audio-visual reasoning benchmark with 684 videos
@@ -7,7 +10,8 @@ Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni
Supports loading QA metadata from:
- Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
-- HuggingFace datasets (``dataset_path``): legacy online mode
+- HuggingFace Hub id (``dataset_path``): resolved with ``snapshot_download`` to the repo's
+ ``qa.json``, which reads the hub cache and therefore also works under ``HF_HUB_OFFLINE``
Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir``
is not set, the first request that needs on-disk media downloads that archive from the Hugging Face
@@ -24,9 +28,9 @@ Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``?
This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``,
``random_seed``, ``self.data``) and implements **two explicit loaders**:
``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface``
- (optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not**
+ (Hub ids, resolved to the repo's ``qa.json`` through the hub cache). The latter is **not**
inheritance; it is the same Hub rows as before, factored into a helper so one class can
- serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``.
+ serve both deployment modes without mandatory ``datasets``.
Usage:
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset
@@ -38,7 +42,7 @@ Usage:
random_seed=42,
)
- # HuggingFace mode (legacy, requires network)
+ # Hub id mode (works offline once the repo is in the HF hub cache)
dataset = DailyOmniDataset(
dataset_path="liarliar/Daily-Omni",
dataset_split="train",
@@ -210,6 +214,68 @@ def _resolve_hf_cache_snapshot(path: Path) -> Path:
return path
+#: QA metadata pulled by :func:`ensure_daily_omni_hub_root`. ``Videos.tar`` is deliberately left
+#: out: it is ~60GB and only :func:`ensure_daily_omni_hub_videos_dir` needs it, lazily, once a
+#: request actually asks for media.
+_DAILY_OMNI_QA_ALLOW_PATTERNS = ["qa.json", "QA.json"]
+
+
+def ensure_daily_omni_hub_root(repo_id: str) -> Path:
+ """Resolve a Daily-Omni Hub id to its snapshot directory, downloading QA metadata if needed.
+
+ Mirrors :func:`...seed_tts_dataset.resolve_seed_tts_root`: ``snapshot_download`` resolves the
+ repo through the **hub cache**, so an air-gapped run (``HF_HUB_OFFLINE=1``) keeps working as
+ long as the dataset was pre-staged there. ``datasets.load_dataset`` cannot do this — it only
+ consults the separate *datasets* cache, which a ``hf download`` never populates.
+
+ Raises:
+ ImportError: if ``huggingface_hub`` is not installed.
+ ValueError: if ``repo_id`` is empty.
+ """
+ rid = (repo_id or "").strip()
+ if not rid:
+ raise ValueError("repo_id is required to download Daily-Omni from Hugging Face")
+
+ try:
+ from huggingface_hub import snapshot_download
+ except ImportError as e:
+ raise ImportError(
+ "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local "
+ "--daily-omni-qa-json / --daily-omni-video-dir mirror."
+ ) from e
+
+ cache = snapshot_download(
+ repo_id=rid,
+ repo_type="dataset",
+ allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS,
+ )
+ root = _resolve_hf_cache_snapshot(Path(cache).resolve())
+ logger.info("Daily-Omni Hub snapshot ready at %s (repo=%s)", root, rid)
+ return root
+
+
+def _daily_omni_qa_load_error(
+ repo_id: str | None,
+ snapshot_error: Exception | None,
+ datasets_error: Exception,
+) -> RuntimeError:
+ """Compose an actionable error after both QA-metadata sources failed.
+
+ Offline CI hits this when the dataset is missing from the hub cache *or* staged under a
+ different repo id, so the message names both attempts and the flags that bypass them.
+ """
+ parts = [f"Could not load Daily-Omni QA metadata for {repo_id!r}."]
+ if snapshot_error is not None:
+ parts.append(f"huggingface_hub snapshot_download failed: {type(snapshot_error).__name__}: {snapshot_error}")
+ parts.append(f"datasets.load_dataset failed: {type(datasets_error).__name__}: {datasets_error}")
+ parts.append(
+ "For offline runs, pre-stage the dataset in the HF hub cache "
+ "(`hf download --repo-type dataset `) and make sure the repo id matches, "
+ "or pass --daily-omni-qa-json / --daily-omni-video-dir pointing at a local mirror."
+ )
+ return RuntimeError(" ".join(parts))
+
+
def resolve_daily_omni_local_root(dataset_path: str | None) -> Path | None:
"""Return the local Daily-Omni mirror root for ``dataset_path``, or ``None`` if not local.
@@ -433,7 +499,8 @@ class DailyOmniDataset(BenchmarkDataset):
QA metadata can be loaded from:
- Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments
- - HuggingFace datasets (``dataset_path``): legacy online mode
+ - HuggingFace Hub id (``dataset_path``): the repo's ``qa.json`` via ``snapshot_download``
+ (hub-cache backed, offline-safe); ``datasets.load_dataset`` only as a parquet fallback
Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set,
the first sample that needs on-disk media downloads that archive from the Hugging Face dataset
@@ -444,7 +511,8 @@ class DailyOmniDataset(BenchmarkDataset):
qa_json_path: Path to local qa.json file (offline mode, preferred). When provided,
``dataset_path`` and ``dataset_split`` are ignored.
dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if
- ``qa_json_path`` is not provided (legacy online mode).
+ ``qa_json_path`` is not provided; resolved through the hub cache, so a pre-staged
+ repo also loads with ``HF_HUB_OFFLINE=1``.
dataset_split: Dataset split to use (default: "train"). Used only in online mode.
random_seed: Random seed for shuffling
video_dir: Directory containing extracted video files (default: None; may be filled lazily
@@ -562,10 +630,13 @@ class DailyOmniDataset(BenchmarkDataset):
def _load_from_local_json(self) -> None:
"""Load QA data from local JSON file."""
- if not self.qa_json_path.exists():
- raise FileNotFoundError(f"QA JSON file not found: {self.qa_json_path}")
+ qa_json_path = self.qa_json_path
+ if qa_json_path is None:
+ raise RuntimeError("qa_json_path is required for local JSON loading")
+ if not qa_json_path.exists():
+ raise FileNotFoundError(f"QA JSON file not found: {qa_json_path}")
- with open(self.qa_json_path, encoding="utf-8") as f:
+ with open(qa_json_path, encoding="utf-8") as f:
data = json.load(f)
# Support both list format and dict with "train"/"test" splits
@@ -602,21 +673,45 @@ class DailyOmniDataset(BenchmarkDataset):
self.data = _ListDatasetIterator(data)
def _load_from_huggingface(self) -> None:
- """Load QA rows via ``datasets.load_dataset`` (legacy / convenience path).
+ """Load QA rows for a Hub id, preferring the repo's ``qa.json`` over ``datasets``.
- Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni``
- and get the same parquet-backed rows as the Hub dataset card, with streaming (or
- non-streaming if ``no_stream=True``) and shuffle.
+ The Hub repo ships ``qa.json`` (1,197 rows) next to ``Videos.tar``, so the primary path is
+ :func:`ensure_daily_omni_hub_root` → :meth:`_load_from_local_json`: ``snapshot_download``
+ reads the hub cache and therefore works under ``HF_HUB_OFFLINE``, the same contract as
+ Seed-TTS. ``datasets.load_dataset`` remains only as a fallback for mirrors
+ that publish parquet rows instead of ``qa.json``; it consults the *datasets* cache rather
+ than the hub cache and streams shards over HTTP, so it cannot serve air-gapped runs.
This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base
always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary
API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing
one ``sample()`` / request-building implementation for both.
"""
+ snapshot_error: Exception | None = None
+ dataset_path = self.dataset_path
+ if dataset_path is None:
+ raise ValueError("dataset_path is required for HuggingFace loading")
+ try:
+ root = ensure_daily_omni_hub_root(dataset_path)
+ except Exception as e: # noqa: BLE001 - any failure just falls through to `datasets`
+ snapshot_error = e
+ else:
+ qa_json = daily_omni_local_qa_json(root)
+ if qa_json is not None:
+ self.qa_json_path = qa_json
+ logger.info("Loading Daily-Omni QA metadata from Hub snapshot: %s", qa_json)
+ self._load_from_local_json()
+ return
+ logger.info(
+ "Hub snapshot %s ships no qa.json; falling back to `datasets.load_dataset`",
+ root,
+ )
+
if load_dataset is None:
- raise ImportError(
- "datasets library is required for HuggingFace mode. "
- "Install with: pip install datasets, or use local JSON mode instead."
+ raise _daily_omni_qa_load_error(
+ self.dataset_path,
+ snapshot_error,
+ ImportError("datasets library is not installed (pip install datasets)"),
)
load_kw: dict[str, Any] = {
@@ -626,7 +721,10 @@ class DailyOmniDataset(BenchmarkDataset):
}
if self.dataset_subset is not None:
load_kw["name"] = self.dataset_subset
- ds = load_dataset(self.dataset_path, **load_kw)
+ try:
+ ds = load_dataset(self.dataset_path, **load_kw)
+ except Exception as e:
+ raise _daily_omni_qa_load_error(self.dataset_path, snapshot_error, e) from e
if not getattr(self, "disable_shuffle", False):
ds = ds.shuffle(seed=self.random_seed)
self.data = ds
@@ -1167,7 +1265,7 @@ class DailyOmniDataset(BenchmarkDataset):
num_video_frames, avg_fps = _probe_video_frames_and_fps(video_path)
duration = num_video_frames / avg_fps
num_seconds = max(1, int(math.ceil(duration)))
- second_timestamps = list(range(num_seconds))
+ second_timestamps = [float(i) for i in range(num_seconds)]
if duration > max_num_frames:
timestamps = [round(i * 0.1, 1) for i in range(int(duration / 0.1))]
diff --git a/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py b/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py
new file mode 100644
index 000000000..04726eee6
--- /dev/null
+++ b/vllm_omni/benchmarks/data_modules/omniinteract_dataset.py
@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniInteract dataset discovery for the standard serving benchmark."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import random
+import shutil
+import tarfile
+import tempfile
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest
+from vllm.tokenizers import TokenizerLike
+from vllm.transformers_utils.repo_utils import hf_fs
+
+OMNIINTERACT_SUBSETS = ("1q1a", "1q1a_math", "1qna")
+DEFAULT_OMNIINTERACT_REPO = "lucky-lance/OmniInteract"
+DEFAULT_MAX_VIDEO_DURATION_S = 3600.0
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniInteractCase:
+ subset: str
+ video_rel: str
+ video_path: Path
+ annotation_path: Path
+ scene_type: str
+
+
+@dataclass(frozen=True)
+class OmniInteractSessionOptions:
+ output_root: Path
+ timeout_s: float
+ media_timeout_s: float
+ ref_audio: str
+ require_response: bool = False
+ max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
+
+
+@dataclass(frozen=True)
+class OmniInteractPreparedInput:
+ duration_s: float
+ pcm16: bytes = field(repr=False)
+ video_frames: tuple[str | None, ...] = field(repr=False)
+ ref_audio_data_url: str = field(repr=False)
+
+
+@dataclass
+class OmniInteractSampleRequest(SampleRequest):
+ omniinteract_case: OmniInteractCase | None = None
+ omniinteract_options: OmniInteractSessionOptions | None = None
+ omniinteract_prepared_input: OmniInteractPreparedInput | None = field(default=None, repr=False)
+
+
+class OmniInteractDataset(BenchmarkDataset):
+ def __init__(
+ self,
+ *,
+ data_root: str | None,
+ dataset_repo: str = DEFAULT_OMNIINTERACT_REPO,
+ subsets: Sequence[str] = OMNIINTERACT_SUBSETS,
+ random_seed: int = 0,
+ disable_shuffle: bool = False,
+ ) -> None:
+ super().__init__(
+ dataset_path=data_root or dataset_repo,
+ random_seed=random_seed,
+ disable_shuffle=disable_shuffle,
+ )
+ self.root = resolve_omniinteract_root(data_root, dataset_repo)
+ self.subsets = tuple(subsets)
+
+ def sample(
+ self,
+ tokenizer: TokenizerLike | None,
+ num_requests: int,
+ *,
+ request_id_prefix: str = "",
+ options: OmniInteractSessionOptions,
+ **_: Any,
+ ) -> list[SampleRequest]:
+ del tokenizer
+ cases = discover_omniinteract_cases(
+ self.root,
+ self.subsets,
+ num_prompts=num_requests,
+ seed=self.random_seed,
+ disable_shuffle=self.disable_shuffle,
+ )
+ return [
+ OmniInteractSampleRequest(
+ prompt="",
+ prompt_len=0,
+ expected_output_len=0,
+ multi_modal_data=None,
+ request_id=f"{request_id_prefix}{index}",
+ omniinteract_case=case,
+ omniinteract_options=options,
+ )
+ for index, case in enumerate(cases)
+ ]
+
+
+def _data_dir(root: Path) -> Path:
+ for candidate in (root, root / "data"):
+ if any((candidate / subset).is_dir() for subset in OMNIINTERACT_SUBSETS):
+ return candidate
+ raise FileNotFoundError(f"OmniInteract data not found under {root}")
+
+
+def _confined_path(root: Path, relative: object, *, field: str) -> Path:
+ text, resolved_root = str(relative or ""), root.resolve()
+ path = Path(text)
+ destination = (root / path).resolve()
+ if not text or path.is_absolute() or ".." in path.parts or not destination.is_relative_to(resolved_root):
+ raise ValueError(f"Unsafe OmniInteract {field} path: {text!r}")
+ return destination
+
+
+def _safe_extract(archive: tarfile.TarFile, target: Path) -> None:
+ root = target.resolve()
+ members = archive.getmembers()
+ for member in members:
+ path = Path(member.name)
+ if (
+ path.is_absolute()
+ or ".." in path.parts
+ or not (root / path).resolve().is_relative_to(root)
+ or not (member.isdir() or member.isfile())
+ ):
+ raise ValueError(f"Unsafe path in OmniInteract archive: {member.name!r}")
+ archive.extractall(target, members=members, filter="data")
+
+
+def _extract_archive(archive: Path, target: Path) -> Path:
+ if target.is_symlink():
+ raise ValueError(f"Refusing to extract through symlink: {target}")
+ stat = archive.stat()
+ published = target / f"{stat.st_size}-{stat.st_mtime_ns}"
+ if published.is_symlink():
+ raise ValueError(f"Refusing to use symlinked extraction: {published}")
+ if published.is_dir():
+ try:
+ return _data_dir(published)
+ except FileNotFoundError:
+ pass
+ target.mkdir(parents=True, exist_ok=True)
+ staging = Path(tempfile.mkdtemp(prefix=".tmp-", dir=target))
+ try:
+ with tarfile.open(archive, "r:*") as handle:
+ _safe_extract(handle, staging)
+ root = _data_dir(staging)
+ try:
+ staging.replace(published)
+ return published / root.relative_to(staging)
+ except OSError:
+ if not published.is_dir():
+ raise
+ return _data_dir(published)
+ finally:
+ shutil.rmtree(staging, ignore_errors=True)
+
+
+def resolve_omniinteract_root(data_root: str | None, dataset_repo: str = DEFAULT_OMNIINTERACT_REPO) -> Path:
+ """Resolve a local tree/archive or download the official archive."""
+
+ if data_root:
+ local = Path(data_root).expanduser().resolve()
+ if local.is_file():
+ return _extract_archive(local, local.parent / f".{local.stem}.vllm_omni_extracted")
+ if not local.is_dir():
+ raise FileNotFoundError(f"--dataset-path does not exist: {local}")
+ try:
+ return _data_dir(local)
+ except FileNotFoundError:
+ for name in ("data.tar.gz", "data.tar"):
+ if (archive := local / name).is_file():
+ return _extract_archive(archive, local / ".vllm_omni_extracted")
+ raise
+
+ cache = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
+ target = cache / "vllm_omni" / "omniinteract" / dataset_repo.replace("/", "__")
+ downloads = target / ".downloads"
+ downloads.mkdir(parents=True, exist_ok=True)
+ filesystem, errors = hf_fs(), []
+ for name in ("data.tar.gz", "data.tar"):
+ archive = downloads / name
+ try:
+ if not archive.is_file():
+ descriptor, temporary_name = tempfile.mkstemp(prefix=f".{name}.", dir=downloads)
+ os.close(descriptor)
+ temporary = Path(temporary_name)
+ try:
+ filesystem.get_file(f"datasets/{dataset_repo}/{name}", str(temporary))
+ if not temporary.stat().st_size:
+ raise ValueError("downloaded archive is empty")
+ temporary.replace(archive)
+ finally:
+ temporary.unlink(missing_ok=True)
+ return _extract_archive(archive, target)
+ except Exception as exc: # noqa: BLE001, PERF203 - Hub errors vary by version
+ errors.append(f"{name}: {exc}")
+ raise FileNotFoundError(f"Could not download OmniInteract from {dataset_repo!r}: {'; '.join(errors)}")
+
+
+def _mapping_cases(root: Path, subset: str) -> list[OmniInteractCase]:
+ mapping_path = root / "video_json_map.json"
+ if not mapping_path.is_file():
+ raise FileNotFoundError(f"Missing OmniInteract mapping: {mapping_path}")
+ mapping = json.loads(mapping_path.read_text())
+ entries = mapping.get("entries") if isinstance(mapping, dict) else None
+ if not isinstance(entries, list):
+ raise ValueError(f"Invalid OmniInteract mapping: {mapping_path}")
+ cases = []
+ for row in entries:
+ if not isinstance(row, dict):
+ raise ValueError(f"Invalid OmniInteract mapping row in {mapping_path}")
+ video_rel, annotation_rel = str(row.get("video") or ""), str(row.get("annotation") or "")
+ video = _confined_path(root, video_rel, field="video")
+ annotation = _confined_path(root, annotation_rel, field="annotation")
+ if not video.is_file() or not annotation.is_file():
+ raise FileNotFoundError(f"OmniInteract mapping references missing files: {video_rel!r}, {annotation_rel!r}")
+ cases.append(
+ OmniInteractCase(
+ subset,
+ video_rel,
+ video,
+ annotation,
+ str(row.get("scene_type") or "multi_turn").lower(),
+ )
+ )
+ return cases
+
+
+def _one_to_many_cases(root: Path) -> list[OmniInteractCase]:
+ videos, annotations = root / "videos_bench", root / "annotations"
+ if not videos.is_dir() or not annotations.is_dir():
+ raise FileNotFoundError(f"Invalid OmniInteract 1qna layout under {root}")
+ cases = []
+ for video in sorted(videos.rglob("*.mp4")):
+ if not video.resolve().is_relative_to(videos.resolve()):
+ raise ValueError(f"Unsafe OmniInteract video path: {video}")
+ relative = video.relative_to(videos)
+ annotation = (annotations / relative).with_suffix(".json").resolve()
+ if not annotation.is_relative_to(annotations.resolve()):
+ raise ValueError(f"Unsafe OmniInteract annotation path: {relative}")
+ if not annotation.is_file():
+ raise FileNotFoundError(f"Missing OmniInteract annotation: {annotation}")
+ cases.append(OmniInteractCase("1qna", f"videos_bench/{relative.as_posix()}", video, annotation, "1qna"))
+ return cases
+
+
+def discover_omniinteract_cases(
+ root: Path,
+ subsets: Sequence[str],
+ *,
+ num_prompts: int,
+ seed: int = 0,
+ disable_shuffle: bool = False,
+) -> list[OmniInteractCase]:
+ invalid = set(subsets) - set(OMNIINTERACT_SUBSETS)
+ if invalid:
+ raise ValueError(f"Unsupported OmniInteract subsets: {sorted(invalid)}")
+ if not subsets:
+ raise ValueError("At least one OmniInteract subset is required")
+ if len(set(subsets)) != len(subsets):
+ raise ValueError("OmniInteract subsets must not contain duplicates")
+ if num_prompts < 0:
+ raise ValueError("num_prompts must be non-negative")
+ data_root, cases = _data_dir(root.resolve()), []
+ for subset in subsets:
+ subset_root = data_root / subset
+ selected = _one_to_many_cases(subset_root) if subset == "1qna" else _mapping_cases(subset_root, subset)
+ if not selected:
+ raise ValueError(f"No OmniInteract sessions found for requested subset {subset!r}")
+ cases.extend(selected)
+ if len({case.video_path for case in cases}) != len(cases):
+ raise ValueError("OmniInteract dataset contains duplicate video paths")
+ if not disable_shuffle:
+ random.Random(seed).shuffle(cases)
+ if num_prompts:
+ if num_prompts > len(cases):
+ logger.warning(
+ "Requested %d OmniInteract prompts but only %d are available; using all cases", num_prompts, len(cases)
+ )
+ cases = cases[:num_prompts]
+ return cases
+
+
+def case_manifest(case: OmniInteractCase, output_dir: Path) -> dict[str, Any]:
+ return {
+ "sample_id": f"{case.subset}__{output_dir.name}",
+ "video": str(case.video_path),
+ "gt_json": str(case.annotation_path.resolve()),
+ "model_json": str((output_dir / "wav_transcript.json").resolve()),
+ "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
+ }
diff --git a/vllm_omni/benchmarks/duplex/__init__.py b/vllm_omni/benchmarks/duplex/__init__.py
new file mode 100644
index 000000000..62d91dd86
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/__init__.py
@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU-friendly Omni-DuplexEval generation and scoring helpers."""
+
+from .omni_duplex_eval_dataset import DuplexSample, load_samples
+from .omni_duplex_eval_metrics import PROTOCOL_PIN
+
+__all__ = ["DuplexSample", "PROTOCOL_PIN", "load_samples"]
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py
new file mode 100644
index 000000000..03997fc8e
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py
@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Media-clock sentence extraction and response JSON normalization."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import regex as re
+
+_TERMINAL = re.compile(r"(.+?[.!?。!?]+)(?:\s+|$)", re.S)
+
+
+@dataclass
+class TimedSentence:
+ sentence: str
+ start: float
+ end: float
+
+ def as_dict(self) -> dict[str, Any]:
+ return {"sentence": self.sentence, "start": round(self.start, 3), "end": round(self.end, 3)}
+
+
+def split_text(text: str) -> list[str]:
+ text = str(text or "").strip()
+ if not text:
+ return []
+ parts: list[str] = []
+ position = 0
+ for match in _TERMINAL.finditer(text):
+ parts.append(match.group(1).strip())
+ position = match.end()
+ tail = text[position:].strip()
+ if tail:
+ parts.append(tail)
+ return [part for part in parts if part]
+
+
+def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "media") -> list[TimedSentence]:
+ """Split output deltas while anchoring text to input ``audio_end_ms``."""
+ if clock != "media":
+ raise ValueError("only clock=media is supported")
+ sentences: list[TimedSentence] = []
+ pending = ""
+ start: float | None = None
+ last: float | None = None
+ for event in events:
+ event_type = event.get("type")
+ if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}:
+ continue
+ delta = event.get("delta")
+ if not isinstance(delta, str) or not delta.strip():
+ continue
+ timestamp = event.get("audio_end_ms")
+ if timestamp is None:
+ timestamp = event.get("_media_clock_ms", event.get("current_time_ms", 0.0))
+ try:
+ # Every accepted event field is explicitly millisecond-valued.
+ now = float(timestamp) / 1000.0
+ except (TypeError, ValueError):
+ now = last or 0.0
+ pending += delta
+ start = now if start is None else start
+ last = now
+ pieces = split_text(pending)
+ complete = pieces if pieces and re.search(r"[.!?。!?]$", pending.strip()) else pieces[:-1]
+ if complete:
+ for piece in complete:
+ sentences.append(TimedSentence(piece, start, max(start, now)))
+ consumed = " ".join(complete)
+ pending = pending[len(consumed) :].strip()
+ start = now if pending else None
+ if pending.strip():
+ sentences.append(TimedSentence(pending.strip(), start or last or 0.0, last or start or 0.0))
+ return sentences
+
+
+def normalize_response_items(raw: Any) -> list[dict[str, Any]]:
+ if isinstance(raw, dict):
+ items = raw.get("sentences", raw.get("chunks"))
+ else:
+ items = raw
+ if not isinstance(items, list):
+ raise ValueError("response JSON must be a list or contain sentences/chunks")
+ normalized = []
+ for item in items:
+ if not isinstance(item, dict):
+ continue
+ text = item.get("sentence", item.get("text", item.get("content", "")))
+ if not isinstance(text, str) or not text.strip():
+ continue
+ try:
+ start = float(item.get("start", item.get("current_time", item.get("time", 0.0))) or 0.0)
+ end = float(item.get("end", item.get("current_time", item.get("time", start))) or start)
+ except (TypeError, ValueError):
+ continue
+ normalized.append({"sentence": text.strip(), "start": start, "end": max(start, end)})
+ return normalized
+
+
+def validate_clock(meta: dict[str, Any], *, allow_invalid: bool = False) -> None:
+ if meta.get("clock") == "invalid" and not allow_invalid:
+ raise ValueError("generate artifact has clock=invalid; pass --allow-invalid-clock to score it")
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
new file mode 100644
index 000000000..9d6013e80
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py
@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Dataset and split normalization for Omni-DuplexEval.
+
+The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an
+already materialized iterable. Media is deliberately resolved at use time so
+the benchmark remains usable in air-gapped environments.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+DEFAULT_DATASET = "Hothan/Omni-DuplexEval"
+RTD_SPLITS = frozenset(
+ {
+ "RTD_world_knowledge",
+ "RTD_counting",
+ "RTD_fine_grained_movement",
+ "RTD_interaction_relation",
+ "RTD_OCR",
+ "RTD_Omni",
+ }
+)
+PR_SPLITS = frozenset({"PR_correction", "PR_event_reminder", "PR_post_event_reminder"})
+_TASK_ALIASES = {
+ "correction": "correction",
+ "pr_correction": "correction",
+ "event_reminder": "proactive_reminder",
+ "pr_event_reminder": "proactive_reminder",
+ "proactive_reminder": "proactive_reminder",
+ "post_event_reminder": "post_event_reminder",
+ "pr_post_event_reminder": "post_event_reminder",
+}
+
+
+def canonical_task_type(value: str) -> str:
+ key = str(value).strip().lower().replace("-", "_").replace(" ", "_")
+ if key not in _TASK_ALIASES:
+ raise ValueError(f"unsupported Omni-DuplexEval task type: {value!r}")
+ return _TASK_ALIASES[key]
+
+
+def task_type_for_split(split: str) -> str:
+ """Match the task routing used by the official HF batch evaluator."""
+ key = str(split).strip().lower()
+ if "correction" in key:
+ return "correction"
+ if "post_event" in key:
+ return "post_event_reminder"
+ if key.startswith("pr"):
+ return "proactive_reminder"
+ raise ValueError(f"cannot infer proactive-reminder task from split {split!r}")
+
+
+def family_for_split(split: str, task_type: str | None = None) -> str:
+ if split in RTD_SPLITS or str(split).upper().startswith("RTD"):
+ return "rtd"
+ if split in PR_SPLITS or str(split).upper().startswith("PR"):
+ return "pr"
+ if task_type:
+ return "rtd" if str(task_type).lower().startswith("rtd") else "pr"
+ raise ValueError(f"cannot infer benchmark family from split {split!r}")
+
+
+def _value(row: dict[str, Any], *keys: str, default: Any = None) -> Any:
+ for key in keys:
+ if key in row and row[key] is not None:
+ return row[key]
+ return default
+
+
+def _float(value: Any, default: float | None = None) -> float | None:
+ try:
+ return None if value is None else float(value)
+ except (TypeError, ValueError):
+ return default
+
+
+@dataclass(frozen=True)
+class DuplexSample:
+ id: str
+ split: str
+ family: str
+ task_type: str | None
+ video: Any
+ question_audio: Any = None
+ question_text: str = ""
+ answer1: str = ""
+ answer2: str = ""
+ reminder1: Any = None
+ reminder2: Any = None
+ video_duration: float | None = None
+ video_type: str = ""
+ raw: dict[str, Any] | None = None
+
+ @classmethod
+ def from_row(cls, row: dict[str, Any], *, split: str | None = None, media_root: Path | None = None) -> DuplexSample:
+ chosen_split = str(split or _value(row, "split", "subset", "config", default=""))
+ task_value = _value(row, "task_type", "task", "type")
+ family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value))
+ task = None
+ if family == "pr":
+ task = canonical_task_type(str(task_value)) if task_value else task_type_for_split(chosen_split)
+ sample_id = str(_value(row, "id", "sample_id", "uid", "name", default=""))
+ video = _value(row, "video", "video_path", "video_file")
+ audio = _value(row, "question_audio", "audio", "question_wav")
+ if media_root:
+
+ def resolve(value: Any) -> Any:
+ if not isinstance(value, str):
+ return value
+ candidate = Path(value).expanduser()
+ return str(candidate if candidate.is_absolute() else media_root / candidate)
+
+ video, audio = resolve(video), resolve(audio)
+ return cls(
+ id=sample_id,
+ split=chosen_split,
+ family=family,
+ task_type=task,
+ video=video,
+ question_audio=audio,
+ question_text=str(_value(row, "question_text", "question", "instruction", default="") or ""),
+ answer1=str(_value(row, "answer1", "answer", "ground_answer", default="") or ""),
+ answer2=str(_value(row, "answer2", default="") or ""),
+ reminder1=_value(row, "reminder1", "reminder_1"),
+ reminder2=_value(row, "reminder2", "reminder_2"),
+ video_duration=_float(_value(row, "video_duration", "duration")),
+ video_type=str(_value(row, "video_type", default="") or ""),
+ raw=dict(row),
+ )
+
+
+def _read_manifest(path: Path) -> list[dict[str, Any]]:
+ if path.suffix.lower() == ".jsonl":
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+ payload = json.loads(path.read_text(encoding="utf-8"))
+ if isinstance(payload, dict):
+ payload = payload.get("data", payload.get("samples", [payload]))
+ if not isinstance(payload, list):
+ raise ValueError("manifest must contain a JSON list")
+ return payload
+
+
+def load_samples(
+ dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET,
+ *,
+ split: str | None = None,
+ family: str = "all",
+ media_root: str | Path | None = None,
+ limit: int | None = None,
+ ids: Iterable[str] | None = None,
+) -> list[DuplexSample]:
+ if isinstance(dataset, (str, Path)) and Path(str(dataset)).exists():
+ rows = _read_manifest(Path(str(dataset)))
+ elif not isinstance(dataset, (str, Path)):
+ rows = list(dataset)
+ else:
+ try:
+ from datasets import load_dataset
+ except ImportError as exc:
+ raise RuntimeError("datasets is required for Hugging Face loading; use a local manifest instead") from exc
+ kwargs: dict[str, Any] = {}
+ if split and split != "all":
+ kwargs["split"] = split
+ loaded = load_dataset(str(dataset), **kwargs)
+ if isinstance(loaded, dict):
+ rows = []
+ for name, table in loaded.items():
+ rows.extend({**dict(row), "split": name} for row in table)
+ else:
+ rows = [dict(row) for row in loaded]
+ wanted = set(str(item) for item in ids) if ids else None
+ root = Path(media_root).expanduser() if media_root else None
+ result = []
+ for row in rows:
+ sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root)
+ if split and split != "all" and sample.split != split:
+ continue
+ if family != "all" and sample.family != family:
+ continue
+ if wanted is not None and sample.id not in wanted:
+ continue
+ result.append(sample)
+ if limit is not None:
+ result = result[: max(0, limit)]
+ return result
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py
new file mode 100644
index 000000000..a92376bd7
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py
@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Evaluate and summarize generated Omni-DuplexEval artifacts."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from .omni_duplex_eval_clock import normalize_response_items, validate_clock
+from .omni_duplex_eval_dataset import DuplexSample
+from .omni_duplex_eval_judge import DuplexJudge
+from .omni_duplex_eval_media import extract_jpeg, materialize_media, video_duration
+from .omni_duplex_eval_metrics import (
+ PROTOCOL_PIN,
+ build_content_prompt,
+ build_reminder_prompt,
+ build_temporal_prompt,
+ parse_judge_json,
+ reminder_window,
+ summarize_pr_results,
+ summarize_temporal_results,
+ temporal_window,
+)
+
+
+def _read(path: Path) -> Any:
+ return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _write(path: Path, value: Any) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def _text(items: list[dict[str, Any]]) -> str:
+ return " ".join(str(item.get("sentence", item.get("text", ""))).strip() for item in items).strip()
+
+
+def evaluate_sample(
+ sample: DuplexSample,
+ response_path: Path,
+ score_path: Path,
+ judge: DuplexJudge,
+ *,
+ judge_fps: int = 2,
+ judge_video_mode: str = "video_url",
+ window_size: float = 10.0,
+ allow_invalid_clock: bool = False,
+) -> dict[str, Any]:
+ raw = _read(response_path)
+ meta = (
+ _read(response_path.with_name(response_path.stem + ".meta.json"))
+ if response_path.with_name(response_path.stem + ".meta.json").exists()
+ else {}
+ )
+ validate_clock(meta, allow_invalid=allow_invalid_clock)
+ items = normalize_response_items(raw)
+ result: dict[str, Any] = {
+ "id": sample.id,
+ "family": sample.family,
+ "task_type": sample.task_type,
+ "protocol_pin": PROTOCOL_PIN,
+ "response_meta": meta,
+ "judge_model": getattr(judge, "model", None),
+ "judge_video_mode": judge_video_mode,
+ }
+ if sample.family == "rtd":
+ video_path = materialize_media(sample.video, score_path.parent / ".media", sample.id, ".mp4")
+ duration = sample.video_duration or video_duration(video_path)
+ temporal_rows = []
+ for item in items:
+ window = temporal_window(item["start"], item["end"], duration)
+ if not window:
+ temporal_rows.append(
+ {
+ **item,
+ "sentence_start": item["start"],
+ "sentence_end": item["end"],
+ "sentence_duration": max(0.0, item["end"] - item["start"]),
+ "error": "No valid temporal window",
+ "temporal_score": 0,
+ "is_relevant": 0,
+ }
+ )
+ continue
+ frames = _extract_frames(video_path, _times(*window, fps=judge_fps))
+ parsed = parse_judge_json(
+ judge.temporal(build_temporal_prompt(*window, item["sentence"], sample.question_text), frames)
+ )
+ temporal_rows.append(
+ {
+ **item,
+ "sentence_start": item["start"],
+ "sentence_end": item["end"],
+ "sentence_duration": max(0.0, item["end"] - item["start"]),
+ "window_start": window[0],
+ "window_end": window[1],
+ "error": None,
+ **parsed,
+ }
+ )
+ content_frames = None
+ if judge_video_mode == "frame-sample":
+ content_frames = _extract_frames(video_path, _content_frame_times(duration))
+ content = parse_judge_json(
+ judge.content(
+ build_content_prompt(_text(items), sample.question_text, [sample.answer1, sample.answer2]),
+ video_path,
+ frames=content_frames,
+ mode=judge_video_mode,
+ )
+ )
+ if content_frames is not None:
+ content["frame_count"] = len(content_frames)
+ result.update(
+ {
+ "temporal": {"sentences": temporal_rows, "summary": summarize_temporal_results(temporal_rows)},
+ "content": content,
+ }
+ )
+ else:
+ if sample.task_type is None:
+ raise ValueError(f"missing proactive-reminder task type for split {sample.split!r}")
+ task = sample.task_type
+ times = _reminder_times(sample)
+ event_rows = []
+ if task == "correction" or not times:
+ response = _text(items)
+ parsed = parse_judge_json(
+ judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
+ )
+ event_rows.append({"response_segment": response, **parsed})
+ else:
+ for start in times:
+ lo, hi = reminder_window(start, window_size)
+ response = " ".join(item["sentence"] for item in items if lo <= item["start"] <= hi)
+ if not response:
+ event_rows.append({"start_time": start, "response_segment": "", "success_score": 0})
+ else:
+ parsed = parse_judge_json(
+ judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
+ )
+ event_rows.append({"start_time": start, "response_segment": response, **parsed})
+ scores = [int(row.get("success_score", 0)) for row in event_rows]
+ result["pr"] = {
+ "events": event_rows,
+ "all_success": bool(scores) and all(score == 1 for score in scores),
+ "total_score": int(bool(scores) and all(score == 1 for score in scores)),
+ }
+ _write(score_path, result)
+ return result
+
+
+def _times(start: float, end: float, *, fps: int) -> list[float]:
+ count = max(1, int((end - start) * fps))
+ return (
+ [start + (end - start) / 2] if count == 1 else [start + i * (end - start) / (count - 1) for i in range(count)]
+ )
+
+
+def _extract_frames(path: str | Path, timestamps: list[float]) -> list[bytes]:
+ """Extract only decodable JPEG frames; clips can end on non-keyframes."""
+ frames = []
+ for timestamp in timestamps:
+ try:
+ frame = extract_jpeg(path, timestamp=timestamp)
+ except (OSError, subprocess.CalledProcessError, ValueError):
+ continue
+ if frame.startswith(b"\xff\xd8"):
+ frames.append(frame)
+ return frames
+
+
+def _content_frame_times(duration: float) -> list[float]:
+ """Sample the full video at approximately one frame every three seconds."""
+ if duration <= 0:
+ return []
+ times = []
+ timestamp = 0.0
+ while timestamp < duration:
+ times.append(min(timestamp, max(0.0, duration - 0.01)))
+ timestamp += 3.0
+ # Qwen2.5-VL has a finite per-request image-token budget. Preserve
+ # coverage while bounding long clips to a safely portable frame count.
+ if len(times) > 16:
+ stride = (len(times) - 1) / 15
+ times = [times[round(i * stride)] for i in range(16)]
+ return times
+
+
+def _reminder_times(sample: DuplexSample) -> list[float]:
+ values = []
+ for value in (sample.reminder1, sample.reminder2):
+ if isinstance(value, dict):
+ value = value.get("start", value.get("time"))
+ try:
+ if value is not None:
+ values.append(float(value))
+ except (TypeError, ValueError):
+ pass
+ return values
+
+
+def summarize_scores(score_root: str | Path) -> dict[str, Any]:
+ rows = []
+ for path in sorted(Path(score_root).rglob("*.json")):
+ if not path.name.endswith("_summary.json"):
+ row = _read(path)
+ row.setdefault("split", path.parent.name)
+ rows.append(row)
+ rtd_temporal = [
+ row["temporal"]["summary"] | {"content_score": row.get("content", {}).get("content_score", 0.0)}
+ for row in rows
+ if "temporal" in row
+ ]
+ pr = [{"task_type": row.get("task_type"), **row.get("pr", {})} for row in rows if "pr" in row]
+ result = {"protocol_pin": PROTOCOL_PIN, "samples": len(rows)}
+ if rtd_temporal:
+ by_task: dict[str, list[dict[str, float]]] = {}
+ for row in rows:
+ if "temporal" in row:
+ by_task.setdefault(str(row.get("split", row.get("task_type", "unknown"))), []).append(
+ {
+ "content": float(row.get("content", {}).get("content_score", 0.0)),
+ "temporal": float(row["temporal"]["summary"].get("avg_temporal_score", 0.0)),
+ }
+ )
+ result["rtd"] = {
+ "mean_content_score": sum(item["content_score"] for item in rtd_temporal) / len(rtd_temporal),
+ "mean_avg_temporal_score": sum(item["avg_temporal_score"] for item in rtd_temporal) / len(rtd_temporal),
+ "by_task": {
+ task: {
+ "mean_content_score": sum(item["content"] for item in values) / len(values),
+ "mean_avg_temporal_score": sum(item["temporal"] for item in values) / len(values),
+ }
+ for task, values in by_task.items()
+ },
+ }
+ if pr:
+ result["pr"] = summarize_pr_results(pr)
+ return result
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py
new file mode 100644
index 000000000..2fa4014d5
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py
@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OpenAI-compatible local judge client."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pybase64 as base64
+import requests
+
+
+def _build_openai_url(base_url: str, api_path: str) -> str:
+ base = base_url.rstrip("/")
+ normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
+ if base.endswith(normalized_path):
+ return base
+ if base.endswith("/v1"):
+ return f"{base}{normalized_path}"
+ return f"{base}/v1{normalized_path}"
+
+
+class DuplexJudge:
+ def __init__(self, base_url: str, model: str, *, api_key: str = "EMPTY", timeout: int = 600) -> None:
+ self.base_url, self.model, self.api_key, self.timeout = base_url, model, api_key, timeout
+
+ def chat(self, content: Any, *, system: str | None = None, max_tokens: int = 1200) -> str:
+ messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": content}]
+ response = requests.post(
+ _build_openai_url(self.base_url, "/chat/completions"),
+ json={
+ "model": self.model,
+ "messages": messages,
+ "temperature": 0.1,
+ "max_tokens": max_tokens,
+ },
+ headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+ timeout=self.timeout,
+ )
+ if not response.ok:
+ raise requests.HTTPError(
+ f"{response.status_code} response from judge: {response.text[:500]}", response=response
+ )
+ return str(response.json()["choices"][0]["message"]["content"])
+
+ def temporal(self, prompt: str, frames: list[bytes]) -> str:
+ content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+ content.extend(
+ {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()}}
+ for frame in frames
+ )
+ return self.chat(content, system="You are a careful evaluator for real-time multimodal systems.")
+
+ def content(
+ self,
+ prompt: str,
+ video: str | Path | None = None,
+ frames: list[bytes] | None = None,
+ *,
+ mode: str = "video_url",
+ ) -> str:
+ content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+ if mode == "video_url" and video is not None:
+ content.append({"type": "video_url", "video_url": {"url": "file://" + str(Path(video).resolve())}})
+ else:
+ content.extend(
+ {
+ "type": "image_url",
+ "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()},
+ }
+ for frame in (frames or [])
+ )
+ return self.chat(content)
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
new file mode 100644
index 000000000..629d8e877
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py
@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Small ffmpeg-backed media helpers used by the duplex runner and judge."""
+
+from __future__ import annotations
+
+import subprocess
+import wave
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+from vllm_omni.experimental.fullduplex.client import PCM16_BYTES_PER_SAMPLE, PCM16_SAMPLE_RATE
+
+
+def _run_ffmpeg(args: list[str]) -> bytes:
+ result = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", *args], capture_output=True, check=True)
+ return result.stdout
+
+
+def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path:
+ """Resolve common Hugging Face media values to a local file."""
+ if isinstance(value, str | Path):
+ return Path(value)
+ path = value.get("path") if isinstance(value, dict) else getattr(value, "path", None)
+ if path:
+ return Path(path)
+ payload = value.get("bytes") if isinstance(value, dict) else None
+ if payload is None and isinstance(value, bytes | bytearray):
+ payload = value
+ if payload is None:
+ raise ValueError(f"cannot materialize media value for {stem!r}")
+ destination = Path(output_dir) / f"{stem}{suffix}"
+ destination.parent.mkdir(parents=True, exist_ok=True)
+ destination.write_bytes(bytes(payload))
+ return destination
+
+
+def video_duration(path: str | Path) -> float:
+ result = subprocess.run(
+ ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)],
+ capture_output=True,
+ text=True,
+ check=True,
+ )
+ return max(0.0, float(result.stdout.strip()))
+
+
+def extract_jpeg(path: str | Path, *, timestamp: float, quality: int = 3) -> bytes:
+ return _run_ffmpeg(
+ [
+ "-ss",
+ f"{max(0.0, timestamp):.3f}",
+ "-i",
+ str(path),
+ "-frames:v",
+ "1",
+ "-q:v",
+ str(quality),
+ "-f",
+ "image2",
+ "pipe:1",
+ ]
+ )
+
+
+def iter_jpegs(
+ path: str | Path, *, fps: float = 1.0, duration: float | None = None, quality: int = 3
+) -> Iterator[tuple[float, bytes]]:
+ end = duration if duration is not None else video_duration(path)
+ step = 1.0 / fps
+ timestamp = 0.0
+ while timestamp < end:
+ try:
+ frame = extract_jpeg(path, timestamp=timestamp, quality=quality)
+ except (OSError, subprocess.CalledProcessError):
+ break
+ if frame:
+ yield timestamp, frame
+ timestamp += step
+
+
+def read_audio_pcm16(path: str | Path) -> bytes:
+ source = Path(path)
+ if source.suffix.lower() == ".wav":
+ try:
+ with wave.open(str(source), "rb") as wav_file:
+ if (
+ wav_file.getnchannels() == 1
+ and wav_file.getsampwidth() == 2
+ and wav_file.getframerate() == PCM16_SAMPLE_RATE
+ and wav_file.getcomptype() == "NONE"
+ ):
+ return wav_file.readframes(wav_file.getnframes())
+ except wave.Error:
+ pass
+ return _run_ffmpeg(["-i", str(source), "-f", "s16le", "-ac", "1", "-ar", str(PCM16_SAMPLE_RATE), "pipe:1"])
+
+
+def iter_av_units(
+ audio_pcm16: bytes, frames: Iterator[tuple[float, bytes]], *, unit_ms: int = 1000
+) -> Iterator[tuple[bytes, bytes | None]]:
+ unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * unit_ms // 1000
+ frame_iter = iter(frames)
+ next_frame = next(frame_iter, None)
+ for offset in range(0, len(audio_pcm16), unit_bytes):
+ timestamp = offset / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+ frame = None
+ if next_frame is not None and next_frame[0] <= timestamp + unit_ms / 1000:
+ frame = next_frame[1]
+ next_frame = next(frame_iter, None)
+ yield audio_pcm16[offset : offset + unit_bytes], frame
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py
new file mode 100644
index 000000000..a373be58d
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py
@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Protocol prompts, windows, parsing, and aggregation.
+
+Prompt wording and window math are copied verbatim from the MIT-licensed
+OpenBMB/Omni-DuplexEval evaluator at :data:`PROTOCOL_PIN`.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from typing import Any
+
+import regex as re
+
+# Copied from OpenBMB/Omni-DuplexEval@PROTOCOL_PIN (MIT).
+# Do not change prompt wording without bumping the pin and recording a score delta.
+PROTOCOL_PIN = "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
+
+
+def temporal_window(sentence_start: float, sentence_end: float, video_duration: float) -> tuple[float, float] | None:
+ start = max(0.0, float(sentence_start) - 2.0)
+ end = max(start + 0.5, float(sentence_end) - 2.0)
+ return (start, end) if end <= float(video_duration) and end - start >= 0.5 else None
+
+
+def reminder_window(start_time: float, window_size: float = 10.0) -> tuple[float, float]:
+ return float(start_time), float(start_time) + float(window_size)
+
+
+def parse_judge_json(text: str) -> dict[str, Any]:
+ for match in re.findall(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", str(text), re.S):
+ try:
+ value = json.loads(re.sub(r",\s*([}\]])", r"\1", match))
+ except (TypeError, ValueError, json.JSONDecodeError):
+ continue
+ if isinstance(value, dict):
+ break
+ else:
+ value = {}
+ patterns = [
+ (r'"temporal_score"\s*:\s*(\d+)', "temporal_score", int),
+ (r'"is_relevant"\s*:\s*(\d+)', "is_relevant", int),
+ (r'"success_score"\s*:\s*([01])', "success_score", int),
+ (r'"content_score"\s*:\s*(\d+(?:\.\d+)?)', "content_score", float),
+ ]
+ for pattern, key, caster in patterns:
+ if match := re.search(pattern, str(text)):
+ parsed = caster(match.group(1))
+ value[key] = round(parsed, 2) if caster is float else parsed
+ for pattern, key in [
+ (r'"temporal_reasoning"\s*:\s*"([^"]*)"', "temporal_reasoning"),
+ (r'"content_reasoning"\s*:\s*"([^"]*)"', "content_reasoning"),
+ (r'"reasoning"\s*:\s*"([^"]*)"', "reasoning"),
+ ]:
+ if match := re.search(pattern, str(text), re.S):
+ value[key] = match.group(1).strip()
+ break
+ for key in ("temporal_score", "is_relevant", "success_score"):
+ if key in value:
+ try:
+ value[key] = int(float(value[key]))
+ except (TypeError, ValueError):
+ value[key] = 0
+ if "content_score" in value:
+ try:
+ value["content_score"] = round(float(value["content_score"]), 2)
+ except (TypeError, ValueError):
+ value["content_score"] = 0.0
+ return value
+
+
+def summarize_temporal_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
+ rows = list(results)
+ valid = [row for row in rows if not row.get("error")]
+ relevant = [row for row in valid if int(row.get("is_relevant", 0)) == 1]
+ irrelevant = [row for row in valid if int(row.get("is_relevant", 0)) == 0]
+ scores = [int(row.get("temporal_score", 0)) for row in relevant]
+ total_valid_duration = sum(float(row.get("sentence_duration", 0.0)) for row in valid)
+ irrelevant_duration = sum(float(row.get("sentence_duration", 0.0)) for row in irrelevant)
+ return {
+ "avg_temporal_score": round(sum(scores) / len(scores), 4) if scores else 0.0,
+ "total_sentences": len(rows),
+ "evaluated_sentences": len(valid),
+ "error_count": len(rows) - len(valid),
+ "relevant_sentences_count": len(relevant),
+ "irrelevant_sentences_count": len(irrelevant),
+ "score_distribution": {
+ "3_points": scores.count(3),
+ "2_points": scores.count(2),
+ "1_point": scores.count(1),
+ "0_points": scores.count(0),
+ },
+ "relevance_stats": {
+ "first_relevant_time": min((row.get("sentence_start") for row in relevant), default=None),
+ "total_relevant_duration": round(sum(float(row.get("sentence_duration", 0.0)) for row in relevant), 4),
+ "total_irrelevant_duration": round(irrelevant_duration, 4),
+ "irrelevant_duration_ratio": round(
+ irrelevant_duration / total_valid_duration if total_valid_duration > 0 else 0.0, 4
+ ),
+ },
+ "excluded_irrelevant_sentences": [
+ {
+ "sentence": row.get("sentence", ""),
+ "start": row.get("sentence_start"),
+ "end": row.get("sentence_end"),
+ "temporal_score": row.get("temporal_score", 0),
+ }
+ for row in irrelevant
+ ],
+ }
+
+
+def summarize_pr_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
+ rows = list(results)
+ successes = [int(row.get("all_success", row.get("total_score", 0))) for row in rows]
+ by_task: dict[str, list[int]] = {}
+ for row, score in zip(rows, successes, strict=False):
+ by_task.setdefault(str(row.get("task_type", "unknown")), []).append(score)
+ return {
+ "samples": len(rows),
+ "mean_all_success": sum(successes) / len(successes) if successes else 0.0,
+ "by_task": {task: sum(values) / len(values) for task, values in by_task.items()},
+ }
+
+
+def build_temporal_prompt(start_time: float, end_time: float, response_text: str, question: str) -> str:
+ return f"""
+You are evaluating a real-time video description system.
+
+Basic information:
+- Video segment: {start_time:.2f}s to {end_time:.2f}s.
+- User instruction: {question}
+- Model response sentence: "{response_text}"
+
+Evaluation principles:
+- Judge only the current video segment.
+- Ignore response content that clearly belongs to earlier or later segments unless it contradicts the current segment.
+- Language does not matter; evaluate semantic correctness and temporal alignment only.
+- Decide whether the sentence is a substantive response or only a filler/polite phrase.
+
+Temporal sensitivity score:
+- 3: Excellent temporal alignment. The sentence accurately describes the current segment and follows the instruction.
+- 2: Mostly aligned. Minor inaccuracies, omissions, or harmless references to nearby segments are allowed.
+- 1: Poor alignment. Major inaccuracies, wrong timing, or mostly irrelevant content.
+- 0: No meaningful alignment with the current segment.
+
+Relevance label:
+- 1: The sentence contains substantive task- or video-related information.
+- 0: The sentence is only filler, acknowledgement, meta-commentary, or generic text without useful content.
+
+Return exactly this JSON object:
+{{
+ "temporal_score": <0, 1, 2, or 3>,
+ "temporal_reasoning": "",
+ "is_relevant": <0 or 1>
+}}
+"""
+
+
+def build_content_prompt(response_text: str, question: str, references: list[str] | None = None) -> str:
+ reference_block = ""
+ if references:
+ reference_block = "Reference annotations, for reference only:\n"
+ for index, reference in enumerate(references, 1):
+ reference_block += f"{index}. {reference}\n"
+
+ return f"""
+You are a precise content-accuracy evaluator for video description.
+
+Basic information:
+- User instruction: {question}
+- Model response: "{response_text}"
+{reference_block}
+
+Evaluation goal:
+Judge whether the model response is factually accurate for the whole video and aligned with the instruction.
+Consider object/action/color/count/spatial/event errors, hallucinations, omissions, and irrelevant content.
+Reference annotations are optional guidance; the primary evidence is the video itself.
+
+Scoring:
+- Use a decimal score from 0.00 to 3.00.
+- Start from 3.00 and deduct for each error.
+- Use 0.00 only when the response is empty, completely irrelevant, or contains no correct video facts.
+- Output exactly two decimal places.
+
+Return exactly this JSON object:
+{{
+ "content_score": ,
+ "content_reasoning": ""
+}}
+"""
+
+
+def build_reminder_prompt(instruction: str, response: str, task_type: str, ground_answer: str = "") -> str:
+ if task_type == "correction":
+ return f"""
+You are judging whether a model successfully completed a correction task.
+
+Task description:
+The user instruction contains an incorrect statement about the video.
+The model should identify the incorrect part and provide the correct information.
+
+Input:
+- User instruction: {instruction}
+- Reference correction: {ground_answer}
+- Model response: "{response}"
+
+Success criteria:
+1. Identify the error implied by the user instruction and the reference correction.
+2. Check whether the model corrected all required error points.
+3. The corrected content must preserve the correct context, including subject, object, action, and attributes.
+4. Ignore extra information that is unrelated to both the instruction and the reference, unless it contradicts them.
+
+Scoring:
+- 1 = all required error points are corrected and the context is consistent.
+- 0 = at least one required correction is missing, wrong, or contextually inconsistent.
+
+Return exactly this JSON object:
+{{
+ "success_score": <0 or 1>,
+ "reasoning": ""
+}}
+"""
+
+ if task_type == "post_event_reminder":
+ timing_note = "The response is evaluated immediately after the target event has occurred."
+ else:
+ timing_note = "The response is evaluated at the target event time or immediately after it."
+
+ return f"""
+You are judging whether a model successfully completed an event-reminder task.
+
+Task description:
+The user gave an instruction asking the system to remind them when a specific event happens.
+{timing_note}
+The provided model text is the response segment generated in the evaluation window.
+
+Input:
+- User instruction: {instruction}
+- Model response segment: "{response}"
+
+Success criteria:
+1. The response clearly refers to the target event in the instruction.
+2. The response communicates a reminder, notification, or confirmation that the event has happened.
+3. Vague narration, unrelated descriptions, or wrong-event reminders are failures.
+4. The response does not need to match the instruction wording exactly if the meaning is clear.
+
+Scoring:
+- 1 = successful reminder.
+- 0 = unsuccessful reminder.
+
+Return exactly this JSON object:
+{{
+ "success_score": <0 or 1>,
+ "reasoning": ""
+}}
+"""
diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py
new file mode 100644
index 000000000..d5a28aa06
--- /dev/null
+++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Single-sample MiniCPM-o native-duplex generation."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import pybase64 as base64
+
+from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, build_realtime_url, wait_for
+
+from .omni_duplex_eval_clock import extract_timed_sentences
+from .omni_duplex_eval_media import iter_av_units, iter_jpegs, materialize_media, read_audio_pcm16, video_duration
+
+
+def _ref_audio(path: str | Path) -> str:
+ value = Path(path).expanduser().read_bytes()
+ return "data:audio/wav;base64," + base64.b64encode(value).decode("ascii")
+
+
+async def generate_sample(
+ sample: Any,
+ *,
+ url: str,
+ model: str,
+ ref_audio: str | Path,
+ output_root: str | Path,
+ fps: float = 1.0,
+ mix: str = "question",
+ pace: str = "realtime",
+ clock: str = "media",
+ overwrite: bool = False,
+ unit_ms: int = 1000,
+) -> Path:
+ output = Path(output_root) / sample.split / f"{sample.id}.json"
+ meta_path = output.with_name(output.stem + ".meta.json")
+ if output.exists() and not overwrite:
+ return output
+ if mix != "question":
+ raise NotImplementedError("v1 supports mix=question; soundtrack mixing is reserved for P1")
+ media_dir = output.parent / ".media"
+ audio_path = materialize_media(sample.question_audio, media_dir, f"{sample.id}_question", ".wav")
+ video_path = materialize_media(sample.video, media_dir, sample.id, ".mp4")
+ pcm = read_audio_pcm16(audio_path)
+ duration = sample.video_duration or video_duration(video_path)
+ frames = iter_jpegs(video_path, fps=fps, duration=duration)
+ realtime = pace == "realtime"
+ if pace not in {"realtime", "as-fast-as-possible"}:
+ raise ValueError("pace must be realtime or as-fast-as-possible")
+ client = RealtimeDuplexClient(build_realtime_url(url, model, autostart=False))
+ response_done = False
+ drain_timeout = None
+ close_timeout = None
+ async with client:
+ await client.configure(model, ref_audio=_ref_audio(ref_audio), instructions="Streaming Omni Conversation.")
+ ack_task = asyncio.create_task(_ack_playback(client))
+ try:
+ await client.stream_av_units(iter_av_units(pcm, frames, unit_ms=unit_ms), realtime=realtime)
+ await client.commit()
+ try:
+ await wait_for(
+ lambda: client.events.count("response.done") > 0,
+ timeout_s=max(20.0, duration + 20.0),
+ label="response.done",
+ )
+ response_done = True
+ except TimeoutError as exc:
+ drain_timeout = str(exc)
+ finally:
+ ack_task.cancel()
+ try:
+ await ack_task
+ except asyncio.CancelledError:
+ pass
+ await client.acknowledge_playback()
+ try:
+ await client.close_session(timeout_s=20.0)
+ except TimeoutError as exc:
+ close_timeout = str(exc)
+ events = list(client.events.events)
+ timed = [sentence.as_dict() for sentence in extract_timed_sentences(events, clock=clock)]
+ output.parent.mkdir(parents=True, exist_ok=True)
+ output.write_text(json.dumps(timed, ensure_ascii=False, indent=2), encoding="utf-8")
+ meta = {
+ "id": sample.id,
+ "split": sample.split,
+ "clock": clock if realtime else "invalid",
+ "pace": pace,
+ "mix": mix,
+ "fps": fps,
+ "unit_ms": unit_ms,
+ "model": model,
+ "response_done": response_done,
+ "drain_timeout": drain_timeout,
+ "close_timeout": close_timeout,
+ "ref_audio_sha256": hashlib.sha256(Path(ref_audio).read_bytes()).hexdigest(),
+ }
+ meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+ return output
+
+
+async def _ack_playback(client: RealtimeDuplexClient) -> None:
+ while True:
+ await asyncio.sleep(0.25)
+ await client.acknowledge_playback()
diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py
index 824b810f3..3d7c983f3 100644
--- a/vllm_omni/benchmarks/metrics/metrics.py
+++ b/vllm_omni/benchmarks/metrics/metrics.py
@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import warnings
from collections import defaultdict
+from collections.abc import Sequence
from dataclasses import field, make_dataclass
import numpy as np
@@ -15,6 +19,11 @@ _FLOAT_LIST_TYPE = list[float]
_INT_LIST_TYPE = list[int]
_MULTIMODAL_BENCHMARK_FIELDS = [
+ ("num_ttft_samples", int, field(default=0)),
+ ("num_tpot_samples", int, field(default=0)),
+ ("num_itl_samples", int, field(default=0)),
+ ("num_audio_ttfp_samples", int, field(default=0)),
+ ("num_audio_rtf_samples", int, field(default=0)),
(defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
(defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
(defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
@@ -176,6 +185,17 @@ def _stage_modality_flags(
)
+def has_metric_samples(metrics: object, metric_name: str) -> bool:
+ sample_count_attr = {
+ "ttft": "num_ttft_samples",
+ "tpot": "num_tpot_samples",
+ "itl": "num_itl_samples",
+ defs.AUDIO_TTFP: "num_audio_ttfp_samples",
+ defs.AUDIO_RTF: "num_audio_rtf_samples",
+ }.get(metric_name)
+ return sample_count_attr is None or getattr(metrics, sample_count_attr, 0) > 0
+
+
def print_metrics(
task_type,
selected_percentile_metrics,
@@ -308,6 +328,9 @@ def process_one_metric(
metric_attribute_name: str,
metrics: MultiModalsBenchmarkMetrics,
):
+ if not has_metric_samples(metrics, metric_attribute_name):
+ return
+
metric_header_map = {
"ttft": "Time to First Token",
"tpot": "Time per Output Token (excl. 1st token)",
@@ -718,10 +741,12 @@ def calculate_metrics(
good_completed = 0
itls: list[float] = []
tpots: list[float] = []
- all_tpots: list[float] = []
+ all_tpots: list[float | None] = []
ttfts: list[float] = []
+ goodput_ttfts: list[float | None] = []
e2els: list[float] = []
audio_ttfps: list[float] = []
+ goodput_audio_ttfps: list[float | None] = []
audio_rtfs: list[float] = []
audio_duration: list[float] = []
audio_frames: list[int] = []
@@ -748,27 +773,49 @@ def calculate_metrics(
output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
actual_output_lens.append(output_len)
total_input += outputs[i].prompt_len
- tpot = 0
+ tpot: float | None = None
if output_len > 1:
if outputs[i].itl:
# Use mean(ITL) directly so per-request TPOT == mean(ITL).
# The ITL list records one entry per SSE chunk; server may
# bundle multiple tokens per chunk, so len(itl)+1 != output_len.
# Using mean(itl) keeps TPOT and ITL on the same footing.
- tpot = sum(outputs[i].itl) / len(outputs[i].itl)
- else:
- try:
- latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft
- except Exception:
- latency_minus_ttft = outputs[i].latency - outputs[i].ttft
- tpot = latency_minus_ttft / (output_len - 1)
- tpots.append(tpot)
- # Note: if output_len <= 1, we regard tpot as 0 for goodput
+ measured_tpot = sum(outputs[i].itl) / len(outputs[i].itl)
+ if np.isfinite(measured_tpot) and measured_tpot > 0:
+ tpot = measured_tpot
+ elif getattr(outputs[i], "tpot_measured", True):
+ text_latency = getattr(outputs[i], "text_latency", None)
+ ttft = getattr(outputs[i], "ttft", None)
+ if (
+ isinstance(text_latency, (int, float))
+ and isinstance(ttft, (int, float))
+ and np.isfinite(text_latency)
+ and np.isfinite(ttft)
+ and text_latency > ttft > 0
+ ):
+ tpot = (text_latency - ttft) / (output_len - 1)
+ if tpot is not None:
+ tpots.append(tpot)
+ # TPOT is not measurable for a single-token response or when the
+ # client did not observe a positive decode interval. In
+ # particular, stage-level output token counts alone must not turn
+ # an unset text_latency into a negative TPOT sample.
all_tpots.append(tpot)
itls += outputs[i].itl
- ttfts.append(outputs[i].ttft)
- audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
- audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
+ session_metrics = getattr(outputs[i], "duplex_session_metrics", None)
+ ttft_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_ttft_ms") is not None
+ audio_ttfp_measured = (
+ not isinstance(session_metrics, dict) or session_metrics.get("mean_ttfp_ms") is not None
+ )
+ audio_rtf_measured = not isinstance(session_metrics, dict) or session_metrics.get("mean_rtf") is not None
+ if ttft_measured:
+ ttfts.append(outputs[i].ttft)
+ if audio_ttfp_measured:
+ audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0))
+ if audio_rtf_measured:
+ audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
+ goodput_ttfts.append(outputs[i].ttft if ttft_measured else None)
+ goodput_audio_ttfps.append(getattr(outputs[i], defs.AUDIO_TTFP, 0.0) if audio_ttfp_measured else None)
audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
@@ -789,14 +836,14 @@ def calculate_metrics(
actual_output_lens.append(0)
if goodput_config_dict:
- valid_metrics = []
+ valid_metrics: list[Sequence[float | None]] = []
slo_values = []
if "ttft" in goodput_config_dict:
- valid_metrics.append(ttfts)
+ valid_metrics.append(goodput_ttfts)
slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
if "audio_ttft" in goodput_config_dict:
- valid_metrics.append(audio_ttfps)
+ valid_metrics.append(goodput_audio_ttfps)
slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
if "tpot" in goodput_config_dict:
valid_metrics.append(all_tpots)
@@ -806,7 +853,7 @@ def calculate_metrics(
slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
for req_metric in zip(*valid_metrics):
- is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
+ is_good_req = all(r is not None and s >= r for s, r in zip(slo_values, req_metric))
if is_good_req:
good_completed += 1
@@ -822,6 +869,7 @@ def calculate_metrics(
# Calculate max output tokens per second metric
max_output_tokens_per_s = 0.0
max_concurrent_requests = 0
+ token_timeline_available = True
# Find the time range across all successful requests
successful_outputs = [output for output in outputs if output.success]
@@ -836,19 +884,26 @@ def calculate_metrics(
concurrent_requests_per_second = np.zeros(duration_seconds)
for i, output in enumerate(successful_outputs):
- # Calculate token generation timestamp using
- # start_time, ttft, and itl
- token_times = [output.start_time + output.ttft]
- current_time = token_times[0]
- for itl_value in output.itl:
- current_time += itl_value
- token_times.append(current_time)
+ session_metrics = getattr(output, "duplex_session_metrics", None)
+ if isinstance(session_metrics, dict):
+ # Duplex TTFT/ITL samples are response-local. Without each
+ # response.created offset they cannot form a session-global
+ # token timeline, so do not publish a misleading peak rate.
+ token_timeline_available = False
+ else:
+ # Calculate token generation timestamp using
+ # start_time, ttft, and itl
+ token_times = [output.start_time + output.ttft]
+ current_time = token_times[0]
+ for itl_value in output.itl:
+ current_time += itl_value
+ token_times.append(current_time)
- # Add tokens to second buckets
- for token_time in token_times:
- second_bucket = int(token_time - min_start_time)
- if 0 <= second_bucket < duration_seconds:
- tokens_per_second[second_bucket] += 1
+ # Add tokens to second buckets
+ for token_time in token_times:
+ second_bucket = int(token_time - min_start_time)
+ if 0 <= second_bucket < duration_seconds:
+ tokens_per_second[second_bucket] += 1
# Track concurrent requests for each second this request was active
request_start_second = int(output.start_time - min_start_time)
@@ -860,18 +915,22 @@ def calculate_metrics(
# Find the maximum tokens per second and corresponding
# concurrent requests
if len(tokens_per_second) > 0:
- max_output_tokens_per_s = float(np.max(tokens_per_second))
+ if token_timeline_available:
+ max_output_tokens_per_s = float(np.max(tokens_per_second))
+ else:
+ max_output_tokens_per_s = float("nan")
max_concurrent_requests = int(np.max(concurrent_requests_per_second))
if TERM_PLOTLIB_AVAILABLE:
import termplotlib as tpl
fig = tpl.figure()
- fig.plot(
- np.arange(len(tokens_per_second)),
- tokens_per_second,
- title="Output tokens per second",
- )
+ if token_timeline_available:
+ fig.plot(
+ np.arange(len(tokens_per_second)),
+ tokens_per_second,
+ title="Output tokens per second",
+ )
fig.plot(
np.arange(len(concurrent_requests_per_second)),
concurrent_requests_per_second,
@@ -881,6 +940,10 @@ def calculate_metrics(
else:
print("tip: install termplotlib and gnuplot to plot the metrics")
+ duplex_metrics_present = any(
+ isinstance(getattr(output, "duplex_session_metrics", None), dict) for output in outputs
+ )
+ missing_duplex_value = float("nan") if duplex_metrics_present else 0
metrics = MultiModalsBenchmarkMetrics(
completed=completed,
failed=len(failed_outputs),
@@ -890,18 +953,18 @@ def calculate_metrics(
request_goodput=good_completed / dur_s,
output_throughput=sum(actual_output_lens) / dur_s,
total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
- mean_ttft_ms=np.mean(ttfts or 0) * 1000, # ttfts is empty if streaming is not supported by the endpoint
- std_ttft_ms=np.std(ttfts or 0) * 1000,
- median_ttft_ms=np.median(ttfts or 0) * 1000,
- percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
- mean_tpot_ms=np.mean(tpots or 0) * 1000,
- std_tpot_ms=np.std(tpots or 0) * 1000,
- median_tpot_ms=np.median(tpots or 0) * 1000,
- percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
- mean_itl_ms=np.mean(itls or 0) * 1000,
- std_itl_ms=np.std(itls or 0) * 1000,
- median_itl_ms=np.median(itls or 0) * 1000,
- percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
+ mean_ttft_ms=np.mean(ttfts or missing_duplex_value) * 1000,
+ std_ttft_ms=np.std(ttfts or missing_duplex_value) * 1000,
+ median_ttft_ms=np.median(ttfts or missing_duplex_value) * 1000,
+ percentiles_ttft_ms=[(p, np.percentile(ttfts or missing_duplex_value, p) * 1000) for p in selected_percentiles],
+ mean_tpot_ms=np.mean(tpots or float("nan")) * 1000,
+ std_tpot_ms=np.std(tpots or float("nan")) * 1000,
+ median_tpot_ms=np.median(tpots or float("nan")) * 1000,
+ percentiles_tpot_ms=[(p, np.percentile(tpots or float("nan"), p) * 1000) for p in selected_percentiles],
+ mean_itl_ms=np.mean(itls or missing_duplex_value) * 1000,
+ std_itl_ms=np.std(itls or missing_duplex_value) * 1000,
+ median_itl_ms=np.median(itls or missing_duplex_value) * 1000,
+ percentiles_itl_ms=[(p, np.percentile(itls or missing_duplex_value, p) * 1000) for p in selected_percentiles],
mean_e2el_ms=np.mean(e2els or 0) * 1000,
std_e2el_ms=np.std(e2els or 0) * 1000,
median_e2el_ms=np.median(e2els or 0) * 1000,
@@ -910,11 +973,16 @@ def calculate_metrics(
max_concurrent_requests=max_concurrent_requests,
rtfx=input_audio_duration / dur_s,
**{
- defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
- defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
- defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
+ "num_ttft_samples": len(ttfts),
+ "num_tpot_samples": len(tpots),
+ "num_itl_samples": len(itls),
+ "num_audio_ttfp_samples": len(audio_ttfps),
+ "num_audio_rtf_samples": len(audio_rtfs),
+ defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or missing_duplex_value) * 1000,
+ defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or missing_duplex_value) * 1000,
+ defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or missing_duplex_value) * 1000,
defs.PERCENTILES_AUDIO_TTFP_MS: [
- (p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles
+ (p, np.percentile(audio_ttfps or missing_duplex_value, p) * 1000) for p in selected_percentiles
],
defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0),
@@ -925,10 +993,12 @@ def calculate_metrics(
defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration),
defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
- defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
- defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
- defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
- defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
+ defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or missing_duplex_value),
+ defs.STD_AUDIO_RTF: np.std(audio_rtfs or missing_duplex_value),
+ defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or missing_duplex_value),
+ defs.PERCENTILES_AUDIO_RTF: [
+ (p, np.percentile(audio_rtfs or missing_duplex_value, p)) for p in selected_percentiles
+ ],
defs.TOTAL_IMAGES: total_images,
defs.IMAGE_THROUGHPUT: total_images / dur_s,
defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py
new file mode 100644
index 000000000..46822b486
--- /dev/null
+++ b/vllm_omni/benchmarks/omniinteract.py
@@ -0,0 +1,933 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniInteract session execution and artifacts for serving benchmarks."""
+
+from __future__ import annotations
+
+import asyncio
+import binascii
+import contextlib
+import copy
+import fcntl
+import hashlib
+import json
+import logging
+import math
+import os
+import subprocess
+import tempfile
+import time
+from collections import Counter
+from collections.abc import Callable, Iterator, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, cast
+from urllib.parse import urljoin, urlsplit
+
+import pybase64 as base64
+
+from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+ DEFAULT_MAX_VIDEO_DURATION_S,
+ OmniInteractCase,
+ OmniInteractPreparedInput,
+ case_manifest,
+)
+from vllm_omni.experimental.fullduplex.client import (
+ PCM16_BYTES_PER_SAMPLE,
+ PCM16_SAMPLE_RATE,
+ RealtimeDuplexClient,
+ RealtimeEventCollector,
+ build_realtime_url,
+ chunk_period_ms,
+ has_residual_model_unit,
+ reference_audio_data_url,
+ summarize_session_request_metrics,
+ write_pcm16_wav,
+)
+
+OUTPUT_SAMPLE_RATE = 24_000
+SUCCESS_ARTIFACTS = (".done", "output.wav", "wav_transcript.json", "events.json", "result.json")
+BATCH_ARTIFACTS = ("batch_summary.json", "official_eval_manifest.jsonl")
+ARTIFACT_LOCK_FILE = ".omniinteract.lock"
+_INPUT_CHUNK_MS = 200
+VIDEO_FPS = 1.0
+_COMPLETION_SETTLE_S = 2.0
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class OmniInteractBenchmarkConfig:
+ model: str
+ base_url: str = "http://127.0.0.1:8000"
+ endpoint: str = "/v1/realtime"
+ output_root: Path = Path("omniinteract-output")
+ timeout_s: float = 900.0
+ media_timeout_s: float = 600.0
+ max_video_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S
+ ref_audio: str | None = None
+ require_response: bool = False
+ extra_headers: dict[str, str] | None = None
+ extra_body: dict[str, object] | None = None
+
+
+@dataclass
+class OmniInteractCaseResult:
+ subset: str
+ video: str
+ output_dir: str
+ success: bool = False
+ error: str = ""
+ session_id: str = ""
+ latency_s: float = 0.0
+ input_audio_chunks: int = 0
+ input_video_frames: int = 0
+ pacing_mean_lag_s: float = 0.0
+ pacing_max_lag_s: float = 0.0
+ responses: int = 0
+ audio_bytes: int = 0
+ audio_clipped_bytes: int = 0
+ transcript: str = ""
+ eligible_for_official_eval: bool = False
+ official_eval_ineligible_reasons: list[str] = field(default_factory=list)
+ artifact_warnings: list[str] = field(default_factory=list)
+ output_tokens: int = 0
+ duplex_request_metrics: list[dict[str, object]] = field(default_factory=list)
+ duplex_session_metrics: dict[str, object] = field(default_factory=dict)
+ _artifact_context: _DeferredArtifactContext | None = field(default=None, repr=False, compare=False)
+
+ def as_dict(self) -> dict[str, Any]:
+ return copy.deepcopy({key: value for key, value in vars(self).items() if not key.startswith("_")})
+
+
+@dataclass(frozen=True)
+class _DeferredArtifactContext:
+ horizon_bytes: int
+ spans: tuple[tuple[int, bytes], ...]
+ chunks: list[dict[str, object]]
+ events: list[dict[str, object]]
+
+
+def benchmark_summary(results: list[OmniInteractCaseResult]) -> dict[str, Any]:
+ succeeded = sum(result.success for result in results)
+ eligible = sum(result.success and result.eligible_for_official_eval for result in results)
+ return {
+ "total": len(results),
+ "success": succeeded,
+ "failed": len(results) - succeeded,
+ "eligible_for_official_eval": eligible,
+ "successful_but_ineligible": succeeded - eligible,
+ "audio_clipped_bytes": sum(result.audio_clipped_bytes for result in results),
+ "results": [result.as_dict() for result in results],
+ }
+
+
+def _run_media_command(
+ command: list[str],
+ *,
+ timeout_s: float,
+ text: bool = False,
+) -> subprocess.CompletedProcess[Any]:
+ try:
+ result = subprocess.run(
+ command,
+ capture_output=True,
+ check=False,
+ text=text,
+ timeout=timeout_s,
+ )
+ except subprocess.TimeoutExpired as exc:
+ raise TimeoutError(f"Media command timed out after {timeout_s:g}s: {command[0]}") from exc
+ except FileNotFoundError as exc:
+ raise RuntimeError(f"Required media command is unavailable: {command[0]}") from exc
+ if result.returncode:
+ error = result.stderr if text else result.stderr.decode("utf-8", "ignore")
+ raise RuntimeError(f"{command[0]} failed: {error.strip()}")
+ return result
+
+
+def prepare_media(
+ video: Path,
+ fps: float,
+ *,
+ timeout_s: float,
+ max_duration_s: float = DEFAULT_MAX_VIDEO_DURATION_S,
+) -> tuple[float, bytes, list[str | None]]:
+ probe = [
+ "ffprobe",
+ *"-v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1".split(),
+ str(video),
+ ]
+ duration = float(_run_media_command(probe, text=True, timeout_s=timeout_s).stdout.strip())
+ if not math.isfinite(duration) or duration <= 0 or duration > max_duration_s:
+ raise ValueError(f"Invalid video duration for {video}: {duration!r}")
+ target_bytes = math.ceil(duration) * PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+ audio = [
+ "ffmpeg",
+ *"-loglevel error -i".split(),
+ str(video),
+ "-t",
+ str(duration),
+ *f"-vn -f s16le -ac 1 -ar {PCM16_SAMPLE_RATE} -fs".split(),
+ str(target_bytes),
+ "pipe:1",
+ ]
+ pcm = _run_media_command(audio, timeout_s=timeout_s).stdout
+ pcm = (pcm + bytes(max(0, target_bytes - len(pcm))))[:target_bytes]
+
+ frame_count = math.ceil(duration * fps)
+ frames: list[str | None] = [None] * frame_count
+ with tempfile.TemporaryDirectory(prefix="vllm-omni-frames-") as temp_dir:
+ output_pattern = Path(temp_dir) / "frame-%06d.jpg"
+ frame_filter = f"select=gte(t\\,(selected_n+0.5)/{fps}),scale=640:640:force_original_aspect_ratio=decrease"
+ frame_command = [
+ "ffmpeg",
+ *"-loglevel error -i".split(),
+ str(video),
+ "-an",
+ "-vf",
+ frame_filter,
+ *f"-frames:v {frame_count} -vsync vfr -q:v 5".split(),
+ str(output_pattern),
+ ]
+ _run_media_command(frame_command, timeout_s=timeout_s)
+ for index, frame_path in enumerate(sorted(Path(temp_dir).glob("frame-*.jpg"))[:frame_count]):
+ frames[index] = base64.b64encode(frame_path.read_bytes()).decode("ascii")
+ return duration, pcm, frames
+
+
+@dataclass(frozen=True)
+class _AudioSegment:
+ event_index: int
+ response_id: str
+ start_s: float
+ end_s: float
+ pcm16: bytes
+
+
+class _Playback:
+ def __init__(self) -> None:
+ self.cursor = 0
+ self.end_s = 0.0
+ self.segments: list[_AudioSegment] = []
+ self.completed: set[str] = set()
+ self.completion_acked: set[str] = set()
+ self.warnings: list[str] = []
+ self._total_samples: dict[str, int] = {}
+ self._response_end_s: dict[str, float] = {}
+
+ def _warn_once(self, warning: str) -> None:
+ if warning not in self.warnings:
+ self.warnings.append(warning)
+
+ def ingest(self, events: RealtimeEventCollector) -> None:
+ while self.cursor < len(events.events):
+ index, self.cursor = self.cursor, self.cursor + 1
+ event = events.events[index]
+ response_id = events.response_id(event)
+ if event.get("type") == "response.done":
+ if response_id:
+ self.completed.add(response_id)
+ continue
+ if event.get("type") != "response.audio.delta":
+ continue
+ encoded = event.get("delta") or event.get("audio")
+ if not response_id:
+ raise ValueError("response audio has no response_id")
+ if not isinstance(encoded, str):
+ raise ValueError("response audio payload is missing")
+ if event.get("format") is None:
+ self._warn_once("response.audio.delta omitted format; assumed pcm16")
+ elif event.get("format") != "pcm16":
+ raise ValueError("OmniInteract output must be pcm16")
+ rate = event.get("sample_rate_hz")
+ if rate is None:
+ rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE
+ self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}")
+ if rate != OUTPUT_SAMPLE_RATE:
+ raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio")
+ try:
+ raw = base64.b64decode(encoded, validate=True)
+ except (ValueError, binascii.Error) as exc:
+ raise ValueError("response audio is not valid base64") from exc
+ if not raw or len(raw) % PCM16_BYTES_PER_SAMPLE:
+ raise ValueError("response audio is empty or not PCM16 aligned")
+ start = max(events.event_received_at_s[index], self.end_s)
+ samples = len(raw) // PCM16_BYTES_PER_SAMPLE
+ segment = _AudioSegment(index, response_id, start, start + samples / rate, raw)
+ self.segments.append(segment)
+ self.end_s = segment.end_s
+ self._total_samples[response_id] = self._total_samples.get(response_id, 0) + samples
+ self._response_end_s[response_id] = segment.end_s
+
+ async def acknowledge(self, client: RealtimeDuplexClient, now: float | None = None) -> None:
+ events = client.events
+ self.ingest(events)
+ now = time.monotonic() if now is None else now
+ for response_id in self.completed - self.completion_acked:
+ total_samples = self._total_samples.get(response_id, 0)
+ if not total_samples or now < self._response_end_s[response_id]:
+ continue
+ played_ms = total_samples * 1000 // OUTPUT_SAMPLE_RATE
+ await client.send_playback_ack(response_id, played_ms)
+ self.completion_acked.add(response_id)
+
+
+async def stream_inputs(
+ client: RealtimeDuplexClient,
+ pcm: bytes,
+ frames: Sequence[str | None],
+ playback: _Playback,
+) -> tuple[int, int, float, float]:
+ bytes_per_second = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+ chunk_bytes = bytes_per_second * _INPUT_CHUNK_MS // 1000
+ started_at = time.monotonic()
+ frame_cursor = sent_frames = 0
+ lags: list[float] = []
+ for chunks, offset in enumerate(range(0, len(pcm), chunk_bytes), start=1):
+ end = min(offset + chunk_bytes, len(pcm))
+ end_ms = end * 1000 // bytes_per_second
+ ready: list[str] = []
+ while frame_cursor < len(frames) and end_ms >= (frame_cursor + 0.5) * 1000 / VIDEO_FPS:
+ if frames[frame_cursor]:
+ ready.append(frames[frame_cursor] or "")
+ frame_cursor += 1
+ sent_frames += len(ready)
+ event: dict[str, object] = {
+ "type": "input_audio_buffer.append",
+ "audio": base64.b64encode(pcm[offset:end]).decode("ascii"),
+ "input_audio_format": "pcm16",
+ "sample_rate_hz": PCM16_SAMPLE_RATE,
+ "duration_ms": (end - offset) * 1000 // bytes_per_second,
+ "audio_end_ms": end_ms,
+ }
+ if ready:
+ event["video_frames"] = ready
+ lags.append(max(0.0, time.monotonic() - started_at - offset / bytes_per_second))
+ await client.send(event)
+ await playback.acknowledge(client)
+ await asyncio.sleep(max(0.0, started_at + end_ms / 1000 - time.monotonic()))
+ return chunks if pcm else 0, sent_frames, sum(lags) / len(lags) if lags else 0.0, max(lags, default=0.0)
+
+
+def _response_status(event: dict[str, object]) -> str | None:
+ response = event.get("response")
+ if isinstance(response, dict) and response.get("status"):
+ return str(response["status"])
+ status = event.get("status")
+ return str(status) if status else None
+
+
+def response_ledger(
+ collector: RealtimeEventCollector,
+ *,
+ end: int | None = None,
+) -> tuple[set[str], set[str]]:
+ created: set[str] = set()
+ done: set[str] = set()
+ for event in collector.events[:end]:
+ event_type = event.get("type")
+ if event_type not in {"response.created", "response.done"}:
+ continue
+ response_id = collector.response_id(event)
+ if not response_id:
+ raise ValueError(f"{event_type} has no response_id")
+ if event_type == "response.created":
+ if response_id in created:
+ raise ValueError(f"duplicate response.created for {response_id}")
+ created.add(response_id)
+ continue
+ if response_id not in created:
+ raise ValueError(f"response.done without response.created for {response_id}")
+ if response_id in done:
+ raise ValueError(f"duplicate response.done for {response_id}")
+ response = event.get("response") if isinstance(event.get("response"), dict) else {}
+ details = event.get("status_details") if isinstance(event.get("status_details"), dict) else {}
+ nested = response.get("status_details") if isinstance(response.get("status_details"), dict) else {}
+ if "failed" in {event.get("status"), response.get("status"), details.get("type"), nested.get("type")}:
+ raise ValueError(f"response.done reports failure for {response_id}")
+ done.add(response_id)
+ return created, done
+
+
+def _is_model_listen_decision(event: dict[str, object]) -> bool:
+ response = event.get("response")
+ metadata = response.get("metadata") if isinstance(response, dict) else event.get("metadata")
+ return isinstance(metadata, dict) and metadata.get("buffering") is not True and metadata.get("model_listen") is True
+
+
+def _has_post_commit_decision(
+ collector: RealtimeEventCollector,
+ events: list[dict[str, object]],
+ *,
+ prior_response_ids: set[str],
+) -> bool:
+ pending_prior = set(prior_response_ids)
+ for event in events:
+ event_type = event.get("type")
+ if event_type == "response.done":
+ response_id = collector.response_id(event)
+ if response_id in pending_prior:
+ pending_prior.remove(response_id)
+ continue
+ if not pending_prior and _response_status(event) != "cancelled":
+ return True
+ elif event_type == "response.listen" and not pending_prior and _is_model_listen_decision(event):
+ return True
+ return False
+
+
+def _raise_if_session_terminated(
+ collector: RealtimeEventCollector,
+ from_index: int,
+ *,
+ explicit_close_from: int | None = None,
+) -> None:
+ errors = collector.errors()
+ if errors:
+ raise RuntimeError(str(errors[-1]))
+ for index, event in enumerate(collector.events[from_index:], start=from_index):
+ event_type = event.get("type")
+ if event_type not in {"session.expired", "session.closed"}:
+ continue
+ reason = event.get("reason")
+ nested = event.get("event")
+ if reason is None and isinstance(nested, dict):
+ reason = nested.get("reason")
+ expected = (
+ event_type == "session.closed"
+ and explicit_close_from is not None
+ and index >= explicit_close_from
+ and reason is None
+ )
+ if expected:
+ continue
+ detail = f": {reason}" if reason else ""
+ prefix = "Unexpected " if explicit_close_from is not None else ""
+ raise RuntimeError(f"{prefix}{event_type}{detail}")
+
+
+def _ensure_final_commit_tail(pcm: bytes, events: list[dict[str, object]]) -> bytes:
+ """Reserve one PCM sample so an exact model unit is flushed by commit.
+
+ A complete unit is otherwise emitted before commit, so its asynchronous
+ decision cannot be correlated with the accepted final input. Keeping an
+ almost-full unit buffered makes the server's final residual flush carry
+ that correlation; the server pads the missing sample back to a full unit.
+ """
+
+ period_ms = chunk_period_ms(events)
+ if len(pcm) >= PCM16_BYTES_PER_SAMPLE and not has_residual_model_unit(pcm, chunk_period_ms=period_ms):
+ return pcm[:-PCM16_BYTES_PER_SAMPLE]
+ return pcm
+
+
+async def wait_for_session_completion(
+ client: RealtimeDuplexClient,
+ playback: _Playback,
+ *,
+ commit_from: int,
+ session_from: int | None = None,
+ timeout_s: float,
+ settle_s: float,
+) -> int:
+ deadline = time.monotonic() + timeout_s
+ committed_index: int | None = None
+ prior_response_ids: set[str] = set()
+ last_event_count = len(client.events.events)
+ stable_since = time.monotonic()
+ while time.monotonic() < deadline:
+ client.raise_if_reader_stopped()
+ _raise_if_session_terminated(client.events, commit_from if session_from is None else session_from)
+ await playback.acknowledge(client)
+ if len(client.events.events) != last_event_count:
+ last_event_count = len(client.events.events)
+ stable_since = time.monotonic()
+ if committed_index is None:
+ committed_index = next(
+ (
+ index
+ for index in range(commit_from, len(client.events.events))
+ if client.events.events[index].get("type") == "input_audio_buffer.committed"
+ ),
+ None,
+ )
+ if committed_index is not None:
+ committed_event = client.events.events[committed_index]
+ committed = committed_event.get("event")
+ if isinstance(committed, dict) and committed.get("overlap_deferred") is True:
+ prior_created, prior_done = response_ledger(client.events, end=committed_index)
+ prior_response_ids = prior_created - prior_done
+ stable_since = time.monotonic()
+ if committed_index is not None:
+ created, done = response_ledger(client.events)
+ post_commit = client.events.events[committed_index + 1 :]
+ decision = _has_post_commit_decision(
+ client.events,
+ post_commit,
+ prior_response_ids=prior_response_ids,
+ )
+ if created == done and decision and time.monotonic() - stable_since >= settle_s:
+ return committed_index
+ await asyncio.sleep(0.05)
+ missing: set[str] = set()
+ with contextlib.suppress(ValueError):
+ created, done = response_ledger(client.events)
+ missing = created - done
+ raise TimeoutError(
+ "Timed out waiting for committed input and stable responses"
+ + (f"; unfinished response_ids={sorted(missing)}" if missing else "")
+ )
+
+
+def _output_dir(root: Path, case: OmniInteractCase) -> Path:
+ relative = case.video_rel.replace("\\", "/")
+ stem = Path(relative).with_suffix("").as_posix().replace("/", "__")
+ digest = hashlib.sha256(f"{case.subset}/{relative}".encode()).hexdigest()[:8]
+ return root / case.subset / f"{stem}--{digest}"
+
+
+@contextlib.contextmanager
+def omniinteract_output_lock(root: Path) -> Iterator[None]:
+ root.mkdir(parents=True, exist_ok=True)
+ with (root / ARTIFACT_LOCK_FILE).open("a+b") as lock_file:
+ fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+ try:
+ yield
+ finally:
+ fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+def _atomic_replace(path: Path, writer: Callable[[Path], None]) -> None:
+ descriptor, temporary_name = tempfile.mkstemp(
+ dir=path.parent,
+ prefix=f".{path.name}.",
+ suffix=".tmp",
+ )
+ os.close(descriptor)
+ temporary = Path(temporary_name)
+ try:
+ writer(temporary)
+ temporary.replace(path)
+ finally:
+ temporary.unlink(missing_ok=True)
+
+
+def _atomic_write_text(path: Path, value: str) -> None:
+ _atomic_replace(path, lambda temporary: temporary.write_text(value, encoding="utf-8"))
+
+
+def _atomic_write_json(path: Path, value: object) -> None:
+ _atomic_write_text(path, json.dumps(value, ensure_ascii=False, indent=2) + "\n")
+
+
+def _atomic_write_wav(path: Path, pcm: bytes | bytearray, rate: int) -> None:
+ _atomic_replace(
+ path,
+ lambda temporary: write_pcm16_wav(
+ temporary,
+ cast(bytes, pcm), # wave.writeframes accepts bytearray.
+ sample_rate_hz=rate,
+ ),
+ )
+
+
+def _clear_artifacts(directory: Path, names: Sequence[str]) -> None:
+ directory.mkdir(parents=True, exist_ok=True)
+ for name in names:
+ (directory / name).unlink(missing_ok=True)
+
+
+def _artifact_summary(case: OmniInteractCase, result: OmniInteractCaseResult) -> dict[str, Any]:
+ return {
+ **result.as_dict(),
+ "annotation": str(case.annotation_path),
+ "scene_type": "1QnA" if case.scene_type == "1qna" else case.scene_type,
+ }
+
+
+def _publish_success_artifacts(
+ directory: Path,
+ result: OmniInteractCaseResult,
+ context: _DeferredArtifactContext,
+ summary: dict[str, Any],
+) -> None:
+ pcm = bytearray(context.horizon_bytes)
+ for offset, chunk in context.spans:
+ pcm[offset : offset + len(chunk)] = chunk
+ try:
+ _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
+ _atomic_write_wav(directory / "output.wav", pcm, OUTPUT_SAMPLE_RATE)
+ _atomic_write_json(
+ directory / "wav_transcript.json",
+ {
+ "text": result.transcript,
+ "chunks": context.chunks,
+ "timestamp_semantics": "serialized playback queue time relative to input streaming start",
+ },
+ )
+ _atomic_write_json(directory / "events.json", context.events)
+ _atomic_write_json(directory / "result.json", result.as_dict())
+ _atomic_write_json(directory / ".done", summary)
+ except Exception:
+ _clear_artifacts(directory, SUCCESS_ARTIFACTS)
+ result.success = result.eligible_for_official_eval = False
+ if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
+ result.official_eval_ineligible_reasons.append("artifact_write_failed")
+ raise
+
+
+def _collect_output(
+ collector: RealtimeEventCollector,
+ playback: _Playback,
+ *,
+ stream_start: float,
+ video_duration_s: float,
+ require_response: bool,
+ retain_events: bool,
+ result: OmniInteractCaseResult,
+) -> _DeferredArtifactContext:
+ playback.ingest(collector)
+ created, done = response_ledger(collector)
+ if created != done:
+ raise ValueError(f"unfinished response_ids: {sorted(created - done)}")
+ done_events = [event for event in collector.events if event.get("type") == "response.done"]
+ cancelled = any(_response_status(event) == "cancelled" for event in done_events)
+ terminal_responses = {
+ response_id
+ for event in done_events
+ if _response_status(event) != "cancelled" and (response_id := collector.response_id(event))
+ }
+ horizon_s = math.ceil(video_duration_s)
+ horizon_bytes = horizon_s * OUTPUT_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE
+ spans: list[tuple[int, bytes]] = []
+ clipped_bytes = 0
+ response_times: dict[str, list[float]] = {}
+ responses_with_audio: set[str] = set()
+ for segment in playback.segments:
+ response_id = segment.response_id
+ if response_id not in created:
+ raise ValueError("response audio has no matching response.created")
+ start_s = max(0.0, segment.start_s - stream_start)
+ end_s = max(start_s, segment.end_s - stream_start)
+ offset = round(start_s * OUTPUT_SAMPLE_RATE) * PCM16_BYTES_PER_SAMPLE
+ writable = max(0, min(len(segment.pcm16), horizon_bytes - offset))
+ clipped_bytes += len(segment.pcm16) - writable
+ if writable:
+ clipped = segment.pcm16[:writable]
+ spans.append((offset, clipped))
+ if any(clipped):
+ responses_with_audio.add(response_id)
+ timing = response_times.setdefault(response_id, [start_s, end_s])
+ timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s)
+
+ text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"}
+ text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types}
+ if None in text_ids or not text_ids <= created:
+ raise ValueError("response transcript has no matching response.created")
+ chunks: list[dict[str, object]] = []
+ texts: list[str] = []
+ for response_id in collector.response_ids:
+ text = collector.response_text(response_id).strip()
+ if not text:
+ continue
+ response_timing = response_times.get(response_id)
+ if response_timing is None or response_timing[0] >= horizon_s:
+ continue
+ texts.append(text)
+ chunks.append(
+ {
+ "response_id": response_id,
+ "text": text,
+ "timestamp": [round(response_timing[0], 6), round(min(response_timing[1], horizon_s), 6)],
+ }
+ )
+ complete_outputs = terminal_responses & responses_with_audio & {str(chunk["response_id"]) for chunk in chunks}
+ if require_response and not complete_outputs:
+ raise ValueError("OmniInteract E2E requires a response with audio and transcript")
+ result.audio_bytes, result.audio_clipped_bytes = (
+ sum(len(segment.pcm16) for segment in playback.segments),
+ clipped_bytes,
+ )
+ result.transcript, result.responses, result.success = " ".join(texts).strip(), len(created), True
+ result.artifact_warnings = list(playback.warnings)
+ result.official_eval_ineligible_reasons = []
+ if clipped_bytes:
+ result.official_eval_ineligible_reasons.append("audio_clipped")
+ if cancelled:
+ result.official_eval_ineligible_reasons.append("cancelled_response")
+ result.eligible_for_official_eval = not result.official_eval_ineligible_reasons
+ audio_sizes = {segment.event_index: len(segment.pcm16) for segment in playback.segments}
+ events = (
+ [
+ {
+ **{key: value for key, value in event.items() if key not in {"delta", "audio"}},
+ "audio_bytes": audio_sizes[index],
+ }
+ if event.get("type") == "response.audio.delta"
+ else dict(event)
+ for index, event in enumerate(collector.events)
+ ]
+ if retain_events
+ else []
+ )
+ return _DeferredArtifactContext(horizon_bytes, tuple(spans), chunks, events)
+
+
+def prepare_success_artifacts(
+ root: Path,
+ case: OmniInteractCase,
+ collector: RealtimeEventCollector,
+ *,
+ playback: _Playback | None = None,
+ stream_start: float,
+ video_duration_s: float,
+ require_response: bool,
+ result: OmniInteractCaseResult,
+ capture_artifacts: bool = True,
+) -> dict[str, Any]:
+ directory = _output_dir(root, case)
+ context = _collect_output(
+ collector,
+ playback or _Playback(),
+ stream_start=stream_start,
+ video_duration_s=video_duration_s,
+ require_response=require_response,
+ retain_events=capture_artifacts,
+ result=result,
+ )
+ result.output_dir = str(directory.resolve())
+ summary = _artifact_summary(case, result)
+ if capture_artifacts:
+ result._artifact_context = context
+ return summary
+
+
+def write_failure_artifacts(root: Path, case: OmniInteractCase, result: OmniInteractCaseResult) -> None:
+ directory = _output_dir(root, case)
+ result.output_dir = str(directory.resolve())
+ result.success = result.eligible_for_official_eval = False
+ _clear_artifacts(directory, SUCCESS_ARTIFACTS)
+ _atomic_write_json(directory / ".failed.json", _artifact_summary(case, result))
+
+
+def publish_deferred_case_artifacts(
+ root: Path,
+ case: OmniInteractCase,
+ result: OmniInteractCaseResult,
+) -> None:
+ context = result._artifact_context
+ result._artifact_context = None
+ if not result.success:
+ write_failure_artifacts(root, case, result)
+ return
+ if context is None:
+ raise RuntimeError("OmniInteract benchmark output lost its deferred artifact context")
+ directory = _output_dir(root, case)
+ _publish_success_artifacts(directory, result, context, _artifact_summary(case, result))
+
+
+def clear_case_artifacts(root: Path, case: OmniInteractCase) -> None:
+ directory = _output_dir(root, case)
+ _clear_artifacts(directory, (*SUCCESS_ARTIFACTS, ".failed.json"))
+
+
+def clear_batch_artifacts(root: Path) -> None:
+ _clear_artifacts(root, BATCH_ARTIFACTS)
+
+
+def write_batch_artifacts(
+ root: Path,
+ cases: list[OmniInteractCase],
+ results: list[OmniInteractCaseResult],
+) -> None:
+ root.mkdir(parents=True, exist_ok=True)
+ rows = [
+ case_manifest(case, _output_dir(root, case))
+ for case, result in zip(cases, results, strict=True)
+ if result.success and result.eligible_for_official_eval
+ ]
+ ineligible = [result for result in results if result.success and not result.eligible_for_official_eval]
+ if ineligible:
+ reasons = Counter(
+ reason for result in ineligible for reason in (result.official_eval_ineligible_reasons or ["unspecified"])
+ )
+ logger.warning(
+ "%d successful OmniInteract cases were excluded from official evaluation: %s",
+ len(ineligible),
+ ", ".join(f"{reason}={count}" for reason, count in sorted(reasons.items())),
+ )
+ _atomic_write_json(root / "batch_summary.json", benchmark_summary(results))
+ _atomic_write_text(
+ root / "official_eval_manifest.jsonl",
+ "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
+ )
+
+
+def _websocket_url(config: OmniInteractBenchmarkConfig, session_id: str) -> str:
+ endpoint = (
+ config.endpoint
+ if urlsplit(config.endpoint).scheme
+ else urljoin(config.base_url.rstrip("/") + "/", config.endpoint.lstrip("/"))
+ )
+ return build_realtime_url(
+ endpoint,
+ config.model,
+ autostart=False,
+ native_duplex=True,
+ session_id=session_id,
+ )
+
+
+def _populate_response_metrics(
+ result: OmniInteractCaseResult,
+ collector: RealtimeEventCollector,
+ *,
+ stream_start: float,
+) -> None:
+ measurement_origin = {
+ "ttft": "response.created client receive to first non-empty text delta",
+ "ttfp": "response.created client receive to first audio packet",
+ "rtf": "response.created client receive to last audio packet divided by emitted audio duration",
+ }
+ request_metrics: list[dict[str, object]] = []
+ output_tokens = 0
+ for request_index, response_id in enumerate(collector.response_ids):
+ timing = collector.timing_summary(
+ after_s=stream_start,
+ input_committed_at_s=None,
+ response_id=response_id,
+ measurement_origin=measurement_origin,
+ )
+ raw_metric = timing.get("request_metrics")
+ stage0 = timing.get("stage0_tokens")
+ metric = {
+ "session_id": result.session_id,
+ "request_index": request_index,
+ "response_id": response_id,
+ **(raw_metric if isinstance(raw_metric, dict) else {}),
+ }
+ if isinstance(stage0, dict):
+ metric["stage0_tokens"] = dict(stage0)
+ output_tokens += int(stage0.get("output_token_count") or 0)
+ if isinstance(raw_metric, dict) or isinstance(stage0, dict):
+ request_metrics.append(metric)
+ result.output_tokens = output_tokens
+ result.duplex_request_metrics = request_metrics
+ result.duplex_session_metrics = summarize_session_request_metrics(
+ request_metrics,
+ session_id=result.session_id,
+ )
+
+
+async def run_omniinteract_case(
+ case: OmniInteractCase,
+ config: OmniInteractBenchmarkConfig,
+ *,
+ request_index: int | str,
+ capture_artifacts: bool = True,
+ prepared_input: OmniInteractPreparedInput | None = None,
+) -> OmniInteractCaseResult:
+ if not config.ref_audio:
+ raise ValueError("ref_audio is required for MiniCPM-o native-duplex audio output")
+ for name, value in (
+ ("timeout_s", config.timeout_s),
+ ("media_timeout_s", config.media_timeout_s),
+ ("max_video_duration_s", config.max_video_duration_s),
+ ):
+ if not math.isfinite(value) or value <= 0:
+ raise ValueError(f"{name} must be finite and positive")
+ session_id = f"omniinteract:{case.subset}:{request_index}:{time.monotonic_ns()}"
+ result = OmniInteractCaseResult(
+ case.subset,
+ str(case.video_path),
+ str(_output_dir(config.output_root, case).resolve()),
+ session_id=session_id,
+ )
+ started_at = time.monotonic()
+ try:
+ if prepared_input is None:
+ reference_audio = reference_audio_data_url(config.ref_audio)
+ duration, pcm, frames = await asyncio.to_thread(
+ prepare_media,
+ case.video_path,
+ VIDEO_FPS,
+ timeout_s=config.media_timeout_s,
+ max_duration_s=config.max_video_duration_s,
+ )
+ else:
+ reference_audio = prepared_input.ref_audio_data_url
+ duration = prepared_input.duration_s
+ pcm = prepared_input.pcm16
+ frames = prepared_input.video_frames
+ if not any(frames):
+ raise ValueError(f"No video frames were decoded from {case.video_path}")
+ async with RealtimeDuplexClient(
+ _websocket_url(config, session_id), additional_headers=config.extra_headers
+ ) as client:
+ session_from = len(client.events.events)
+ await client.configure(
+ config.model,
+ ref_audio=reference_audio,
+ session_id=session_id,
+ extra_body=config.extra_body,
+ idle_timeout_s=config.timeout_s,
+ timeout_s=min(config.timeout_s, 20.0),
+ )
+ pcm = _ensure_final_commit_tail(pcm, client.events.events)
+ playback = _Playback()
+ stream_start = time.monotonic()
+ try:
+ input_duration_s = len(pcm) / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+ upload_timeout_s = config.timeout_s + input_duration_s
+ try:
+ chunks, frame_count, mean_lag, max_lag = await asyncio.wait_for(
+ stream_inputs(client, pcm, frames, playback), timeout=upload_timeout_s
+ )
+ except asyncio.TimeoutError as exc:
+ raise TimeoutError(f"Realtime upload timed out after {upload_timeout_s:g}s") from exc
+ commit_from = len(client.events.events)
+ await client.commit()
+ await wait_for_session_completion(
+ client,
+ playback,
+ commit_from=commit_from,
+ session_from=session_from,
+ timeout_s=config.timeout_s,
+ settle_s=_COMPLETION_SETTLE_S,
+ )
+ await playback.acknowledge(client, playback.end_s)
+ _raise_if_session_terminated(client.events, session_from)
+ close_from = len(client.events.events)
+ await client.close_session(timeout_s=min(config.timeout_s, 20.0))
+ except Exception:
+ with contextlib.suppress(Exception):
+ await client.close_session(timeout_s=min(config.timeout_s, 20.0))
+ raise
+ errors = client.events.errors()
+ if errors:
+ raise RuntimeError(str(errors[-1]))
+ _raise_if_session_terminated(client.events, session_from, explicit_close_from=close_from)
+ result.latency_s = time.monotonic() - started_at
+ result.input_audio_chunks, result.input_video_frames = chunks, frame_count
+ result.pacing_mean_lag_s, result.pacing_max_lag_s = mean_lag, max_lag
+ _populate_response_metrics(result, client.events, stream_start=stream_start)
+ await asyncio.to_thread(
+ prepare_success_artifacts,
+ config.output_root,
+ case,
+ client.events,
+ playback=playback,
+ stream_start=stream_start,
+ video_duration_s=duration,
+ require_response=config.require_response,
+ result=result,
+ capture_artifacts=capture_artifacts,
+ )
+ except Exception as exc:
+ result.success = result.eligible_for_official_eval = False
+ if "case_failed" not in result.official_eval_ineligible_reasons:
+ result.official_eval_ineligible_reasons.append("case_failed")
+ result.error = str(exc)
+ result.latency_s = max(0.0, time.monotonic() - started_at)
+ return result
diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py
index 3f83a1b38..fd52c197b 100644
--- a/vllm_omni/benchmarks/patch/patch.py
+++ b/vllm_omni/benchmarks/patch/patch.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
import asyncio
import contextlib
import io
@@ -16,7 +19,6 @@ from dataclasses import dataclass, replace
from datetime import datetime
from pathlib import Path
from typing import Any, Literal
-from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import aiohttp
import numpy as np
@@ -31,6 +33,7 @@ from vllm.benchmarks.lib.endpoint_request_func import (
RequestFuncOutput,
StreamedResponseHandler,
_get_chat_content,
+ _get_headers,
_update_headers_common,
_update_payload_common,
_validate_api_url,
@@ -46,6 +49,13 @@ from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
daily_omni_local_videos_dir,
resolve_daily_omni_local_root,
)
+from vllm_omni.benchmarks.data_modules.omniinteract_dataset import (
+ DEFAULT_OMNIINTERACT_REPO,
+ OmniInteractDataset,
+ OmniInteractPreparedInput,
+ OmniInteractSampleRequest,
+ OmniInteractSessionOptions,
+)
from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT,
@@ -56,8 +66,27 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
)
from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
+from vllm_omni.benchmarks.omniinteract import (
+ VIDEO_FPS,
+ OmniInteractBenchmarkConfig,
+ OmniInteractCaseResult,
+ clear_batch_artifacts,
+ clear_case_artifacts,
+ prepare_media,
+ publish_deferred_case_artifacts,
+ run_omniinteract_case,
+ write_failure_artifacts,
+)
+from vllm_omni.benchmarks.omniinteract import (
+ benchmark_summary as omniinteract_benchmark_summary,
+)
+from vllm_omni.benchmarks.omniinteract import (
+ write_batch_artifacts as write_omniinteract_batch_artifacts,
+)
from vllm_omni.experimental.fullduplex.client import (
RealtimeDuplexClient,
+ build_realtime_url,
+ reference_audio_data_url,
summarize_session_request_metrics,
wait_for,
)
@@ -90,6 +119,17 @@ def should_request_stage_metrics(args: Any) -> bool:
if backend in _IMAGE_STAGE_METRICS_BACKENDS:
return True
+ percentile_metrics = getattr(args, "percentile_metrics", ()) or ()
+ if isinstance(percentile_metrics, str):
+ percentile_metrics = percentile_metrics.split(",")
+ selected_metrics = {str(metric).strip().lower() for metric in percentile_metrics}
+ if backend == "openai-chat-omni" and selected_metrics.intersection({"tpot", "itl"}):
+ # Client receive timestamps normally provide these metrics. Under
+ # enough load, however, every text token can be coalesced into one
+ # event-loop read, leaving no positive interval. Request the engine's
+ # Stage 0 timings so that path still has an authoritative measurement.
+ return True
+
extra_body = getattr(args, "extra_body", None) or {}
modalities = extra_body.get("modalities") if isinstance(extra_body, dict) else None
return backend == "openai-chat-omni" and "image" in (modalities or [])
@@ -221,6 +261,91 @@ def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFu
rfi.extra_body = base
+def _attach_omniinteract_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None:
+ if not isinstance(sample, OmniInteractSampleRequest):
+ return
+ setattr(rfi, "omniinteract_case", sample.omniinteract_case)
+ setattr(rfi, "omniinteract_options", sample.omniinteract_options)
+ setattr(rfi, "omniinteract_prepared_input", sample.omniinteract_prepared_input)
+
+
+def _append_error(existing: str, message: str) -> str:
+ return f"{existing}\n{message}" if existing else message
+
+
+def _finalize_omniinteract_batch(
+ input_requests: list[SampleRequest],
+ outputs: list[RequestFuncOutput],
+) -> dict[str, object] | None:
+ rows = [
+ (sample, output)
+ for sample, output in zip(input_requests, outputs, strict=True)
+ if isinstance(sample, OmniInteractSampleRequest)
+ ]
+ if not rows:
+ return None
+ options = rows[0][0].omniinteract_options
+ if not isinstance(options, OmniInteractSessionOptions):
+ raise RuntimeError("OmniInteract benchmark output lost its artifact options")
+ cases, results = [], []
+ artifact_errors: list[str] = []
+ for sample, output in rows:
+ case = sample.omniinteract_case
+ result = getattr(output, "omniinteract_case_result", None)
+ if case is None or not isinstance(result, OmniInteractCaseResult):
+ raise RuntimeError("OmniInteract benchmark output lost its dataset identity")
+ cases.append(case)
+ results.append(result)
+ try:
+ publish_deferred_case_artifacts(options.output_root, case, result)
+ except Exception as exc: # noqa: BLE001 - artifact backends raise heterogeneous errors
+ result.success = result.eligible_for_official_eval = False
+ if "artifact_write_failed" not in result.official_eval_ineligible_reasons:
+ result.official_eval_ineligible_reasons.append("artifact_write_failed")
+ artifact_error = f"Artifact publication failed: {exc}"
+ result.error = _append_error(result.error, artifact_error)
+ output.error = _append_error(output.error, artifact_error)
+ message = f"{case.subset}/{case.video_rel}: {artifact_error}"
+ artifact_errors.append(message)
+ logger.exception(message)
+ try:
+ write_failure_artifacts(options.output_root, case, result)
+ except Exception as recovery_exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure
+ message = f"{case.subset}/{case.video_rel}: failure artifact publication failed: {recovery_exc}"
+ artifact_errors.append(message)
+ logger.exception(message)
+ try:
+ write_omniinteract_batch_artifacts(options.output_root, cases, results)
+ except Exception as exc: # noqa: BLE001 - preserve benchmark metrics after I/O failure
+ message = f"Batch artifact publication failed: {exc}"
+ artifact_errors.append(message)
+ logger.exception(message)
+ try:
+ clear_batch_artifacts(options.output_root)
+ except Exception as cleanup_exc: # noqa: BLE001 - best-effort removal of partial batch artifacts
+ message = f"Partial batch artifact cleanup failed: {cleanup_exc}"
+ artifact_errors.append(message)
+ logger.exception(message)
+ summary = omniinteract_benchmark_summary(results)
+ compact_summary = {key: value for key, value in summary.items() if key != "results"}
+ compact_summary["artifacts_complete"] = not artifact_errors
+ if artifact_errors:
+ compact_summary["artifact_errors"] = artifact_errors
+ return compact_summary
+
+
+def _prepare_omniinteract_batch(input_requests: list[SampleRequest]) -> None:
+ roots: set[Path] = set()
+ for sample in input_requests:
+ if not isinstance(sample, OmniInteractSampleRequest):
+ continue
+ root = sample.omniinteract_options.output_root.resolve()
+ roots.add(root)
+ clear_case_artifacts(sample.omniinteract_options.output_root, sample.omniinteract_case)
+ for root in roots:
+ clear_batch_artifacts(root)
+
+
def _daily_omni_repo_from_args(args) -> str | None:
"""Resolve HuggingFace repo id for Daily-Omni from CLI args.
@@ -250,6 +375,7 @@ def get_samples(args, tokenizer):
"ttsd",
"sound-effect",
)
+ is_omniinteract = args.dataset_name == "omniinteract"
# Check if we need to handle omni-related backends/datasets
is_omni_backend = args.backend in [
@@ -258,12 +384,68 @@ def get_samples(args, tokenizer):
"openai-realtime-duplex",
"daily-omni",
]
- is_omni_dataset = is_daily_omni or is_seed_tts or args.dataset_name == "random-mm"
+ is_omni_dataset = is_daily_omni or is_seed_tts or is_omniinteract or args.dataset_name == "random-mm"
if not is_omni_backend and not is_omni_dataset:
# Not an omni-related request, delegate to original implementation
return get_samples_old(args, tokenizer)
+ if is_omniinteract:
+ dataset_path = getattr(args, "dataset_path", None)
+ data_root = None
+ dataset_repo = DEFAULT_OMNIINTERACT_REPO
+ if dataset_path:
+ candidate = Path(dataset_path).expanduser()
+ if candidate.exists() or candidate.is_absolute():
+ data_root = str(candidate)
+ else:
+ dataset_repo = str(dataset_path)
+ dataset = OmniInteractDataset(
+ data_root=data_root,
+ dataset_repo=dataset_repo,
+ subsets=tuple(getattr(args, "omniinteract_subsets")),
+ random_seed=args.seed,
+ disable_shuffle=getattr(args, "disable_shuffle", False),
+ )
+ options = OmniInteractSessionOptions(
+ output_root=Path(getattr(args, "omniinteract_output_dir")),
+ timeout_s=float(getattr(args, "omniinteract_timeout_s")),
+ media_timeout_s=float(getattr(args, "omniinteract_media_timeout_s")),
+ ref_audio=str(getattr(args, "omniinteract_ref_audio")),
+ require_response=bool(getattr(args, "omniinteract_require_response")),
+ max_video_duration_s=float(getattr(args, "omniinteract_max_video_duration_s")),
+ )
+ requests = dataset.sample(
+ tokenizer,
+ args.num_prompts,
+ request_id_prefix=args.request_id_prefix,
+ options=options,
+ )
+ if not requests:
+ raise ValueError("No OmniInteract sessions were selected")
+ encoded_ref_audio = reference_audio_data_url(options.ref_audio)
+ assert encoded_ref_audio is not None
+ for request in requests:
+ case = request.omniinteract_case
+ assert isinstance(request, OmniInteractSampleRequest) and case is not None
+ duration, pcm, frames = prepare_media(
+ case.video_path,
+ VIDEO_FPS,
+ timeout_s=options.media_timeout_s,
+ max_duration_s=options.max_video_duration_s,
+ )
+ if not any(frames):
+ raise ValueError(f"No video frames were decoded from {case.video_path}")
+ request.omniinteract_prepared_input = OmniInteractPreparedInput(
+ duration_s=duration,
+ pcm16=pcm,
+ video_frames=tuple(frames),
+ ref_audio_data_url=encoded_ref_audio,
+ )
+ # Replace OmniInteract's ``0`` (all) with the measured request count.
+ args.num_prompts = len(requests)
+ return requests
+
# Handle Daily-Omni dataset
if is_daily_omni:
# Support:
@@ -476,6 +658,7 @@ class MixRequestFuncOutput(RequestFuncOutput):
image_pixels: int = 0
denoise_step_latency_ms: float = 0.0
text_latency: float = 0.0
+ tpot_measured: bool = True
#: Worst-case streaming-audio underrun (wall-clock seconds the player
#: would have been starved). Populated by the audio-speech backend; ``0.0``
#: for backends that do not run continuity analysis.
@@ -682,6 +865,33 @@ def _update_output_stage_metrics_from_payload(
output.stage_metrics.update(stage_snapshot)
+def _apply_chat_stage0_token_timings(output: MixRequestFuncOutput) -> bool:
+ """Apply native Stage 0 timings from a chat response snapshot."""
+ stage_metrics = output.stage_metrics
+ if not isinstance(stage_metrics, dict):
+ return False
+ stage0 = stage_metrics.get("0")
+ if not isinstance(stage0, dict):
+ return False
+
+ output_tokens = coerce_positive_int_scalar(stage0.get(defs.NUM_TOKENS_OUT))
+ expected_output_tokens = coerce_positive_int_scalar(output.output_tokens)
+ if output_tokens is None or expected_output_tokens is None or output_tokens != expected_output_tokens:
+ return False
+
+ return _apply_stage0_token_timings(
+ output,
+ [
+ {
+ "output_token_count": output_tokens,
+ "itls_ms": stage0.get(defs.VLLM_ITLS_MS),
+ "tpot_ms": stage0.get(defs.VLLM_TPOT_MS),
+ }
+ ],
+ expected_output_tokens=expected_output_tokens,
+ )
+
+
def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]:
if not isinstance(metrics, dict):
return 0, 0.0, 0, 0.0
@@ -949,6 +1159,14 @@ async def async_request_openai_chat_omni_completions(
output.latency = timestamp - st
output.generated_text = generated_text
+ if output.output_tokens > 1 and not any(
+ isinstance(value, int | float)
+ and not isinstance(value, bool)
+ and np.isfinite(value)
+ and value > 0
+ for value in output.itl
+ ):
+ _apply_chat_stage0_token_timings(output)
if output.itl:
# Align text_latency with ITL so TPOT formula and
# mean(ITL) are consistent. Do NOT infer output_tokens
@@ -1313,11 +1531,158 @@ async def async_request_openai_audio_speech(
def _realtime_websocket_url(api_url: str) -> str:
- parts = urlsplit(api_url)
- scheme = "wss" if parts.scheme == "https" else "ws"
- query = dict(parse_qsl(parts.query, keep_blank_values=True))
- query.setdefault("duplex", "1")
- return urlunsplit((scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
+ return build_realtime_url(api_url, None, native_duplex=None)
+
+
+def _nonnegative_number(value: object) -> bool:
+ return isinstance(value, int | float) and not isinstance(value, bool) and np.isfinite(value) and value >= 0
+
+
+def _apply_stage0_token_timings(
+ output: MixRequestFuncOutput,
+ raw_stages: list[object],
+ *,
+ expected_output_tokens: int,
+) -> bool:
+ """Populate aggregate token timing from complete per-response Stage 0 metrics."""
+ output.itl = []
+ output.text_latency = output.ttft
+ output.tpot_measured = False
+
+ stages: list[tuple[int, dict[str, object]]] = []
+ for stage in raw_stages:
+ if not isinstance(stage, dict):
+ return False
+ token_count = stage.get("output_token_count")
+ if not isinstance(token_count, int) or isinstance(token_count, bool) or token_count < 0:
+ return False
+ stages.append((token_count, stage))
+
+ if not stages or sum(count for count, _ in stages) != expected_output_tokens:
+ return False
+
+ timed = [(count - 1, stage) for count, stage in stages if count > 1]
+ if not timed:
+ return False
+
+ measured_itls: list[float] = []
+ exact_itls = True
+ for intervals, stage in timed:
+ values = stage.get("itls_ms")
+ if (
+ not isinstance(values, list)
+ or len(values) != intervals
+ or not all(_nonnegative_number(value) for value in values)
+ ):
+ exact_itls = False
+ break
+ measured_itls.extend(float(value) / 1000.0 for value in values if isinstance(value, int | float))
+ if exact_itls:
+ if sum(measured_itls) > 0:
+ output.itl = measured_itls
+ output.text_latency = output.ttft + sum(output.itl)
+ output.tpot_measured = True
+ return True
+
+ weighted_tpot_ms = 0.0
+ for intervals, stage in timed:
+ tpot_ms = stage.get("tpot_ms")
+ if not _nonnegative_number(tpot_ms) or not isinstance(tpot_ms, int | float) or tpot_ms <= 0:
+ break
+ weighted_tpot_ms += float(tpot_ms) * intervals
+ else:
+ decode_intervals = sum(intervals for intervals, _ in timed)
+ weighted_tpot = weighted_tpot_ms / decode_intervals / 1000.0
+ # ``calculate_metrics`` consumes a request-shaped latency. Preserve
+ # the weighted per-response TPOT when projecting it onto the combined
+ # output token count; gaps and TTFTs between turns are intentionally
+ # excluded.
+ output.text_latency = output.ttft + weighted_tpot * (expected_output_tokens - 1)
+ output.tpot_measured = True
+ return True
+
+ return False
+
+
+async def _async_request_omniinteract(
+ request_func_input: RequestFuncInput,
+ *,
+ pbar: tqdm | None,
+) -> MixRequestFuncOutput:
+ case = getattr(request_func_input, "omniinteract_case", None)
+ options = getattr(request_func_input, "omniinteract_options", None)
+ output = MixRequestFuncOutput()
+ output.prompt_len = request_func_input.prompt_len
+ output.start_time = time.perf_counter()
+ try:
+ if case is None or not isinstance(options, OmniInteractSessionOptions):
+ raise ValueError("OmniInteract RequestFuncInput is missing its dataset session layout")
+ prepared_input = getattr(request_func_input, "omniinteract_prepared_input", None)
+ if not isinstance(prepared_input, OmniInteractPreparedInput):
+ raise ValueError("OmniInteract media must be prepared before benchmark request timing")
+ headers = _get_headers()
+ _update_headers_common(headers, request_func_input)
+ config = OmniInteractBenchmarkConfig(
+ endpoint=request_func_input.api_url,
+ model=request_func_input.model_name or request_func_input.model,
+ output_root=options.output_root,
+ timeout_s=options.timeout_s,
+ media_timeout_s=options.media_timeout_s,
+ max_video_duration_s=options.max_video_duration_s,
+ ref_audio=options.ref_audio,
+ require_response=options.require_response,
+ extra_headers=headers or None,
+ extra_body=dict(request_func_input.extra_body or {}) or None,
+ )
+ case_result = await run_omniinteract_case(
+ case,
+ config,
+ request_index=request_func_input.request_id or uuid.uuid4().hex,
+ capture_artifacts=request_func_input.request_id is not None,
+ prepared_input=prepared_input,
+ )
+ output.latency = case_result.latency_s
+ output.generated_text = case_result.transcript
+ output.output_tokens = case_result.output_tokens
+ output.audio_duration = case_result.audio_bytes / (24_000 * 2)
+ output.audio_frames = case_result.audio_bytes // 2
+ session_metrics = case_result.duplex_session_metrics
+ output.ttft = float(session_metrics.get("mean_ttft_ms") or 0.0) / 1000.0
+ output.audio_ttfp = float(session_metrics.get("mean_ttfp_ms") or 0.0) / 1000.0
+ output.audio_rtf = float(session_metrics.get("mean_rtf") or 0.0)
+ token_timing_measured = _apply_stage0_token_timings(
+ output,
+ [request_metric.get("stage0_tokens") for request_metric in case_result.duplex_request_metrics],
+ expected_output_tokens=output.output_tokens,
+ )
+ if not token_timing_measured:
+ if output.output_tokens > 1 or (output.output_tokens == 0 and output.generated_text):
+ logger.warning(
+ "OmniInteract session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
+ case_result.session_id,
+ )
+ output.duplex_request_metrics = case_result.duplex_request_metrics
+ output.duplex_session_metrics = session_metrics
+ output.success = case_result.success
+ output.error = case_result.error
+ setattr(output, "omniinteract_case", case)
+ setattr(output, "omniinteract_case_result", case_result)
+ except Exception:
+ output.success = False
+ output.error = traceback.format_exc()
+ logger.error("OmniInteract Realtime request failed: %s", output.error)
+ if case is not None:
+ case_result = OmniInteractCaseResult(
+ subset=case.subset,
+ video=str(case.video_path),
+ output_dir="",
+ error=output.error,
+ )
+ setattr(output, "omniinteract_case", case)
+ setattr(output, "omniinteract_case_result", case_result)
+ if pbar:
+ pbar.update(1)
+ return output
async def async_request_openai_realtime_duplex(
@@ -1325,14 +1690,18 @@ async def async_request_openai_realtime_duplex(
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> MixRequestFuncOutput:
- """Run one or more Seed-TTS turns through one explicit Realtime TTS session."""
del session
+ if getattr(request_func_input, "omniinteract_case", None) is not None:
+ return await _async_request_omniinteract(request_func_input, pbar=pbar)
output = MixRequestFuncOutput()
output.prompt_len = request_func_input.prompt_len
output.start_time = time.perf_counter()
speech_extra = getattr(request_func_input, "seed_tts_speech_extra", None)
if not isinstance(speech_extra, dict):
speech_extra = {}
+ else:
+ speech_extra = dict(speech_extra)
+ speech_extra["return_stage_metrics"] = True
configured_turns = getattr(request_func_input, "seed_tts_turns", ())
turn_prompts = [
(
@@ -1387,8 +1756,10 @@ async def async_request_openai_realtime_duplex(
)
await client.send({"type": "response.create"})
await wait_for(
- lambda: client.events.count("response.done") > done_before
- or len(client.events.errors()) > errors_before,
+ lambda: (
+ client.events.count("response.done") > done_before
+ or len(client.events.errors()) > errors_before
+ ),
timeout_s=180.0,
label=f"Seed-TTS Realtime TTS turn {request_index} response.done",
)
@@ -1433,19 +1804,7 @@ async def async_request_openai_realtime_duplex(
channels=1,
)
)
- turn_transcripts.append(
- "".join(
- str(event.get("delta") or "")
- for event in client.events.events
- if client.events.response_id(event) == response_id
- and event.get("type")
- in {
- "response.audio_transcript.delta",
- "response.output_text.delta",
- "response.text.delta",
- }
- )
- )
+ turn_transcripts.append(client.events.response_text(response_id))
await client.acknowledge_playback()
request_finished_at = time.perf_counter()
session_metrics = summarize_session_request_metrics(
@@ -1473,6 +1832,16 @@ async def async_request_openai_realtime_duplex(
for timing in turn_timings
if isinstance((stage0 := timing.get("stage0_tokens")), dict)
)
+ token_timing_measured = _apply_stage0_token_timings(
+ output,
+ [timing.get("stage0_tokens") for timing in turn_timings],
+ expected_output_tokens=output.output_tokens,
+ )
+ if not token_timing_measured and output.output_tokens > 1:
+ logger.warning(
+ "Realtime TTS session %s omitted complete engine token timing; standard TPOT/ITL are unavailable",
+ session_id,
+ )
output.success = True
except Exception:
output.success = False
@@ -1520,6 +1889,7 @@ from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, ge
from vllm_omni.benchmarks.metrics.metrics import (
MultiModalsBenchmarkMetrics,
calculate_metrics,
+ has_metric_samples,
)
# ruff: noqa: E402
@@ -1626,6 +1996,7 @@ async def benchmark(
)
_attach_daily_omni_to_request_func_input(input_requests[0], test_input)
_attach_seed_tts_to_request_func_input(input_requests[0], test_input)
+ _attach_omniinteract_to_request_func_input(input_requests[0], test_input)
if ready_check_timeout_sec > 0:
test_output = await wait_for_endpoint(
@@ -1741,6 +2112,7 @@ async def benchmark(
print(f"Probe request rate: {probe_request_rate} req/s")
probe_task = asyncio.create_task(probe_loop())
+ _prepare_omniinteract_batch(input_requests)
benchmark_start_time = time.perf_counter()
tasks: list[asyncio.Task] = []
@@ -1801,6 +2173,7 @@ async def benchmark(
)
_attach_daily_omni_to_request_func_input(request, request_func_input)
_attach_seed_tts_to_request_func_input(request, request_func_input)
+ _attach_omniinteract_to_request_func_input(request, request_func_input)
tasks.append(
asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar))
)
@@ -1815,6 +2188,8 @@ async def benchmark(
benchmark_duration = time.perf_counter() - benchmark_start_time
+ omniinteract_summary = _finalize_omniinteract_batch(input_requests, outputs)
+
if task_type == TaskType.GENERATION:
metrics, actual_output_lens = calculate_metrics(
input_requests=input_requests,
@@ -1839,6 +2214,13 @@ async def benchmark(
actual_output_lens = 0
if isinstance(metrics, MultiModalsBenchmarkMetrics):
+
+ def measured_ttft(output: RequestFuncOutput) -> float | None:
+ session_metrics = getattr(output, "duplex_session_metrics", None)
+ if isinstance(session_metrics, dict) and session_metrics.get("mean_ttft_ms") is None:
+ return None
+ return output.ttft
+
result = {
"duration": benchmark_duration,
"completed": metrics.completed,
@@ -1859,7 +2241,7 @@ async def benchmark(
"input_lens": [output.prompt_len for output in outputs],
"start_times": [output.start_time for output in outputs],
"output_lens": actual_output_lens,
- "ttfts": [output.ttft for output in outputs],
+ "ttfts": [measured_ttft(output) for output in outputs],
"itls": [output.itl for output in outputs],
"generated_texts": [output.generated_text for output in outputs],
"errors": [output.error for output in outputs],
@@ -1893,6 +2275,8 @@ async def benchmark(
]
if duplex_session_metrics:
result["duplex_session_metrics"] = duplex_session_metrics
+ if omniinteract_summary is not None:
+ result["omniinteract"] = omniinteract_summary
from vllm_omni.benchmarks.data_modules.daily_omni_eval import (
compute_daily_omni_accuracy_metrics,
@@ -1949,6 +2333,8 @@ async def benchmark(
# metric.
if metric_attribute_name not in result_percentile_metrics:
return
+ if not has_metric_samples(metrics, metric_attribute_name):
+ return
# No text tokens generated (e.g. pure TTS speech endpoint): per-token
# latency metrics (ttft/tpot/itl) are undefined, so skip them.
is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio"))
diff --git a/vllm_omni/benchmarks/serve.py b/vllm_omni/benchmarks/serve.py
index 8e3b19ea9..2c4f8d130 100644
--- a/vllm_omni/benchmarks/serve.py
+++ b/vllm_omni/benchmarks/serve.py
@@ -1,6 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
import argparse
import asyncio
+import contextlib
import os
+from pathlib import Path
from typing import Any
from vllm.benchmarks.serve import main_async
@@ -9,6 +16,7 @@ from vllm.benchmarks.serve import main_async
# This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
# Must be imported before any vllm.benchmarks module usage
import vllm_omni.benchmarks.patch.patch # noqa: F401
+from vllm_omni.benchmarks.omniinteract import omniinteract_output_lock
from vllm_omni.benchmarks.patch.patch import (
maybe_enable_stage_metrics,
set_print_stage,
@@ -28,4 +36,10 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
getattr(args, "extra_body", None),
enabled=should_request_stage_metrics(args),
)
- return asyncio.run(main_async(args))
+ lock = (
+ omniinteract_output_lock(Path(args.omniinteract_output_dir))
+ if getattr(args, "dataset_name", None) == "omniinteract"
+ else contextlib.nullcontext()
+ )
+ with lock:
+ return asyncio.run(main_async(args))
diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py
index 3069025be..f53a1ef74 100644
--- a/vllm_omni/config/__init__.py
+++ b/vllm_omni/config/__init__.py
@@ -2,7 +2,6 @@
Configuration module for vLLM-Omni.
"""
-from vllm_omni.config.config_factory import StageConfigFactory
from vllm_omni.config.lora import LoRAConfig
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.config.omni_config import (
@@ -22,7 +21,6 @@ from vllm_omni.config.omni_config import (
VllmOmniGenerationStageConfig,
VllmOmniOrchestratorConfig,
)
-from vllm_omni.config.pipeline_registry import register_pipeline
from vllm_omni.config.stage_config import (
PIPELINE_WIDE_ENGINE_FIELDS,
DeployConfig,
@@ -42,6 +40,15 @@ from vllm_omni.config.yaml_util import (
to_dict,
)
+# StageConfigFactory / register_pipeline pull pipeline_registry, which eagerly
+# imports PI0_PIPELINE → diffusion.data. Keep those lazy so
+# `from vllm_omni.config.lora import LoRAConfig` (used while data.py is still
+# loading) cannot close a circular import through DiffusionOutput.
+_LAZY_ATTRS = {
+ "StageConfigFactory": ("vllm_omni.config.config_factory", "StageConfigFactory"),
+ "register_pipeline": ("vllm_omni.config.pipeline_registry", "register_pipeline"),
+}
+
__all__ = [
# Legacy model-level configs.
"LoRAConfig",
@@ -82,3 +89,19 @@ __all__ = [
"merge_configs",
"to_dict",
]
+
+
+def __getattr__(name: str):
+ spec = _LAZY_ATTRS.get(name)
+ if spec is None:
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+ module_name, attr_name = spec
+ import importlib
+
+ value = getattr(importlib.import_module(module_name), attr_name)
+ globals()[name] = value
+ return value
+
+
+def __dir__() -> list[str]:
+ return sorted(set(__all__) | set(globals()))
diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py
index a1778eda3..b2eb64bca 100644
--- a/vllm_omni/config/config_factory.py
+++ b/vllm_omni/config/config_factory.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Config factories for vllm-omni, e.g., StageConfigFactory."""
from __future__ import annotations
@@ -15,6 +15,7 @@ from transformers import PretrainedConfig
from vllm.logger import init_logger
from vllm.transformers_utils.config import get_config
from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
+from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
from vllm_omni.config.endpoint_policy import EndpointRestriction
from vllm_omni.config.omni_config import VllmOmniConfig
@@ -28,8 +29,10 @@ from vllm_omni.config.stage_config import (
build_stage_runtime_overrides,
load_deploy_config,
merge_pipeline_deploy,
+ normalize_pipeline_cli_overrides,
)
from vllm_omni.config.yaml_util import create_config
+from vllm_omni.diffusion.io_support import get_diffusion_output_type
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
logger = init_logger(__name__)
@@ -45,6 +48,39 @@ logger = init_logger(__name__)
_DEFAULT_PARALLEL_DEGREE = 1
+@functools.cache
+def _materialize_object_storage_configs(model: str) -> str:
+ """Materialize an object-storage model URI's config files locally.
+
+ vLLM's Run:AI streamer keeps ``s3://``/``gs://``/``az://`` URIs opaque until
+ each stage builds its ``ModelConfig``; parent-process resolution (HF config
+ lookup, pipeline/pipeline-key matching) would instead hand the URI to
+ ``huggingface_hub`` helpers, which reject it with ``HFValidationError``.
+ Pull the lightweight files once into vLLM's deterministic
+ ``model_streamer/`` directory so config reads work here, and so the
+ stage processes' own pull lands in that same directory.
+
+ Returns the input unchanged for non object-storage paths.
+ """
+ if not is_runai_obj_uri(model):
+ return model
+ object_storage_model = ObjectStorageModel(url=model)
+ object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"])
+ logger.info("Materialized object-storage configs for %s at %s", model, object_storage_model.dir)
+ return object_storage_model.dir
+
+
+def _name_match_candidate(model: str) -> str:
+ """Last path component of a model reference, used for name-based matching.
+
+ Object-storage URIs and HF repo ids carry non-model segments (bucket name,
+ organization) that must not participate in substring matching; e.g. a
+ bucket named ``qwen3-tts-models`` holding a ``Qwen3-Omni`` checkpoint must
+ not resolve to the ``qwen3_tts`` pipeline.
+ """
+ return model.rstrip("/").rsplit("/", 1)[-1]
+
+
def with_trust_remote_code_override(
overrides: Mapping[str, Any],
trust_remote_code: bool | None,
@@ -117,7 +153,7 @@ class StageConfigFactory:
"""
hf_config = None
try:
- return get_config(model, trust_remote_code=trust_remote_code)
+ return get_config(_materialize_object_storage_configs(model), trust_remote_code=trust_remote_code)
except Exception as e:
logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
return hf_config
@@ -163,10 +199,12 @@ class StageConfigFactory:
if hf_config is not None:
return hf_config.model_type
+ config_source = _materialize_object_storage_configs(model)
+
# Fallback: read config.json directly for custom model types that
# are not registered with transformers (e.g. qwen3_tts).
try:
- config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+ config_dict = get_hf_file_to_dict("config.json", config_source, revision=None)
if config_dict:
if "model_type" in config_dict:
return config_dict["model_type"]
@@ -183,14 +221,17 @@ class StageConfigFactory:
# model_index.json with _class_name that maps to a pipeline key via
# PipelineConfig.diffusers_class_name.
try:
- model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
+ model_index = get_hf_file_to_dict("model_index.json", config_source, revision=None)
if model_index and "_class_name" in model_index:
class_name = model_index["_class_name"]
for obj in OMNI_PIPELINES.values():
# If we have a resolver, call it with the optional hf_config
# to get the default pipeline config for this key
pipeline_cfg = obj(hf_config) if callable(obj) else obj
- if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
+ if pipeline_cfg is not None and class_name in (
+ pipeline_cfg.diffusers_class_name,
+ *pipeline_cfg.diffusers_class_aliases,
+ ):
logger.info(
"Detected pipeline %r from model_index.json (_class_name=%r)",
pipeline_cfg.model_type,
@@ -203,8 +244,10 @@ class StageConfigFactory:
# Final fallback: some models (e.g. CosyVoice3) ship an empty
# config.json and rely on naming conventions. Match the model path
# basename against registered pipeline keys — longest match wins
- # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
- model_lower = model.lower().replace("-", "").replace("_", "")
+ # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). Only
+ # the basename is scanned so URI segments such as the bucket name
+ # cannot select an unrelated pipeline.
+ model_lower = _name_match_candidate(model).lower().replace("-", "").replace("_", "")
best: str | None = None
best_len = 0
for registered_key in OMNI_PIPELINES.keys():
@@ -395,6 +438,8 @@ class StageConfigFactory:
load-balancer policy (``None`` when no strategy set one) travels with the
stages instead of through a mutable out-param.
"""
+ cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
+ deploy_cfg: DeployConfig | None
if user_deploy_config is not None:
deploy_cfg = user_deploy_config
elif deploy_config_path is not None:
@@ -402,9 +447,12 @@ class StageConfigFactory:
assert deploy_cfg is not None
elif pipeline_cfg.default_deploy_config_name is not None:
deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name)
+ assert deploy_cfg is not None
else:
deploy_cfg = DeployConfig()
+ assert deploy_cfg is not None
+
cli_async_chunk = cli_overrides.get("async_chunk")
if cli_async_chunk is not None:
deploy_cfg.async_chunk = bool(cli_async_chunk)
@@ -583,6 +631,7 @@ class StageConfigFactory:
engine_args["dtype"] = str(engine_args["dtype"])
engine_args.setdefault("max_num_seqs", 1)
+ model_class_name = engine_args.get("model_class_name")
config_dict: dict[str, Any] = {
"stage_id": 0,
@@ -593,7 +642,7 @@ class StageConfigFactory:
},
"engine_args": create_config(engine_args),
"final_output": True,
- "final_output_type": "image",
+ "final_output_type": get_diffusion_output_type(model_class_name),
}
return [config_dict]
diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py
index d22c0acc1..af53b8691 100644
--- a/vllm_omni/config/endpoint_policy.py
+++ b/vllm_omni/config/endpoint_policy.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Endpoint restriction policy for omni pipelines."""
from dataclasses import dataclass
@@ -8,7 +8,14 @@ from typing import NamedTuple
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
-from vllm.entrypoints.serve.utils.error_response import create_error_response
+
+# vLLM < 0.28 keeps create_error_response under serve.utils (it is not
+# re-exported from vllm.entrypoints.serve); 0.28+ moved it under
+# serve.exception_handling and re-exports it from the package root.
+try:
+ from vllm.entrypoints.serve import create_error_response
+except ImportError:
+ from vllm.entrypoints.serve.utils.error_response import create_error_response
class RouteTarget(NamedTuple):
diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py
new file mode 100644
index 000000000..2a7385c1d
--- /dev/null
+++ b/vllm_omni/config/environment_variable_inventory.py
@@ -0,0 +1,292 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Reviewed inventory of environment variables used by vLLM-Omni.
+
+This module classifies environment-variable names; it does not read them or
+make every classified name part of the public configuration contract. In
+particular, model-specific variables remain transitional until their recorded
+disposition is implemented.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import Enum
+from types import MappingProxyType
+
+
+class EnvironmentVariableCategory(str, Enum):
+ """Ownership and documentation boundary for an environment variable."""
+
+ PUBLIC_OMNI = "public_omni"
+ INHERITED_VLLM = "inherited_vllm"
+ PLATFORM_EXTERNAL = "platform_external"
+ MODEL_SPECIFIC = "model_specific"
+ BENCHMARK_TRANSITIONAL = "benchmark_transitional"
+ INTERNAL = "internal"
+
+
+class ModelEnvironmentVariableDisposition(str, Enum):
+ """Required destination for a model-specific environment variable."""
+
+ PROMOTE = "promote"
+ REQUEST_SCOPE = "request_scope"
+ EXTERNAL = "external"
+ INTERNALIZE = "internalize"
+ DEPRECATE_REMOVE = "deprecate_remove"
+
+
+@dataclass(frozen=True)
+class EnvironmentVariableClassification:
+ """Classification metadata for one statically identifiable name."""
+
+ name: str
+ category: EnvironmentVariableCategory
+ model_disposition: ModelEnvironmentVariableDisposition | None = None
+ redact_value: bool = False
+
+ @property
+ def is_public_omni(self) -> bool:
+ return self.category is EnvironmentVariableCategory.PUBLIC_OMNI
+
+
+# New entries must use the VLLM_OMNI_ prefix. The regression test explicitly
+# grandfathers the older public names below.
+_PUBLIC_OMNI = (
+ "DIFFUSION_ATTENTION_BACKEND",
+ "DIFFUSION_CACHE_ADAPTER",
+ "DIFFUSION_CACHE_BACKEND",
+ "OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
+ "OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
+ "OMNI_DIFFUSION_SESSION_STATE_MANAGER",
+ "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
+ "SPEAKER_MAX_UPLOADED",
+ "SPEAKER_SAMPLES_DIR",
+ "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT",
+ "VLLM_OMNI_EVENT_DRIVEN_ORCH",
+ "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
+ "VLLM_OMNI_ORCH_MONITOR_PATH",
+ "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",
+ "VLLM_OMNI_SERVER_STORAGE__FILE_TTL",
+ "VLLM_OMNI_SERVER_STORAGE__PATH",
+ "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL",
+ "VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP",
+ "VLLM_OMNI_STORAGE_MAX_CONCURRENCY",
+ "VLLM_OMNI_STORAGE_PATH",
+ "VLLM_OMNI_USE_QUACK_FP8",
+ "VLLM_OMNI_VIDEO_SYNC_TIMEOUT",
+ "VLLM_VIDEO_ASYNC_CHUNK",
+ "VLLM_VIDEO_AUDIO_DELTA_MODE",
+)
+
+_INHERITED_VLLM = (
+ "CUDA_HOME",
+ "CUDA_VISIBLE_DEVICES",
+ "LOCAL_RANK",
+ "NO_COLOR",
+ "VLLM_ALLOW_LONG_MAX_MODEL_LEN",
+ "VLLM_BATCH_INVARIANT",
+ "VLLM_CACHE_ROOT",
+ "VLLM_DISABLE_LOG_LOGO",
+ "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE",
+ "VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
+ "VLLM_LOGGING_COLOR",
+ "VLLM_MOONCAKE_BOOTSTRAP_PORT",
+ "VLLM_PLUGINS",
+ "VLLM_ROCM_USE_AITER",
+ "VLLM_ROCM_USE_AITER_RMSNORM",
+ "VLLM_TUNED_CONFIG_FOLDER",
+ "VLLM_USE_MODELSCOPE",
+ "VLLM_USE_OINK_OPS",
+ "VLLM_WORKER_MULTIPROC_METHOD",
+ "VLLM_XPU_USE_SAMPLER_KERNEL",
+)
+
+_PLATFORM_EXTERNAL = (
+ "ASCEND_LAUNCH_BLOCKING",
+ "ASCEND_RT_VISIBLE_DEVICES",
+ "FLASHINFER_DISABLE_VERSION_CHECK",
+ "HF_HOME",
+ "HF_HUB_DISABLE_XET",
+ "HF_TOKEN",
+ "HIP_VISIBLE_DEVICES",
+ "HUGGINGFACE_HUB_TOKEN",
+ "MASTER_ADDR",
+ "MASTER_PORT",
+ "MINDIE_SD_FA_TYPE",
+ "MORI_RDMA_DEVICES",
+ "MUSA_VISIBLE_DEVICES",
+ "MV2_COMM_WORLD_LOCAL_RANK",
+ "NCCL_ASYNC_ERROR_HANDLING",
+ "OMPI_COMM_WORLD_LOCAL_RANK",
+ "OPENAI_API_KEY",
+ "PYTHONPATH",
+ "QUACK_CACHE_DIR",
+ "RANK",
+ "RANK_ID",
+ "RAY_RAYLET_PID",
+ "RDMA_DEVICE_NAME",
+ "SAGE_ATTN_TENSOR_LAYOUT",
+ "VLLM_LOCAL_RANK",
+ "WORLD_SIZE",
+ "XDG_CACHE_HOME",
+)
+
+_MODEL_PROMOTE = (
+ "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH",
+ "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
+ "COSMOS3_SOUND_TOKENIZER_PATH",
+ "COSYVOICE3_ESTIMATOR_ONNX",
+ "COSYVOICE3_TRT",
+ "COSYVOICE3_TRT_CACHE",
+ "DYNIN_CONFIG_PATH",
+ "DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY",
+ "DYNIN_MAGVIT_REMOTE_CODE_REPO_ID",
+ "DYNIN_MAGVIT_REMOTE_CODE_REVISION",
+ "DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY",
+ "DYNIN_REMOTE_CODE_REPO_ID",
+ "DYNIN_REMOTE_CODE_REVISION",
+ "HIGGS_AUDIO_TOKENIZER_PATH",
+ "INTERNVLA_A1_COSMOS_DECODER_PATH",
+ "INTERNVLA_A1_COSMOS_DIR",
+ "INTERNVLA_A1_COSMOS_ENCODER_PATH",
+ "INTERNVLA_A1_PROCESSOR_DIR",
+ "MIMO_AUDIO_TOKENIZER_CUDA_GRAPH",
+ "MIMO_AUDIO_TOKENIZER_DEVICE",
+ "MIMO_AUDIO_TOKENIZER_PATH",
+ "MING_CFM_CUDAGRAPH",
+ "MINICPMO_TOKEN2WAV_TRT",
+ "MINICPMO_TOKEN2WAV_TRT_CACHE",
+ "MINICPMO_TOKEN2WAV_TRT_DTYPE",
+ "NEMOTRON_VOICECHAT_LLM_PATH",
+ "OMNIVOICE_CUDA_GRAPH",
+ "OMNIVOICE_TF32",
+ "VLLM_GEPARD_CHUNK_FRAMES",
+ "VLLM_GEPARD_FIRST_CHUNK_FRAMES",
+ "VLLM_GEPARD_LOOKBACK_FRAMES",
+ "VLLM_OMNI_FISH_KVCACHE_ATTN",
+ "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER",
+)
+
+_MODEL_REQUEST_SCOPE = (
+ "DOTS_TTS_DIT_NUM_STEPS",
+ "GR00T_NOISE_SEED",
+ "STEP_AUDIO2_DEFAULT_PROMPT_WAV",
+ "VLLM_GEPARD_END_FADE_MS",
+ "VLLM_GEPARD_END_SILENCE_MS",
+ "VLLM_GEPARD_GREEDY",
+)
+
+# No audited model variable is currently owned by a supported third-party
+# contract. Keep the disposition explicit so future reviews cannot silently
+# fold an externally owned name into "promote".
+_MODEL_EXTERNAL: tuple[str, ...] = ()
+
+_MODEL_INTERNALIZE = (
+ "DOTS_TTS_BETA_TRACE",
+ "DYNIN_S2U_VENDOR_ROOT",
+ "DZ_PHASE_TIMING",
+ "MING_TTS_STAGE1_FINAL_LOG",
+ "MOSS_TTS_DEBUG_STOP",
+ "NEMOTRON_VOICECHAT_DEBUG_FUNCTION_TIMELINE",
+ "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS",
+ "VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS_LOG_EVERY",
+ "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS",
+ "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_FILE",
+ "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY",
+)
+
+_MODEL_DEPRECATE_REMOVE = (
+ "HIGGS_AUDIO_V2_TOKENIZER_PATH",
+ "VLLM_OMNI_LINGBOT_ACTION_ROOT",
+ "VLLM_OMNI_VOXCPM_CODE_PATH",
+ "VLLM_USE_FLASHINFER_MOE_FP16",
+ "XCODEC1_PATH",
+ "model_stage",
+)
+
+_BENCHMARK_TRANSITIONAL = (
+ "DAILY_OMNI_EXTRACT_MODE",
+ "DAILY_OMNI_SAVE_EVAL_ITEMS",
+ "MAX_NUM_FRAMES",
+ "SEED_TTS_EVAL_DEVICE",
+ "SEED_TTS_HF_WHISPER_MODEL",
+ "SEED_TTS_SIM_DEVICE",
+ "SEED_TTS_SIM_EVAL",
+ "SEED_TTS_UTMOS_EVAL",
+ "SEED_TTS_UTMOS_HF_REPO",
+ "SEED_TTS_UTMOS_JIT_FILE",
+ "SEED_TTS_WAVLM_MAX_SECONDS",
+ "SEED_TTS_WAVLM_MIN_SAMPLES",
+ "SEED_TTS_WAVLM_MODEL",
+ "SEED_TTS_WER_EVAL",
+ "SEED_TTS_WER_SAVE_AUDIO_DIR",
+ "SEED_TTS_WER_SAVE_ITEMS",
+ "VLLM_DAILY_OMNI_MEDIA_REPO",
+ "VLLM_OMNI_BENCH_AUDIO_CHANNELS",
+ "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S",
+ "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
+)
+
+_INTERNAL = (
+ "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT",
+ "VLLM_OMNI_REPLICA_ID",
+)
+
+_REDACTED = frozenset(
+ {
+ "HF_TOKEN",
+ "HUGGINGFACE_HUB_TOKEN",
+ "OPENAI_API_KEY",
+ }
+)
+
+
+def _build_inventory() -> Mapping[str, EnvironmentVariableClassification]:
+ inventory: dict[str, EnvironmentVariableClassification] = {}
+
+ def add(
+ names: tuple[str, ...],
+ category: EnvironmentVariableCategory,
+ disposition: ModelEnvironmentVariableDisposition | None = None,
+ ) -> None:
+ for name in names:
+ if name in inventory:
+ raise ValueError(f"Environment variable {name!r} is classified more than once")
+ inventory[name] = EnvironmentVariableClassification(
+ name=name,
+ category=category,
+ model_disposition=disposition,
+ redact_value=name in _REDACTED,
+ )
+
+ add(_PUBLIC_OMNI, EnvironmentVariableCategory.PUBLIC_OMNI)
+ add(_INHERITED_VLLM, EnvironmentVariableCategory.INHERITED_VLLM)
+ add(_PLATFORM_EXTERNAL, EnvironmentVariableCategory.PLATFORM_EXTERNAL)
+ add(_MODEL_PROMOTE, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.PROMOTE)
+ add(
+ _MODEL_REQUEST_SCOPE,
+ EnvironmentVariableCategory.MODEL_SPECIFIC,
+ ModelEnvironmentVariableDisposition.REQUEST_SCOPE,
+ )
+ add(_MODEL_EXTERNAL, EnvironmentVariableCategory.MODEL_SPECIFIC, ModelEnvironmentVariableDisposition.EXTERNAL)
+ add(
+ _MODEL_INTERNALIZE,
+ EnvironmentVariableCategory.MODEL_SPECIFIC,
+ ModelEnvironmentVariableDisposition.INTERNALIZE,
+ )
+ add(
+ _MODEL_DEPRECATE_REMOVE,
+ EnvironmentVariableCategory.MODEL_SPECIFIC,
+ ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE,
+ )
+ add(_BENCHMARK_TRANSITIONAL, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL)
+ add(_INTERNAL, EnvironmentVariableCategory.INTERNAL)
+ return MappingProxyType(dict(sorted(inventory.items())))
+
+
+# All reviewed environment-variable names, indexed by exact name. This is
+# classification metadata, not the executable value registry used by
+# ``vllm.envs``.
+ENVIRONMENT_VARIABLE_INVENTORY: Mapping[str, EnvironmentVariableClassification] = _build_inventory()
diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py
index 9164e1352..7be569f47 100644
--- a/vllm_omni/config/model.py
+++ b/vllm_omni/config/model.py
@@ -93,6 +93,7 @@ class OmniModelConfig(ModelConfig):
hf_text_config: The sub text_config of the model's hf_config (default: None)
stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
async_chunk: If set to True, perform async chunk
+ session_mode: Request lifecycle mode, either turn-based or duplex
model_stage: Stage type identifier, e.g., "thinker" or "talker"
(default: "thinker")
model_arch: Model architecture name
@@ -121,6 +122,7 @@ class OmniModelConfig(ModelConfig):
stage_id: int = 0
async_chunk: bool = False
+ session_mode: str = "turn"
retains_state_across_chunks: bool = False
# Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
active_stream_window: int = 0
@@ -141,11 +143,18 @@ class OmniModelConfig(ModelConfig):
}
)
subtalker_sampling_params: dict[str, Any] | None = None
+ silence_ban_frames: int = 0
omni_kv_config: dict | None = None
codec_frame_rate_hz: float | None = None
task_type: str | None = None
enable_sleep_mode: bool = False
has_sampling_extra_args: bool = False
+ # Key names (not values) of the stage's default sampling ``extra_args``.
+ # Engine-core code runs before any request arrives, so this is the only
+ # place it can learn which request-shaping conventions a stage uses (e.g.
+ # ``cfg_role`` for classifier-free-guidance request pairs).
+ sampling_extra_args_keys: tuple[str, ...] = ()
+ requires_full_payload_input: bool = False
@property
def registry(self):
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 4b5c73a39..58dc55fb6 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Structured vLLM-Omni configuration classes.
This module is additive for Phase 2 of RFC #4021.
@@ -12,12 +12,22 @@ from __future__ import annotations
import copy
from collections.abc import Mapping
-from dataclasses import dataclass, field, fields
+from dataclasses import InitVar, dataclass, field, fields
+from functools import wraps
+from inspect import Parameter, signature
from pathlib import Path
from typing import Any, Literal, TypeAlias, TypedDict, cast
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict, Field, model_validator
+from typing_extensions import Self
+from vllm.config import CacheConfig as VllmCacheConfig
+from vllm.config import CompilationConfig as VllmCompilationConfig
+from vllm.config import LoadConfig as VllmLoadConfig
+from vllm.config import ParallelConfig as VllmParallelConfig
+from vllm.config import ProfilerConfig as VllmProfilerConfig
+from vllm.config import SchedulerConfig as VllmSchedulerConfig
from vllm.config.utils import config
+from vllm.engine.arg_utils import EngineArgs as VllmEngineArgs
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_omni.config.stage_config import (
@@ -36,6 +46,7 @@ from vllm_omni.config.stage_config import (
_select_processor_funcs,
build_stage_runtime_overrides,
load_deploy_config,
+ normalize_pipeline_cli_overrides,
)
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -53,7 +64,6 @@ _NON_STAGE_ENGINE_CLI_FIELDS = frozenset(
"model",
"omni",
"output_modalities",
- "stage_configs_path",
"stage_id",
"tokenizer",
}
@@ -78,6 +88,49 @@ class _QuantizationEngineOverrides(TypedDict, total=False):
quantization: str
+class _TrackExplicitConfigFields:
+ """Record constructor inputs without adding a serialized config field."""
+
+ @model_validator(mode="wrap")
+ @classmethod
+ def _record_explicit_config_fields(cls, value: Any, handler: Any) -> Any:
+ result = handler(value)
+ kwargs = getattr(value, "kwargs", None)
+ if kwargs is not None:
+ explicit_fields = frozenset(kwargs)
+ elif isinstance(value, cls):
+ explicit_fields = getattr(value, "_omni_explicit_fields", frozenset())
+ else:
+ explicit_fields = frozenset()
+ object.__setattr__(result, "_omni_explicit_fields", explicit_fields)
+ return result
+
+
+def _enforce_keyword_only_init(cls: type[Any]) -> type[Any]:
+ """Make inherited Pydantic dataclass fields keyword-only as well."""
+ generated_init = cls.__init__
+ generated_signature = signature(cls)
+ keyword_only_signature = generated_signature.replace(
+ parameters=[
+ parameter.replace(kind=Parameter.KEYWORD_ONLY)
+ if parameter.kind in {Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD}
+ else parameter
+ for parameter in generated_signature.parameters.values()
+ ]
+ )
+
+ @wraps(generated_init)
+ def keyword_only_init(self: Any, *args: Any, **kwargs: Any) -> None:
+ if args:
+ raise TypeError(f"{cls.__name__}() accepts keyword arguments only; got {len(args)} positional argument(s)")
+ generated_init(self, **kwargs)
+
+ cls.__init__ = keyword_only_init
+ cls.__signature__ = keyword_only_signature
+ setattr(cls, "__match_args__", ())
+ return cls
+
+
class _ModelEngineOverrides(TypedDict, total=False):
model: str
model_arch: str
@@ -91,6 +144,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
trust_remote_code: bool
dtype: Any
attention_backend: Any
+ attention_config: Any
moe_backend: str
hf_overrides: Any
limit_mm_per_prompt: dict[str, Any]
@@ -99,6 +153,7 @@ class _ModelEngineOverrides(TypedDict, total=False):
active_stream_window: int
enable_sleep_mode: bool
subtalker_sampling_params: dict[str, Any]
+ silence_ban_frames: int
has_sampling_extra_args: bool
custom_voice_dir: str
task_type: str
@@ -106,7 +161,6 @@ class _ModelEngineOverrides(TypedDict, total=False):
enforce_eager: bool
max_cudagraph_capture_size: int
enable_flashinfer_autotune: bool
- compilation_config: dict[str, Any]
enable_multithread_weight_load: bool
num_weight_load_threads: int
disable_autocast: bool
@@ -148,7 +202,6 @@ class _RuntimeEngineOverrides(TypedDict, total=False):
num_gpus: int
log_level: str
log_stats: bool
- profiler_config: dict[str, Any]
class _ParallelConfigEngineOverrides(TypedDict, total=False):
@@ -160,6 +213,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False):
ring_degree: int
allgather_degree: int
ulysses_mode: str
+ ulysses_a2a_permute: bool
cfg_parallel_size: int
vae_patch_parallel_size: int
vae_parallel_mode: str
@@ -192,6 +246,8 @@ class _StageEngineValues:
runtime: _RuntimeEngineOverrides
parallel: _ParallelEngineOverrides
diffusion: _DiffusionEngineOverrides
+ compilation_config: Mapping[str, Any] | VllmCompilationConfig | None
+ profiler_config: Mapping[str, Any] | VllmProfilerConfig | None
@dataclass(frozen=True)
@@ -322,6 +378,7 @@ class OmniStageModelConfig:
trust_remote_code: bool = False
dtype: Any = "auto"
attention_backend: Any = None
+ attention_config: Any = None
moe_backend: str = "auto"
hf_overrides: Any = None
limit_mm_per_prompt: dict[str, Any] | None = None
@@ -329,10 +386,12 @@ class OmniStageModelConfig:
interleave_mm_strings: bool | None = None
media_io_kwargs: dict[str, Any] | None = None
active_stream_window: int = Field(default=0, ge=0)
+ session_mode: str = "turn"
duplex_max_sessions: int = Field(default=1, ge=1)
enable_sleep_mode: bool = False
default_sampling_params: dict[str, Any] | None = None
subtalker_sampling_params: dict[str, Any] | None = None
+ silence_ban_frames: int = 0
has_sampling_extra_args: bool = False
custom_voice_dir: str | None = None
task_type: str | None = None
@@ -340,7 +399,6 @@ class OmniStageModelConfig:
enforce_eager: bool = False
max_cudagraph_capture_size: int | None = Field(default=None, ge=0)
enable_flashinfer_autotune: bool | None = None
- compilation_config: dict[str, Any] | None = None
enable_multithread_weight_load: bool = True
num_weight_load_threads: int = Field(default=4, ge=1)
disable_autocast: bool = False
@@ -349,23 +407,24 @@ class OmniStageModelConfig:
# StagePipelineConfig.model_subdir/tokenizer_subdir on the legacy path.
model_subdir: str | None = None
tokenizer_subdir: str | None = None
+ requires_full_payload_input: bool = False
-@config
-class OmniStageLoadConfig:
- """Per-stage loading behavior and resolved tokenizer input."""
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageLoadConfig(_TrackExplicitConfigFields, VllmLoadConfig):
+ """vLLM loading behavior plus Omni stage-specific tokenizer inputs."""
tokenizer: str | None = None
- download_dir: str | None = None
skip_tokenizer_init: bool = False
- load_format: str = "auto"
tokenizer_mode: str = "auto"
config_format: str | None = None
skip_mm_profiling: bool | None = None
-@config
-class OmniStageCacheConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageCacheConfig(_TrackExplicitConfigFields, VllmCacheConfig):
"""Per-stage engine cache and memory behavior.
This is separate from ``_DiffusionConfigProjection.cache_config``, which configures
@@ -382,17 +441,26 @@ class OmniStageCacheConfig:
mamba_ssm_cache_dtype: str | None = None
-@config
-class OmniStageSchedulerConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageSchedulerConfig(_TrackExplicitConfigFields, VllmSchedulerConfig):
"""Per-stage request scheduling behavior."""
+ # Upstream receives max_model_len only while materializing SchedulerConfig.
+ # Omni retains it as unresolved stage input until the owning engine process.
max_num_seqs: int | None = Field(default=None, ge=1)
max_num_batched_tokens: int | None = Field(default=None, ge=1)
max_model_len: int | None = Field(default=None, ge=-1)
+ is_encoder_decoder: InitVar[bool] = False # type: ignore[assignment]
enable_chunked_prefill: bool | None = None
async_scheduling: bool | None = None
- def __post_init__(self) -> None:
+ def __post_init__(self, is_encoder_decoder: bool = False) -> None:
+ # Upstream initializes these derived fields in its terminal post-init.
+ # Keep them serializable here without running model-dependent checks.
+ self.max_num_encoder_input_tokens = self.max_num_batched_tokens
+ self.encoder_cache_size = self.max_num_batched_tokens
+
if (
self.max_num_batched_tokens is not None
and self.max_num_seqs is not None
@@ -431,24 +499,73 @@ class OmniStageRuntimeConfig:
num_gpus: int = Field(default=1, ge=1)
log_level: str = "info"
log_stats: bool = False
- profiler_config: dict[str, Any] | None = None
-@config
-class OmniStageParallelConfig:
+@_enforce_keyword_only_init
+@config(kw_only=True)
+class OmniStageParallelConfig(_TrackExplicitConfigFields, VllmParallelConfig):
"""Common per-stage distributed parallelism behavior."""
- pipeline_parallel_size: int = Field(default=1, ge=1)
- data_parallel_size: int = Field(default=1, ge=1)
- tensor_parallel_size: int = Field(default=1, ge=1)
- enable_expert_parallel: bool = False
- world_size: int = Field(default=1, ge=1, init=False)
+ # EngineArgs intentionally leaves these unresolved. The upstream terminal
+ # ParallelConfig defaults (rank 0, local size 1, port 29550, worker
+ # ``"auto"``) must not turn into explicit engine inputs in the head
+ # process merely because this transport class inherits ParallelConfig.
+ data_parallel_size_local: int | None = Field(default=None, ge=0)
+ data_parallel_rank: int | None = Field(default=None, ge=0)
+ data_parallel_rpc_port: int | None = None
+ worker_cls: str | None = None
+
+ @model_validator(mode="after")
+ def _validate_parallel_config(self) -> Self:
+ """Run upstream validation without resolving deferred engine inputs."""
+ deferred_values = (
+ self.data_parallel_size_local,
+ self.data_parallel_rank,
+ self.data_parallel_rpc_port,
+ self.worker_cls,
+ getattr(self, "all2all_backend", None),
+ getattr(self, "disable_custom_all_reduce", None),
+ )
+ if self.data_parallel_size_local is not None and self.data_parallel_size_local > self.data_parallel_size:
+ raise ValueError(
+ f"data_parallel_size_local ({self.data_parallel_size_local}) "
+ f"must be <= data_parallel_size ({self.data_parallel_size})"
+ )
+ if self.data_parallel_rank is not None and not 0 <= self.data_parallel_rank < self.data_parallel_size:
+ raise ValueError(
+ f"data_parallel_rank ({self.data_parallel_rank}) must be in the range [0, {self.data_parallel_size})"
+ )
+ self.data_parallel_size_local = self.data_parallel_size
+ self.data_parallel_rank = 0
+ self.data_parallel_rpc_port = VllmParallelConfig.data_parallel_rpc_port
+ self.worker_cls = VllmParallelConfig.worker_cls
+ try:
+ VllmParallelConfig._validate_parallel_config(self)
+ finally:
+ (
+ self.data_parallel_size_local,
+ self.data_parallel_rank,
+ self.data_parallel_rpc_port,
+ self.worker_cls,
+ self.all2all_backend,
+ self.disable_custom_all_reduce,
+ ) = deferred_values
+ return self
def __post_init__(self) -> None:
+ # Keep config construction transport-safe. Upstream runtime backend,
+ # rank, port, and platform resolution happens in the owning process.
+ self.data_parallel_index = self.data_parallel_rank
self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size
+ @property
+ def world_size_across_dp(self) -> int:
+ # Omni's public world_size has historically included DP.
+ return self.world_size
-@config
+
+@_enforce_keyword_only_init
+@config(kw_only=True)
class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
"""Diffusion-stage distributed parallelism behavior."""
@@ -457,6 +574,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
ring_degree: int = Field(default=1, ge=1)
allgather_degree: int = Field(default=1, ge=1)
ulysses_mode: str = "strict"
+ ulysses_a2a_permute: bool = False
cfg_parallel_size: int = Field(default=1, ge=1)
vae_patch_parallel_size: int = Field(default=1, ge=1)
text_encoder_tp_size: int = Field(default=1, ge=1)
@@ -467,6 +585,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
hsdp_replicate_size: int = Field(default=1, ge=1)
def __post_init__(self) -> None:
+ self.data_parallel_index = self.data_parallel_rank
self.sequence_parallel_size = (
self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
)
@@ -488,12 +607,17 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
* self.cfg_parallel_size
)
if self.use_hsdp:
- if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
- raise ValueError(
- "HSDP (use_hsdp=True) cannot be used with TP or DP "
- f"(tensor_parallel_size={self.tensor_parallel_size}, "
- f"data_parallel_size={self.data_parallel_size})"
- )
+ incompatible = []
+ if self.tensor_parallel_size > 1:
+ incompatible.append("TP")
+ if self.data_parallel_size > 1:
+ incompatible.append("DP")
+ if self.pipeline_parallel_size > 1:
+ incompatible.append("PP")
+ if self.enable_expert_parallel:
+ incompatible.append("EP")
+ if incompatible:
+ raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
if self.hsdp_shard_size == -1:
if other_parallel_world_size == 1:
raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1")
@@ -544,7 +668,7 @@ class _DiffusionConfigProjection:
dtype: Any = "auto"
trust_remote_code: bool = False
revision: str | None = None
- distributed_executor_backend: str = "mp"
+ distributed_executor_backend: str | None = None
dist_timeout: int | None = None
nccl_port: int | None = None
master_port: int | None = None
@@ -559,6 +683,7 @@ class _DiffusionConfigProjection:
cache_config: Any = field(default_factory=dict)
enable_cache_dit_summary: bool = False
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+ diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True)
enable_prompt_embed_cache: bool = False
prompt_embed_cache_size: int = Field(default=32, ge=1)
enable_session_state_manager: bool = False
@@ -576,6 +701,9 @@ class _DiffusionConfigProjection:
enable_distributed_layerwise_offload: bool = False
dlo_use_allgather: bool = True
dlo_resident_layers: int = Field(default=0, ge=0)
+ host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled"
+ host_weight_runtime_root: str | None = None
+ dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0)
pin_cpu_memory: bool = True
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -605,6 +733,7 @@ class _DiffusionConfigProjection:
default_factory=lambda: {
"transformer": True,
"vae": True,
+ "text_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -621,7 +750,7 @@ class _DiffusionConfigProjection:
from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs)
- valid_fields = {f.name for f in fields(cls)}
+ valid_fields = {f.name for f in fields(cast(Any, cls))}
return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields})
def __post_init__(self) -> None:
@@ -633,6 +762,8 @@ class _DiffusionConfigProjection:
TransformerConfig,
build_attention_config,
parse_kv_cache_skip_selector,
+ validate_dlo_host_registration_options,
+ validate_host_weight_runtime_options,
)
from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
from vllm_omni.quantization import build_quant_config
@@ -697,6 +828,11 @@ class _DiffusionConfigProjection:
)
self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+ if (
+ self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+ and self.diffusion_kv_max_rows_per_request is None
+ ):
+ raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set")
self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
@@ -705,6 +841,17 @@ class _DiffusionConfigProjection:
elif self.max_cpu_loras < 1:
raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
+ validate_host_weight_runtime_options(
+ mode=self.host_weight_runtime_mode,
+ root=self.host_weight_runtime_root,
+ )
+ self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+ limit_gib=self.dlo_host_registration_limit_gib,
+ enable_dlo=self.enable_distributed_layerwise_offload,
+ use_allgather=self.dlo_use_allgather,
+ hwr_mode=self.host_weight_runtime_mode,
+ )
+
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
raise ValueError(
"diffusers_load_kwargs and diffusers_call_kwargs are only "
@@ -768,7 +915,7 @@ class _DiffusionConfigProjection:
setattr(self, name, _copy_value(getattr(omni_diffusion_config, name)))
-_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection))
+_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(cast(Any, _DiffusionConfigProjection)))
# Current OmniDiffusionConfig still contains a flat mix of shared engine,
# runtime, parallel, and diffusion-specific knobs. Keep this classification
@@ -813,6 +960,10 @@ _DIFFUSION_MOVED_SHARED_FIELDS = frozenset(
"omni_kv_config",
"cfg_kv_collect_func",
"max_num_seqs",
+ "kv_cache_memory_bytes",
+ "gpu_memory_utilization",
+ "max_num_batched_tokens",
+ "max_model_len",
"enable_sleep_mode",
"enforce_eager",
"enable_multithread_weight_load",
@@ -827,6 +978,7 @@ _STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS)
_DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset(
{
"diffusion_attention_backend",
+ "fastvideo_vsa_topk",
"kv_cache_dtype",
"kv_cache_skip_layers",
"kv_cache_skip_steps",
@@ -838,6 +990,58 @@ _DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMP
"stage_id",
}
+
+def _upstream_engine_field_map(
+ config_cls: type[Any],
+ *,
+ aliases: Mapping[str, str] = {},
+ exclude: frozenset[str] = frozenset(),
+) -> dict[str, str]:
+ """Map reusable upstream config fields to their EngineArgs inputs."""
+ engine_fields = frozenset(config_field.name for config_field in fields(VllmEngineArgs))
+ return {
+ config_field.name: engine_name
+ for config_field in fields(config_cls)
+ if config_field.init
+ and config_field.name not in exclude
+ and (engine_name := aliases.get(config_field.name, config_field.name)) in engine_fields
+ }
+
+
+_LOAD_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(VllmLoadConfig)
+_CACHE_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+ VllmCacheConfig,
+ aliases={"cache_dtype": "kv_cache_dtype"},
+)
+_SCHEDULER_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+ VllmSchedulerConfig,
+ aliases={"policy": "scheduling_policy"},
+ # ``scheduler_cls`` is selected by the immutable stage topology, while
+ # ``disable_hybrid_kv_cache_manager`` belongs to the cache concern in the
+ # Omni schema. Keep both out of the dynamic SchedulerConfig projection so
+ # one input cannot acquire two owners.
+ exclude=frozenset(
+ {
+ "scheduler_cls",
+ "disable_hybrid_kv_cache_manager",
+ }
+ ),
+)
+_PARALLEL_CONFIG_ENGINE_FIELD_MAP = _upstream_engine_field_map(
+ VllmParallelConfig,
+ aliases={"data_parallel_master_ip": "data_parallel_address"},
+ # Runtime owns worker selection. The private API-process fields are
+ # terminal vLLM internals rather than per-stage user inputs.
+ exclude=frozenset(
+ {
+ "distributed_executor_backend",
+ "worker_cls",
+ "_api_process_count",
+ "_api_process_rank",
+ }
+ ),
+)
+
_QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__)
_MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__)
_LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__)
@@ -845,7 +1049,17 @@ _CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__)
_SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__)
_CONNECTOR_ENGINE_FIELDS = frozenset(_ConnectorEngineOverrides.__annotations__)
_RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__)
-_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIRECT_VLLM_CONFIG_ENGINE_FIELDS = frozenset({"compilation_config", "profiler_config"})
+_LLM_LOAD_ENGINE_FIELDS = _LOAD_ENGINE_FIELDS | frozenset(_LOAD_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_CACHE_ENGINE_FIELDS = _CACHE_ENGINE_FIELDS | frozenset(_CACHE_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_SCHEDULER_ENGINE_FIELDS = _SCHEDULER_ENGINE_FIELDS | frozenset(_SCHEDULER_CONFIG_ENGINE_FIELD_MAP.values())
+_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP.values())
+_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(
+ f.name for f in fields(OmniStageDiffusionParallelConfig)
+) & frozenset(_ParallelConfigEngineOverrides.__annotations__)
+_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP = {name: name for name in _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS}
+_LLM_PARALLEL_CONFIG_FIELDS = frozenset(_PARALLEL_CONFIG_ENGINE_FIELD_MAP)
+_PARALLEL_CONFIG_ENGINE_FIELDS = _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
_PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
_COMMON_STAGE_ENGINE_FIELDS = (
_QUANTIZATION_ENGINE_FIELDS
@@ -855,12 +1069,16 @@ _COMMON_STAGE_ENGINE_FIELDS = (
| _SCHEDULER_ENGINE_FIELDS
| _CONNECTOR_ENGINE_FIELDS
| _RUNTIME_ENGINE_FIELDS
+ | _DIRECT_VLLM_CONFIG_ENGINE_FIELDS
+)
+_LLM_STAGE_ENGINE_FIELDS = (
+ _COMMON_STAGE_ENGINE_FIELDS
+ | _LLM_LOAD_ENGINE_FIELDS
+ | _LLM_CACHE_ENGINE_FIELDS
+ | _LLM_SCHEDULER_ENGINE_FIELDS
+ | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS
+ | {"parallel_config"}
)
-_LLM_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageParallelConfig)) - {"world_size"}
-_DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(f.name for f in fields(OmniStageDiffusionParallelConfig)) - {
- "world_size"
-}
-_LLM_STAGE_ENGINE_FIELDS = _COMMON_STAGE_ENGINE_FIELDS | _LLM_PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"}
_DIFFUSION_OWNED_STAGE_ENGINE_FIELDS = (
_COMMON_STAGE_ENGINE_FIELDS
| _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS
@@ -881,6 +1099,27 @@ _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE = {
StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_ENGINE_FIELDS,
StageExecutionType.DIFFUSION: _DIFFUSION_PARALLEL_CONFIG_ENGINE_FIELDS,
}
+_PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE = {
+ StageExecutionType.LLM_AR: _LLM_PARALLEL_CONFIG_FIELDS,
+ StageExecutionType.LLM_GENERATION: _LLM_PARALLEL_CONFIG_FIELDS,
+ StageExecutionType.DIFFUSION: frozenset(_DIFFUSION_PARALLEL_CONFIG_FIELD_MAP),
+}
+
+_LOAD_STAGE_ENGINE_FIELD_MAP = {
+ **{name: name for name in _LOAD_ENGINE_FIELDS},
+ **_LOAD_CONFIG_ENGINE_FIELD_MAP,
+}
+_CACHE_STAGE_ENGINE_FIELD_MAP = {
+ **{name: name for name in _CACHE_ENGINE_FIELDS},
+ **_CACHE_CONFIG_ENGINE_FIELD_MAP,
+}
+_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {
+ **{name: name for name in _SCHEDULER_ENGINE_FIELDS},
+ **_SCHEDULER_CONFIG_ENGINE_FIELD_MAP,
+}
+_DIFFUSION_LOAD_STAGE_ENGINE_FIELD_MAP = {name: name for name in _LOAD_ENGINE_FIELDS}
+_DIFFUSION_CACHE_STAGE_ENGINE_FIELD_MAP = {name: name for name in _CACHE_ENGINE_FIELDS}
+_DIFFUSION_SCHEDULER_STAGE_ENGINE_FIELD_MAP = {name: name for name in _SCHEDULER_ENGINE_FIELDS}
def _validate_stage_engine_override_ownership(
@@ -891,13 +1130,19 @@ def _validate_stage_engine_override_ownership(
try:
owner_fields = _STAGE_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
parallel_owner_fields = _PARALLEL_CONFIG_ENGINE_FIELDS_BY_EXECUTION_TYPE[execution_type]
+ parallel_config_fields = _PARALLEL_CONFIG_FIELDS_BY_EXECUTION_TYPE[execution_type]
except KeyError as exc:
raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc
unowned_fields = set(overrides) - owner_fields
parallel_config = overrides.get("parallel_config")
if isinstance(parallel_config, Mapping):
- unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - parallel_owner_fields)
+ # Nested values traditionally use upstream config names, while flat
+ # EngineArgs uses aliases such as ``data_parallel_address``. Accept
+ # either spelling at the boundary; the builder canonicalizes aliases
+ # before instantiating the inherited config.
+ accepted_parallel_names = parallel_config_fields | parallel_owner_fields
+ unowned_fields.update(f"parallel_config.{name}" for name in set(parallel_config) - accepted_parallel_names)
if unowned_fields:
names = ", ".join(sorted(unowned_fields))
raise ValueError(
@@ -918,7 +1163,7 @@ def _global_stage_cli_fields() -> frozenset[str]:
)
externally_consumed = (
_NON_STAGE_ENGINE_CLI_FIELDS
- | frozenset(f.name for f in fields(VllmOmniOrchestratorConfig))
+ | frozenset(f.name for f in fields(cast(Any, VllmOmniOrchestratorConfig)))
| (orchestrator_field_names() - _STAGE_ENGINE_FIELDS)
)
return candidates - externally_consumed
@@ -968,17 +1213,28 @@ def _stage_engine_values(
topology.execution_type,
engine,
)
+ if topology.execution_type in {
+ StageExecutionType.LLM_AR,
+ StageExecutionType.LLM_GENERATION,
+ }:
+ load_engine_fields = _LLM_LOAD_ENGINE_FIELDS
+ cache_engine_fields = _LLM_CACHE_ENGINE_FIELDS
+ scheduler_engine_fields = _LLM_SCHEDULER_ENGINE_FIELDS
+ else:
+ load_engine_fields = _LOAD_ENGINE_FIELDS
+ cache_engine_fields = _CACHE_ENGINE_FIELDS
+ scheduler_engine_fields = _SCHEDULER_ENGINE_FIELDS
return _StageEngineValues(
quantization=cast(
_QuantizationEngineOverrides,
_select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS),
),
model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)),
- load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)),
- cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)),
+ load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, load_engine_fields)),
+ cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, cache_engine_fields)),
scheduler=cast(
_SchedulerEngineOverrides,
- _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS),
+ _select_engine_overrides(engine, scheduler_engine_fields),
),
connector=cast(
_ConnectorEngineOverrides,
@@ -987,6 +1243,8 @@ def _stage_engine_values(
runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)),
parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)),
diffusion=_DiffusionEngineOverrides.from_engine(engine),
+ compilation_config=_copy_value(engine.get("compilation_config")),
+ profiler_config=_copy_value(engine.get("profiler_config")),
)
@@ -1003,7 +1261,7 @@ def _stage_sampling_params(
def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]:
overrides: dict[str, Any] = {}
- for config_field in fields(VllmOmniOrchestratorConfig):
+ for config_field in fields(cast(Any, VllmOmniOrchestratorConfig)):
name = config_field.name
if name == "deploy_config_path":
continue
@@ -1041,6 +1299,8 @@ class BaseVllmOmniStageConfig:
connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig)
runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig)
parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig)
+ compilation_config: VllmCompilationConfig | None = None
+ profiler_config: VllmProfilerConfig | None = None
quantization_config: _QuantizationConfigType = None
@property
@@ -1117,6 +1377,10 @@ class BaseVllmOmniStageConfig:
def prompt_expand_func(self) -> str | None:
return self.stage_pipeline_config.prompt_expand_func
+ @property
+ def prompt_transform_func(self) -> str | None:
+ return self.stage_pipeline_config.prompt_transform_func
+
@property
def cfg_kv_collect_func(self) -> str | None:
return self.stage_pipeline_config.cfg_kv_collect_func
@@ -1192,6 +1456,8 @@ def _build_common_stage_config_kwargs(
parallel_config,
),
"parallel_config": parallel_config,
+ "compilation_config": _copy_value(engine.compilation_config),
+ "profiler_config": _copy_value(engine.profiler_config),
"quantization_config": _copy_value(quantization_config),
},
input_proc,
@@ -1204,8 +1470,8 @@ def _with_resolved_processors(
input_proc: str | None,
next_stage_proc: str | None,
) -> StageConfigType:
- stage_config._resolved_custom_process_input_func = input_proc
- stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc
+ setattr(stage_config, "_resolved_custom_process_input_func", input_proc)
+ setattr(stage_config, "_resolved_custom_process_next_stage_input_func", next_stage_proc)
return stage_config
@@ -1319,13 +1585,16 @@ def _build_stage_config(
builder = _STAGE_CONFIG_BUILDERS[topology.execution_type]
except KeyError as exc:
raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc
- return builder(
- pipeline,
- deploy,
- topology,
- stage_deploy,
- engine,
- model=model,
+ return cast(
+ StageConfigType,
+ builder(
+ pipeline,
+ deploy,
+ topology,
+ stage_deploy,
+ engine,
+ model=model,
+ ),
)
@@ -1352,6 +1621,7 @@ def _build_model_config(
) -> OmniStageModelConfig:
default_sampling_params = _stage_sampling_params(stage_deploy, topology)
kwargs = _config_kwargs(engine)
+ kwargs["requires_full_payload_input"] = topology.requires_full_payload_input
kwargs["model"] = _first_defined(kwargs.get("model"), model)
if "model_arch" not in kwargs:
kwargs["model_arch"] = topology.model_arch or pipeline.model_arch or None
@@ -1369,8 +1639,9 @@ def _build_model_config(
kwargs["model_subdir"] = topology.model_subdir
if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None:
kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
- return OmniStageModelConfig(
+ return cast(Any, OmniStageModelConfig)(
default_sampling_params=default_sampling_params,
+ session_mode=deploy.session_mode,
duplex_max_sessions=duplex_max_sessions,
**kwargs,
)
@@ -1386,12 +1657,51 @@ def _build_load_config(
return OmniStageLoadConfig(**kwargs)
+def _config_kwargs_from_engine_args(
+ engine: Mapping[str, Any],
+ field_map: Mapping[str, str],
+) -> dict[str, Any]:
+ """Convert EngineArgs names back to their upstream config field names."""
+ return {
+ config_name: _copy_value(engine[engine_name])
+ for config_name, engine_name in field_map.items()
+ if engine_name in engine and engine[engine_name] is not None
+ }
+
+
+def _normalize_config_mapping(
+ values: Mapping[str, Any],
+ field_map: Mapping[str, str],
+) -> dict[str, Any]:
+ """Canonicalize config and EngineArgs spellings to config field names."""
+ engine_to_config = {engine_name: config_name for config_name, engine_name in field_map.items()}
+ normalized: dict[str, Any] = {}
+ source_names: dict[str, str] = {}
+ for name, value in values.items():
+ if value is None:
+ continue
+ config_name = name if name in field_map else engine_to_config.get(name, name)
+ if config_name in normalized:
+ raise ValueError(
+ f"Config mapping specifies both {source_names[config_name]!r} and {name!r} "
+ f"for upstream field {config_name!r}"
+ )
+ normalized[config_name] = _copy_value(value)
+ source_names[config_name] = name
+ return normalized
+
+
+def _omni_config_kwargs(engine: Mapping[str, Any], fields_: frozenset[str]) -> dict[str, Any]:
+ return {name: _copy_value(engine[name]) for name in fields_ if name in engine and engine[name] is not None}
+
+
def _build_cache_config(
deploy: DeployConfig,
engine: _CacheEngineOverrides,
execution_type: StageExecutionType,
) -> OmniStageCacheConfig:
- kwargs = _config_kwargs(engine)
+ kwargs = _config_kwargs_from_engine_args(engine, _CACHE_CONFIG_ENGINE_FIELD_MAP)
+ kwargs.update(_omni_config_kwargs(engine, _CACHE_ENGINE_FIELDS))
if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None:
kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching)
if "disable_hybrid_kv_cache_manager" not in kwargs and execution_type == StageExecutionType.LLM_GENERATION:
@@ -1405,7 +1715,8 @@ def _build_scheduler_config(
engine: _SchedulerEngineOverrides,
execution_type: StageExecutionType,
) -> OmniStageSchedulerConfig:
- kwargs = _config_kwargs(engine)
+ kwargs = _config_kwargs_from_engine_args(engine, _SCHEDULER_CONFIG_ENGINE_FIELD_MAP)
+ kwargs.update(_omni_config_kwargs(engine, _SCHEDULER_ENGINE_FIELDS))
if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None:
kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill)
if "async_scheduling" not in kwargs and execution_type == StageExecutionType.LLM_AR:
@@ -1420,7 +1731,7 @@ def _build_connector_config(
) -> OmniStageConnectorConfig:
output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
- return OmniStageConnectorConfig(
+ return cast(Any, OmniStageConnectorConfig)(
async_chunk=bool(deploy.async_chunk),
omni_kv_config=_copy_value(engine.get("omni_kv_config")),
output_connectors=_copy_value(output_connectors) if output_connectors else None,
@@ -1453,16 +1764,21 @@ def _build_parallel_config(
config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig,
) -> OmniStageParallelConfig:
parallel_config = _mapping_or_empty(engine.get("parallel_config"))
- config_fields = {
- config_field.name
- for config_field in fields(config_cls)
- if getattr(config_field.default, "init", None) is not False
- }
- kwargs = {
- name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None
- }
+ config_fields = frozenset(signature(config_cls).parameters)
+ is_diffusion = issubclass(config_cls, OmniStageDiffusionParallelConfig)
+ field_map = _DIFFUSION_PARALLEL_CONFIG_FIELD_MAP if is_diffusion else _PARALLEL_CONFIG_ENGINE_FIELD_MAP
+ owned_parallel_fields = frozenset(field_map)
+ kwargs = _config_kwargs(engine) if is_diffusion else _config_kwargs_from_engine_args(engine, field_map)
+ kwargs = {name: value for name, value in kwargs.items() if name in owned_parallel_fields}
kwargs = {name: value for name, value in kwargs.items() if name in config_fields}
- kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields})
+ nested_kwargs = _normalize_config_mapping(parallel_config, field_map)
+ kwargs.update(
+ {
+ name: value
+ for name, value in nested_kwargs.items()
+ if name in owned_parallel_fields and name in config_fields
+ }
+ )
if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None:
kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size)
if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None:
@@ -1528,6 +1844,7 @@ class VllmOmniConfig:
"""Create a structured config from a resolved pipeline and deploy YAML."""
if cli_overrides is None:
cli_overrides = {}
+ cli_overrides = normalize_pipeline_cli_overrides(pipeline_cfg, cli_overrides)
deploy, loaded_deploy_config_path = _get_deploy_config(
pipeline_cfg,
@@ -1564,11 +1881,11 @@ class VllmOmniConfig:
for topology in pipeline_cfg.stages
)
- orchestrator_config = VllmOmniOrchestratorConfig(
+ orchestrator_config = cast(Any, VllmOmniOrchestratorConfig)(
deploy_config_path=loaded_deploy_config_path,
**_orchestrator_cli_overrides(cli_overrides),
)
- return cls(
+ return cast(Any, cls)(
pipeline_config=pipeline_cfg,
stage_configs=stage_configs,
orchestrator_config=orchestrator_config,
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index f02a3578e..83f8de383 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Pipeline registry and factory for vllm-omni.
``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig``
@@ -33,6 +33,7 @@ from vllm.logger import init_logger
from vllm_omni.config.stage_config import (
PipelineConfig,
)
+from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE
from vllm_omni.model_executor.models.audex.pipeline import (
AUDEX_S2S_PIPELINE,
AUDEX_THINKER_ONLY_PIPELINE,
@@ -47,9 +48,11 @@ from vllm_omni.model_executor.models.bagel.pipeline import (
)
from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
+from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE
from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE
from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
+from vllm_omni.model_executor.models.gepard.pipeline import GEPARD_PIPELINE
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
from vllm_omni.model_executor.models.glm_tts.pipeline import GLM_TTS_PIPELINE
from vllm_omni.model_executor.models.gr00t.pipeline import GR00T_N1D7_PIPELINE
@@ -82,6 +85,8 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import (
MING_TTS_PIPELINE,
)
from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE
+from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE
+from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE
from vllm_omni.model_executor.models.moss_tts.pipeline import (
MOSS_TTS_LOCAL_PIPELINE,
MOSS_TTS_PIPELINE,
@@ -97,12 +102,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
QWEN2_5_OMNI_PIPELINE,
QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
)
-from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline
-from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
-from vllm_omni.model_executor.models.soulx_singer.pipeline import (
- SOULXSINGER_SVC_PIPELINE,
- SOULXSINGER_SVS_PIPELINE,
+from vllm_omni.model_executor.models.qwen3_omni.pipeline import (
+ QWEN3_OMNI_THINKER_ONLY_PIPELINE,
+ resolve_qwen3_omni_pipeline,
)
+from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
from vllm_omni.model_executor.models.step_audio2.pipeline import (
STEP_AUDIO2_ASR_PIPELINE,
STEP_AUDIO2_PIPELINE,
@@ -127,6 +131,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
# has no model_type key).
"nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
"qwen3_omni_moe": resolve_qwen3_omni_pipeline,
+ "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE,
"qwen3_tts": QWEN3_TTS_PIPELINE,
"step_audio_2": STEP_AUDIO2_PIPELINE,
"step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
@@ -137,6 +142,8 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"lance": LANCE_PIPELINE,
"dreamzero": DREAMZERO_PIPELINE,
"Gr00tN1d7": GR00T_N1D7_PIPELINE,
+ "pi0": PI0_PIPELINE,
+ "gepard": GEPARD_PIPELINE,
"glm_image": GLM_IMAGE_PIPELINE,
"hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,
"hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE,
@@ -144,6 +151,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE,
"wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
"voxcpm2": VOXCPM2_PIPELINE,
+ "dots_tts": DOTS_TTS_PIPELINE,
"cosyvoice3": COSYVOICE3_PIPELINE,
"audex_tts": AUDEX_TTS_PIPELINE,
"audex_tta": AUDEX_TTA_PIPELINE,
@@ -164,6 +172,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
"ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
"moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
+ "minimax_h3_disaggregated": MINIMAX_H3_PIPELINE,
"omnivoice": OMNIVOICE_PIPELINE,
"mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
"mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
@@ -171,13 +180,12 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
"moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
"moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
"minicpmo_4_5": MINICPMO_4_5_PIPELINE,
+ "minimax_music3": MINIMAX_MUSIC3_PIPELINE,
"higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE,
"higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE,
"dynin_omni": DYNIN_OMNI_PIPELINE,
"indextts2": INDEXTTS2_PIPELINE,
"indextts2_5": INDEXTTS25_PIPELINE,
- "soulxsinger_svc": SOULXSINGER_SVC_PIPELINE,
- "soulxsinger_svs": SOULXSINGER_SVS_PIPELINE,
}
@@ -188,18 +196,13 @@ def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_typ
since resolvers can return multiple different PipelineConfigs depending on the
consumed config.
"""
- errors: list[str] = []
if isinstance(pipeline, PipelineConfig):
- errors = pipeline.validate()
model_type = model_type if model_type is not None else pipeline.model_type
- else:
- if model_type is None:
- raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
+ elif model_type is None:
+ raise ValueError("Model type must be explicitly provided when registering a pipeline resolver")
if model_type in OMNI_PIPELINES:
- errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered")
- if errors:
- logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors)
+ logger.warning(f"Model type {model_type} is already registered; the old mapping will be clobbered")
OMNI_PIPELINES[model_type] = pipeline
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index 496aea208..2c1dc670b 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -1,10 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Stage configuration system for vLLM-Omni."""
from __future__ import annotations
-import dataclasses
import functools
import re
import warnings
@@ -56,7 +55,7 @@ def build_stage_runtime_overrides(
``internal_keys`` defaults to the union of
``arg_utils.internal_blacklist_keys()`` and ``arg_utils.SHARED_FIELDS``
so that neither orchestrator-only fields nor shared-pipeline fields
- (``model`` / ``stage_configs_path`` / ``log_stats`` / ``stage_id``) leak
+ (``model`` / ``log_stats`` / ``stage_id``) leak
into a stage's per-stage runtime overrides — the orchestrator sets those
uniformly for every stage, they are not per-stage knobs. Callers can
pass an explicit set for tests or specialized flows.
@@ -90,44 +89,41 @@ def build_stage_runtime_overrides(
return result
-def strip_parent_engine_args(
- kwargs: dict[str, Any],
- *,
- parent_fields: dict[str, dataclasses.Field],
- keep_keys: set[str] | frozenset[str] = frozenset(),
- strip_keys: set[str] | frozenset[str] = frozenset(),
- no_warn_keys: set[str] | frozenset[str] = frozenset(),
-) -> tuple[dict[str, Any], list[str]]:
- """Strip parent ``EngineArgs`` fields before merging into stage YAML."""
- overridden: list[str] = []
- result: dict[str, Any] = {}
-
- for key, value in kwargs.items():
- if key in strip_keys:
+def normalize_pipeline_cli_overrides(
+ pipeline: PipelineConfig,
+ cli_overrides: dict[str, Any],
+) -> dict[str, Any]:
+ """Translate pipeline-owned global CLI aliases into stage-scoped overrides."""
+ normalized = dict(cli_overrides)
+ for source, (stage_id, target) in pipeline.stage_cli_aliases.items():
+ invalid_stage_keys = [
+ key
+ for key, value in normalized.items()
+ if value is not None
+ and (match := _STAGE_OVERRIDE_PATTERN.match(key)) is not None
+ and match.group(2) == source
+ and int(match.group(1)) != stage_id
+ ]
+ if invalid_stage_keys:
+ invalid = ", ".join(sorted(invalid_stage_keys))
+ raise ValueError(
+ f"{invalid} cannot be set for pipeline {pipeline.model_type!r}; "
+ f"{source} belongs to stage {stage_id} as {target}."
+ )
+ value = normalized.pop(source, None)
+ if value is None:
continue
-
- if key not in parent_fields or key in keep_keys:
- result[key] = value
+ stage_key = f"stage_{stage_id}_{target}"
+ stage_value = normalized.get(stage_key)
+ if stage_value is not None and stage_value != value:
+ warnings.warn(
+ f"Ignoring {source}={value!r} because {stage_key}={stage_value!r} takes precedence.",
+ UserWarning,
+ stacklevel=2,
+ )
continue
-
- field_def = parent_fields[key]
- if field_def.default is not dataclasses.MISSING:
- default = field_def.default
- elif field_def.default_factory is not dataclasses.MISSING:
- default = field_def.default_factory()
- else:
- default = dataclasses.MISSING
-
- if default is dataclasses.MISSING or value is None:
- continue
-
- if dataclasses.is_dataclass(default) and not isinstance(default, type):
- default = asdict(default)
-
- if value != default and key not in no_warn_keys:
- overridden.append(key)
-
- return result, sorted(overridden)
+ normalized[stage_key] = value
+ return normalized
def _apply_diffusion_parallel_runtime_overrides(
@@ -232,6 +228,10 @@ class StagePipelineConfig:
# Alternates picked by ``merge_pipeline_deploy`` based on ``deploy.async_chunk``.
async_chunk_process_next_stage_input_func: str | None = None
sync_process_input_func: str | None = None
+ # Rewrites the Stage-0 view of a raw prompt before vLLM input processing.
+ # The callable receives ``(prompt, sampling_params_list)``; downstream
+ # stages continue to receive the original prompt.
+ prompt_transform_func: str | None = None
prompt_expand_func: str | None = None
cfg_kv_collect_func: str | None = None
omni_kv_config: dict[str, Any] | None = None
@@ -242,6 +242,16 @@ class StagePipelineConfig:
# by ``stage_init_utils._resolve_model_tokenizer_paths``.
model_subdir: str | None = None
tokenizer_subdir: str | None = None
+ # Model-owned hook that resolves a pipeline root to this stage's checkpoint.
+ # Consumed and removed before backend engine args are constructed.
+ model_path_resolver: str | None = None
+ # Keep a single-replica diffusion stage in the orchestrator process.
+ # Disabled by default so existing multi-stage pipelines retain subprocess
+ # isolation unless their topology explicitly opts in.
+ inline_diffusion: bool = False
+ # Whether the non-async path waits for a complete upstream payload from
+ # the model-runner connector before scheduling this stage.
+ requires_full_payload_input: bool = False
extras: dict[str, Any] = field(default_factory=dict)
@@ -275,6 +285,7 @@ class PipelineConfig:
# this value to auto-detect the pipeline. Only needed for diffusers-style
# multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model.
diffusers_class_name: str | None = None
+ diffusers_class_aliases: tuple[str, ...] = ()
endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
# Optional model-owned duplex planner loaded by the stable engine runtime.
duplex_runtime_extension: str | None = None
@@ -288,6 +299,19 @@ class PipelineConfig:
# Bundled deploy defaults for this concrete pipeline topology. The file is
# loaded from vllm_omni/deploy; None uses DeployConfig defaults.
default_deploy_config_name: str | None = None
+ # Global CLI spelling -> (stage id, stage-local spelling).
+ stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ errors = self.get_validation_errors()
+
+ if errors:
+ # If we have multiple errors, put them on separate indented lines for readability
+ multi_sep = "\n\t"
+ error_str = multi_sep.join(errors)
+ if len(errors) > 1:
+ error_str = f"{multi_sep}{error_str}"
+ raise ValueError(f"PipelineConfig initialization failed with the following error(s): {error_str}")
def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
"""Look up a stage by its ID."""
@@ -296,7 +320,7 @@ class PipelineConfig:
return stage
return None
- def validate(self) -> list[str]:
+ def get_validation_errors(self) -> list[str]:
"""Return list of topology errors (empty if valid)."""
errors: list[str] = []
if not self.stages:
@@ -314,6 +338,11 @@ class PipelineConfig:
errors.append(f"Stage {stage.stage_id} references itself")
if not any(not s.input_sources for s in self.stages):
errors.append("No entry point (stage with empty input_sources)")
+ # Request completion and client-visible output are both driven by stages
+ # that declare ``final_output`` (see ``Orchestrator._route_output``). A
+ # pipeline without one can never emit a result, so every request hangs.
+ if not any(s.final_output for s in self.stages):
+ errors.append("No terminal stage (stage with final_output=True)")
return errors
@@ -342,6 +371,7 @@ class StageDeployConfig:
default_sampling_params: dict[str, Any] | None = None
default_pooling_params: dict[str, Any] | None = None
subtalker_sampling_params: dict[str, Any] | None = None
+ silence_ban_frames: int = 0
# === Generic stage engine fields ===
# Parallelism, scheduler, and memory-capacity controls.
@@ -375,6 +405,7 @@ class StageDeployConfig:
# Diffusion parallel_config deploy/runtime override fields.
ulysses_degree: int | None = None
ulysses_mode: str | None = None
+ ulysses_a2a_permute: bool | None = None
ring_degree: int | None = None
allgather_degree: int | None = None
sequence_parallel_size: int | None = None
@@ -396,6 +427,7 @@ class StageDeployConfig:
diffusers_call_kwargs: dict[str, Any] | None = None
diffusion_quantization_config: str | None = None
diffusion_attention_backend: str | None = None
+ fastvideo_vsa_topk: int | None = None
diffusion_attention_config: dict[str, Any] | None = None
# Diffusion execution, cache, and VAE behavior.
@@ -424,6 +456,9 @@ class StageDeployConfig:
enable_distributed_layerwise_offload: bool | None = None
dlo_use_allgather: bool | None = None
dlo_resident_layers: int | None = None
+ host_weight_runtime_mode: str | None = None
+ host_weight_runtime_root: str | None = None
+ dlo_host_registration_limit_gib: float | None = None
# Diffusion-specific debug and observability knobs.
enable_diffusion_pipeline_profiler: bool | None = None
@@ -573,7 +608,13 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
_DEEP_MERGE_KEYS = frozenset(
- {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"}
+ {
+ "default_sampling_params",
+ "default_pooling_params",
+ "subtalker_sampling_params",
+ "engine_extras",
+ "engine_args",
+ }
)
@@ -673,6 +714,11 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]:
def load_deploy_config(path: str | Path) -> DeployConfig:
"""Load a deploy YAML (with optional base_config inheritance)."""
raw_dict = resolve_deploy_yaml(path)
+ if "stage_args" in raw_dict:
+ raise ValueError(
+ f"Deploy config {path} uses the removed `stage_args` schema; "
+ "define topology in PipelineConfig and deployment overrides under `stages`."
+ )
stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])]
@@ -855,6 +901,9 @@ def _build_engine_args(
engine_args["model_subdir"] = ps.model_subdir
if ps.tokenizer_subdir:
engine_args["tokenizer_subdir"] = ps.tokenizer_subdir
+ if ps.model_path_resolver:
+ engine_args["model_path_resolver"] = ps.model_path_resolver
+ engine_args["inline_diffusion"] = ps.inline_diffusion
# Pipeline-wide top-level DeployConfig settings, applied to every stage.
for name in _PIPELINE_WIDE_ENGINE_FIELDS:
@@ -872,6 +921,7 @@ def _build_engine_args(
# Materialize the resolved pipeline-wide async_chunk value into every
# stage so explicit False overrides do not get lost downstream.
engine_args["async_chunk"] = bool(deploy.async_chunk)
+ engine_args["session_mode"] = deploy.session_mode
if deploy.session_mode == "duplex":
# The engine admission limit is also the authoritative capacity for
# model-owned streaming state. Propagate it to every stage instead of
@@ -879,6 +929,7 @@ def _build_engine_args(
engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions
if ps.omni_kv_config:
engine_args["omni_kv_config"] = dict(ps.omni_kv_config)
+ engine_args["requires_full_payload_input"] = ps.requires_full_payload_input
return engine_args
@@ -900,6 +951,8 @@ def _build_extras(
extras["output_connectors"] = dict(ds.output_connectors)
if ds is not None and ds.input_connectors:
extras["input_connectors"] = dict(ds.input_connectors)
+ if ps.prompt_transform_func:
+ extras["prompt_transform_func"] = ps.prompt_transform_func
if ps.prompt_expand_func:
extras["prompt_expand_func"] = ps.prompt_expand_func
if ps.cfg_kv_collect_func:
@@ -988,6 +1041,7 @@ def merge_pipeline_deploy(
scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls),
hf_config_name=ps.hf_config_name,
is_comprehension=ps.owns_tokenizer,
+ sampling_constraints=dict(ps.sampling_constraints),
yaml_engine_args=engine_args,
yaml_runtime=runtime,
yaml_extras=extras,
@@ -1015,6 +1069,7 @@ class StageConfig:
scheduler_cls: str | None = None
hf_config_name: str | None = None
is_comprehension: bool = False
+ sampling_constraints: dict[str, Any] = field(default_factory=dict)
yaml_engine_args: dict[str, Any] = field(default_factory=dict)
yaml_runtime: dict[str, Any] = field(default_factory=dict)
yaml_extras: dict[str, Any] = field(default_factory=dict)
@@ -1075,6 +1130,7 @@ class StageConfig:
"final_output": self.final_output,
"final_output_type": self.final_output_type,
"is_comprehension": self.is_comprehension,
+ "sampling_constraints": dict(self.sampling_constraints),
}
if self.custom_process_input_func:
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 66da231a6..ad7baf764 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -1,7 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
+import time
from collections import defaultdict
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
from typing import Any
import numpy as np
@@ -104,6 +108,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# Cache per-request flag to avoid repeated deserialization of additional_information
self._omits_kv_transfer_cache: dict[str, bool] = {}
+
+ # KV-wait start ts for the vllm_omni:kv_wait_s metric; see
+ # _emit_kv_wait_output for the engine-core → orchestrator carry.
+ self._kv_wait_start_ts: dict[str, float] = {}
self._init_omni_io_scheduling_state()
# Snapshot prompt length for each streaming input update
self._new_prompt_len_snapshot: dict[str, int] = {}
@@ -113,6 +121,70 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# Output placeholders are zero when async scheduling isn't used
return request.num_computed_tokens - request.num_output_placeholders
+ def _resolve_kv_connector_type(self) -> str:
+ """Connector backend name for the ``kv_wait_s`` label, or ``unknown``."""
+ connector_cfg = self._get_omni_kv_config_value("connector_config")
+ if isinstance(connector_cfg, dict):
+ c_type = connector_cfg.get("type")
+ if isinstance(c_type, str) and c_type:
+ return c_type
+ return "unknown"
+
+ def _emit_kv_wait_output(
+ self,
+ outputs: dict[int, list[OmniEngineCoreOutput]],
+ req_id: str,
+ req: Request,
+ ) -> None:
+ """Carry the KV-wait duration for ``req_id`` across the process boundary.
+
+ Pops the start ts recorded at ENTER (``_free_request``); skips silently
+ when none was recorded. The wait rides ``kv_transfer_params`` to the
+ orchestrator, which calls ``observe_kv_wait``.
+ """
+ start_ts = self._kv_wait_start_ts.pop(req_id, None)
+ if start_ts is None:
+ return
+ kv_wait_s = max(time.monotonic() - start_ts, 0.0)
+ outputs.setdefault(req.client_index, []).append(
+ OmniEngineCoreOutput(
+ request_id=req_id,
+ new_token_ids=[],
+ kv_transfer_params={
+ "kv_wait_s": kv_wait_s,
+ "connector_type": self._resolve_kv_connector_type(),
+ },
+ )
+ )
+
+ def _clear_kv_wait_starts(self, request_ids: Iterable[str]) -> None:
+ """Drop incomplete KV-wait measurements for terminal requests."""
+ wait_starts = getattr(self, "_kv_wait_start_ts", None)
+ if wait_starts is None:
+ return
+ for request_id in request_ids:
+ wait_starts.pop(request_id, None)
+
+ def finish_requests(
+ self,
+ request_ids: str | Iterable[str] | None,
+ finished_status: RequestStatus,
+ ) -> list[Request]:
+ """Finish requests and discard any incomplete KV-wait timing."""
+ if isinstance(request_ids, str):
+ cleanup_ids = (request_ids,)
+ finish_request_ids: str | tuple[str, ...] | None = request_ids
+ elif request_ids is None:
+ cleanup_ids = ()
+ finish_request_ids = None
+ else:
+ finish_request_ids = list(request_ids) if isinstance(request_ids, Iterator) else request_ids
+ cleanup_ids = tuple(finish_request_ids)
+
+ finished = super().finish_requests(finish_request_ids, finished_status)
+ self._clear_kv_wait_starts(cleanup_ids)
+ return finished
+
def _get_kv_transfer_criteria(self) -> dict | None:
return self._get_omni_kv_config_value("kv_transfer_criteria")
@@ -223,15 +295,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# them. Upstream vllm raises RuntimeError on this status; omni allows
# async abort (e.g. client disconnect during TTS streaming) to leave
# requests in the waiting/running queues temporarily.
- for queue in (self.waiting, self.running):
- for req in list(queue):
- if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
- queue.remove(req)
+ waiting = getattr(self, "waiting")
+ self._drop_aborted_queued_requests()
self._process_pending_omni_inputs(model_mode="ar")
+ self._drop_aborted_queued_requests()
+ self._resync_streaming_input_counter()
original_waiting = None
if self._should_defer_waiting_admission():
- original_waiting = self.waiting
+ original_waiting = waiting
self.waiting = create_request_queue(self.policy)
try:
@@ -270,8 +342,25 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None)
num_nans_in_logits = model_runner_output.num_nans_in_logits
kv_connector_output = model_runner_output.kv_connector_output
+ ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
+ # Every GPU write enqueued by this and earlier steps has completed, so
+ # it is safe to return deferred-free blocks to the pool. This is the
+ # update-side half of the upstream v0.28 deferred-free fence; the schedule
+ # half advances sched_step_seq inside super().schedule().
+ # getattr: __new__-constructed test schedulers carry no
+ # defer_block_free attribute.
+ if (
+ getattr(self, "defer_block_free", False)
+ # getattr again: SchedulerOutput is mocked field-by-field in the
+ # scheduler unit tests, and dataclass fields without defaults are
+ # invisible to MagicMock(spec=...).
+ and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0
+ ):
+ self.processed_step_seq += 1
+ self._drain_deferred_frees()
+
perf_stats: PerfStats | None = None
if self.perf_metrics and self.perf_metrics.is_enabled():
perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
@@ -306,6 +395,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
kv_transfer_params={"kv_ready": True},
)
)
+ # EXIT: extraction ack ends the wait — carry the duration.
+ # req is still alive here (the del happens in the cleanup
+ # loop below); skip if it was already freed upstream.
+ if req is not None:
+ self._emit_kv_wait_output(outputs, req_id, req)
except Exception:
init_logger(__name__).exception("Failed to pre-process KV extraction for %s", req_id)
@@ -345,15 +439,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# request is aborted while the model is executing it (e.g.,
# in pipeline parallelism or async scheduling).
continue
- if output_is_stale:
+ req_index = model_runner_output.req_id_to_index[req_id]
+ generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+
+ stale_async_tokens = int(getattr(request, "async_tokens_to_discard", 0) or 0)
+ async_output_is_stale = bool(generated_token_ids and stale_async_tokens > 0)
+ if async_output_is_stale:
+ # Drain this marker even when the same frame also belongs to
+ # the scheduled-token stale window below. Both accounting
+ # domains must consume the old frame before new output passes.
+ request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids))
+
+ if output_is_stale or async_output_is_stale:
# Output of a step scheduled before the request's in-flight
# tokens were discarded (segment stop / session replacement).
# num_computed_tokens was rolled back at the discard site, so
# this output must not be appended or emitted.
continue
- req_index = model_runner_output.req_id_to_index[req_id]
- generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
status_before_stop = request.status
new_logprobs = None
logprob_validation_failed = False
@@ -412,6 +515,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
mm_output = mm_outputs[req_index] if mm_outputs else None
inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None
kv_transfer_params = None
+ ec_transfer_params = None
finish_reason = None
routed_experts = None
@@ -463,13 +567,33 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
request.resumable = False
stopped = True
+ # Finalize prefill stats BEFORE stop handling (upstream v0.28
+ # order): _free_request below releases the KV blocks, after which
+ # estimate_cached_tokens(request) reports 0. kv_transfer_params is
+ # omitted from the emission predicate here because it only becomes
+ # non-None when stopped is already True.
+ prefill_stats = None
+ if new_token_ids or mm_output is not None or pooler_output is not None or stopped:
+ prefill_stats = request.take_prefill_stats()
+ if prefill_stats is not None:
+ prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request))
+
+ confirmed_num_computed_tokens = None
+ boundary_generation = None
if stopped:
+ if self.chunk_transfer_adapter is not None:
+ confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request)
if model_runner_output.routed_experts is not None:
routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
# Capture finish_reason BEFORE _handle_stopped_request, which may
# reset the status to WAITING for streaming requests that continue.
finish_reason = request.get_finished_reason()
+ if self.chunk_transfer_adapter is not None:
+ try:
+ boundary_generation = int(getattr(request, "_omni_segment_generation", 0) or 0)
+ except (TypeError, ValueError):
+ boundary_generation = 0
finished = self._handle_stopped_request(request)
is_segment_finished = not finished
if finished:
@@ -478,9 +602,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# for streaming input request only
if self.chunk_transfer_adapter:
if self.vllm_config.model_config.stage_id != 0:
- # Downstream async-chunk stages receive real payloads from the
- # connector. This update only resumes polling for the next segment.
- self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
+ # Only a connector-fed receiver in native duplex can
+ # poll the next segment without an external update.
+ # Sender-only and turn-mode stages remain parked.
+ self._resume_downstream_chunk_receiver(request)
outstanding_async_tokens = request.num_output_placeholders
# Always record the discard signal (0 when nothing is in
# flight). Upstream a0c092ee72 removed the
@@ -499,8 +624,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# valid new-segment frames or underflowed the drain
# assert whenever placeholder counts diverged from
# scheduled counts (spec drafts, in-flight prefill).
+ # Assign, never accumulate: _handle_stopped_request above
+ # may already have fenced the same in-flight tokens through
+ # a queued streaming update, and adding twice swallows the
+ # next duplex unit's listen/speak under async scheduling.
if request.num_in_flight_tokens > 0:
- request.num_stale_output_tokens += request.num_in_flight_tokens
+ request.num_stale_output_tokens = request.num_in_flight_tokens
if outstanding_async_tokens > 0:
# Discard only outputs that are already in flight and
# roll back their optimistic computed-token accounting.
@@ -509,7 +638,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
request.spec_token_ids = []
request._output_token_ids.clear()
if finished:
- kv_transfer_params, _ = self._free_request(request)
+ kv_transfer_params, ec_transfer_params = self._free_request(request)
if status_before_stop == RequestStatus.RUNNING:
stopped_running_reqs.add(request)
elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:
@@ -537,8 +666,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
pooling_output=pooling_output_payload,
multimodal_output=mm_output,
stop_reason=request.stop_reason,
- prefill_stats=request.take_prefill_stats(),
+ prefill_stats=prefill_stats,
kv_transfer_params=kv_transfer_params,
+ ec_transfer_params=ec_transfer_params,
routed_experts=routed_experts,
num_nans_in_logits=request.num_nans_in_logits,
is_segment_finished=is_segment_finished,
@@ -551,10 +681,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if self.chunk_transfer_adapter is not None and (
inter_stage_output is not None or is_segment_finished or finished
):
+ save_kwargs = {
+ "new_token_ids": new_token_ids,
+ "confirmed_num_computed_tokens": confirmed_num_computed_tokens,
+ }
+ if is_segment_finished:
+ save_kwargs["segment_generation"] = boundary_generation
self.chunk_transfer_adapter.save_async(
inter_stage_output,
request,
is_segment_finished,
+ **save_kwargs,
)
self._remove_stopped_requests_from_queues(
@@ -576,6 +713,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if kv_connector_output:
self._update_from_kv_xfer_finished(kv_connector_output)
+ # EC Connector: update state from worker-side EC connector output.
+ # Use getattr for safety with test __new__/SimpleNamespace code paths.
+ if getattr(self, "ec_connector", None) is not None and ec_connector_output:
+ self.ec_connector.update_connector_output(ec_connector_output)
+
kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
self._publish_kv_cache_events()
@@ -635,10 +777,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# pre-replacement frame will drain, so the counter reaches exactly
# zero and the new segment's frames are never swallowed. This also
# covers an in-flight prefill chunk, which carries no placeholders
- # but must still have its late output dropped.
+ # but must still have its late output dropped. Assign, never
+ # accumulate: update_from_output fences the same in-flight tokens when
+ # a resumable stop applies a queued update through this helper.
in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
if in_flight_tokens > 0:
- session.num_stale_output_tokens += in_flight_tokens
+ session.num_stale_output_tokens = in_flight_tokens
if outstanding_async_tokens > 0:
# Async scheduling may already have sampled the previous
# segment's next token. Drop that late token instead of
@@ -648,9 +792,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
session.num_output_placeholders = 0
session.spec_token_ids = []
stage_id = self.vllm_config.model_config.stage_id
+
+ update_infos = (
+ getattr(update, "model_intermediate_buffer", None),
+ getattr(update, "additional_information", None),
+ )
if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks:
self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
if stage_id != 0:
+ session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
# Downstream async-chunk stages receive real payloads from the
# connector. This update only resumes polling for the next segment.
self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
@@ -668,10 +818,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if self.log_stats:
session.record_event(EngineCoreEventType.QUEUED)
return
- update_infos = (
- getattr(update, "model_intermediate_buffer", None),
- getattr(update, "additional_information", None),
- )
replace_streaming_prompt = any(
isinstance(info, dict)
and isinstance(info.get("meta"), dict)
@@ -679,8 +825,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
for info in update_infos
)
if replace_streaming_prompt:
+ self._release_replaced_streaming_prompt_cache(session)
self._replace_streaming_session(session, update)
return
+ session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
super()._update_request_as_session(session, update)
if hasattr(update, "model_intermediate_buffer"):
session.model_intermediate_buffer = update.model_intermediate_buffer
@@ -702,6 +850,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# 1. Standard cleanup parts from base _free_request
connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)
+ # EC Connector: mirror the KV hook (upstream v0.28 _free_request).
+ # The contract requires firing before the encoder cache is freed so
+ # the connector can inspect per-request state (e.g. which mm_hashes
+ # it recorded during save_caches()) and emit ec_transfer_params for
+ # the response body. getattr: __new__-constructed test schedulers
+ # carry no ec_connector attribute.
+ ec_xfer_params: dict[str, Any] | None = None
+ if getattr(self, "ec_connector", None) is not None:
+ ec_delay_free, ec_xfer_params = self.ec_connector.request_finished(request)
+ connector_delay_free_blocks |= ec_delay_free
+
self.encoder_cache_manager.free(request)
request_id = request.request_id
self.finished_req_ids.add(request_id)
@@ -725,11 +884,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
# It triggered but hasn't finished yet. We MUST wait.
logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
self.waiting_for_transfer_free.add(request_id)
+ self._kv_wait_start_ts[request_id] = time.monotonic()
kv_xfer_params = None
- return kv_xfer_params, None
+ return kv_xfer_params, ec_xfer_params
elif request_id in self.waiting_for_transfer_free:
# Blocks held until KV extraction completes in a future step.
- return None, None
+ return None, ec_xfer_params
else:
logger.debug(
f"[Omni] Request {request_id} finished and transfer no longer ACTIVE (extracted/acked). "
@@ -737,6 +897,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
)
else:
self.waiting_for_transfer_free.add(request_id)
+ self._kv_wait_start_ts[request_id] = time.monotonic()
confirmed_computed = self._get_confirmed_num_computed_tokens(request)
self._mark_request_for_kv_transfer(request_id, confirmed_computed)
# Return KV transfer metadata so it propagates to RequestOutput
@@ -766,14 +927,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if isinstance(add_info, dict):
add_info.update(kv_xfer_params)
- return kv_xfer_params, None
+ return kv_xfer_params, ec_xfer_params
# 3. Standard Freeing
delay_free_blocks |= connector_delay_free_blocks
if not delay_free_blocks:
self._free_blocks(request)
- return kv_xfer_params, None
+ return kv_xfer_params, ec_xfer_params
finally:
self._free_input_coordinator_request(request_id)
# Normal completion runs through here, not finish_requests()
diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py
index ae19a41a3..db0eedcf7 100644
--- a/vllm_omni/core/sched/omni_generation_scheduler.py
+++ b/vllm_omni/core/sched/omni_generation_scheduler.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
import time
@@ -98,7 +101,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
# Temporary queue: preserve waiting order while requests await input.
skipped_waiting_requests = create_request_queue(self.policy)
req_index = 0
+ self._drop_aborted_queued_requests()
self._process_pending_omni_inputs(model_mode="generation")
+ self._drop_aborted_queued_requests()
+ self._resync_streaming_input_counter()
# OMNI: Track requests that are already finished (e.g., marked by connector)
# These should be removed from running and not scheduled
@@ -309,6 +315,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
ec_meta = self.ec_connector.build_connector_meta(scheduler_output)
scheduler_output.ec_connector_metadata = ec_meta
+ # Advance the fence only for non-empty steps (those that actually
+ # write KV and have their output processed later in
+ # update_from_output). Must precede _update_after_schedule, which
+ # stamps request.last_sched_seq from the advanced value; the other
+ # half drains in update_from_output. The fallback path above gets
+ # both halves from super().schedule(). getattr: __new__-constructed
+ # test schedulers carry no defer_block_free attribute.
+ if getattr(self, "defer_block_free", False) and total_num_scheduled_tokens > 0:
+ self.sched_step_seq += 1
+
# Update internal state (advance num_computed_tokens, free encoder inputs,
# etc.)
self._update_after_schedule(scheduler_output)
@@ -345,8 +361,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
num_nans_in_logits = model_runner_output.num_nans_in_logits
kv_connector_output = model_runner_output.kv_connector_output
+ ec_connector_output = getattr(model_runner_output, "ec_connector_output", None)
cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
+
+ # Every GPU write enqueued by this and earlier steps has completed, so
+ # it is safe to return deferred-free blocks to the pool. This is the
+ # update-side half of the upstream v0.28 deferred-free fence; the schedule
+ # half advances sched_step_seq in the fast path above.
+ # getattr: __new__-constructed test schedulers carry no
+ # defer_block_free attribute.
+ if (
+ getattr(self, "defer_block_free", False)
+ # getattr again: SchedulerOutput is mocked field-by-field in the
+ # scheduler unit tests, and dataclass fields without defaults are
+ # invisible to MagicMock(spec=...).
+ and getattr(scheduler_output, "total_num_scheduled_tokens", 0) > 0
+ ):
+ self.processed_step_seq += 1
+ self._drain_deferred_frees()
+
perf_stats: PerfStats | None = None
if self.perf_metrics and self.perf_metrics.is_enabled():
perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
@@ -438,6 +472,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
new_logprobs = None
new_token_ids = generated_token_ids
kv_transfer_params = None
+ ec_transfer_params = None
pooler_output = pooler_outputs[req_index] if pooler_outputs else None
mm_output = mm_outputs[req_index] if mm_outputs else None
status_before_stop = request.status
@@ -474,6 +509,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
# (does not affect protocol)
stopped = True
+ # Finalize prefill stats BEFORE stop handling (upstream v0.28
+ # order): _free_request below releases the KV blocks, after which
+ # estimate_cached_tokens(request) reports 0. kv_transfer_params is
+ # omitted from the emission predicate here because it only becomes
+ # non-None when stopped is already True.
+ prefill_stats = None
+ if new_token_ids or mm_output is not None or pooler_output is not None or stopped:
+ prefill_stats = request.take_prefill_stats()
+ if prefill_stats is not None:
+ prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request))
+
if stopped:
if model_runner_output.routed_experts is not None:
routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
@@ -485,7 +531,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
if self.chunk_transfer_adapter:
self.chunk_transfer_adapter.segment_finished_requests.discard(req_id)
if finished:
- kv_transfer_params, _ = self._free_request(request)
+ kv_transfer_params, ec_transfer_params = self._free_request(request)
if self.chunk_transfer_adapter is not None:
self.chunk_transfer_adapter.cleanup(
request.request_id,
@@ -527,8 +573,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
pooling_output=pooling_output_payload,
multimodal_output=mm_output,
stop_reason=request.stop_reason,
- prefill_stats=request.take_prefill_stats(),
+ prefill_stats=prefill_stats,
kv_transfer_params=kv_transfer_params,
+ ec_transfer_params=ec_transfer_params,
routed_experts=routed_experts,
num_nans_in_logits=request.num_nans_in_logits,
is_segment_finished=is_segment_finished,
@@ -547,8 +594,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
finished = self._handle_stopped_request(request)
is_segment_finished = not finished
kv_transfer_params = None
+ ec_transfer_params = None
if finished:
- kv_transfer_params, _ = self._free_request(request)
+ kv_transfer_params, ec_transfer_params = self._free_request(request)
if self.chunk_transfer_adapter is not None:
self.chunk_transfer_adapter.cleanup(
request.request_id,
@@ -562,6 +610,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
finish_reason=finish_reason,
stop_reason=request.stop_reason,
kv_transfer_params=kv_transfer_params,
+ ec_transfer_params=ec_transfer_params,
is_segment_finished=is_segment_finished,
)
stopped_running_reqs.add(request)
@@ -587,6 +636,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
if kv_connector_output:
self._update_from_kv_xfer_finished(kv_connector_output)
+ # EC Connector: update state from worker-side EC connector output.
+ # Use getattr for safety with test __new__/SimpleNamespace code paths.
+ if getattr(self, "ec_connector", None) is not None and ec_connector_output:
+ self.ec_connector.update_connector_output(ec_connector_output)
+
kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
self._publish_kv_cache_events()
diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py
index e0df774c0..58c3546cd 100644
--- a/vllm_omni/core/sched/omni_scheduler_mixin.py
+++ b/vllm_omni/core/sched/omni_scheduler_mixin.py
@@ -1,8 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
from __future__ import annotations
+import math
import os
import time
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
from typing import Any
import torch
@@ -42,25 +46,56 @@ logger = init_logger(__name__)
_STATS_INTERVAL_S = 1.0
-# Upper bound on how long a request may sit in full-payload-input wait
-# (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
-# before the scheduler force-fails it. Defends against stuck consumer-side
-# requests when the producer drops a full-payload, send fails, or recv
-# never arrives. Override per-deployment via
-# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net.
+# Upper bound on how long a request may wait for stage input before the
+# scheduler force-fails it. Defends against stuck consumer-side requests when
+# the producer drops a payload, send fails, or recv never arrives. Override
+# per-deployment via VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set exactly 0 to disable
+# the safety net. Negative and non-finite values are rejected at startup.
#
-# Scope: this constant only covers the full-payload coordinator path
-# (``input_coordinator``). The async-chunk path uses
-# ``chunk_transfer_adapter`` and is not affected by this constant.
+# Scope: both transports. The full-payload path measures time parked in
+# WAITING_FOR_INPUT (``OmniSchedulingCoordinator._waiting_since``); the
+# async-chunk path measures time stalled in WAITING_FOR_CHUNK
+# (``OmniChunkTransferAdapter._waiting_since``, reset on each chunk arrival).
+# One knob, because the question it answers -- "how long may a request wait for
+# stage input?" -- does not depend on which transport carries it.
+_INPUT_WAIT_TIMEOUT_DEFAULT = 600.0
_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")
try:
DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
except ValueError:
logger.warning(
- "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.",
+ "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to %s seconds.",
_INPUT_WAIT_TIMEOUT_RAW,
+ _INPUT_WAIT_TIMEOUT_DEFAULT,
+ )
+ DEFAULT_INPUT_WAIT_TIMEOUT_S = _INPUT_WAIT_TIMEOUT_DEFAULT
+
+if not math.isfinite(DEFAULT_INPUT_WAIT_TIMEOUT_S):
+ # nan compares false against everything, so every deadline check silently
+ # passes; inf makes the deadline unreachable. Both disable the net while
+ # looking like a number, which is the failure this item exists to stop.
+ raise ValueError(
+ f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is not finite. "
+ "Use a positive number of seconds, or exactly 0 to disable the deadline."
+ )
+if DEFAULT_INPUT_WAIT_TIMEOUT_S < 0:
+ # Fail at startup rather than substituting a default. `-1` is a common idiom
+ # for "no limit", so silently turning it into 600 would give an operator who
+ # meant "never time out" mysterious failures ten minutes in.
+ raise ValueError(
+ f"VLLM_OMNI_INPUT_WAIT_TIMEOUT_S={_INPUT_WAIT_TIMEOUT_RAW!r} is negative. "
+ "Use a positive number of seconds, or exactly 0 to disable the deadline."
+ )
+elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0:
+ # Zero stays a supported opt-out, but it is no longer silent: it disables
+ # the deadline on BOTH transports, so a request whose stage input never
+ # arrives waits forever. Say which nets just went away.
+ logger.warning(
+ "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=0: stage-input deadlines are DISABLED for "
+ "both the full-payload path (WAITING_FOR_INPUT) and the async-chunk path "
+ "(WAITING_FOR_CHUNK). A request whose input never arrives will wait "
+ "indefinitely instead of failing."
)
- DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0
class OmniSchedulerMixin:
@@ -114,9 +149,13 @@ class OmniSchedulerMixin:
def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None:
"""Replace a downstream stage's placeholder with its next payload."""
+ session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1
adapter = getattr(self, "chunk_transfer_adapter", None)
if adapter is not None:
adapter.segment_finished_requests.discard(session.request_id)
+ watermark = getattr(adapter, "requests_num_chunks_sent", None)
+ if watermark is not None:
+ watermark.pop(session.external_req_id, None)
session._output_token_ids.clear()
session._all_token_ids.clear()
# In-flight outputs from the previous segment were optimistically
@@ -127,7 +166,10 @@ class OmniSchedulerMixin:
# pre-replacement frame will drain, so the counter reaches exactly
# zero; a placeholder-based seed swallowed valid new-segment frames
# whenever placeholder counts diverged from scheduled counts.
- session.num_stale_output_tokens += int(getattr(session, "num_in_flight_tokens", 0) or 0)
+ # num_in_flight_tokens already includes any undrained stale share.
+ # Assign instead of accumulating so callers that fenced the same
+ # rollover before entering this helper do not count it twice.
+ session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
session.num_output_placeholders = 0
session.spec_token_ids = []
new_prompt = update.prompt_token_ids or ()
@@ -153,6 +195,44 @@ class OmniSchedulerMixin:
if self.log_stats:
session.record_event(EngineCoreEventType.QUEUED)
+ def _release_replaced_streaming_prompt_cache(self, session: Request) -> None:
+ """Discard cache state that belongs to a replaced prompt."""
+ # A prompt replacement is not a normal streaming extension: none of
+ # the old KV blocks or encoder state is valid for the new prompt. Use
+ # the scheduler's block-free path so an in-flight GPU step is fenced
+ # correctly before the blocks return to the pool.
+ self._free_request_blocks(session)
+ self.encoder_cache_manager.free(session)
+ getattr(self, "_inflight_prefills", set()).discard(session)
+
+ def _reset_ready_async_chunk_replacements(self) -> None:
+ """Release stale cache state after an async-chunk prompt rollover."""
+ adapter = getattr(self, "chunk_transfer_adapter", None)
+ if adapter is None:
+ return
+ replaced_ids = getattr(adapter, "replaced_streaming_prompt_ids", None)
+ ready_ids = getattr(adapter, "requests_with_ready_chunks", None)
+ if not replaced_ids or not ready_ids:
+ return
+
+ for request_id in tuple(replaced_ids & ready_ids):
+ request = self.requests.get(request_id)
+ if request is None:
+ replaced_ids.discard(request_id)
+ continue
+ # The streaming update may already have fenced this same in-flight
+ # frame. Seed idempotently so the replacement does not count it twice.
+ request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0)
+ request.num_output_placeholders = 0
+ request.spec_token_ids = []
+ self._release_replaced_streaming_prompt_cache(request)
+ watermark = getattr(adapter, "requests_num_chunks_sent", None)
+ if watermark is not None:
+ watermark.pop(request.external_req_id, None)
+ # Consume this marker after the one-time cache reset. The separate
+ # ready-chunk marker remains until scheduler admission succeeds.
+ replaced_ids.discard(request_id)
+
def _consume_pending_connector_output(self, model_mode: str) -> None:
"""Drain ``self._latest_omni_connector_output`` into the coordinator.
@@ -184,6 +264,29 @@ class OmniSchedulerMixin:
self.running,
scheduler_requests=self.requests,
)
+ # Prompt-window contracts are finalized on this scheduler thread
+ # and may surface a permanent receive failure.
+ self._process_chunk_receive_failures()
+ self._reset_ready_async_chunk_replacements()
+ self._process_pending_chunk_timeouts()
+ self._log_failed_chunk_sends()
+
+ def _process_chunk_receive_failures(self) -> None:
+ """Fail requests whose consumed async chunk violated its contract."""
+ adapter = getattr(self, "chunk_transfer_adapter", None)
+ collector = getattr(adapter, "collect_failed_receive_request_ids", None)
+ if collector is None:
+ return
+ failures = collector()
+ present_ids = {request_id for request_id in failures if request_id in self.requests}
+ if not present_ids:
+ return
+ logger.error(
+ "Marking %d request(s) as FINISHED_ERROR after invalid connector input: %s",
+ len(present_ids),
+ {request_id: failures[request_id] for request_id in sorted(present_ids)},
+ )
+ self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
def _restore_omni_wait_queues(self) -> None:
"""Restore requests temporarily parked by Omni input gates."""
@@ -211,9 +314,9 @@ class OmniSchedulerMixin:
Disabled when ``DEFAULT_INPUT_WAIT_TIMEOUT_S`` is <= 0.
Scope: only covers ``input_coordinator`` (full-payload path).
- Async-chunk requests park in ``chunk_transfer_adapter`` instead
- and are not handled here -- if a similar safety net is needed
- for the chunk path, it belongs in the chunk adapter.
+ Async-chunk requests park in ``chunk_transfer_adapter`` instead and
+ are handled by ``_process_pending_chunk_timeouts`` below, which shares
+ this timeout.
"""
if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
return
@@ -234,6 +337,84 @@ class OmniSchedulerMixin:
)
self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
+ def _log_failed_chunk_sends(self) -> None:
+ """Surface chunks the sender gave up on (R1.2 of #4855).
+
+ ``save_loop`` pops a task with ``popleft()`` and never re-queues it, and
+ ``connector.put`` can report failure without raising, so a dropped chunk
+ is a give-up rather than a retry. Until this change the only trace was a
+ warning that printed ``None`` for the request id, because the task dict
+ has no ``request_id`` key.
+
+ This only reports. Failing the request from here does not work: the
+ producer is not the ``final_output`` stage, so an abort synthesized by
+ ``finish_requests`` is not turned into a client-visible termination by
+ ``Orchestrator._route_output`` (see the ``final_output`` guards there).
+ Making a producer-side drop client-visible needs an orchestrator-side
+ path like ``_fail_request_dead_stage``; the consumer's R1.1 deadline is
+ what actually ends the request today.
+ """
+ adapter = getattr(self, "chunk_transfer_adapter", None)
+ if adapter is None:
+ return
+ failures = adapter.collect_failed_send_request_ids()
+ for request_id, reason in failures.items():
+ if request_id in self.requests:
+ logger.error(
+ "[OmniScheduler] req=%s: chunk send gave up (%s); the consumer will "
+ "fail on the stage-input deadline",
+ request_id,
+ reason,
+ )
+ else:
+ # Drained but no longer scheduled here: the request finished or
+ # was aborted between the failed send and this sweep. Still say
+ # so -- collect_* empties the map, so staying quiet would make
+ # this the one drop path in a change whose point is visibility.
+ logger.warning(
+ "[OmniScheduler] req=%s: chunk send gave up (%s) after the request "
+ "left this scheduler (finished or aborted)",
+ request_id,
+ reason,
+ )
+
+ def _process_pending_chunk_timeouts(self) -> None:
+ """Force-fail requests stalled in ``WAITING_FOR_CHUNK`` too long.
+
+ The async-chunk counterpart of ``_process_pending_input_timeouts``,
+ called right after ``chunk_transfer_adapter.process_pending_chunks``.
+ Until now the chunk path had no safety net at all: the scheduler-side
+ comment on the full-payload net said a similar guard "belongs in the
+ chunk adapter", while the worker mixin assumed the full-payload net
+ already covered it. Only the first was true, so a dropped terminal
+ chunk or a producer that exhausted its send retries parked the request
+ forever (vllm-project/vllm-omni#3833). Since ``async_chunk: true`` is
+ the default for the flagship TTS pipelines, that was the default
+ streaming path.
+
+ Shares ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` with the full-payload net:
+ one knob for "how long may a request wait for stage input", regardless
+ of which transport carries it. Disabled when the value is <= 0.
+ """
+ if DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0:
+ return
+ adapter = getattr(self, "chunk_transfer_adapter", None)
+ if adapter is None or not getattr(adapter, "receives_chunks", False):
+ return
+ timed_out_ids = adapter.collect_timed_out_request_ids(timeout_s=DEFAULT_INPUT_WAIT_TIMEOUT_S)
+ if not timed_out_ids:
+ return
+ present_ids = {req_id for req_id in timed_out_ids if req_id in self.requests}
+ if not present_ids:
+ return
+ logger.warning(
+ "Marking %d request(s) as FINISHED_ERROR after stalling > %.0fs waiting for a chunk: %s",
+ len(present_ids),
+ DEFAULT_INPUT_WAIT_TIMEOUT_S,
+ sorted(present_ids),
+ )
+ self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
+
def _capture_omni_connector_output(self, model_runner_output: Any) -> None:
"""Stash the model runner's omni_connector_output for next schedule().
@@ -315,6 +496,7 @@ class OmniSchedulerMixin:
stop_reason: Any = None,
prefill_stats: Any = None,
kv_transfer_params: Any = None,
+ ec_transfer_params: Any = None,
routed_experts: Any = None,
num_nans_in_logits: int = 0,
is_segment_finished: bool | None = False,
@@ -333,6 +515,7 @@ class OmniSchedulerMixin:
events=request.take_events(),
prefill_stats=prefill_stats,
kv_transfer_params=kv_transfer_params,
+ ec_transfer_params=ec_transfer_params,
trace_headers=request.trace_headers,
routed_experts=routed_experts,
num_nans_in_logits=num_nans_in_logits,
@@ -407,8 +590,41 @@ class OmniSchedulerMixin:
if stopped_running_reqs:
self.running = remove_all(self.running, stopped_running_reqs)
if stopped_preempted_reqs:
- self.waiting.remove_requests(stopped_preempted_reqs)
- self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+ # ``_handle_stopped_request`` re-enqueues a resumable segment as
+ # WAITING_FOR_STREAMING_REQ before this cleanup runs. A
+ # downstream async-chunk request can have entered the update with
+ # the older WAITING_FOR_CHUNK status, which puts it in
+ # ``stopped_preempted_reqs``; deleting it here would immediately
+ # undo the requeue and the receiver would never poll chunk 1.
+ removable = {
+ request
+ for request in stopped_preempted_reqs
+ if not (
+ request.resumable
+ and request.status
+ in (
+ RequestStatus.WAITING,
+ RequestStatus.WAITING_FOR_STREAMING_REQ,
+ )
+ )
+ }
+ self.waiting.remove_requests(removable)
+ self.skipped_waiting.remove_requests(removable)
+
+ def _resume_downstream_chunk_receiver(self, request: Request) -> None:
+ """Resume duplex connector polling without an external update."""
+ adapter = self.chunk_transfer_adapter
+ adapter.segment_finished_requests.discard(request.request_id)
+ if (
+ not adapter.receives_chunks
+ or getattr(self.vllm_config.model_config, "session_mode", "turn") != "duplex"
+ or request.status != RequestStatus.WAITING_FOR_STREAMING_REQ
+ ):
+ return
+ self.num_waiting_for_streaming_input -= 1
+ request.status = RequestStatus.WAITING
+ self.skipped_waiting.remove_requests((request,))
+ self._enqueue_waiting_request(request)
def _aggregate_kv_connector_stats(
self,
@@ -459,12 +675,36 @@ class OmniSchedulerMixin:
finished_status: RequestStatus,
) -> list[Request]:
"""Finish requests and clean all Omni-owned queue/coordinator state."""
- if self.chunk_transfer_adapter:
- self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
+ if isinstance(request_ids, str):
+ target_request_ids = {request_ids}
+ elif request_ids is None:
+ target_request_ids = set(self.requests)
+ else:
+ if isinstance(request_ids, Iterator):
+ request_ids = tuple(request_ids)
+ target_request_ids = set(request_ids)
- self._realign_request_status_to_queues(request_ids)
+ skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
+ pre_adapter_streaming_wait_ids = {
+ rid
+ for rid in target_request_ids & skipped_waiting_ids
+ if (req := self.requests.get(rid)) is not None
+ and (
+ req.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+ or (getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED)
+ )
+ }
+
+ if chunk_transfer_adapter := getattr(self, "chunk_transfer_adapter", None):
+ chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
+
+ self._realign_request_status_to_queues(
+ request_ids,
+ pre_adapter_streaming_wait_ids=pre_adapter_streaming_wait_ids,
+ )
finished = super().finish_requests(request_ids, finished_status)
- self._purge_finished_from_running()
+ self._purge_finished_from_running(target_request_ids)
+ self._resync_streaming_input_counter()
for request in finished:
self._free_input_coordinator_request(request.request_id)
@@ -480,6 +720,8 @@ class OmniSchedulerMixin:
def _realign_request_status_to_queues(
self,
request_ids: str | Iterable[str] | None,
+ *,
+ pre_adapter_streaming_wait_ids: set[str] | None = None,
) -> None:
"""Realign ``request.status`` to actual queue membership.
@@ -506,7 +748,13 @@ class OmniSchedulerMixin:
Realign here: if a request lives in ``self.running`` but its
status is not ``RUNNING``, set it to ``RUNNING``; symmetrically
flip ``RUNNING → WAITING`` when the request is actually in
- ``self.waiting``. This is a localized safety net for
+ ``self.waiting``. A resumable segment stop still held in
+ ``self.skipped_waiting`` is restored to
+ ``WAITING_FOR_STREAMING_REQ`` so upstream also balances its
+ paused-session counter. Because adapter cleanup may restore a
+ connector-owned request's prior status first, ``finish_requests``
+ snapshots these counter-bearing rows before invoking the adapter.
+ This is a localized safety net for
``requests_origin_status`` staleness on the admit transition;
it does not touch the adapter's invariants and is complementary
to the chunk-transfer-adapter deque purge that already runs
@@ -539,17 +787,36 @@ class OmniSchedulerMixin:
running_ids = {r.request_id for r in self.running}
waiting_ids = {r.request_id for r in self.waiting}
+ skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
for rid in ids_to_align:
req = self.requests.get(rid)
- if req is None or req.is_finished():
+ if req is None:
continue
- if rid in running_ids and req.status != RequestStatus.RUNNING:
+ # A persistent Session may be closed after its current segment
+ # reached FINISHED_STOPPED but while it is still owned by a live
+ # scheduler/connector queue. vLLM skips already-finished requests
+ # in finish_requests(), leaking the KV and worker slot. Only
+ # recover requests with positive queue ownership; an off-queue
+ # terminal can legitimately be waiting for deferred block free.
+ resumable_segment_stop = bool(
+ getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED
+ )
+ streaming_wait = (
+ rid in pre_adapter_streaming_wait_ids
+ if pre_adapter_streaming_wait_ids is not None
+ else resumable_segment_stop
+ )
+ if req.is_finished() and not (resumable_segment_stop or streaming_wait):
+ continue
+ if rid in skipped_waiting_ids and streaming_wait:
+ req.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+ elif rid in running_ids and req.status != RequestStatus.RUNNING:
req.status = RequestStatus.RUNNING
- elif rid in waiting_ids and req.status == RequestStatus.RUNNING:
+ elif rid in waiting_ids and (req.status == RequestStatus.RUNNING or resumable_segment_stop):
req.status = RequestStatus.WAITING
- def _purge_finished_from_running(self) -> None:
+ def _purge_finished_from_running(self, target_request_ids: set[str] | None = None) -> None:
"""Defensive post-finish sweep of ``self.running``.
Belt-and-suspenders to ``_realign_request_status_to_queues``:
@@ -566,18 +833,14 @@ class OmniSchedulerMixin:
residue after ``super().finish_requests`` so any stale entries
are reclaimed).
- Scope of the predicate. ``is_finished()`` covers entries the
- upstream ``finish_requests`` already drained from ``self.requests``
- but failed to remove from ``self.running``; the
- ``request_id not in self.requests`` arm catches the same surface
- from a different angle and is the post-cleanup mirror of the
- deque purge ``_purge_untracked_chunk_requests`` already runs at
- the chunk-transfer-adapter layer. It does **not** by itself make
- arbitrary direct deletions of ``self.requests`` safe -- callers
- that pop ``self.requests`` outside the standard finish path
- still have to go through ``_free_request`` (or equivalent) for
- block / connector / coordinator cleanup. This sweep only
- reclaims the ``self.running`` slot reference.
+ A resumable ``FINISHED_STOPPED`` request may legitimately remain
+ in ``self.running`` between realtime segments. Preserve such a
+ request unless it belongs to this finish call. Other finished or
+ untracked entries are stale and can be swept defensively.
+
+ When ``target_request_ids`` is ``None`` or empty, every resumable
+ ``FINISHED_STOPPED`` entry is preserved. Production
+ ``finish_requests`` always passes its resolved finish set.
In-place via ``self.running[:] = ...`` for minor consistency
with idiomatic vLLM scheduler mutation; upstream
@@ -594,4 +857,38 @@ class OmniSchedulerMixin:
"""
if not self.running:
return
- self.running[:] = [req for req in self.running if not req.is_finished() and req.request_id in self.requests]
+ target_request_ids = target_request_ids or set()
+
+ def keep_running(req: Request) -> bool:
+ if req.request_id not in self.requests:
+ return False
+ if not req.is_finished():
+ return True
+ resumable_segment_stop = bool(
+ getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED
+ )
+ return resumable_segment_stop and req.request_id not in target_request_ids
+
+ self.running[:] = [req for req in self.running if keep_running(req)]
+
+ def _drop_aborted_queued_requests(self) -> None:
+ for queue in (self.waiting, self.skipped_waiting):
+ aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED]
+ if aborted:
+ queue.remove_requests(aborted)
+ self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
+
+ def _resync_streaming_input_counter(self) -> None:
+ counter = getattr(self, "num_waiting_for_streaming_input", None)
+ if counter is None:
+ return
+ parked = sum(
+ 1
+ for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None))
+ if queue
+ for request in queue
+ if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+ )
+ if parked != counter:
+ logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked)
+ self.num_waiting_for_streaming_input = parked
diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py
index 75f234d8c..93f40daa0 100644
--- a/vllm_omni/core/sched/omni_scheduling_coordinator.py
+++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py
@@ -22,75 +22,19 @@ from vllm_omni.core.sched.output import OmniChunkRecvHandle
logger = init_logger(__name__)
-# (arch, model_stage) pairs that route their full_payload stage input via
-# the worker connector and therefore need the scheduler-side coordinator to
-# park requests in WAITING_FOR_INPUT until the recv side delivers. This set
-# must stay aligned with the arch scope of `init_omni_connectors` in
-# gpu_ar_model_runner.py and gpu_generation_model_runner.py. Adding a stage
-# here without also wiring its worker connector init produces a permanent
-# Stage 1 hang (gate parks the request, no transport ever releases it).
-#
-_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
- {
- ("Qwen3OmniMoeForConditionalGeneration", "talker"),
- ("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
- # qwen2_5_omni thinker->talker uses the real full-payload
- # producer builder (text_hidden_states routed via
- # pooler_output["hidden"] -> accumulator -> connector). Both
- # stages of qwen2_5_omni are enabled.
- ("Qwen2_5OmniForConditionalGeneration", "talker"),
- ("Qwen2_5OmniForConditionalGeneration", "code2wav"),
- # covo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
- ("CovoAudioForConditionalGeneration", "code2wav"),
- # mimo_audio: fused_thinker_talker (Stage 0) -> code2wav (Stage 1).
- ("MiMoAudioModel", "code2wav"),
- # qwen3_tts: Qwen3TTSTalkerForConditionalGeneration (Stage 0)
- # -> Qwen3TTSCode2Wav (Stage 1). Stage 1 is the consumer.
- ("Qwen3TTSCode2Wav", "code2wav"),
- # minicpmo_4_5: Talker (Stage 1) -> Code2Wav (Stage 2).
- ("MiniCPMO45Code2Wav", "code2wav"),
- # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
- ("CosyVoice3Model", "cosyvoice3_code2wav"),
- # nemotron_voicechat: talker (Stage 1) -> code2wav (Stage 2). Stage 2
- # waits for the talker's full-payload code stacks; the thinker (Stage 0)
- # -> talker (Stage 1) hop is token-path only and must NOT be listed.
- ("NemotronVoiceChatCode2Wav", "code2wav"),
- # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
- # The default deploy is async_chunk; this covers async_chunk: false.
- ("AudexCode2Wav", "audex_code2wav"),
- # audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync
- # full-payload — CNN codec decoded over the full sequence).
- ("AudexXCodec1", "audex_xcodec"),
- # indextts2 / indextts2_5: talker (Stage 0) -> s2mel decoder
- # (Stage 1). Stage 1 consumes the complete mel/optional-latent payload.
- ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
- ("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
- # dynin: token2text (Stage 0) -> token2image (Stage 1) ->
- # token2audio (Stage 2). Producer wires via
- # custom_process_next_stage_input_func: *_full_payload in deploy yaml.
- ("DyninOmniForConditionalGeneration", "token2image"),
- ("DyninOmniForConditionalGeneration", "token2audio"),
- }
-)
-
-
def uses_full_payload_input_coordinator(model_config: Any) -> bool:
"""Returns True if this stage parks pending requests in
WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector.
- Gated by (model_arch, model_stage) — see _FULL_PAYLOAD_INPUT_STAGES for the
- rationale on why this is a whitelist instead of a marker-driven structural
- gate.
+ Gated by the topology-declared ``requires_full_payload_input`` capability on
+ downstream (stage_id > 0) stages, and only on the non-async-chunk path
+ (async-chunk stages are fed through the streamed connector instead).
"""
if getattr(model_config, "stage_id", 0) <= 0:
return False
if getattr(model_config, "async_chunk", False):
return False
- key = (
- getattr(model_config, "model_arch", None),
- getattr(model_config, "model_stage", None),
- )
- return key in _FULL_PAYLOAD_INPUT_STAGES
+ return bool(getattr(model_config, "requires_full_payload_input", False))
class OmniSchedulingCoordinator:
diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py
index db2cf15ea..2ce1eed70 100644
--- a/vllm_omni/data_entry_keys.py
+++ b/vllm_omni/data_entry_keys.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Structured payload types for inter-stage communication.
Categories under ``OmniPayload``:
@@ -12,12 +12,16 @@ Categories under ``OmniPayload``:
from __future__ import annotations
+from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, TypedDict
import msgspec
import numpy as np
import torch
+REQUEST_ARTIFACT_DIRS_KEY = "_omni_request_artifact_dirs"
+TRANSFORM_OWNED_META_KEYS = frozenset({"minimax_h3_prepared_reference_videos"})
+
if TYPE_CHECKING:
from vllm_omni.engine import AdditionalInformationEntry, AdditionalInformationPayload
@@ -55,6 +59,7 @@ class Ids(TypedDict, total=False):
output: list[int]
speech_token: list[int]
prior_image: list[int]
+ streaming_prompt_previous_codes: list[int]
class OmniPayloadMeta(TypedDict, total=False):
@@ -74,7 +79,11 @@ class OmniPayloadMeta(TypedDict, total=False):
override_keys: list[tuple[str, str]]
num_processed_tokens: int
next_stage_prompt_len: int
+ next_stage_generation_tokens: int
replace_streaming_prompt: bool
+ streaming_prompt_recompute: bool
+ streaming_condition_seq: int
+ replace_runtime_additional_information: bool
ar_width: int
eol_token_id: int
visual_token_start_id: int
@@ -91,6 +100,13 @@ class OmniPayloadMeta(TypedDict, total=False):
ref_context_included: bool
talker_prefill_offset: int
omni_final_stage_id: int
+ # Per-request audio seed. Stages that draw their own noise (flow-matching
+ # / diffusion decoders) need the producing stage's seed to stay
+ # reproducible, and the producing stage's SamplingParams do not travel
+ # with the payload.
+ audio_seed: int
+ token_role_ids: torch.Tensor
+ minimax_h3_prepared_reference_videos: str
class OmniPayload(TypedDict, total=False):
@@ -111,7 +127,9 @@ class OmniPayload(TypedDict, total=False):
# ── msgspec.Struct mirror of the TypedDicts (runtime-validated) ──
-class _StructBase(msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True):
+class _StructBase( # type: ignore[call-arg]
+ msgspec.Struct, omit_defaults=True, kw_only=True, forbid_unknown_fields=True
+):
pass
@@ -153,6 +171,7 @@ class IdsStruct(_StructBase):
output: list[int] | None = None
speech_token: list[int] | None = None
prior_image: list[int] | None = None
+ streaming_prompt_previous_codes: list[int] | None = None
class MetaStruct(_StructBase):
@@ -169,7 +188,11 @@ class MetaStruct(_StructBase):
override_keys: list[tuple[str, str]] | None = None
num_processed_tokens: int | None = None
next_stage_prompt_len: int | None = None
+ next_stage_generation_tokens: int | None = None
replace_streaming_prompt: bool | None = None
+ streaming_prompt_recompute: bool | None = None
+ streaming_condition_seq: int | None = None
+ replace_runtime_additional_information: bool | None = None
ar_width: int | None = None
eol_token_id: int | None = None
visual_token_start_id: int | None = None
@@ -189,6 +212,9 @@ class MetaStruct(_StructBase):
codec_left_context_frames: int | None = None
code_flat_numel: int | None = None
omni_final_stage_id: int | None = None
+ audio_seed: int | None = None
+ token_role_ids: torch.Tensor | None = None
+ minimax_h3_prepared_reference_videos: str | None = None
class OmniPayloadStruct(_StructBase):
@@ -302,7 +328,7 @@ def _dtype_to_name(dtype: torch.dtype) -> str:
_NESTED_KEYS = frozenset({"hidden_states", "embed", "ids", "codes", "meta"})
-def flatten_payload(payload: dict[str, Any]) -> dict[str, Any]:
+def flatten_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
"""Flatten a nested ``OmniPayload`` to dotted keys.
Nested sub-dicts under ``_NESTED_KEYS`` are expanded:
@@ -353,17 +379,25 @@ def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry:
t_cpu = t.detach().to("cpu").contiguous()
return AdditionalInformationEntry(
- tensor_data=t_cpu.numpy().tobytes(),
+ tensor_data=t_cpu.flatten().view(torch.uint8).numpy().tobytes(),
tensor_shape=list(t_cpu.shape),
tensor_dtype=_dtype_to_name(t_cpu.dtype),
)
def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor:
- dt = np.dtype(entry.tensor_dtype or "float32")
- arr = np.frombuffer(entry.tensor_data, dtype=dt) # type: ignore[arg-type]
- arr = arr.reshape(entry.tensor_shape)
- return torch.from_numpy(arr.copy())
+ dtype_name = entry.tensor_dtype or "float32"
+ dtype = getattr(torch, dtype_name, None)
+ if not isinstance(dtype, torch.dtype):
+ raise ValueError(f"Unsupported tensor dtype: {dtype_name}")
+
+ if entry.tensor_shape is None:
+ raise ValueError("Tensor shape is required")
+ if not entry.tensor_data:
+ return torch.empty(0, dtype=dtype).reshape(entry.tensor_shape)
+
+ data = torch.frombuffer(bytearray(entry.tensor_data), dtype=torch.uint8)
+ return data.view(dtype).reshape(entry.tensor_shape)
def serialize_payload(
diff --git a/vllm_omni/deploy/README_higgs_audio_v3.md b/vllm_omni/deploy/README_higgs_audio_v3.md
index ea1363069..44f7f75c7 100644
--- a/vllm_omni/deploy/README_higgs_audio_v3.md
+++ b/vllm_omni/deploy/README_higgs_audio_v3.md
@@ -44,6 +44,65 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
```
+## Why Stage 0 Pins Native FlashInfer
+
+The three Stage 0 profiles set `attention_backend: FLASHINFER` and also set:
+
+```yaml
+attention_config:
+ use_trtllm_attention: false
+```
+
+`FLASHINFER` names the overall vLLM attention backend. Starting with vLLM
+0.25, that backend can internally route Hopper decode to the TRT-LLM XQA
+kernel. The XQA eligibility check answers whether the kernel can run, not
+whether it is the fastest kernel for the current workload. In particular, the
+decode builder can select XQA statically when all of the following hold:
+
+- the explicit `use_trtllm_attention` override is not `false`;
+- the device is Hopper/SM90 and the XQA cubin is available; and
+- the query-head count is divisible by the KV-head count.
+
+Higgs has 32 query heads and 8 KV heads, so it passes the head-layout check.
+The static gate does not include a model dtype, KV-cache dtype, context-length,
+batch-shape, or measured-performance comparison. This matters because kernel
+support is not a performance guarantee.
+
+Higgs also has an unusual output shape that must not be confused with its
+attention shape. One autoregressive audio step advances one token per request.
+The resulting hidden row is then projected into eight codebook distributions
+with shape `[1, 8, codebook_size]`. The eight codebooks therefore do not turn
+the attention query into `q_len=8`; at concurrency one, XQA still receives a
+single BF16 query row. At this small per-step shape, any XQA advantage for
+other KV dtypes, batch sizes, or context lengths does not necessarily apply.
+The difference is paid in every transformer layer and every generated audio
+step, so even a small per-call disadvantage accumulates in TTFP and RTF.
+
+The upstream XQA integration also showed that the result is dtype-sensitive:
+its FP8 benchmark favored XQA, while its BF16 benchmark reported 21.16 ms mean
+TPOT for XQA versus 19.47 ms for FA3. That is not a direct native-FlashInfer
+comparison, but it demonstrates why SM90 compatibility alone is insufficient
+to choose the fastest BF16 kernel. See the
+[vLLM XQA benchmark](https://github.com/vllm-project/vllm/pull/43232#issuecomment-4501176335).
+
+The end-to-end H800 A/B in
+[issue #5584](https://github.com/vllm-project/vllm-omni/issues/5584#issuecomment-5188700218)
+isolated the routing decision on the same vLLM 0.25 stack:
+
+| Decode path | TTFP (ms) | RTF |
+|---|---:|---:|
+| XQA (automatic default) | 175.76 | 0.3405 |
+| Native FlashInfer (`use_trtllm_attention: false`) | 166.39 | 0.3233 |
+| vLLM 0.24 reference | 166.68 | 0.3209 |
+
+This proves that XQA selection causes the observed local regression and that
+native FlashInfer recovers the pre-v0.25 range. It does not by itself identify
+one defective XQA instruction or memory transaction. Establishing that deeper
+kernel-level cause requires an attention-only microbenchmark and GPU profiler
+trace for the exact Higgs batch/context distribution. Until vLLM has a
+workload-aware selector or XQA is faster for this BF16 shape, the model-level
+pin is the deterministic choice.
+
## Notes
- Stage 1 remains `enforce_eager: true` in both profiles.
diff --git a/vllm_omni/deploy/bagel_single_stage.yaml b/vllm_omni/deploy/bagel_single_stage.yaml
index bcfbad253..fb0872af8 100644
--- a/vllm_omni/deploy/bagel_single_stage.yaml
+++ b/vllm_omni/deploy/bagel_single_stage.yaml
@@ -9,7 +9,7 @@
# --deploy-config vllm_omni/deploy/bagel_single_stage.yaml
#
# Or programmatically:
-# Omni(model="...", deploy_config_path="vllm_omni/deploy/bagel_single_stage.yaml")
+# Omni(model="...", deploy_config="vllm_omni/deploy/bagel_single_stage.yaml")
pipeline: bagel_single_stage
async_chunk: false
diff --git a/vllm_omni/deploy/dots_tts.yaml b/vllm_omni/deploy/dots_tts.yaml
new file mode 100644
index 000000000..acd28de52
--- /dev/null
+++ b/vllm_omni/deploy/dots_tts.yaml
@@ -0,0 +1,28 @@
+# dots.tts deployment config (single-stage AR talker).
+#
+# Conservative first-integration settings: eager execution, no CUDA
+# Graphs, no batched side-path optimizations yet — see voxcpm2.yaml for
+# the runtime knobs a mature deployment of this model family grows.
+#
+# Verified locally on RTX 5080 16GB.
+
+async_chunk: false
+dtype: bfloat16
+
+stages:
+ - stage_id: 0
+ max_num_seqs: 4
+ gpu_memory_utilization: 0.8
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ async_scheduling: true
+ max_num_batched_tokens: 4096
+ max_model_len: 4096
+ devices: "0"
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 4096
+ repetition_penalty: 1.0
diff --git a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
index 360b72974..b6e5bc7fe 100644
--- a/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
+++ b/vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml
@@ -11,7 +11,7 @@
#
# Or programmatically:
# Omni(model="fishaudio/s2-pro",
-# deploy_config_path="vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml")
+# deploy_config="vllm_omni/deploy/fish_qwen3_omni_2gpu.yaml")
async_chunk: true
enable_chunked_prefill: false
diff --git a/vllm_omni/deploy/gepard.yaml b/vllm_omni/deploy/gepard.yaml
new file mode 100644
index 000000000..26000d616
--- /dev/null
+++ b/vllm_omni/deploy/gepard.yaml
@@ -0,0 +1,40 @@
+# Gepard-1.0: single-stage native-AR FSQ/NanoCodec TTS (~0.5B).
+# Target 1x L4 24GB. Zero-shot only; enforce_eager until the perf PR.
+#
+# The pipeline pin is required, not cosmetic: the checkpoint's config.json
+# self-identifies as qwen3_5_text, so model_type inference cannot reach the
+# gepard registry entry, and the architectures fallback would match
+# Qwen3_5ForCausalLM against the diffusion pipeline.
+pipeline: gepard
+async_chunk: false
+dtype: bfloat16
+trust_remote_code: true
+
+stages:
+ - stage_id: 0
+ max_num_seqs: 4
+ gpu_memory_utilization: 0.3 # AR LM + NanoCodec fit in ~2 GiB
+ enforce_eager: true
+ enable_prefix_caching: false
+ async_scheduling: true
+ max_num_batched_tokens: 4096
+ max_model_len: 4096
+ devices: "0"
+ # Override architectures so vLLM resolves our talker rather than the native
+ # Qwen3.5 text model. Do NOT override model_type: the backbone's own
+ # model_type check needs it to stay qwen3_5_text.
+ hf_overrides:
+ architectures: ["GepardTalkerForConditionalGeneration"]
+ default_sampling_params:
+ # All 32 heads are sampled in-model, so vLLM's sampler is pinned
+ # deterministic and just argmaxes what forward chose.
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ # One output token is one frame, so this is the frame budget: 1000 frames
+ # is ~46 s at 21.5 fps. A backstop, not a target -- normal utterances are
+ # 40-420 frames -- but a load-bearing one, because this checkpoint's stop
+ # head does not fire on every batched trajectory.
+ max_tokens: 1000
+ seed: 42
+ detokenize: false
diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml
index da299415b..41e1a9677 100644
--- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml
+++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml
@@ -38,6 +38,11 @@ stages:
gpu_memory_utilization: 0.6
enforce_eager: true
attention_backend: FLASHINFER
+ # Eight codebooks share one attention row, so decode remains BF16 q_len=1
+ # per request. vLLM's SM90 XQA capability gate does not benchmark this
+ # shape; keep the faster native path (README_higgs_audio_v3.md, #5584).
+ attention_config:
+ use_trtllm_attention: false
trust_remote_code: true
# Stage-0 prefix caching is on by default. The talker's HS opt-out and
# deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
@@ -78,3 +83,10 @@ stages:
max_tokens: 65536
seed: 42
repetition_penalty: 1.0
+
+platforms:
+ rocm:
+ stages:
+ - stage_id: 0
+ # FlashInfer is not supported on ROCm.
+ attention_backend: TRITON_ATTN
diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml
index 8a751483c..bb37bdb4c 100644
--- a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml
+++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml
@@ -37,6 +37,11 @@ stages:
gpu_memory_utilization: 0.6
enforce_eager: true
attention_backend: FLASHINFER
+ # Eight codebooks share one attention row, so decode remains BF16 q_len=1
+ # per request. vLLM's SM90 XQA capability gate does not benchmark this
+ # shape; keep the faster native path (README_higgs_audio_v3.md, #5584).
+ attention_config:
+ use_trtllm_attention: false
trust_remote_code: true
# Stage-0 prefix caching is on by default. The talker's HS opt-out and
# deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
index 0d6ae8218..13c6961e6 100644
--- a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
+++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml
@@ -43,6 +43,11 @@ stages:
gpu_memory_utilization: 0.6
enforce_eager: false
attention_backend: FLASHINFER
+ # Eight codebooks share one attention row, so decode remains BF16 q_len=1
+ # per request. vLLM's SM90 XQA capability gate does not benchmark this
+ # shape; keep the faster native path (README_higgs_audio_v3.md, #5584).
+ attention_config:
+ use_trtllm_attention: false
trust_remote_code: true
enable_prefix_caching: true
async_scheduling: false
diff --git a/vllm_omni/deploy/mimo_audio_5090d.yaml b/vllm_omni/deploy/mimo_audio_5090d.yaml
new file mode 100644
index 000000000..5e6b7b035
--- /dev/null
+++ b/vllm_omni/deploy/mimo_audio_5090d.yaml
@@ -0,0 +1,70 @@
+# MiMo-Audio deploy for 1x RTX 5090 / 5090D 32GB (Blackwell, SM 12.x).
+#
+# Verified offline TTS (`tts_sft`) with async-chunk streaming on a single GPU.
+# Stage 0 uses TRITON_ATTN to avoid FlashAttention-2 PTX/toolchain issues on
+# Blackwell. Stage 1 keeps a small VRAM budget; pair with
+# MIMO_AUDIO_TOKENIZER_DEVICE=cpu in the recipe when tokenizer VRAM is tight.
+#
+# Default async-chunk layout matches vllm_omni/deploy/mimo_audio.yaml (both
+# stages on device 0). For the legacy 2-GPU sync pipeline, see the header
+# comment in mimo_audio.yaml.
+async_chunk: true
+dtype: bfloat16
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ codec_streaming: true
+ connector_get_sleep_s: 0.001
+ connector_get_max_wait_first_chunk: 3000
+ connector_get_max_wait: 300
+ codec_chunk_frames: 30
+ codec_left_context_frames: 40
+
+stages:
+ - stage_id: 0
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.78
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ max_num_batched_tokens: 4096
+ max_model_len: 4096
+ devices: "0"
+ attention_backend: TRITON_ATTN
+ enable_flashinfer_autotune: false
+ skip_mm_profiling: true
+ mm_processor_cache_gb: 0
+ output_connectors:
+ to_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.6
+ top_p: 0.95
+ top_k: 50
+ max_tokens: 18192
+ seed: 42
+ repetition_penalty: 1.1
+
+ - stage_id: 1
+ max_num_seqs: 1
+ gpu_memory_utilization: 0.12
+ enforce_eager: true
+ trust_remote_code: true
+ enable_prefix_caching: false
+ async_scheduling: false
+ max_num_batched_tokens: 4096
+ max_model_len: 4096
+ devices: "0"
+ attention_backend: TRITON_ATTN
+ enable_flashinfer_autotune: false
+ skip_mm_profiling: true
+ mm_processor_cache_gb: 0
+ input_connectors:
+ from_stage_0: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ top_p: 1.0
+ top_k: -1
+ max_tokens: 18192
+ seed: 42
diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml
index bc34bfe9f..023101223 100644
--- a/vllm_omni/deploy/minicpmo_4_5.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5.yaml
@@ -1,6 +1,15 @@
# MiniCPM-o 4.5 three-stage deploy on one large-memory GPU.
pipeline: minicpmo_4_5
async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+# Keep the active window aligned with the session and stage capacities below.
+# A smaller window serializes first-packet generation for concurrent chat turns.
+active_stream_window: 4
+duplex_session:
+ # Matches Stage 0/1 max_num_seqs below.
+ max_sessions: 4
trust_remote_code: true
enable_prefix_caching: false
@@ -32,7 +41,7 @@ stages:
max_num_batched_tokens: 16384
devices: "0"
# Required for MiniCPM interleaved image/audio Daily-Omni packs (~78% recipe).
- # Keep these in YAML so they survive both --deploy-config and --stage-configs-path.
+ # Keep these in YAML so they are applied through --deploy-config.
interleave_mm_strings: true
media_io_kwargs:
video:
@@ -55,6 +64,13 @@ stages:
- stage_id: 1
max_num_seqs: 4
+ # Native duplex uses MiniCPMTTS.generate_with_buffer's supported
+ # sliding_recompute policy: retain one completed condition/codec chunk
+ # when the next condition arrives. Other deploy profiles continue to use
+ # the checkpoint's full_attention policy.
+ hf_overrides:
+ tts_config:
+ attention_type: sliding_recompute
# Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace.
gpu_memory_utilization: 0.15
enforce_eager: false
@@ -64,15 +80,20 @@ stages:
devices: "0"
output_connectors:
to_stage_2: connector_of_shared_memory
- # Codec sampling comes from the checkpoint's tts_config. These parameters
- # govern only vLLM's binary continue/stop token.
+ # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler
+ # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1).
+ # Values track upstream utils.TTSSamplingParams (min_p/win_size/tau_r are
+ # declared there but never read, so they stay out). repetition_penalty is
+ # consumed by the Talker, which rescores it over upstream's 16-frame window
+ # instead of the Sampler's whole-stream presence penalty.
+ # min_tokens is upstream's min_new_token.
default_sampling_params:
- temperature: 0.0
- top_p: 1.0
- top_k: -1
+ temperature: 0.8
+ top_p: 0.85
+ top_k: 25
+ repetition_penalty: 1.05
min_tokens: 50
- max_tokens: 2048
- stop_token_ids: [1]
+ max_tokens: 4096
seed: 42
detokenize: false
@@ -111,3 +132,9 @@ platforms:
max_num_batched_tokens: 8192
compilation_config:
cudagraph_mode: PIECEWISE
+ - stage_id: 2
+ # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+ # estimator kernels through the NPU platform patch.
+ additional_config:
+ code2wav_enable_npu_graph: true
+ code2wav_max_npu_graphs: 32
diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
index 2a4abed42..9c3818813 100644
--- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
@@ -4,6 +4,15 @@
# isolated from codec/vocoder contention while preserving stage-2 batching.
pipeline: minicpmo_4_5
async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+# Keep the active window aligned with the session and stage capacities below.
+# A smaller window serializes first-packet generation for concurrent chat turns.
+active_stream_window: 4
+duplex_session:
+ # Matches Stage 0/1 max_num_seqs below.
+ max_sessions: 4
trust_remote_code: true
enable_prefix_caching: false
@@ -52,16 +61,14 @@ stages:
devices: "1"
output_connectors:
to_stage_2: connector_of_shared_memory
- # Codec sampling is model-owned by the checkpoint's tts_config (defaults:
- # T=0.8, top-k=25, top-p=0.85, repetition penalty=1.05, seed=42).
- # These parameters only govern vLLM's binary continue/stop token.
default_sampling_params:
- temperature: 0.0
- top_p: 1.0
- top_k: -1
+ temperature: 0.8
+ top_p: 0.85
+ top_k: 25
+ # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window.
+ repetition_penalty: 1.05
min_tokens: 50
- max_tokens: 2048
- stop_token_ids: [1]
+ max_tokens: 4096
seed: 42
detokenize: false
@@ -93,3 +100,9 @@ platforms:
max_num_batched_tokens: 8192
compilation_config:
cudagraph_mode: PIECEWISE
+ - stage_id: 2
+ # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+ # estimator kernels through the NPU platform patch.
+ additional_config:
+ code2wav_enable_npu_graph: true
+ code2wav_max_npu_graphs: 32
diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml
index b34673910..5eddbb6f4 100644
--- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml
@@ -4,6 +4,15 @@
# Code2Wav run independently on GPUs 0, 1, and 2.
pipeline: minicpmo_4_5
async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. Every other
+# duplex_session knob keeps its runtime default.
+session_mode: duplex
+# Keep the active window aligned with the session and stage capacities below.
+# A smaller window serializes first-packet generation for concurrent chat turns.
+active_stream_window: 4
+duplex_session:
+ # Matches Stage 0/1 max_num_seqs below.
+ max_sessions: 4
trust_remote_code: true
enable_prefix_caching: false
@@ -57,15 +66,14 @@ stages:
devices: "1"
output_connectors:
to_stage_2: connector_of_shared_memory
- # vLLM samples the binary continue/stop row; codec sampling comes from
- # the checkpoint's tts_config and currently uses a deterministic seed.
default_sampling_params:
- temperature: 0.0
- top_p: 1.0
- top_k: -1
+ temperature: 0.8
+ top_p: 0.85
+ top_k: 25
+ # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window.
+ repetition_penalty: 1.05
min_tokens: 50
- max_tokens: 2048
- stop_token_ids: [1]
+ max_tokens: 4096
seed: 42
detokenize: false
@@ -86,4 +94,20 @@ stages:
max_tokens: 65536
detokenize: true
-platforms: {}
+platforms:
+ npu:
+ stages:
+ - stage_id: 0
+ max_num_batched_tokens: 8192
+ compilation_config:
+ cudagraph_mode: PIECEWISE
+ - stage_id: 1
+ max_num_batched_tokens: 8192
+ compilation_config:
+ cudagraph_mode: PIECEWISE
+ - stage_id: 2
+ # Stage 2 stays outer-eager and captures only exact-shape CFM DiT
+ # estimator kernels through the NPU platform patch.
+ additional_config:
+ code2wav_enable_npu_graph: true
+ code2wav_max_npu_graphs: 32
diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml
index 7a0513167..806f3330e 100644
--- a/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_3gpu_stage1_replicas.yaml
@@ -5,7 +5,6 @@
# Stage 1: replicas on GPU 1 and 2
base_config: minicpmo_4_5.yaml
pipeline: minicpmo_4_5
-session_mode: duplex
stages:
- stage_id: 0
diff --git a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml
index da28456fe..16d9f79da 100644
--- a/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_4gpu_stage1_replicas.yaml
@@ -5,7 +5,6 @@
# Stage 1: replicas on GPU 1, 2, and 3
base_config: minicpmo_4_5.yaml
pipeline: minicpmo_4_5
-session_mode: duplex
stages:
- stage_id: 0
diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml
index 813a2e2eb..eacc069d1 100644
--- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml
@@ -15,6 +15,10 @@
# so 4096 is the empirically-verified safe ceiling.
pipeline: minicpmo_4_5
async_chunk: true
+# Native full-duplex shares this profile with /v1/chat/completions. The default
+# duplex_session caps sessions at 1, matching Stage 0/1 max_num_seqs below.
+session_mode: duplex
+active_stream_window: 1
trust_remote_code: true
enable_prefix_caching: false
@@ -69,15 +73,14 @@ stages:
devices: "4"
output_connectors:
to_stage_2: connector_of_shared_memory
- # vLLM samples the binary continue/stop row; codec sampling comes from
- # the checkpoint's tts_config and currently uses a deterministic seed.
default_sampling_params:
- temperature: 0.0
- top_p: 1.0
- top_k: -1
+ temperature: 0.8
+ top_p: 0.85
+ top_k: 25
+ # See minicpmo_4_5.yaml: the Talker rescores this over a 16-frame window.
+ repetition_penalty: 1.05
min_tokens: 50
- max_tokens: 2048
- stop_token_ids: [1]
+ max_tokens: 4096
seed: 42
detokenize: false
diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml
index f4ad2c6b4..f377ac040 100644
--- a/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml
+++ b/vllm_omni/deploy/minicpmo_4_5_8x4090_stage1_replicas.yaml
@@ -5,7 +5,6 @@
# Stage 1: replicas on GPU 4, 5, 6, and 7
base_config: minicpmo_4_5_8x4090.yaml
pipeline: minicpmo_4_5
-session_mode: duplex
stages:
- stage_id: 0
diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml
deleted file mode 100644
index df0e40be9..000000000
--- a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml
+++ /dev/null
@@ -1,41 +0,0 @@
-# MiniCPM-o 4.5 native duplex deployment.
-#
-# The registered minicpmo_4_5 pipeline owns the fixed
-# Thinker -> Talker -> Code2Wav topology. This overlay changes session
-# lifetime and bounds all three stages to two live duplex sessions.
-base_config: minicpmo_4_5.yaml
-pipeline: minicpmo_4_5
-session_mode: duplex
-active_stream_window: 1
-duplex_session:
- idle_ttl_s: 300
- disconnect_grace_s: 30
- reaper_interval_s: 5
- resume_replay_ttl_s: 60
- resume_replay_max_bytes_per_session: 8388608
- max_pending_input_bytes_per_session: 16777216
- max_pending_turns_per_session: 4
- # Matches the configured Stage0/Stage1 sequence capacity below.
- max_sessions: 2
- completed_append_cache_size: 256
-
-stages:
- - stage_id: 0
- max_num_seqs: 2
- async_scheduling: false
-
- - stage_id: 1
- max_num_seqs: 2
- async_scheduling: false
- # Talker context is 4096 (tts_config.max_position_embeddings).
- max_model_len: 4096
- # The split Talker exposes a binary continue/stop vocabulary; native
- # segment and turn boundaries travel as explicit output metadata.
- default_sampling_params:
- max_tokens: 2048
- min_tokens: 0
- stop_token_ids: [1]
-
- - stage_id: 2
- max_num_seqs: 2
- async_scheduling: false
diff --git a/vllm_omni/deploy/minimax_h3_disaggregated.yaml b/vllm_omni/deploy/minimax_h3_disaggregated.yaml
new file mode 100644
index 000000000..1c1ee6c40
--- /dev/null
+++ b/vllm_omni/deploy/minimax_h3_disaggregated.yaml
@@ -0,0 +1,28 @@
+# MiniMax H3: vLLM Qwen3-VL text encoder + diffusion video/audio generation.
+pipeline: minimax_h3_disaggregated
+async_chunk: false
+trust_remote_code: true
+distributed_executor_backend: mp
+
+stages:
+ - stage_id: 0
+ devices: "0,1"
+ max_num_seqs: 1
+ tensor_parallel_size: 2
+ default_sampling_params:
+ max_tokens: 1
+ temperature: 0.0
+ detokenize: false
+
+ - stage_id: 1
+ devices: "2,3,4,5"
+ max_num_seqs: 1
+ model_loaded:
+ text_encoder: false
+ parallel_config:
+ tensor_parallel_size: 1
+ ulysses_degree: 4
+ vae_patch_parallel_size: 4
+ default_sampling_params:
+ seed: 42
+ num_inference_steps: 50
diff --git a/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
new file mode 100644
index 000000000..d7e84db42
--- /dev/null
+++ b/vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
@@ -0,0 +1,31 @@
+# MiniMax H3 Turbo: five-sigma defaults for the disaggregated topology.
+pipeline: minimax_h3_disaggregated
+async_chunk: false
+trust_remote_code: true
+distributed_executor_backend: mp
+
+stages:
+ - stage_id: 0
+ devices: "0,1"
+ max_num_seqs: 1
+ tensor_parallel_size: 2
+ default_sampling_params:
+ max_tokens: 1
+ temperature: 0.0
+ detokenize: false
+
+ - stage_id: 1
+ devices: "2,3,4,5"
+ max_num_seqs: 1
+ model_loaded:
+ text_encoder: false
+ parallel_config:
+ tensor_parallel_size: 1
+ ulysses_degree: 4
+ vae_patch_parallel_size: 4
+ default_sampling_params:
+ seed: 42
+ num_inference_steps: 5
+ extra_args:
+ flow_shift: 6.0
+ audio_flow_shift: 3.0
diff --git a/vllm_omni/deploy/minimax_music3.yaml b/vllm_omni/deploy/minimax_music3.yaml
new file mode 100644
index 000000000..6e8a13d9d
--- /dev/null
+++ b/vllm_omni/deploy/minimax_music3.yaml
@@ -0,0 +1,130 @@
+# MiniMax-Music3 deploy: Stage 0 (AR talker) -> Stage 1 (acoustic decoder),
+# both on one 140 GB card, streaming 200-frame windows over shared memory.
+#
+# Auto-discovered default for model_type=minimax_music3.
+#
+# Classifier-free guidance is mandatory for this checkpoint: every request
+# decodes as a cond/uncond pair in stage 0, so stage-0 KV, max_num_seqs and
+# prefill budget are all sized for TWO rows per user request. The stage-0
+# extra_args block below is what turns on has_sampling_extra_args (so per-row
+# sampling extra args reach the model's forward) and is also how the CFG
+# scheduler patches learn this stage is paired.
+pipeline: minimax_music3
+async_chunk: true
+
+# The repo-root config.json reports model_type=minimax_music3, which upstream
+# transformers does not ship a config class for.
+trust_remote_code: true
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ # Stream stereo PCM back to the client as stage 1 finishes each window.
+ codec_streaming: true
+ # Window geometry is fixed by the checkpoint (200-frame windows on a
+ # 100-frame hop, models/minimax_music3/constants.py), so there are no
+ # codec_chunk_frames knobs here: changing the window changes the audio.
+ shm_threshold_bytes: 65536
+ connector_get_sleep_s: 0.01
+ # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
+ # for a chunk than a per-step codec stage would. First chunk additionally
+ # covers stage-0 prefill of both CFG rows.
+ connector_get_max_wait_first_chunk: 30000
+ connector_get_max_wait: 10000
+
+stages:
+ - stage_id: 0
+ # 16 rows = 8 concurrent guided requests, since CFG doubles every request.
+ max_num_seqs: 16
+ # ~16 GB of bf16 Qwen3-8B weights plus KV for 16 full-length rows
+ # (36 layers x 8 KV heads x 128 head_dim ~= 144 KB/token, so ~1.4 GB per
+ # 10240-token row) fits comfortably inside 0.6 of a 140 GB card.
+ gpu_memory_utilization: 0.6
+ dtype: bfloat16
+ # The checkpoint's max_position_embeddings is 10240 and the whole song is
+ # decoded in-context. Note this cannot hold a max-length prompt AND a
+ # max-length song at the same time: the prompt cap is 5000 tokens and the
+ # frame budget is 9000, which overruns 10240. That is a property of the
+ # checkpoint, not of this file; long captions shorten the maximum song.
+ max_model_len: 10240
+ # One scheduler step must be able to prefill both rows of a pair
+ # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
+ max_num_batched_tokens: 10240
+ # Two rows at 8B is a launch-bound decode step, so the backbone is
+ # captured. FULL_DECODE_ONLY only ever dispatches a graph for a uniform
+ # decode batch; prefill and mixed steps stay eager, which is what the
+ # talker's frame-feedback fallback is for. The talker's own sampler and
+ # depth decoder run after the graph, not inside it.
+ #
+ # What makes the forward capturable: the frame feedback writes a
+ # fixed-size slab in prepare_runner_inputs (outside the captured region)
+ # and the forward selects over it with torch.where, instead of an
+ # index_copy against an index tensor whose address and length changed
+ # every step. The row-to-request binding the sampler reads is likewise
+ # rebuilt in prepare_runner_inputs, because a replayed step never runs
+ # the forward's Python body.
+ enforce_eager: false
+ compilation_config:
+ # No compilation. The talker walks the backbone's decoder layers itself,
+ # so the compiled Qwen3Model.forward is never entered and inductor has
+ # nothing to do; what the mode still decides is which kernels vLLM
+ # defaults to. Left at its non-eager default it picks the elementwise
+ # PyTorch norms over the fused vllm_c ones, which is both slower and
+ # numerically different: measured here, that alone moved the sampled
+ # codes off the eager trajectory by frame 4. Graph capture is
+ # independent of this and still happens.
+ mode: 0
+ # 16 rows is the stage ceiling, i.e. 8 guided requests.
+ cudagraph_capture_sizes: [1, 2, 4, 8, 16]
+ cudagraph_mode: FULL_DECODE_ONLY
+ cudagraph_num_of_warmups: 1
+ # Belt and braces on the same point: the non-eager default is "none",
+ # which hands the norm and activation ops to elementwise PyTorch.
+ custom_ops: ["all"]
+ # A CFG pair must be sampled inside one engine step for the logits blend
+ # and the cond->uncond token copy to line up.
+ async_scheduling: false
+ # Off during bring-up: the uncond row is a long run of identical
+ # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
+ # its cond partner, which is exactly the num_computed_tokens skew the CFG
+ # pair patches then have to undo.
+ enable_prefix_caching: false
+ devices: "0"
+ output_connectors:
+ to_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 1.0
+ # The frame budget is enforced by the model. This ceiling only has to
+ # leave room for the extra step that delivers the final window.
+ max_tokens: 9002
+ extra_args:
+ cfg_role: cond
+ cfg_scale: 1.5
+
+ - stage_id: 1
+ # One row per user request: the uncond companion is stage-0-final and
+ # never reaches the acoustic stage.
+ max_num_seqs: 8
+ # No KV cache on this stage; the budget only covers the DiT and DAV
+ # weights plus per-window latents and activations.
+ gpu_memory_utilization: 0.2
+ # The flow-matching solver degrades measurably in bfloat16
+ # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
+ dtype: float32
+ enforce_eager: true
+ # Async scheduling builds per-request payloads against the NEXT step's
+ # input batch, which breaks when chunk arrival staggers requests.
+ async_scheduling: false
+ enable_prefix_caching: false
+ # One placeholder prompt token per arriving window, so a six-minute song
+ # (90 windows) needs a tiny context.
+ max_model_len: 4096
+ max_num_batched_tokens: 4096
+ devices: "0"
+ input_connectors:
+ from_stage_0: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 4096
+ seed: 0
diff --git a/vllm_omni/deploy/minimax_music3_2gpu.yaml b/vllm_omni/deploy/minimax_music3_2gpu.yaml
new file mode 100644
index 000000000..8c463e0e7
--- /dev/null
+++ b/vllm_omni/deploy/minimax_music3_2gpu.yaml
@@ -0,0 +1,97 @@
+# MiniMax-Music3 two-GPU layout: AR talker on device 0, acoustic decoder on
+# device 1.
+#
+# Splitting the stages removes the memory-budget compromise of the single-card
+# file: the talker's CFG-doubled KV no longer competes with the DiT and DAV
+# weights, and the flow-matching solver no longer competes with AR decode for
+# SMs. Everything else (window geometry, CFG contract, sampling defaults)
+# matches minimax_music3.yaml, so the two files stay behaviorally identical.
+pipeline: minimax_music3
+async_chunk: true
+
+# The repo-root config.json reports model_type=minimax_music3, which upstream
+# transformers does not ship a config class for.
+trust_remote_code: true
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ # Stream stereo PCM back to the client as stage 1 finishes each window.
+ codec_streaming: true
+ # Window geometry is fixed by the checkpoint (200-frame windows on a
+ # 100-frame hop, models/minimax_music3/constants.py), so there are no
+ # codec_chunk_frames knobs here: changing the window changes the audio.
+ shm_threshold_bytes: 65536
+ connector_get_sleep_s: 0.01
+ # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer
+ # for a chunk than a per-step codec stage would. First chunk additionally
+ # covers stage-0 prefill of both CFG rows.
+ connector_get_max_wait_first_chunk: 30000
+ connector_get_max_wait: 10000
+
+stages:
+ - stage_id: 0
+ # 32 rows = 16 concurrent guided requests, since CFG doubles every request.
+ # Twice the single-card figure: this stage now owns its whole device.
+ max_num_seqs: 32
+ gpu_memory_utilization: 0.85
+ dtype: bfloat16
+ # The checkpoint's max_position_embeddings is 10240 and the whole song is
+ # decoded in-context. Note this cannot hold a max-length prompt AND a
+ # max-length song at the same time: the prompt cap is 5000 tokens and the
+ # frame budget is 9000, which overruns 10240. That is a property of the
+ # checkpoint, not of this file; long captions shorten the maximum song.
+ max_model_len: 10240
+ # One scheduler step must be able to prefill both rows of a pair
+ # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked.
+ max_num_batched_tokens: 10240
+ # The talker owns its sampler and drives a depth decoder per AR step;
+ # leave graph capture off until that path is verified graph-safe.
+ enforce_eager: true
+ # A CFG pair must be sampled inside one engine step for the logits blend
+ # and the cond->uncond token copy to line up.
+ async_scheduling: false
+ # Off during bring-up: the uncond row is a long run of identical
+ # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than
+ # its cond partner, which is exactly the num_computed_tokens skew the CFG
+ # pair patches then have to undo.
+ enable_prefix_caching: false
+ devices: "0"
+ output_connectors:
+ to_stage_1: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 1.0
+ # The frame budget is enforced by the model. This ceiling only has to
+ # leave room for the extra step that delivers the final window.
+ max_tokens: 9002
+ extra_args:
+ cfg_role: cond
+ cfg_scale: 1.5
+
+ - stage_id: 1
+ # One row per user request: the uncond companion is stage-0-final and
+ # never reaches the acoustic stage.
+ max_num_seqs: 16
+ # No KV cache on this stage, but it now has a device to itself, so the
+ # solver can hold more windows in flight than the shared-card profile.
+ gpu_memory_utilization: 0.6
+ # The flow-matching solver degrades measurably in bfloat16
+ # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py).
+ dtype: float32
+ enforce_eager: true
+ # Async scheduling builds per-request payloads against the NEXT step's
+ # input batch, which breaks when chunk arrival staggers requests.
+ async_scheduling: false
+ enable_prefix_caching: false
+ # One placeholder prompt token per arriving window, so a six-minute song
+ # (90 windows) needs a tiny context.
+ max_model_len: 4096
+ max_num_batched_tokens: 4096
+ devices: "1"
+ input_connectors:
+ from_stage_0: connector_of_shared_memory
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 4096
+ seed: 0
diff --git a/vllm_omni/deploy/moss_tts_nano.yaml b/vllm_omni/deploy/moss_tts_nano.yaml
index 2c8fc54c0..e9eab184b 100644
--- a/vllm_omni/deploy/moss_tts_nano.yaml
+++ b/vllm_omni/deploy/moss_tts_nano.yaml
@@ -6,8 +6,8 @@
# * enforce_eager=true — the model uses upstream inference_stream() with
# trust_remote_code; CUDA graph capture is not wired up.
# * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB.
-# * max_num_seqs=4 — mirrors the legacy stage_configs contract; per-request
-# streaming generators are stored in self._stream_gens.
+# * max_num_seqs=4 — supports per-request streaming generators stored in
+# self._stream_gens.
#
# MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling
# bypasses the dummy-MM pass at profile time.
diff --git a/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
new file mode 100644
index 000000000..1a69b24e6
--- /dev/null
+++ b/vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml
@@ -0,0 +1,47 @@
+# Nemotron VoiceChat native Full-Duplex deployment.
+#
+# One scheduler segment consumes one 1280-sample (80 ms) microphone frame.
+# Stage 0 remains one resumable request, so its NemotronH KV/Mamba state is
+# preserved between frame appends. Stage 1 and the codec overlap per frame.
+base_config: nemotron_labs_voicechat_streaming.yaml
+pipeline: nemotron_voicechat
+session_mode: duplex
+active_stream_window: 1
+
+duplex_session:
+ idle_ttl_s: 300
+ disconnect_grace_s: 30
+ reaper_interval_s: 5
+ resume_replay_ttl_s: 60
+ resume_replay_max_bytes_per_session: 8388608
+ max_pending_input_bytes_per_session: 16777216
+ max_pending_turns_per_session: 4
+ max_sessions: 1
+ completed_append_cache_size: 256
+
+async_chunk: true
+
+connectors:
+ connector_of_shared_memory:
+ name: SharedMemoryConnector
+ extra:
+ shm_threshold_bytes: 65536
+ codec_chunk_frames: 1
+
+stages:
+ - stage_id: 0
+ max_num_seqs: 1
+ async_scheduling: false
+ default_sampling_params:
+ temperature: 0.0
+ max_tokens: 1
+ ignore_eos: true
+ seed: 0
+
+ - stage_id: 1
+ max_num_seqs: 1
+ async_scheduling: false
+
+ - stage_id: 2
+ max_num_seqs: 1
+ async_scheduling: false
diff --git a/vllm_omni/deploy/pi0.yaml b/vllm_omni/deploy/pi0.yaml
new file mode 100644
index 000000000..8d93855d8
--- /dev/null
+++ b/vllm_omni/deploy/pi0.yaml
@@ -0,0 +1,52 @@
+# π0 (Pi-Zero) VLA deploy: single diffusion stage.
+#
+# Topology (PI0_PIPELINE) is declared in vllm_omni/diffusion/models/pi0/pipeline_pi0.py.
+# π0 takes multi-camera images + a language instruction + robot state and emits
+# a continuous action chunk [chunk_size, max_action_dim] via flow matching.
+# Actions are exposed through the OpenPI realtime serving API
+# (/v1/realtime/robot/openpi); they ride in multimodal_output["actions"].
+#
+# Default: one GPU, TP=1, eager (flow matching needs no CUDA graph), bsz 1
+# (OpenPI policy serving is one observation -> one action reply). float32 to
+# match the LeRobot pi0_base reference numerics.
+
+pipeline: pi0
+async_chunk: false
+distributed_executor_backend: mp
+dtype: float32
+
+stages:
+ - stage_id: 0
+ devices: "0"
+ max_num_seqs: 1
+ enforce_eager: true
+ model_class_name: Pi0Pipeline
+ model_config:
+ paligemma_variant: gemma_2b
+ action_expert_variant: gemma_300m
+ chunk_size: 50
+ max_action_dim: 32
+ max_state_dim: 32
+ num_inference_steps: 10
+ image_resolution: [224, 224]
+ tokenizer_max_length: 48
+ max_cameras: 3
+ # Camera identities (the order the model attends to). Serving maps raw obs
+ # keys onto these via image_key_map; leave image_key_map empty to match
+ # incoming keys verbatim.
+ image_feature_keys:
+ - observation.images.base_0_rgb
+ - observation.images.left_wrist_0_rgb
+ - observation.images.right_wrist_0_rgb
+ image_key_map: {}
+ # OpenPI handshake metadata advertised to the robot client.
+ policy_server_config:
+ image_resolution: [224, 224]
+ n_external_cameras: 1
+ needs_wrist_camera: true
+ needs_stereo_camera: false
+ needs_session_id: false
+ action_space: joint_position
+ action_horizon: 50
+ action_dim: 32
+ max_cameras: 3
diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml
index 76878804d..59378d06b 100644
--- a/vllm_omni/deploy/qwen3_omni_moe.yaml
+++ b/vllm_omni/deploy/qwen3_omni_moe.yaml
@@ -69,6 +69,15 @@ stages:
repetition_penalty: 1.1
platforms:
+ # vLLM 0.27's compiled native MRoPE path can vertically fuse the thinker's
+ # Q/K RMSNorm reductions into the rotary operations and materially change
+ # Qwen3-Omni audio output (#6090). Keep MRoPE on its CUDA path instead.
+ cuda:
+ stages:
+ - stage_id: 0
+ compilation_config:
+ custom_ops: ["+rotary_embedding"]
+
# The supported Qwen3-Omni ModelOpt checkpoint quantizes the thinker only. Clear the
# checkpoint-level metadata for the BF16 audio stages on the verified MUSA
# profile before vLLM performs its normal quantization auto-detection.
diff --git a/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
index 7ccc64750..d0b3d5ac7 100644
--- a/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+++ b/vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
@@ -1,6 +1,8 @@
-# Deploy config for Qwen3-Omni-MoE thinker-only checkpoints.
+# Deploy config for Qwen3-Omni-MoE thinker-only (text output, no talker/code2wav).
+# Works for Captioner/Thinking checkpoints and for Instruct weights when this
+# YAML is passed via --deploy-config (pipeline key forces the single-stage topology).
# Verified on 2x H100-80G GPUs.
-pipeline: qwen3_omni_moe
+pipeline: qwen3_omni_moe_thinker_only
distributed_executor_backend: mp
enable_prefix_caching: true
diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml
index 08531869a..747b68852 100644
--- a/vllm_omni/deploy/qwen3_tts.yaml
+++ b/vllm_omni/deploy/qwen3_tts.yaml
@@ -33,6 +33,17 @@ connectors:
# initial_codec_chunk_frames (including the per-request API field).
# Example:
# codec_chunk_ramp: [4, 4, 8, 16, 25]
+ # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp):
+ # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback.
+ # Self-tunes across hardware (NPU/NVIDIA) and load conditions.
+ # codec_chunk_adaptive: true
+ # codec_chunk_min_frames: 2
+ # Code2Wav decode + transfer margin
+ # codec_chunk_safety_margin_ms: 50
+ # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp
+ # codec_chunk_ramp_divisor: 15
+ # min chunk-size growth per emit
+ # codec_chunk_ramp_delta_min: 2
# Decoder CUDA Graph sizing:
# * async_chunk=true derives all stateful ICL/xvec shapes from
# codec_chunk_ramp when enabled, otherwise from
@@ -73,6 +84,10 @@ stages:
# audio (HTTP 400). See #4962.
min_tokens: 2
repetition_penalty: 1.05
+ # Suppress codebook-0 silence tokens for the first N decode frames of a
+ # request (#4966). Applies to every request on this stage, so intended
+ # for x-vector-only (Base voice-clone) deploys. 0 = off.
+ silence_ban_frames: 0
subtalker_sampling_params:
do_sample: true
temperature: 0.9
diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
index ff086beb5..7a73dd50a 100644
--- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
+++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
@@ -38,6 +38,17 @@ connectors:
# initial_codec_chunk_frames (including the per-request API field).
# Example:
# codec_chunk_ramp: [4, 4, 8, 16, 25]
+ # Adaptive chunk ramp (opt-in, takes precedence over fixed codec_chunk_ramp):
+ # Chunk 0 uses dynamic IC; chunk 1+ sizes from live buffer-depth feedback.
+ # Self-tunes across hardware (NPU/NVIDIA) and load conditions.
+ # codec_chunk_adaptive: true
+ # codec_chunk_min_frames: 2
+ # Code2Wav decode + transfer margin
+ # codec_chunk_safety_margin_ms: 50
+ # EWMA->delta divisor (delta=max(delta_min, int(ewma/divisor))); higher=gentler ramp
+ # codec_chunk_ramp_divisor: 15
+ # min chunk-size growth per emit
+ # codec_chunk_ramp_delta_min: 2
# async_chunk derives stateful ICL/xvec graph shapes from the ramp when
# enabled, otherwise from the fixed chunk settings above, and ignores
# decode_cudagraph_capture_sizes. Keep B>1
diff --git a/vllm_omni/deploy/soulxsinger_svc.yaml b/vllm_omni/deploy/soulxsinger_svc.yaml
deleted file mode 100644
index a9d815806..000000000
--- a/vllm_omni/deploy/soulxsinger_svc.yaml
+++ /dev/null
@@ -1,19 +0,0 @@
-# SoulX-Singer SVC: single-stage DiT (preprocess inline in pre_process_func).
-#
-# Use: vllm serve --omni
-pipeline: soulxsinger_svc
-async_chunk: false
-dtype: float16
-trust_remote_code: true
-distributed_executor_backend: "mp"
-
-stages:
- - stage_id: 0
- max_num_seqs: 1
- gpu_memory_utilization: 0.5
- enforce_eager: false
- devices: "0"
- default_sampling_params:
- num_inference_steps: 32
- guidance_scale: 3.0
- seed: 42
diff --git a/vllm_omni/deploy/soulxsinger_svs.yaml b/vllm_omni/deploy/soulxsinger_svs.yaml
deleted file mode 100644
index 3f021cb53..000000000
--- a/vllm_omni/deploy/soulxsinger_svs.yaml
+++ /dev/null
@@ -1,19 +0,0 @@
-# SoulX-Singer SVS: single-stage DiT (preprocess inline in pre_process_func).
-#
-# Use: vllm serve --omni
-pipeline: soulxsinger_svs
-async_chunk: false
-dtype: float16
-trust_remote_code: true
-distributed_executor_backend: "mp"
-
-stages:
- - stage_id: 0
- max_num_seqs: 1
- gpu_memory_utilization: 0.5
- enforce_eager: false
- devices: "0"
- default_sampling_params:
- num_inference_steps: 32
- guidance_scale: 3.0
- seed: 42
diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml
index bdb5cfecb..0d0a9124c 100644
--- a/vllm_omni/deploy/wan2_2_ti2v.yaml
+++ b/vllm_omni/deploy/wan2_2_ti2v.yaml
@@ -8,7 +8,7 @@ distributed_executor_backend: mp
stages:
- stage_id: 0
devices: "0"
- max_num_seqs: 1
+ max_num_seqs: 4 # in-flight request batch width (CLI: --max-num-seqs)
tensor_parallel_size: 1
enforce_eager: true
# Keep VAE tiling on by default for higher resolutions / longer clips.
diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py
index f3a5f99ef..c0b9a738e 100644
--- a/vllm_omni/diffusion/attention/backends/abstract.py
+++ b/vllm_omni/diffusion/attention/backends/abstract.py
@@ -1,9 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Literal, TypeVar
import torch
@@ -15,6 +15,11 @@ class AttentionBackend(ABC):
accept_output_buffer: bool = False
supports_piecewise_spans: bool = False
+ # A backend that supports this capability can consume the opaque paged-KV
+ # context prepared by the diffusion Worker data plane. Keeping the
+ # capability on the backend class prevents a paged request from silently
+ # falling back to dense attention on an incompatible implementation.
+ supports_paged_kv: bool = False
# The backend can represent a contiguous valid K/V prefix by slicing the
# tensors instead of materializing a padding mask. Models may use this to
# avoid a slower masked-attention plan when tail padding is not semantic.
@@ -22,12 +27,30 @@ class AttentionBackend(ABC):
@classmethod
def supports_packed_mask_free(cls) -> bool:
- """Whether packed attention never reads attn_mask on this platform.
+ """Whether [real, pad] packed layouts can run without attn_mask.
When True, models that pack a [real, pad] two-document layout and
- carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
- constructing the padding mask entirely. Backends whose mask-free
- behavior is platform-dependent must check current_omni_platform.
+ provide ``AttentionMetadata.packed_padding`` alongside the packed
+ cu_seqlens/max_seqlen metadata may skip constructing the padding mask
+ entirely. Backends whose mask-free behavior is platform-dependent must
+ check current_omni_platform.
+ """
+ return False
+
+ @classmethod
+ def supports_multi_doc_packed_varlen(cls) -> bool:
+ """Whether this backend keeps N-document packed boundaries isolated.
+
+ When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens
+ as a genuine block-diagonal attention plan (a dedicated varlen kernel,
+ not a padding-mask rebuild), so a caller may pack multiple real
+ requests into one forward without attention crossing document
+ boundaries. When False, callers packing more than one real document
+ must run one forward per document; otherwise a backend that only
+ supports a ``[real, pad]`` two-document contract, or that ignores
+ cu_seqlens outright, will silently attend across request boundaries.
+ Backends whose kernel selection is platform-dependent must consult
+ ``current_omni_platform``.
"""
return False
@@ -81,6 +104,19 @@ class AttentionBackend(ABC):
supported_head_sizes = cls.get_supported_head_sizes()
return (not supported_head_sizes) or head_size in supported_head_sizes
+ @classmethod
+ def indexes_kv_by_block_stride(cls) -> bool:
+ """Whether this backend reads K/V pages by the runtime block stride.
+
+ Returning ``True`` means the physical cache layout has ``num_blocks``
+ as its outer stride, so native vLLM may safely use page-size padding
+ when it unifies cache layouts across layers. Dense diffusion backends
+ conservatively keep the default ``False``; a paged backend should
+ override this only when its kernel actually follows that layout.
+ """
+
+ return False
+
@dataclass(frozen=True, slots=True)
class QueryRange:
@@ -89,9 +125,22 @@ class QueryRange:
global_start: int
+@dataclass(frozen=True, slots=True)
+class VideoTokenSpan:
+ """One physical video-grid slice in packed document 0."""
+
+ start: int
+ latent_grid: tuple[int, int, int]
+ role: Literal["reference", "target"]
+
+ @property
+ def length(self) -> int:
+ return self.latent_grid[0] * self.latent_grid[1] * self.latent_grid[2]
+
+
@dataclass(frozen=True, slots=True)
class VideoTokenLayout:
- """Where the video segment sits inside a packed multimodal sequence.
+ """Video-grid slices in a packed multimodal sequence.
A model that packs its sequence as ``[prefix | t*h*w video rows | padding]``
publishes this so backends can recover spatiotemporal locality; the prefix
@@ -99,11 +148,32 @@ class VideoTokenLayout:
Publishing it also asserts that any ``attn_mask`` masks only the trailing
padding, so ``prefix_len + t*h*w`` is the used length of the sequence.
+ ``prefix_len``/``latent_grid`` retain the original one-tail contract. A
+ Ref2VA layout instead publishes ``used_len`` and every physical video
+ span, so audio and image rows between videos remain dense.
+
Plain ints, so reading it never forces a device-to-host sync.
"""
- prefix_len: int
- latent_grid: tuple[int, int, int]
+ prefix_len: int | None = None
+ latent_grid: tuple[int, int, int] | None = None
+ used_len: int | None = None
+ video_spans: tuple[VideoTokenSpan, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class PackedPaddingMetadata:
+ """Producer-validated mask-free view of padding in a [real, pad] packing.
+
+ The cumulative-length tensors are canonical two-element ``[0, length]``
+ views. Consumers may use them without reading device scalars because the
+ producer owns the packing.
+ """
+
+ q_length: int
+ kv_length: int
+ cu_seqlens_q: torch.Tensor
+ cu_seqlens_k: torch.Tensor
@dataclass
@@ -153,6 +223,10 @@ class AttentionMetadata:
# locality (block-sparse selection, tiled masks). Dense backends ignore it.
video_layout: VideoTokenLayout | None = None
+ # Canonical mask-free view of structural suffix padding. Backends that do
+ # not advertise supports_packed_mask_free ignore it.
+ packed_padding: PackedPaddingMetadata | None = None
+
T = TypeVar("T", bound=AttentionMetadata)
@@ -208,6 +282,18 @@ class AttentionImpl(ABC, Generic[T]):
else:
raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
+ def forward_paged(self, paged_kv_context: Any) -> torch.Tensor:
+ """Execute one Worker-prepared paged-KV attention call.
+
+ The context is intentionally opaque to the common attention layer.
+ Backends opt in by setting ``supports_paged_kv`` on their backend
+ class and implementing this method. Dense callers continue to use
+ ``forward`` unchanged.
+ """
+
+ del paged_kv_context
+ raise NotImplementedError(f"{type(self).__name__} does not support paged KV attention")
+
def forward_cuda(
self,
query: torch.Tensor,
diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py
index ffcc2fb44..ab814f713 100644
--- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py
+++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py
@@ -85,6 +85,24 @@ class CuDNNAttentionImpl(AttentionImpl):
enable_gqa = query.shape[2] != key.shape[2]
query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+ # cuDNN has no kernel for degenerate attention shapes when either
+ # sequence has one token. Route those degenerate shapes before entering the
+ # cuDNN-only context: under torch.compile, FakeTensor evaluation raises
+ # before the eager RuntimeError fallback below can run.
+ if query.shape[-2] == 1 or key.shape[-2] == 1:
+ with sdpa_kernel([SDPBackend.MATH]):
+ output = torch.nn.functional.scaled_dot_product_attention(
+ query,
+ key,
+ value,
+ attn_mask=attention_mask,
+ dropout_p=0.0,
+ is_causal=self.causal,
+ scale=self.softmax_scale,
+ enable_gqa=enable_gqa,
+ )
+ return output.permute(0, 2, 1, 3)
+
# Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
# PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
# cuDNN gets runtime-disabled in the same call and the dispatcher falls
@@ -93,7 +111,7 @@ class CuDNNAttentionImpl(AttentionImpl):
#
# Fall back to the default SDPA dispatcher if cuDNN rejects the shape,
# e.g. under torch.compile where Dynamo sees a symbolic head_dim and
- # cuDNN's kernel selection fails (observed in LTX-2 audio attention).
+ # cuDNN's kernel selection fails for some symbolic attention shapes.
# The unpinned dispatcher then picks EFFICIENT/MATH instead of raising.
try:
with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py
new file mode 100644
index 000000000..fcd592dd8
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py
@@ -0,0 +1,431 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import functools
+import math
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+ AttentionBackend,
+ AttentionImpl,
+ AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+
+logger = init_logger(__name__)
+
+
+# Keep the external FastVideo pybind/CUDA kernel opaque to torch.compile.
+# This mirrors the SageAttention3 backend pattern: tracing the raw extension
+# through Dynamo can reach Inductor scheduling with unstable internal op names
+# (e.g. KeyError: "op12"). The custom op gives Dynamo a single Tensor->Tensor
+# boundary and lets Inductor schedule the surrounding Wan block normally.
+if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"):
+
+ @torch.library.custom_op("vllm_omni::fastvideo_vsa_bshd", mutates_args=())
+ def _fastvideo_vsa_bshd_op(
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ variable_block_sizes: torch.Tensor,
+ q_variable_block_sizes: torch.Tensor,
+ compress_attn_weight: torch.Tensor,
+ topk: int,
+ block_t: int,
+ block_h: int,
+ block_w: int,
+ ) -> torch.Tensor:
+ from fastvideo_kernel import video_sparse_attn_bshd
+
+ return video_sparse_attn_bshd(
+ query,
+ key,
+ value,
+ variable_block_sizes=variable_block_sizes,
+ q_variable_block_sizes=q_variable_block_sizes,
+ topk=topk,
+ block_size=(block_t, block_h, block_w),
+ compress_attn_weight=compress_attn_weight if compress_attn_weight.numel() else None,
+ )
+
+ @_fastvideo_vsa_bshd_op.register_fake
+ def _(
+ query,
+ key,
+ value,
+ variable_block_sizes,
+ q_variable_block_sizes,
+ compress_attn_weight,
+ topk,
+ block_t,
+ block_h,
+ block_w,
+ ):
+ del (
+ key,
+ value,
+ variable_block_sizes,
+ q_variable_block_sizes,
+ compress_attn_weight,
+ topk,
+ block_t,
+ block_h,
+ block_w,
+ )
+ return torch.empty_like(query)
+
+
+_fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd
+
+
+@functools.lru_cache(maxsize=32)
+def _get_tile_partition_indices(
+ dit_seq_shape: tuple[int, int, int],
+ tile_size: tuple[int, int, int],
+ device: torch.device,
+) -> torch.Tensor:
+ t_size, h_size, w_size = dit_seq_shape
+ tile_t, tile_h, tile_w = tile_size
+ indices = torch.arange(t_size * h_size * w_size, device=device, dtype=torch.long).reshape(t_size, h_size, w_size)
+ tiles = []
+ for tile_t_idx in range(math.ceil(t_size / tile_t)):
+ for tile_h_idx in range(math.ceil(h_size / tile_h)):
+ for tile_w_idx in range(math.ceil(w_size / tile_w)):
+ tiles.append(
+ indices[
+ tile_t_idx * tile_t : min((tile_t_idx + 1) * tile_t, t_size),
+ tile_h_idx * tile_h : min((tile_h_idx + 1) * tile_h, h_size),
+ tile_w_idx * tile_w : min((tile_w_idx + 1) * tile_w, w_size),
+ ].flatten()
+ )
+ return torch.cat(tiles, dim=0)
+
+
+@functools.lru_cache(maxsize=32)
+def _construct_variable_block_sizes(
+ dit_seq_shape: tuple[int, int, int],
+ tile_size: tuple[int, int, int],
+ device: torch.device,
+) -> torch.Tensor:
+ num_tiles = tuple(math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, tile_size))
+
+ def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.Tensor:
+ sizes = torch.full((n_tiles,), tile, dtype=torch.int32, device=device)
+ remainder = dim_len - (n_tiles - 1) * tile
+ sizes[-1] = remainder if remainder > 0 else tile
+ return sizes
+
+ t_sizes = _sizes(dit_seq_shape[0], tile_size[0], num_tiles[0])
+ h_sizes = _sizes(dit_seq_shape[1], tile_size[1], num_tiles[1])
+ w_sizes = _sizes(dit_seq_shape[2], tile_size[2], num_tiles[2])
+ return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1)
+
+
+@functools.lru_cache(maxsize=32)
+def _get_non_pad_index(variable_block_sizes: torch.Tensor, max_block_size: int) -> torch.Tensor:
+ num_blocks = variable_block_sizes.shape[0]
+ device = variable_block_sizes.device
+ starts = torch.arange(num_blocks, device=device) * max_block_size
+ padded_index = starts[:, None] + torch.arange(max_block_size, device=device)[None, :]
+ valid = torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None]
+ return padded_index[valid]
+
+
+@torch.compiler.disable
+def _get_tile_metadata(
+ dit_seq_shape: tuple[int, int, int],
+ tile_size: tuple[int, int, int],
+ block_elements: int,
+ device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, tile_size, device)
+ variable_block_sizes = _construct_variable_block_sizes(dit_seq_shape, tile_size, device)
+ non_pad_index = _get_non_pad_index(variable_block_sizes, block_elements)
+ untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
+ return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index
+
+
+def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None:
+ if attn_metadata is None:
+ return None
+ value = attn_metadata.extra.get("vsa_dit_seq_shape")
+ if value is None:
+ return None
+ if not isinstance(value, (list, tuple)) or len(value) != 3:
+ return None
+ return (int(value[0]), int(value[1]), int(value[2]))
+
+
+def _get_gate_compress(attn_metadata: AttentionMetadata | None) -> torch.Tensor | None:
+ if attn_metadata is None:
+ return None
+ value = attn_metadata.extra.get("gate_compress")
+ return value if isinstance(value, torch.Tensor) else None
+
+
+def _preserve_vsa_all_blocks(attn_metadata: AttentionMetadata | None) -> bool:
+ if attn_metadata is None:
+ return False
+ return attn_metadata.extra.get("preserve_vsa_all_blocks") is True
+
+
+class FastVideoVSABackend(AttentionBackend):
+ accept_output_buffer: bool = True
+
+ @staticmethod
+ def get_supported_head_sizes() -> list[int]:
+ # FastVideo VSA is intended for video DiT head sizes such as 64/128.
+ # Keep this permissive and let the runtime fallback handle unsupported
+ # cases from the installed fastvideo-kernel build.
+ return []
+
+ @staticmethod
+ def get_name() -> str:
+ return "FASTVIDEO_VSA"
+
+ @staticmethod
+ def get_impl_cls() -> type[FastVideoVSAImpl]:
+ return FastVideoVSAImpl
+
+
+class FastVideoVSAImpl(AttentionImpl):
+ def __init__(
+ self,
+ num_heads: int,
+ head_size: int,
+ softmax_scale: float,
+ causal: bool = False,
+ num_kv_heads: int | None = None,
+ prefix: str = "",
+ qkv_layout: str | None = None,
+ backend_kwargs: Mapping[str, Any] | None = None,
+ **extra_impl_args,
+ ) -> None:
+ backend_kwargs = backend_kwargs or {}
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+ self.head_size = head_size
+ self.softmax_scale = softmax_scale
+ self.causal = causal
+ self.qkv_layout = qkv_layout
+
+ self.topk = int(backend_kwargs.get("topk", 64))
+ self.block_size = self._parse_block_size(backend_kwargs.get("block_size", (4, 8, 8)))
+ self.block_elements = self.block_size[0] * self.block_size[1] * self.block_size[2]
+ self.min_seq_len = int(backend_kwargs.get("min_seq_len", self.block_elements * 2))
+ self.fallback_on_error = bool(backend_kwargs.get("fallback_on_error", True))
+ self.disable_when_sp_active = bool(backend_kwargs.get("disable_when_sp_active", True))
+
+ self.sdpa_fallback = SDPAImpl(
+ num_heads=num_heads,
+ head_size=head_size,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ num_kv_heads=num_kv_heads,
+ qkv_layout=qkv_layout,
+ )
+
+ if self.block_elements != 256:
+ logger.warning(
+ "FASTVIDEO_VSA currently uses fastvideo_kernel.video_sparse_attn_bshd, "
+ "which supports only 256-token blocks. Configured block_size=%s "
+ "(product=%d) will fall back to SDPA.",
+ self.block_size,
+ self.block_elements,
+ )
+
+ @staticmethod
+ def _parse_block_size(value: Any) -> tuple[int, int, int]:
+ if isinstance(value, int):
+ return (value, value, value)
+ if isinstance(value, (list, tuple)) and len(value) == 3:
+ return (int(value[0]), int(value[1]), int(value[2]))
+ raise ValueError(f"FASTVIDEO_VSA block_size must be an int or length-3 tuple/list, got {value!r}")
+
+ def _fallback(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None,
+ reason: str,
+ ) -> torch.Tensor:
+ logger.warning_once("FASTVIDEO_VSA falling back to SDPA: %s", reason)
+ return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
+ def _fallback_reason(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None,
+ ) -> str | None:
+ if self.causal:
+ return "causal attention is not supported"
+ if self.block_elements != 256:
+ return f"block_elements must be 256, got {self.block_elements}"
+ if self.topk <= 0:
+ return f"topk must be positive, got {self.topk}"
+ if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+ return f"expected [B, S, H, D] tensors, got {query.shape}, {key.shape}, {value.shape}"
+ if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
+ return "batch dimensions must match"
+ if query.shape[2:] != key.shape[2:] or query.shape[2:] != value.shape[2:]:
+ return "head/head_dim dimensions must match"
+ if query.shape[1] != key.shape[1] or query.shape[1] != value.shape[1]:
+ return "initial VSA backend supports self-attention with Sq == Skv only"
+ if query.shape[1] < self.min_seq_len:
+ return f"sequence length {query.shape[1]} is below min_seq_len {self.min_seq_len}"
+ dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
+ if dit_seq_shape is None:
+ return "vsa_dit_seq_shape metadata is required"
+ if math.prod(dit_seq_shape) != query.shape[1]:
+ return f"vsa_dit_seq_shape product {math.prod(dit_seq_shape)} != sequence length {query.shape[1]}"
+ num_blocks = math.prod(
+ math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
+ )
+ if self.topk > num_blocks:
+ return f"topk {self.topk} > num_blocks {num_blocks}"
+ if query.dtype not in (torch.float16, torch.bfloat16):
+ return f"dtype {query.dtype} is not supported"
+ if key.dtype != query.dtype or value.dtype != query.dtype:
+ return "q/k/v dtypes must match"
+ if query.device.type != "cuda" or key.device.type != "cuda" or value.device.type != "cuda":
+ return "q/k/v must be CUDA tensors"
+ expected_scale = self.head_size**-0.5
+ if abs(float(self.softmax_scale) - float(expected_scale)) > 1e-6:
+ return f"softmax_scale {self.softmax_scale} differs from FastVideo VSA scale {expected_scale}"
+ if attn_metadata is not None and attn_metadata.attn_mask is not None:
+ return "attention masks are not supported"
+ if attn_metadata is not None and attn_metadata.full_attn_spans is not None:
+ return "piecewise/full attention spans are not supported"
+ if self.num_heads != self.num_kv_heads:
+ return "GQA/MQA is not supported"
+ if self.disable_when_sp_active and is_forward_context_available():
+ ctx = get_forward_context()
+ if getattr(ctx, "sp_active", False):
+ return "sequence parallel context is active"
+ return None
+
+ def forward_cuda(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None = None,
+ ) -> torch.Tensor:
+ reason = self._fallback_reason(query, key, value, attn_metadata)
+ if reason is not None:
+ return self._fallback(query, key, value, attn_metadata, reason)
+
+ seq_len = query.shape[1]
+ dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
+ assert dit_seq_shape is not None
+ num_blocks = math.prod(
+ math.ceil(seq_dim / tile_dim) for seq_dim, tile_dim in zip(dit_seq_shape, self.block_size)
+ )
+ preserve_all_blocks = _preserve_vsa_all_blocks(attn_metadata)
+ use_native_sdpa = self.topk == num_blocks and not preserve_all_blocks
+ route = "SDPA" if use_native_sdpa else "VSA_ALL_BLOCKS" if self.topk == num_blocks else "VSA"
+ checkpoint_mode = "fastvideo_dmd" if preserve_all_blocks else "native"
+ logger.info_once(
+ "FASTVIDEO_VSA routing: seq_len=%d, dit_seq_shape=%s, block_size=%s, num_blocks=%d, "
+ "topk=%d, keep_ratio=%.1f%%, checkpoint_mode=%s, route=%s",
+ seq_len,
+ dit_seq_shape,
+ self.block_size,
+ num_blocks,
+ self.topk,
+ 100.0 * self.topk / num_blocks,
+ checkpoint_mode,
+ route,
+ )
+ if use_native_sdpa:
+ return self._fallback(
+ query,
+ key,
+ value,
+ attn_metadata,
+ f"topk {self.topk} selects all blocks for a native checkpoint",
+ )
+
+ try:
+ tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index = _get_tile_metadata(
+ dit_seq_shape,
+ self.block_size,
+ self.block_elements,
+ query.device,
+ )
+
+ padded_len = variable_block_sizes.numel() * self.block_elements
+ target_shape = (query.shape[0], padded_len, query.shape[2], query.shape[3])
+ query_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype)
+ key_tiled = torch.zeros_like(query_tiled)
+ value_tiled = torch.zeros_like(query_tiled)
+ query_tiled[:, non_pad_index] = query[:, tile_partition_indices]
+ key_tiled[:, non_pad_index] = key[:, tile_partition_indices]
+ value_tiled[:, non_pad_index] = value[:, tile_partition_indices]
+ # Gate behavior is checkpoint-driven, not user-configured.
+ # Wan VSA layers always provide a gate projection. Its zero
+ # initialization makes checkpoints without gate weights sparse-only;
+ # checkpoints containing to_gate_compress weights use the learned gate.
+ gate_compress = _get_gate_compress(attn_metadata)
+ if gate_compress is None:
+ gate_compress = torch.zeros_like(query)
+ elif gate_compress.shape != query.shape:
+ raise ValueError(f"gate_compress shape {gate_compress.shape} must match query shape {query.shape}")
+ gate_tiled = torch.zeros_like(query_tiled)
+ gate_tiled[:, non_pad_index] = gate_compress[:, tile_partition_indices]
+ compress_attn_weight = gate_tiled
+
+ output = _fastvideo_vsa_bshd_op(
+ query_tiled.contiguous(),
+ key_tiled.contiguous(),
+ value_tiled.contiguous(),
+ variable_block_sizes,
+ variable_block_sizes,
+ compress_attn_weight,
+ self.topk,
+ self.block_size[0],
+ self.block_size[1],
+ self.block_size[2],
+ )
+ return output[:, untile_combined_index].contiguous()
+ except Exception as exc:
+ if not self.fallback_on_error:
+ raise
+ return self._fallback(query, key, value, attn_metadata, f"VSA kernel failed: {exc}")
+
+ def forward_npu(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None = None,
+ ) -> torch.Tensor:
+ return self.sdpa_fallback.forward_npu(query, key, value, attn_metadata)
+
+ def forward_xpu(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None = None,
+ ) -> torch.Tensor:
+ return self.sdpa_fallback.forward_xpu(query, key, value, attn_metadata)
+
+ def forward_musa(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attn_metadata: AttentionMetadata | None = None,
+ ) -> torch.Tensor:
+ return self.sdpa_fallback.forward_musa(query, key, value, attn_metadata)
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py
index bdef1dd65..a376216fa 100644
--- a/vllm_omni/diffusion/attention/backends/flash_attn.py
+++ b/vllm_omni/diffusion/attention/backends/flash_attn.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import os
from functools import partial
@@ -7,14 +7,11 @@ from functools import partial
import torch
from vllm.logger import init_logger
-from vllm_omni.diffusion.attention.backends.abstract import (
- AttentionBackend,
- AttentionImpl,
- AttentionMetadata,
-)
+from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
piecewise_attn,
+ run_paged_piecewise_plan,
)
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
from vllm_omni.platforms import current_omni_platform
@@ -25,6 +22,7 @@ logger = init_logger(__name__)
class FlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
supports_piecewise_spans: bool = True
+ supports_paged_kv: bool = True
@classmethod
def supports_packed_mask_free(cls) -> bool:
@@ -35,6 +33,20 @@ class FlashAttentionBackend(AttentionBackend):
# constructing it there.
return current_omni_platform.is_cuda() or current_omni_platform.is_npu()
+ @classmethod
+ def supports_multi_doc_packed_varlen(cls) -> bool:
+ # CUDA / ROCm / MUSA all route through ``forward_cuda``, which
+ # dispatches ``_forward_varlen_packed`` -> ``flash_attn_varlen_func``
+ # over the caller's cu_seqlens without a mask, so an arbitrary
+ # N-document packing keeps its boundaries. NPU's forward path
+ # (``_resolve_packed_seq_npu``) only accepts a ``[real, pad]``
+ # two-document layout (cu_seqlens shape <= 3) and otherwise falls
+ # back to a padding-mask rebuild that spans the whole packed row,
+ # so N>=2 real documents would silently attend across request
+ # boundaries. XPU never consumes cu_seqlens from ``extra``. Both
+ # therefore stay False here.
+ return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa()
+
@classmethod
def supports_attention_mask(cls) -> bool:
return True
@@ -77,12 +89,14 @@ class FlashAttentionImpl(AttentionImpl):
prefix: str = "",
qkv_layout: str | None = None,
backend_kwargs: dict | None = None,
+ role: str = "self",
**extra_impl_args,
) -> None:
self.num_heads = num_heads
self.causal = causal
self.softmax_scale = softmax_scale
self.qkv_layout = qkv_layout
+ self.is_cross_attn = role == "cross"
cfg = get_current_diffusion_config_or_none()
self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
if backend_kwargs:
@@ -136,6 +150,7 @@ class FlashAttentionImpl(AttentionImpl):
attention_mask: torch.Tensor,
) -> torch.Tensor:
from vllm_omni.diffusion.attention.backends.utils.fa import (
+ _index_first_axis,
_pad_input,
_unpad_input,
_upad_input,
@@ -143,10 +158,21 @@ class FlashAttentionImpl(AttentionImpl):
)
assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
- query_length = query.size(1)
- q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
- query, key, value, attention_mask, query_length, _unpad_input
- )
+ batch_size, query_length = query.shape[:2]
+ if not self.is_cross_attn and query_length == key.size(1):
+ q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+ query, key, value, attention_mask, query_length, _unpad_input
+ )
+ else:
+ # Cross-attention: the mask covers keys only, so keep every query row.
+ k, indices_k, cu_seq_lens_k, max_length_k, _ = _unpad_input(key, attention_mask)
+ v = _index_first_axis(value, indices_k)
+ q = query.flatten(0, 1)
+ cu_seq_lens_q = torch.arange(
+ 0, (batch_size + 1) * query_length, query_length, dtype=torch.int32, device=query.device
+ )
+ max_length_q = query_length
+ indices_q = None
out_unpad = flash_attn_varlen_func(
q,
@@ -162,7 +188,9 @@ class FlashAttentionImpl(AttentionImpl):
},
)
out_unpad = self._unwrap_flash_output(out_unpad)
- return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+ if indices_q is None:
+ return out_unpad.reshape(batch_size, query_length, *out_unpad.shape[1:])
+ return _pad_input(out_unpad, indices_q, batch_size, query_length)
def _forward_varlen_packed(
self,
@@ -246,6 +274,105 @@ class FlashAttentionImpl(AttentionImpl):
# (b s) h d -> b s h d
return out.reshape(batch_size, q_len, *out.shape[1:])
+ def forward_paged(self, paged_kv_context) -> torch.Tensor:
+ """Run platform-native paged attention through the Omni backend.
+
+ ``DiffusionPagedAttentionAdapter`` owns BlockTables and prepares the
+ rank-local native cache context. The adapter no longer performs the
+ attention call itself; this method is the backend-owned execution
+ boundary. Its native layer wrapper keeps vLLM version-specific cache
+ and kernel details out of Omni's common ``Attention`` layer. CUDA uses
+ the native vLLM FlashAttention writer/kernel contract. Ascend writes
+ the complete K/V span once before its piecewise FIA calls.
+ """
+
+ layer = getattr(paged_kv_context, "layer", None)
+ if layer is None:
+ raise TypeError("paged_kv_context must expose a native diffusion attention layer")
+ kv_cache = getattr(layer, "kv_cache", None)
+ if kv_cache is None:
+ raise RuntimeError(f"Native KV cache is not bound for diffusion layer {layer.layer_name!r}")
+ native_impl = getattr(layer, "impl", None)
+ if native_impl is None:
+ raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}")
+ prewrite_kv = current_omni_platform.requires_diffusion_paged_kv_prewrite()
+ read_kv_from_cache = prewrite_kv and layer.attn_backend.forward_includes_kv_cache_update
+ # The GPU/default path preserves native cache-update ownership. Ascend
+ # prewrites through vLLM-Ascend's normal-layout writer so piecewise FIA
+ # segments do not repeatedly scatter the same layer K/V.
+ if not layer.attn_backend.forward_includes_kv_cache_update or prewrite_kv:
+ cache_update = getattr(native_impl, "do_kv_cache_update", None)
+ if not callable(cache_update):
+ raise RuntimeError(
+ f"Native attention implementation for {layer.layer_name!r} cannot update the KV cache"
+ )
+ cache_update(
+ layer,
+ paged_kv_context.key_write,
+ paged_kv_context.value_write,
+ kv_cache,
+ paged_kv_context.slot_mapping,
+ )
+
+ def run_native_attention(
+ query: torch.Tensor,
+ key: torch.Tensor | None,
+ value: torch.Tensor | None,
+ native_metadata,
+ output: torch.Tensor | None,
+ ) -> torch.Tensor:
+ if output is None:
+ output = torch.empty(
+ (query.shape[0], layer.num_heads, layer.head_size_v),
+ dtype=query.dtype,
+ device=query.device,
+ )
+ return native_impl.forward(
+ layer,
+ query,
+ None if read_kv_from_cache else key,
+ None if read_kv_from_cache else value,
+ kv_cache,
+ native_metadata,
+ output,
+ )
+
+ if paged_kv_context.piecewise_plan is not None:
+ # Ascend FIA can execute identical CFG rows as one batched call
+ # per piece. Keep that batch layout through the piecewise runner
+ # so it can concatenate row-major results instead of emitting an
+ # indexed ScatterUpdate for every segment. The CUDA path keeps
+ # its existing output-buffer contract (including graph capture).
+ use_homogeneous_batch = (
+ current_omni_platform.is_npu() and paged_kv_context.piecewise_plan.homogeneous_batch_shape is not None
+ )
+ output = None
+ if not use_homogeneous_batch:
+ output = torch.empty(
+ (paged_kv_context.query.shape[0], layer.num_heads, layer.head_size_v),
+ dtype=paged_kv_context.query.dtype,
+ device=paged_kv_context.query.device,
+ )
+ output = run_paged_piecewise_plan(
+ paged_kv_context.query,
+ None if read_kv_from_cache else paged_kv_context.key_write,
+ None if read_kv_from_cache else paged_kv_context.value_write,
+ paged_kv_context.piecewise_plan,
+ paged_kv_context.piecewise_native_metadata,
+ run_native_attention,
+ output_buffer=output,
+ use_homogeneous_batch=use_homogeneous_batch,
+ )
+ else:
+ output = run_native_attention(
+ paged_kv_context.query,
+ paged_kv_context.key_write,
+ paged_kv_context.value_write,
+ paged_kv_context.native_metadata,
+ None,
+ )
+ return paged_kv_context.restore_output(output)
+
def forward_cuda(
self,
query: torch.Tensor,
@@ -417,15 +544,8 @@ class FlashAttentionImpl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata = None,
) -> torch.Tensor:
- try:
- from mindiesd import attention_forward
- except ImportError:
- raise ImportError(
- "FlashAttentionBackend NPU implementation requires MindIE-SD. "
- "Please install MindIE-SD to enable NPU attention support. "
- "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
- "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
- )
+ from mindiesd import attention_forward
+
# Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the
# model marks extra["npu_attn_varlen"] and carries packed metadata, so
# the padding document is excluded without reading or materializing
@@ -534,15 +654,8 @@ class FlashAttentionImpl(AttentionImpl):
return None
seq_q, seq_k = resolved
- try:
- from mindiesd import attention_forward_varlen
- except ImportError:
- raise ImportError(
- "FlashAttentionBackend NPU implementation requires MindIE-SD. "
- "Please install MindIE-SD to enable NPU attention support. "
- "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
- "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
- )
+ from mindiesd import attention_forward_varlen
+
q = query.squeeze(0) # [T, N, D] == TND
k = key.squeeze(0)
v = value.squeeze(0)
@@ -590,15 +703,8 @@ class FlashAttentionImpl(AttentionImpl):
_, seq_k = resolved
used_k = seq_k[0] # real document length (first cumulative end)
- try:
- from mindiesd import attention_forward
- except ImportError:
- raise ImportError(
- "FlashAttentionBackend NPU implementation requires MindIE-SD. "
- "Please install MindIE-SD to enable NPU attention support. "
- "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
- "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
- )
+ from mindiesd import attention_forward
+
# mindiesd always takes BSND input regardless of `layout`; the arg
# selects the op-internal layout and laser only supports BNSD, so do
# NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here.
diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py
index 28b032712..9bbde02fa 100644
--- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py
+++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py
@@ -1,8 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
+import functools
+import inspect
import math
from dataclasses import dataclass
from typing import Any
@@ -50,6 +52,27 @@ _MISSING_MINDIESD = (
"Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN"
)
+_INCOMPATIBLE_MINDIESD = (
+ "RAINFUSION_ATTN requires a MindIE-SD build whose sparse_attention supports the video_spans "
+ "argument. Please upgrade MindIE-SD or select FLASH_ATTN."
+)
+
+
+# Whether the installed mindiesd ``sparse_attention`` accepts ``precision=``.
+# Releases without it accept the kwarg through ``**kwargs`` but silently ignore
+# it, so a requested mix/fp8 mode would silently run the BF16 path. Cached
+# because forwards run per layer per denoise step.
+@functools.cache
+def _mindiesd_supports_precision() -> bool:
+ try:
+ from inspect import signature
+
+ from mindiesd import sparse_attention
+
+ return "precision" in signature(sparse_attention).parameters
+ except Exception:
+ return False
+
def _try_extract_layer_index(prefix: str) -> int | None:
if not prefix:
@@ -60,6 +83,13 @@ def _try_extract_layer_index(prefix: str) -> int | None:
return None
+def _supports_video_spans(sparse_attention: Any) -> bool:
+ try:
+ return "video_spans" in inspect.signature(sparse_attention).parameters
+ except (TypeError, ValueError):
+ return False
+
+
@dataclass(frozen=True)
class RainFusionConfig:
"""Resolved RainFusion controls for one attention layer.
@@ -72,6 +102,8 @@ class RainFusionConfig:
sparsity: float = 0.0
start_step: int = 0
+ end_step: int = 0
+ precision: str = "bf16"
skip_layers: frozenset[int] = frozenset()
@classmethod
@@ -80,6 +112,8 @@ class RainFusionConfig:
return cls(
sparsity=float(bk.get("sparsity", 0.0)),
start_step=int(bk.get("start_step", 0)),
+ end_step=int(bk.get("end_step", 0)),
+ precision=str(bk.get("precision", "bf16")),
skip_layers=frozenset(bk.get("skip_layers") or ()),
)
@@ -92,9 +126,10 @@ class RainFusionConfig:
class RainFusionPlan:
"""Per-forward geometry handed to the rf_v2 kernel."""
- prefix_len: int
used_len: int
- latent_shape: list[int]
+ prefix_len: int | None = None
+ latent_shape: list[int] | None = None
+ video_spans: list[dict[str, object]] | None = None
class RainFusionAttentionBackend(AttentionBackend):
@@ -107,6 +142,12 @@ class RainFusionAttentionBackend(AttentionBackend):
if find_spec("mindiesd") is None:
raise ValueError(_MISSING_MINDIESD)
+ try:
+ from mindiesd import sparse_attention
+ except ImportError as exc:
+ raise ValueError(_MISSING_MINDIESD) from exc
+ if not _supports_video_spans(sparse_attention):
+ raise ValueError(_INCOMPATIBLE_MINDIESD)
@staticmethod
def get_supported_head_sizes() -> list[int]:
@@ -232,8 +273,17 @@ class RainFusionAttentionImpl(AttentionImpl):
return None
if is_forward_context_available():
step_idx = get_forward_context().denoise_step_idx
+ total_steps = get_forward_context().total_denoise_steps
if step_idx is not None and step_idx < rf.start_step:
return None
+ # Tail fallback: keep the last ``end_step`` denoise steps dense.
+ if (
+ rf.end_step > 0
+ and step_idx is not None
+ and total_steps is not None
+ and step_idx >= total_steps - rf.end_step
+ ):
+ return None
if self.qkv_layout is None:
# The sparse path reads the sequence off dim 1, which the tensors alone
# do not establish, and the dense fallback resolves an absent layout its
@@ -266,6 +316,14 @@ class RainFusionAttentionImpl(AttentionImpl):
)
return None
+ if layout.video_spans:
+ return self._resolve_multi_span_plan(layout, max_seqlen_q)
+
+ if layout.prefix_len is None or layout.latent_grid is None:
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: video layout has neither a legacy video tail nor multi-video spans."
+ )
+ return None
prefix_len = int(layout.prefix_len)
latent_shape = [int(dim) for dim in layout.latent_grid]
# rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of
@@ -304,11 +362,97 @@ class RainFusionAttentionImpl(AttentionImpl):
video_len,
)
return RainFusionPlan(
- prefix_len=prefix_len,
used_len=used_len,
+ prefix_len=prefix_len,
latent_shape=latent_shape,
)
+ def _resolve_multi_span_plan(self, layout, max_seqlen_q: int) -> RainFusionPlan | None:
+ """Validate Ref2VA's non-contiguous video grids before sparse dispatch."""
+ used_len = layout.used_len
+ if used_len is None or int(used_len) != int(max_seqlen_q):
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: multi-video layout used_len=%r does not match packed document 0 (%d).",
+ used_len,
+ int(max_seqlen_q),
+ )
+ return None
+
+ spans: list[dict[str, object]] = []
+ span_summaries: list[str] = []
+ previous_end = 0
+ target_count = 0
+ video_seqlen = 0
+ boundary_dense_seqlen = 0
+ previous_video_length: int | None = None
+ for span in sorted(layout.video_spans, key=lambda item: item.start):
+ grid = tuple(int(dim) for dim in span.latent_grid)
+ length = math.prod(grid)
+ start = int(span.start)
+ if (
+ len(grid) != 3
+ or any(dim <= 0 for dim in grid)
+ or start < previous_end
+ or start + length > int(used_len)
+ ):
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: invalid multi-video span start=%d grid=%s used_len=%d.",
+ start,
+ grid,
+ int(used_len),
+ )
+ return None
+ role = span.role
+ if role not in ("reference", "target"):
+ logger.warning_once("RAINFUSION_ATTN staying dense: unsupported multi-video span role %r.", role)
+ return None
+ if role == "target":
+ target_count += 1
+ if previous_video_length is not None:
+ # rf_v2 works on fixed 128-token blocks. The preceding clip
+ # needs these dense rows to complete its tail block before
+ # this clip begins, otherwise one sparse block would cross a
+ # clip boundary.
+ boundary_dense_seqlen += (-previous_video_length) % _BLOCK_SIZE
+ spans.append({"start": start, "latent_shape": list(grid)})
+ span_summaries.append(f"role={role}, start={start}, seqlen={length}, latent_shape={grid}")
+ previous_end = start + length
+ video_seqlen += length
+ previous_video_length = length
+
+ if target_count != 1:
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: Ref2VA layout must contain exactly one target video span, got %d.",
+ target_count,
+ )
+ return None
+ if video_seqlen < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: multi-video seqlen=%d is under the sparse threshold seqlen=%d.",
+ video_seqlen,
+ _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
+ )
+ return None
+ dense_context_seqlen = int(used_len) - video_seqlen
+ if boundary_dense_seqlen > dense_context_seqlen:
+ logger.warning_once(
+ "RAINFUSION_ATTN staying dense: multi-video spans need dense_context_seqlen=%d "
+ "to isolate clip block boundaries, but this layout has only %d.",
+ boundary_dense_seqlen,
+ dense_context_seqlen,
+ )
+ return None
+ logger.info_once(
+ "RAINFUSION_ATTN multi-video active: sparsity=%.2f, spans=[%s], "
+ "video_seqlen=%d, dense_context_seqlen=%d, valid_packed_seqlen=%d.",
+ self.rainfusion.sparsity,
+ "; ".join(span_summaries),
+ video_seqlen,
+ dense_context_seqlen,
+ int(used_len),
+ )
+ return RainFusionPlan(used_len=int(used_len), video_spans=spans)
+
def _forward_sparse_npu(
self,
query: torch.Tensor,
@@ -320,26 +464,45 @@ class RainFusionAttentionImpl(AttentionImpl):
from mindiesd import sparse_attention
except ImportError:
raise ImportError(_MISSING_MINDIESD)
+ if self.rainfusion.precision != "bf16" and not _mindiesd_supports_precision():
+ raise RuntimeError(
+ f"block_sparse.precision={self.rainfusion.precision!r} requires MindIE-SD "
+ "with sparse_attention(precision=...) support; the installed mindiesd "
+ "silently ignores it and would run the BF16 path. Install a compatible "
+ "MindIE-SD release or use precision='bf16'."
+ )
used = plan.used_len
q, k, v = (tensor[:, :used] for tensor in (query, key, value))
# Ulysses has already gathered the full sequence onto this rank and split
# the heads, so read the head count off the tensor rather than num_heads.
- out = sparse_attention(
- q,
- k,
- v,
- scale=self.softmax_scale,
- head_num=query.shape[-2],
- input_layout=_INPUT_LAYOUT,
- inner_precise=_INNER_PRECISE,
- sparse_type="rf_v2",
- txt_len=plan.prefix_len,
- block_size=_BLOCK_SIZE,
- latent_shape_q=plan.latent_shape,
- latent_shape_k=plan.latent_shape,
- sparsity=self.rainfusion.sparsity,
- )
+ common_kwargs: dict[str, object] = {
+ "scale": self.softmax_scale,
+ "head_num": query.shape[-2],
+ "input_layout": _INPUT_LAYOUT,
+ "inner_precise": _INNER_PRECISE,
+ "block_size": _BLOCK_SIZE,
+ "sparsity": self.rainfusion.sparsity,
+ "precision": self.rainfusion.precision,
+ }
+ if plan.video_spans is not None:
+ out = sparse_attention(
+ q,
+ k,
+ v,
+ sparse_type="rf_v2",
+ video_spans=plan.video_spans,
+ **common_kwargs,
+ )
+ else:
+ assert plan.prefix_len is not None and plan.latent_shape is not None
+ common_kwargs.update(
+ sparse_type="rf_v2",
+ txt_len=plan.prefix_len,
+ latent_shape_q=plan.latent_shape,
+ latent_shape_k=plan.latent_shape,
+ )
+ out = sparse_attention(q, k, v, **common_kwargs)
if used == query.shape[1]:
return out
padded = torch.zeros_like(query)
diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py
index fdfaf5f9c..1af079696 100644
--- a/vllm_omni/diffusion/attention/backends/registry.py
+++ b/vllm_omni/diffusion/attention/backends/registry.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Diffusion attention backend registry.
This module provides an enum-based registry for diffusion attention backends,
@@ -23,7 +23,7 @@ logger = init_logger(__name__)
class _DiffusionBackendEnumMeta(EnumMeta):
"""Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
- def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
+ def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum": # type: ignore[override]
"""Get backend by name with helpful error messages."""
try:
return super().__getitem__(name) # type: ignore[return-value]
@@ -62,6 +62,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
+ FASTVIDEO_VSA = "vllm_omni.diffusion.attention.backends.fastvideo_vsa.FastVideoVSABackend"
FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend"
diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py
index c63280fc8..967e81c6e 100644
--- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py
+++ b/vllm_omni/diffusion/attention/backends/trtllm_attn.py
@@ -1,10 +1,11 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import functools
import inspect
import math
from dataclasses import dataclass, replace
+from typing import NamedTuple, cast
import torch
from vllm.logger import init_logger
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
AttentionBackend,
AttentionImpl,
AttentionMetadata,
+ PackedPaddingMetadata,
)
logger = init_logger(__name__)
@@ -42,8 +44,9 @@ class SkipSoftmaxConfig:
return cls(
threshold=_validate_control(bk.get("skip_softmax_threshold"), "skip_softmax_threshold", 0.0, None),
target_sparsity=_validate_control(bk.get("target_sparsity"), "target_sparsity", 0.0, 1.0),
- disabled_until_timestep=_validate_control(
- bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0
+ disabled_until_timestep=cast(
+ float,
+ _validate_control(bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0),
),
)
@@ -131,7 +134,7 @@ class QuantConfig:
return self.dtype_qk is not None
def quantize(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, quantize_fn):
- qk_quant_dtype = _QK_QUANT_DTYPES[self.dtype_qk]
+ qk_quant_dtype = _QK_QUANT_DTYPES[cast(str, self.dtype_qk)]
q_q, k_q, v_q, q_sfs, k_sfs, v_sfs = quantize_fn(
q,
k,
@@ -145,6 +148,18 @@ class QuantConfig:
return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk
+class _PackedLayout(NamedTuple):
+ """Active token ranges and kernel metadata for a flattened ragged batch."""
+
+ q_tokens: int
+ kv_tokens: int
+ cu_seqlens_q: torch.Tensor
+ cu_seqlens_kv: torch.Tensor
+ batch_size: int
+ seq_lens: torch.Tensor
+ known_min_kv_length: int | None = None
+
+
def _workspace_bytes() -> int:
import vllm.envs as envs
@@ -154,6 +169,14 @@ def _workspace_bytes() -> int:
class TrtllmAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True
+ @classmethod
+ def supports_packed_mask_free(cls) -> bool:
+ return True
+
+ @classmethod
+ def supports_multi_doc_packed_varlen(cls) -> bool:
+ return True
+
@staticmethod
def get_supported_head_sizes() -> list[int]:
return [128]
@@ -250,6 +273,97 @@ class TrtllmAttentionImpl(AttentionImpl):
cls._workspace = ws
return ws
+ @staticmethod
+ def _prepare_packed_padding_layout(
+ physical_batch: int,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ packed_padding: PackedPaddingMetadata,
+ extra: dict,
+ ) -> _PackedLayout:
+ """Prepare a [real, pad] layout from an explicit host-side boundary."""
+ valid_q_tokens = packed_padding.q_length
+ valid_kv_tokens = packed_padding.kv_length
+ max_q_len = extra["max_seqlen_q"]
+ max_kv_len = extra["max_seqlen_k"]
+ lengths = (valid_q_tokens, valid_kv_tokens, max_q_len, max_kv_len)
+ if any(isinstance(value, bool) or not isinstance(value, int) for value in lengths):
+ raise ValueError("Mask-free packed TRTLLM attention lengths must be Python integers")
+ if physical_batch != 1:
+ raise ValueError("Mask-free packed TRTLLM attention requires a single packed batch")
+
+ total_q_tokens = q.shape[0]
+ total_kv_tokens = k.shape[0]
+ if not 0 < valid_q_tokens <= total_q_tokens:
+ raise ValueError(
+ "PackedPaddingMetadata.q_length must be within the packed Q sequence, "
+ f"got {valid_q_tokens} for length {total_q_tokens}"
+ )
+ if not 0 < valid_kv_tokens <= total_kv_tokens:
+ raise ValueError(
+ "PackedPaddingMetadata.kv_length must be within the packed K/V sequence, "
+ f"got {valid_kv_tokens} for length {total_kv_tokens}"
+ )
+ if max_q_len != valid_q_tokens or max_kv_len != valid_kv_tokens:
+ raise ValueError("Mask-free packed TRTLLM attention lengths must match max_seqlen_q/k")
+ published_kv_length = extra.get("valid_kv_length")
+ if published_kv_length is not None:
+ if isinstance(published_kv_length, bool) or not isinstance(published_kv_length, int):
+ raise ValueError("valid_kv_length must be a Python integer")
+ if published_kv_length != valid_kv_tokens:
+ raise ValueError("PackedPaddingMetadata.kv_length must match valid_kv_length")
+
+ cu_seq_lens_q = packed_padding.cu_seqlens_q
+ cu_seq_lens_kv = packed_padding.cu_seqlens_k
+ if cu_seq_lens_q.dtype != torch.int32 or cu_seq_lens_kv.dtype != torch.int32:
+ raise ValueError("Mask-free packed TRTLLM attention requires int32 cu_seqlens")
+ if cu_seq_lens_q.device != q.device or cu_seq_lens_kv.device != k.device:
+ raise ValueError("Mask-free packed TRTLLM attention metadata must be on the Q/K device")
+ if cu_seq_lens_q.shape != (2,) or cu_seq_lens_kv.shape != (2,):
+ raise ValueError("Mask-free packed TRTLLM attention requires canonical two-element cu_seqlens")
+ return _PackedLayout(
+ q_tokens=valid_q_tokens,
+ kv_tokens=valid_kv_tokens,
+ cu_seqlens_q=cu_seq_lens_q,
+ cu_seqlens_kv=cu_seq_lens_kv,
+ batch_size=1,
+ # Canonical packed-padding metadata is [0, valid_kv_tokens], so this slice
+ # is the one-element sequence-length view expected by the kernel.
+ seq_lens=cu_seq_lens_kv[1:],
+ known_min_kv_length=valid_kv_tokens,
+ )
+
+ @staticmethod
+ def _prepare_generic_packed_layout(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cu_seq_lens_q: torch.Tensor,
+ cu_seq_lens_kv: torch.Tensor,
+ ) -> _PackedLayout:
+ """Validate and normalize a general ragged batch."""
+ if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
+ raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")
+
+ # Packed producers may retain a zero-length trailing segment when the
+ # input is already aligned. It is not a real sequence and is removed
+ # before deriving the batch size and per-sequence lengths.
+ while (
+ cu_seq_lens_q.numel() > 2
+ and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
+ and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
+ ):
+ cu_seq_lens_q = cu_seq_lens_q[:-1]
+ cu_seq_lens_kv = cu_seq_lens_kv[:-1]
+
+ return _PackedLayout(
+ q_tokens=q.shape[0],
+ kv_tokens=k.shape[0],
+ cu_seqlens_q=cu_seq_lens_q,
+ cu_seqlens_kv=cu_seq_lens_kv,
+ batch_size=cu_seq_lens_q.numel() - 1,
+ seq_lens=(cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous(),
+ )
+
def forward_cuda(
self,
query: torch.Tensor,
@@ -266,12 +380,16 @@ class TrtllmAttentionImpl(AttentionImpl):
has_packed_metadata = len(present_packed_keys) == len(packed_keys)
attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
- if attn_mask is not None and not has_packed_metadata:
+ packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None
+ if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata):
+ raise ValueError("packed_padding must be PackedPaddingMetadata")
+ if packed_padding is not None and not has_packed_metadata:
+ raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata")
+ if attn_mask is not None:
raise ValueError(
- "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with "
- "mask_sp_padding=True and a non-divisible seq len). Either set "
- "parallel_config.mask_sp_padding=False, or select a mask-capable backend "
- "(e.g. CUDNN_ATTN / TORCH_SDPA)."
+ "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding "
+ "with packed-padding metadata, or select a mask-capable backend such as "
+ "CUDNN_ATTN or TORCH_SDPA."
)
if not HAS_FLASHINFER:
@@ -282,7 +400,6 @@ class TrtllmAttentionImpl(AttentionImpl):
physical_batch, q_len, num_q_heads, head_dim = query.shape
kv_len, num_kv_heads = key.shape[1], key.shape[2]
-
device = query.device
q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous()
@@ -290,6 +407,7 @@ class TrtllmAttentionImpl(AttentionImpl):
v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
output_tokens = q.shape[0]
+ known_min_kv_length: int | None = kv_len
if has_packed_metadata:
cu_seq_lens_q = extra["cu_seqlens_q"]
cu_seq_lens_kv = extra["cu_seqlens_k"]
@@ -297,48 +415,36 @@ class TrtllmAttentionImpl(AttentionImpl):
raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional")
if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2:
raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches")
- if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
- raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")
- # Drop trailing empty sequences retained by packed producers when
- # alignment adds no padding tokens.
- while (
- cu_seq_lens_q.numel() > 2
- and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
- and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
- ):
- cu_seq_lens_q = cu_seq_lens_q[:-1]
- cu_seq_lens_kv = cu_seq_lens_kv[:-1]
+ # PackedPaddingMetadata is an optional producer-published shortcut
+ # for structural suffix padding. Inputs without it retain the
+ # general cu_seqlens path for arbitrary ragged Q/KV batches.
+ if packed_padding is not None:
+ packed_layout = self._prepare_packed_padding_layout(
+ physical_batch,
+ q,
+ k,
+ packed_padding,
+ extra,
+ )
+ else:
+ # Generic metadata can describe any number of real sequences,
+ # not only a single [real, pad] pair.
+ packed_layout = self._prepare_generic_packed_layout(
+ q,
+ k,
+ cu_seq_lens_q,
+ cu_seq_lens_kv,
+ )
- # A packed prefix mask denotes structural suffix padding, not another
- # document. Exclude it before quantization and attention; in particular,
- # SAGE block quantization must not span the valid/padding boundary.
- if attn_mask is not None:
- if q.shape[0] != k.shape[0] or attn_mask.numel() != q.shape[0]:
- raise ValueError("Packed TRTLLM prefix masks require equal Q/K lengths")
- flat_mask = attn_mask.reshape(-1).to(dtype=torch.bool)
- valid_tokens = int(flat_mask.sum().item())
- expected_mask = torch.arange(flat_mask.numel(), device=device) < valid_tokens
- if not torch.equal(flat_mask, expected_mask):
- raise ValueError("TRTLLM_ATTN only supports prefix-valid packed attention masks")
-
- q_boundary = (cu_seq_lens_q == valid_tokens).nonzero(as_tuple=False)
- kv_boundary = (cu_seq_lens_kv == valid_tokens).nonzero(as_tuple=False)
- if q_boundary.numel() != 1 or kv_boundary.numel() != 1:
- raise ValueError("Packed TRTLLM prefix length must be a Q/K sequence boundary")
- q_end = int(q_boundary.item()) + 1
- kv_end = int(kv_boundary.item()) + 1
- if q_end != kv_end:
- raise ValueError("Packed TRTLLM prefix masks require matching Q/K sequence batches")
-
- q = q[:valid_tokens]
- k = k[:valid_tokens]
- v = v[:valid_tokens]
- cu_seq_lens_q = cu_seq_lens_q[:q_end].contiguous()
- cu_seq_lens_kv = cu_seq_lens_kv[:kv_end].contiguous()
-
- batch = cu_seq_lens_q.numel() - 1
- seq_lens = (cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous()
+ q = q[: packed_layout.q_tokens]
+ k = k[: packed_layout.kv_tokens]
+ v = v[: packed_layout.kv_tokens]
+ cu_seq_lens_q = packed_layout.cu_seqlens_q
+ cu_seq_lens_kv = packed_layout.cu_seqlens_kv
+ batch = packed_layout.batch_size
+ seq_lens = packed_layout.seq_lens
+ known_min_kv_length = packed_layout.known_min_kv_length
max_q_len = int(extra["max_seqlen_q"])
max_kv_len = int(extra["max_seqlen_k"])
else:
@@ -361,7 +467,13 @@ class TrtllmAttentionImpl(AttentionImpl):
sage_kwargs: dict = {}
# The SAGE kernel requires every KV sequence to contain at least one full
# quantization block. Small auxiliary attention sites use the dense kernel.
- use_sage = self.quant.enabled and bool(torch.all(seq_lens >= self.quant.k_block_size).item())
+ use_sage = False
+ if self.quant.enabled:
+ if known_min_kv_length is not None:
+ sage_lengths_supported = known_min_kv_length >= self.quant.k_block_size
+ else:
+ sage_lengths_supported = bool(torch.all(seq_lens >= self.quant.k_block_size).item())
+ use_sage = sage_lengths_supported
if self.quant.enabled and not use_sage:
message = (
f"TRTLLM_ATTN SAGE quantization is configured for attention role {self.role!r}, but at least one "
diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py
index 5d83d8dbf..a8676c34a 100644
--- a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py
+++ b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Piecewise attention for mixed causal / full (bidirectional) masks.
@@ -15,6 +15,8 @@ Per segment:
from __future__ import annotations
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal, NamedTuple
import torch
@@ -90,6 +92,7 @@ def piecewise_attn(
):
_check_homogeneous(full_attn_spans)
spans = full_attn_spans[0]
+ ranges: tuple[tuple[int, int, int], ...]
if query_ranges is None:
query_len = query.shape[1]
ranges = ((0, query_len, key.shape[1] - query_len),)
@@ -123,3 +126,282 @@ def piecewise_attn(
if len(outputs) == 1:
return outputs[0]
return torch.cat(outputs, dim=1)
+
+
+@dataclass(frozen=True, slots=True)
+class PagedPiecewiseSegment:
+ """One aligned segment across rows in a packed paged batch."""
+
+ row_segments: tuple[Segment, ...]
+ query_indices: torch.Tensor
+ query_range: tuple[int, int] | None
+ # Local [start, end) range shared by every row when the batch can retain
+ # its original [B, S] layout. The packed query_indices remain available
+ # for metadata construction and the heterogeneous fallback.
+ local_query_range: tuple[int, int] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class PagedPiecewisePlan:
+ """Piecewise segment batches for one packed paged-attention batch."""
+
+ spans: tuple[tuple[tuple[int, int], ...], ...]
+ segments: tuple[PagedPiecewiseSegment, ...]
+ num_query_tokens: int
+ segments_cover_query_contiguously: bool
+ # (batch_size, query_len) when all rows have the same local segment
+ # layout. Native FIA still receives flattened row-major tokens, but the
+ # runner can restore outputs with one batch concat instead of one indexed
+ # scatter per segment.
+ homogeneous_batch_shape: tuple[int, int] | None = None
+
+
+PagedPiecewiseRunner = Callable[
+ [torch.Tensor, torch.Tensor | None, torch.Tensor | None, object, torch.Tensor | None],
+ torch.Tensor,
+]
+
+
+def build_paged_piecewise_plan(
+ full_attn_spans: Sequence[Sequence[tuple[int, int]]],
+ query_offsets: Sequence[int],
+ query_lens: Sequence[int],
+ seq_lens: Sequence[int],
+ *,
+ device: torch.device | str | None = None,
+) -> PagedPiecewisePlan:
+ """Build packed indices for corresponding piecewise segments in each row."""
+
+ row_count = len(full_attn_spans)
+ if row_count == 0 or not (row_count == len(query_offsets) == len(query_lens) == len(seq_lens)):
+ raise ValueError("Paged piecewise inputs must have one entry per row")
+
+ spans: tuple[tuple[tuple[int, int], ...], ...] = tuple(
+ tuple((span_start, span_end) for span_start, span_end in row_spans) for row_spans in full_attn_spans
+ )
+
+ packed_offsets = [0]
+ for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True):
+ previous_end = 0
+ for start, end in row_spans:
+ if start < previous_end or not 0 <= start < end <= seq_len:
+ raise ValueError("Paged piecewise spans must be sorted, non-overlapping, and within the sequence")
+ previous_end = end
+ if query_offset < 0 or query_len <= 0 or query_offset + query_len > seq_len:
+ raise ValueError("Paged piecewise query range must be within the sequence")
+ packed_offsets.append(packed_offsets[-1] + query_len)
+
+ segments_by_row = tuple(
+ tuple(build_segments(row_spans, query_offset, query_len))
+ for row_spans, query_offset, query_len in zip(spans, query_offsets, query_lens, strict=True)
+ )
+ if len({len(row_segments) for row_segments in segments_by_row}) != 1:
+ raise ValueError("Paged piecewise rows must produce aligned segments")
+
+ packed_segments = []
+ for row_segments in zip(*segments_by_row, strict=True):
+ if len({segment.mode for segment in row_segments}) != 1:
+ raise ValueError("Paged piecewise rows must use the same attention mode per segment")
+ query_indices: list[int] = []
+ local_ranges: list[tuple[int, int]] = []
+ for row_index, (segment, query_offset) in enumerate(zip(row_segments, query_offsets, strict=True)):
+ local_start = segment.q_start - query_offset
+ local_end = segment.q_end - query_offset
+ local_ranges.append((local_start, local_end))
+ query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end))
+ query_range = None
+ if query_indices and query_indices == list(range(query_indices[0], query_indices[-1] + 1)):
+ query_range = (query_indices[0], query_indices[-1] + 1)
+ local_query_range = local_ranges[0] if len(set(local_ranges)) == 1 else None
+ packed_segments.append(
+ PagedPiecewiseSegment(
+ row_segments=tuple(row_segments),
+ query_indices=torch.tensor(query_indices, dtype=torch.long, device=device),
+ query_range=query_range,
+ local_query_range=local_query_range,
+ )
+ )
+
+ covered = 0
+ for segment in packed_segments:
+ if segment.query_range is None or segment.query_range[0] != covered:
+ break
+ covered = segment.query_range[1]
+ segments_cover_query_contiguously = covered == packed_offsets[-1]
+ homogeneous_batch_shape: tuple[int, int] | None = None
+ if (
+ row_count > 1
+ and len(set(query_lens)) == 1
+ and all(segment.local_query_range is not None for segment in packed_segments)
+ ):
+ # A common local range is sufficient even when rows have different
+ # global offsets or KV lengths; those values remain row-specific in
+ # native metadata.
+ homogeneous_batch_shape = (row_count, int(query_lens[0]))
+ return PagedPiecewisePlan(
+ spans=spans,
+ segments=tuple(packed_segments),
+ num_query_tokens=packed_offsets[-1],
+ segments_cover_query_contiguously=segments_cover_query_contiguously,
+ homogeneous_batch_shape=homogeneous_batch_shape,
+ )
+
+
+def _run_homogeneous_paged_piecewise_plan(
+ query: torch.Tensor,
+ key: torch.Tensor | None,
+ value: torch.Tensor | None,
+ plan: PagedPiecewisePlan,
+ segment_metadata: Sequence[object],
+ segment_runner: PagedPiecewiseRunner,
+ output_buffer: torch.Tensor | None,
+) -> torch.Tensor:
+ """Run a homogeneous batch while retaining the old [B, S] layout.
+
+ The native FIA contract is flattened TND, so each segment is flattened
+ only for the kernel call. Results are reshaped back to [B, L] and
+ concatenated along the sequence dimension, which avoids the per-segment
+ ``index_copy_`` used by heterogeneous packed rows.
+ """
+
+ batch_shape = plan.homogeneous_batch_shape
+ assert batch_shape is not None
+ batch_size, query_len = batch_shape
+ if query.shape[0] != batch_size * query_len:
+ raise ValueError(
+ f"Homogeneous paged piecewise plan expects {batch_size * query_len} query tokens, got {query.shape[0]}"
+ )
+ query_rows = query.reshape(batch_size, query_len, *query.shape[1:])
+ key_rows = None if key is None else key.reshape(batch_size, query_len, *key.shape[1:])
+ value_rows = None if value is None else value.reshape(batch_size, query_len, *value.shape[1:])
+
+ segment_outputs: list[torch.Tensor] = []
+ for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
+ local_range = segment.local_query_range
+ if local_range is None:
+ raise ValueError("Homogeneous paged piecewise execution requires a shared local segment range")
+ start, end = local_range
+ segment_query = query_rows[:, start:end].contiguous().reshape(-1, *query.shape[1:])
+ segment_key = None if key_rows is None else key_rows[:, start:end].contiguous().reshape(-1, *key.shape[1:])
+ segment_value = (
+ None if value_rows is None else value_rows[:, start:end].contiguous().reshape(-1, *value.shape[1:])
+ )
+ segment_output = segment_runner(segment_query, segment_key, segment_value, metadata, None)
+ expected_tokens = batch_size * (end - start)
+ if segment_output.shape[0] != expected_tokens:
+ raise ValueError(
+ f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+ f"for a {expected_tokens}-token homogeneous segment"
+ )
+ segment_outputs.append(segment_output.reshape(batch_size, end - start, *segment_output.shape[1:]))
+
+ if not segment_outputs:
+ result = torch.empty_like(query)
+ elif len(segment_outputs) == 1:
+ result = segment_outputs[0].reshape(plan.num_query_tokens, *segment_outputs[0].shape[2:])
+ else:
+ result_rows = torch.cat(segment_outputs, dim=1)
+ result = result_rows.reshape(plan.num_query_tokens, *result_rows.shape[2:])
+ if output_buffer is not None:
+ if output_buffer.shape != result.shape:
+ raise ValueError(
+ f"Paged piecewise output buffer shape {tuple(output_buffer.shape)} does not match "
+ f"homogeneous result shape {tuple(result.shape)}"
+ )
+ output_buffer.copy_(result)
+ return output_buffer
+ return result
+
+
+def run_paged_piecewise_plan(
+ query: torch.Tensor,
+ key: torch.Tensor | None,
+ value: torch.Tensor | None,
+ plan: PagedPiecewisePlan,
+ segment_metadata: Sequence[object],
+ segment_runner: PagedPiecewiseRunner,
+ output_buffer: torch.Tensor | None = None,
+ *,
+ use_homogeneous_batch: bool = False,
+) -> torch.Tensor:
+ """Run piecewise attention with view fast paths for contiguous segments.
+
+ ``use_homogeneous_batch`` is an opt-in for native backends that can keep
+ identical rows in the legacy batch layout. The default remains the
+ indexed packed path so existing GPU/heterogeneous contracts are unchanged.
+ """
+
+ if query.shape[0] != plan.num_query_tokens:
+ raise ValueError(f"Paged piecewise plan has {plan.num_query_tokens} query tokens, got {query.shape[0]}")
+ if output_buffer is not None and output_buffer.shape[0] != plan.num_query_tokens:
+ raise ValueError(
+ f"Paged piecewise output buffer has {output_buffer.shape[0]} tokens; expected {plan.num_query_tokens}"
+ )
+
+ if use_homogeneous_batch and plan.homogeneous_batch_shape is not None:
+ return _run_homogeneous_paged_piecewise_plan(
+ query,
+ key,
+ value,
+ plan,
+ segment_metadata,
+ segment_runner,
+ output_buffer,
+ )
+
+ output = output_buffer
+ contiguous_outputs = []
+ for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
+ indices = segment.query_indices
+ segment_output_buffer = None
+ if segment.query_range is None:
+ segment_query = query.index_select(0, indices)
+ segment_key = None if key is None else key.index_select(0, indices)
+ segment_value = None if value is None else value.index_select(0, indices)
+ else:
+ start, end = segment.query_range
+ segment_query = query[start:end]
+ segment_key = None if key is None else key[start:end]
+ segment_value = None if value is None else value[start:end]
+ if output is not None:
+ segment_output_buffer = output[start:end]
+ segment_output = segment_runner(
+ segment_query,
+ segment_key,
+ segment_value,
+ metadata,
+ segment_output_buffer,
+ )
+ if segment_output.shape[0] != indices.shape[0]:
+ raise ValueError(
+ f"Paged piecewise runner returned {segment_output.shape[0]} tokens "
+ f"for a {indices.shape[0]}-token segment"
+ )
+ if segment_output_buffer is not None and segment_output is not segment_output_buffer:
+ if segment_output.shape != segment_output_buffer.shape:
+ raise ValueError(
+ f"Paged piecewise runner returned shape {tuple(segment_output.shape)} "
+ f"for output buffer shape {tuple(segment_output_buffer.shape)}"
+ )
+ segment_output_buffer.copy_(segment_output)
+ segment_output = segment_output_buffer
+ if plan.segments_cover_query_contiguously:
+ if output_buffer is None:
+ contiguous_outputs.append(segment_output)
+ continue
+ if output is None:
+ output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:]))
+ if segment.query_range is None:
+ output.index_copy_(0, indices, segment_output)
+ else:
+ start, end = segment.query_range
+ output[start:end].copy_(segment_output)
+
+ if output_buffer is not None and plan.segments_cover_query_contiguously:
+ return output_buffer
+ if contiguous_outputs:
+ if len(contiguous_outputs) == 1:
+ return contiguous_outputs[0]
+ return torch.cat(contiguous_outputs, dim=0)
+ assert output is not None
+ return output
diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py
index 7eb74a941..7dca3c1e6 100644
--- a/vllm_omni/diffusion/attention/layer.py
+++ b/vllm_omni/diffusion/attention/layer.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) Microsoft Corporation and Jiarui Fang
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team & Jiarui Fang
@@ -11,8 +11,10 @@ from dataclasses import replace
import torch
import torch.nn as nn
+from vllm.config import VllmConfig, set_current_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.models.utils import extract_layer_index
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
@@ -21,8 +23,13 @@ from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
-from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.forward_context import (
+ get_forward_context,
+ get_ulysses_mode,
+ is_forward_context_available,
+)
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
@@ -38,6 +45,8 @@ def _try_extract_layer_index(prefix: str) -> int | None:
class Attention(nn.Module):
+ _scheduler_paged_kv = False
+
def __init__(
self,
num_heads: int,
@@ -61,12 +70,28 @@ class Attention(nn.Module):
# perf for this layer (e.g. Wan2.2 cross-attn has short sequences and
# block-FP8 quant offers no win). Default False = follow global config.
disable_kv_quant: bool = False,
+ # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion
+ # attention remains dense and contributes no native KVCacheSpec.
+ paged_kv_cache_role: str | None = None,
+ paged_kv_cache_dtype: torch.dtype | None = None,
):
super().__init__()
self.role = role
self.role_category = role_category
self.qkv_layout = qkv_layout
+ # ``prefix`` is also the stable layer identity used by vLLM's native
+ # KV-cache metadata. Keep it on the Omni layer so the active paged
+ # adapter can dispatch the already-resharded Q/K/V to the matching
+ # native cache tensor without replacing this Omni execution path.
+ self.prefix = prefix
+ if paged_kv_cache_role == "":
+ raise ValueError("paged_kv_cache_role must be non-empty when provided")
+ self.paged_kv_cache_role = paged_kv_cache_role
+ self.paged_kv_cache_dtype = paged_kv_cache_dtype
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+ self.head_size = head_size
# Resolve backend via role-aware config.
# The global diffusion config is set during model init via
@@ -89,6 +114,44 @@ class Attention(nn.Module):
role_category=role_category,
allow_trtllm_default=allow_trtllm_default,
)
+ scheduler_paged_kv = (
+ config is not None
+ and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ is DiffusionKVCacheMode.PAGED_SCHEDULER
+ )
+ self._scheduler_paged_kv = scheduler_paged_kv
+ if (
+ scheduler_paged_kv
+ and paged_kv_cache_role is not None
+ and spec is None
+ and not attn_backend_cls.supports_paged_kv
+ ):
+ # FLASH_ATTN is an Omni selector, not a device-specific kernel.
+ # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
+ # attention on NPU. A platform may legitimately default dense
+ # attention to SDPA (for example when an optional diffusion FA
+ # package is absent), but a paged layer must advertise the
+ # capability so the Worker can register its native cache view.
+ # Resolve the selector only for this marked layer; unmarked dense
+ # attention keeps the platform default unchanged.
+ from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+ dense_backend_name = attn_backend_cls.get_name()
+ paged_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+ attn_backend_cls, spec = get_attn_backend_for_role(
+ role=role,
+ head_size=head_size,
+ attention_config=paged_attention_config,
+ role_category=role_category,
+ allow_trtllm_default=allow_trtllm_default,
+ )
+ logger.info(
+ "Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
+ "does not support Scheduler-owned KV",
+ role,
+ attn_backend_cls.get_name(),
+ dense_backend_name,
+ )
parallel_config = getattr(config, "parallel_config", None)
allgather_degree = getattr(parallel_config, "allgather_degree", 1)
# TODO: Move AllGather-KV compatibility into an AttentionBackend capability
@@ -171,6 +234,25 @@ class Attention(nn.Module):
self._disable_kv_quant: bool = disable_kv_quant
self._init_kv_cache_quantization(config)
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
+ """Return native rank-local geometry for an opted-in paged cache."""
+
+ if self.paged_kv_cache_role is None:
+ return None
+ dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
+ # Keep backend layout discovery under the same config context used by
+ # upstream vLLM's attention-spec collector.
+ with set_current_vllm_config(vllm_config):
+ indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride()
+ return FullAttentionSpec(
+ block_size=vllm_config.cache_config.block_size,
+ num_kv_heads=self.num_kv_heads,
+ head_size=self.head_size,
+ dtype=dtype,
+ indexes_kv_by_block_stride=indexes_kv_by_block_stride,
+ non_causal=not self.causal,
+ )
+
def _get_active_parallel_strategy(self):
"""Get the parallel strategy based on current SP active state.
@@ -279,19 +361,94 @@ class Attention(nn.Module):
) -> torch.Tensor:
# Get the appropriate parallel strategy based on SP active state
strategy = self._get_active_parallel_strategy()
+ paged_adapter = self._active_paged_kv_adapter()
+ in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+ if (
+ self._scheduler_paged_kv
+ and self.paged_kv_cache_role is not None
+ and paged_adapter is None
+ and not in_kv_memory_profile
+ ):
+ raise RuntimeError(
+ "Scheduler-paged diffusion attention reached model forward without an active Worker adapter. "
+ "Only the startup KV memory profile may execute before paged KV initialization."
+ )
+ use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
+ if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
+ raise NotImplementedError(
+ f"Diffusion paged KV requires an Omni backend with paged support; "
+ f"selected {self.attn_backend.get_name()}"
+ )
+ if use_paged_attention and strategy is not self._no_parallel_strategy:
+ strategy_name = strategy.name
+ if self.use_ring or strategy_name == "ring":
+ raise NotImplementedError(
+ "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP"
+ )
+ if strategy_name == "allgather_kv":
+ raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism")
+ if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict":
+ raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses")
# 1. Prepare inputs (Communication / Resharding)
# For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v
# For Ring: Concat joint_q
query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)
- attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+ # Scheduler rows describe the logical sequence, while strict Ulysses
+ # may append synthetic tokens solely to make the image shard divisible.
+ # Remove those tokens after the all-to-all and put zero placeholders
+ # back before the reverse all-to-all.
+ paged_sp_padding = 0
+ paged_sp_padding_offset = query.shape[1]
+ if use_paged_attention and strategy is not self._no_parallel_strategy and strategy.name == "ulysses":
+ forward_ctx = get_forward_context() if is_forward_context_available() else None
+ paged_sp_padding = int(getattr(forward_ctx, "sp_padding_size", 0))
+ if paged_sp_padding:
+ joint_len = int(getattr(ctx, "joint_len", 0))
+ joint_strategy = str(getattr(ctx, "joint_strategy", "front"))
+ paged_sp_padding_offset = query.shape[1] - (joint_len if joint_strategy == "rear" else 0)
+ padding_start = paged_sp_padding_offset - paged_sp_padding
+ if padding_start < 0:
+ raise ValueError(
+ "Paged Ulysses padding exceeds the post-all-to-all sequence: "
+ f"padding={paged_sp_padding}, sequence={query.shape[1]}"
+ )
- # 2. Kernel Execution (Computation)
- if self.use_ring and strategy is not self._no_parallel_strategy:
- out = self._run_ring_attention(query, key, value, attn_metadata)
+ def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor:
+ return torch.cat(
+ (tensor[:, :padding_start], tensor[:, paged_sp_padding_offset:]),
+ dim=1,
+ ).contiguous()
+
+ query = _remove_paged_sp_padding(query)
+ key = _remove_paged_sp_padding(key)
+ value = _remove_paged_sp_padding(value)
+
+ # 2. This is the shared GPU/NPU boundary. The Worker adapter prepares
+ # the native page-table context after SP has produced rank-local Q/K/V;
+ # backend resolution below selects CUDA or Ascend execution.
+ if use_paged_attention:
+ assert paged_adapter is not None
+ paged_kv_context = paged_adapter.prepare_layer_context(
+ self.prefix,
+ query,
+ key,
+ value,
+ omni_attn_metadata=attn_metadata,
+ )
+ out = self.attention.forward_paged(paged_kv_context)
else:
- out = self._run_local_attention(query, key, value, attn_metadata)
+ attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+ if self.use_ring and strategy is not self._no_parallel_strategy:
+ out = self._run_ring_attention(query, key, value, attn_metadata)
+ else:
+ out = self._run_local_attention(query, key, value, attn_metadata)
+
+ if paged_sp_padding:
+ padding_start = paged_sp_padding_offset - paged_sp_padding
+ output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:]))
+ out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1)
# 3. Post-processing (Reverse Communication)
# For Ulysses: AllToAll Output, and AllGather Joint Output
@@ -299,6 +456,19 @@ class Attention(nn.Module):
return out
+ @staticmethod
+ def _active_paged_kv_adapter():
+ """Return the Worker adapter selected by Runner-owned metadata."""
+
+ if not is_forward_context_available():
+ return None
+ return getattr(get_forward_context(), "paged_kv_adapter", None)
+
+ def is_paged_kv_active(self) -> bool:
+ """Return whether this layer will use Scheduler-managed paged KV."""
+
+ return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None
+
def _run_local_attention(self, query, key, value, attn_metadata):
self._assert_piecewise_compatible(attn_metadata)
@@ -309,7 +479,25 @@ class Attention(nn.Module):
)
return self.sdpa_fallback.forward(query, key, value, attn_metadata)
- # Fallback to standard attention
+ in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
+ # NPU only: the startup KV-capacity profile needs tensor shapes, not a
+ # paged attention result. If dense MindIE-SD is absent, SDPA provides
+ # that profile forward. Formal GPU/NPU paged requests never use this
+ # branch because their Worker adapter is active.
+ if (
+ self._scheduler_paged_kv
+ and self.paged_kv_cache_role is not None
+ and in_kv_memory_profile
+ and current_omni_platform.is_npu()
+ and self.attn_backend.get_name() == "FLASH_ATTN"
+ and not current_omni_platform.supports_diffusion_dense_flash_attention()
+ ):
+ logger.warning_once(
+ "The startup KV memory profile is using SDPA because MindIE-SD is unavailable. "
+ "Formal paged requests still use the platform-native paged attention backend."
+ )
+ return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+
return self.attention.forward(query, key, value, attn_metadata)
def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py
new file mode 100644
index 000000000..08e9e0e4a
--- /dev/null
+++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused MiniMax H3 modulation with FP32 accumulation."""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _indexed_scale_shift_kernel(
+ output_ptr,
+ x_ptr,
+ shift_ptr,
+ scale_ptr,
+ indices_ptr,
+ hidden_size,
+ stride_x_row,
+ stride_shift_row,
+ stride_scale_row,
+ stride_indices,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, block_n)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+ shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+ scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+
+ tl.store(
+ output_ptr + row * stride_x_row + columns,
+ x * (1.0 + scale) + shift,
+ mask=mask,
+ )
+
+
+@triton.jit
+def _indexed_gate_kernel(
+ output_ptr,
+ x_ptr,
+ gate_ptr,
+ other_ptr,
+ indices_ptr,
+ hidden_size,
+ stride_output_row,
+ stride_x_row,
+ stride_gate_row,
+ stride_other_row,
+ stride_indices,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, block_n)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+ gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
+ other = tl.load(other_ptr + row * stride_other_row + columns, mask=mask, other=0.0).to(tl.float32)
+
+ tl.store(
+ output_ptr + row * stride_output_row + columns,
+ x + gate * other,
+ mask=mask,
+ )
+
+
+@triton.jit
+def _rms_norm_indexed_scale_shift_kernel(
+ output_ptr,
+ x_ptr,
+ weight_ptr,
+ shift_ptr,
+ scale_ptr,
+ indices_ptr,
+ hidden_size: tl.constexpr,
+ eps: tl.constexpr,
+ stride_x_row,
+ stride_shift_row,
+ stride_scale_row,
+ stride_indices,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, block_n)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
+ weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
+ variance = tl.sum(x * x, axis=0) / hidden_size
+ normalized = x * tl.rsqrt(variance + eps) * weight
+
+ shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+ scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+ tl.store(
+ output_ptr + row * hidden_size + columns,
+ normalized * (1.0 + scale) + shift,
+ mask=mask,
+ )
+
+
+@triton.jit
+def _indexed_gate_rms_norm_scale_shift_kernel(
+ residual_out_ptr,
+ modulated_out_ptr,
+ residual_ptr,
+ gate_ptr,
+ branch_ptr,
+ weight_ptr,
+ shift_ptr,
+ scale_ptr,
+ indices_ptr,
+ hidden_size: tl.constexpr,
+ eps: tl.constexpr,
+ stride_residual_row,
+ stride_gate_row,
+ stride_branch_row,
+ stride_shift_row,
+ stride_scale_row,
+ stride_indices,
+ block_n: tl.constexpr,
+):
+ row = tl.program_id(0)
+ columns = tl.arange(0, block_n)
+ mask = columns < hidden_size
+ index = tl.load(indices_ptr + row * stride_indices)
+
+ residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32)
+ gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
+ branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32)
+ updated = residual + gate * branch
+ tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask)
+
+ weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
+ variance = tl.sum(updated * updated, axis=0) / hidden_size
+ normalized = updated * tl.rsqrt(variance + eps) * weight
+ shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
+ scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
+ tl.store(
+ modulated_out_ptr + row * hidden_size + columns,
+ normalized * (1.0 + scale) + shift,
+ mask=mask,
+ )
+
+
+def indexed_scale_shift_(
+ x: torch.Tensor,
+ shift: torch.Tensor,
+ scale: torch.Tensor,
+ indices: torch.Tensor,
+) -> torch.Tensor:
+ """Apply indexed scale/shift in-place to a disposable contiguous input."""
+ if x.is_cpu:
+ x.copy_((x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype))
+ return x
+ rows, hidden_size = x.shape
+ if rows == 0:
+ return x
+ _indexed_scale_shift_kernel[(rows,)](
+ x,
+ x,
+ shift,
+ scale,
+ indices,
+ hidden_size,
+ x.stride(0),
+ shift.stride(0),
+ scale.stride(0),
+ indices.stride(0),
+ block_n=triton.next_power_of_2(hidden_size),
+ num_warps=8,
+ )
+ return x
+
+
+def indexed_gate(
+ x: torch.Tensor,
+ gate: torch.Tensor,
+ other: torch.Tensor,
+ indices: torch.Tensor,
+) -> torch.Tensor:
+ """Return ``x + gate[indices] * other`` without indexed temporaries."""
+ if x.is_cpu:
+ return (x + gate.index_select(0, indices) * other).to(x.dtype)
+ output = torch.empty_like(x)
+ rows, hidden_size = x.shape
+ if rows == 0:
+ return output
+ _indexed_gate_kernel[(rows,)](
+ output,
+ x,
+ gate,
+ other,
+ indices,
+ hidden_size,
+ output.stride(0),
+ x.stride(0),
+ gate.stride(0),
+ other.stride(0),
+ indices.stride(0),
+ block_n=triton.next_power_of_2(hidden_size),
+ num_warps=8,
+ )
+ return output
+
+
+def rms_norm_indexed_scale_shift(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ shift: torch.Tensor,
+ scale: torch.Tensor,
+ indices: torch.Tensor,
+ eps: float,
+) -> torch.Tensor:
+ """Fuse H3 RMSNorm with its indexed AdaLN affine transform."""
+ if x.is_cpu:
+ input_dtype = x.dtype
+ normalized = x.float()
+ variance = normalized.pow(2).mean(-1, keepdim=True)
+ normalized = normalized * torch.rsqrt(variance + eps)
+ normalized = (weight.float() * normalized).to(input_dtype)
+ return (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(input_dtype)
+ output = torch.empty_like(x)
+ rows, hidden_size = x.shape
+ if rows:
+ _rms_norm_indexed_scale_shift_kernel[(rows,)](
+ output,
+ x,
+ weight,
+ shift,
+ scale,
+ indices,
+ hidden_size,
+ eps,
+ x.stride(0),
+ shift.stride(0),
+ scale.stride(0),
+ indices.stride(0),
+ block_n=triton.next_power_of_2(hidden_size),
+ num_warps=8,
+ )
+ return output
+
+
+def indexed_gate_rms_norm_scale_shift(
+ residual: torch.Tensor,
+ gate: torch.Tensor,
+ branch: torch.Tensor,
+ weight: torch.Tensor,
+ shift: torch.Tensor,
+ scale: torch.Tensor,
+ indices: torch.Tensor,
+ eps: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Fuse gated residual, the following RMSNorm, and AdaLN affine."""
+ if residual.is_cpu:
+ input_dtype = residual.dtype
+ residual_out = (residual + gate.index_select(0, indices) * branch).to(input_dtype)
+ normalized = residual_out.float()
+ variance = normalized.pow(2).mean(-1, keepdim=True)
+ normalized = normalized * torch.rsqrt(variance + eps)
+ normalized = (weight.float() * normalized).to(input_dtype)
+ modulated_out = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
+ input_dtype
+ )
+ return residual_out, modulated_out
+ residual_out = torch.empty_like(residual)
+ modulated_out = torch.empty_like(residual)
+ rows, hidden_size = residual.shape
+ if rows:
+ _indexed_gate_rms_norm_scale_shift_kernel[(rows,)](
+ residual_out,
+ modulated_out,
+ residual,
+ gate,
+ branch,
+ weight,
+ shift,
+ scale,
+ indices,
+ hidden_size,
+ eps,
+ residual.stride(0),
+ gate.stride(0),
+ branch.stride(0),
+ shift.stride(0),
+ scale.stride(0),
+ indices.stride(0),
+ block_n=triton.next_power_of_2(hidden_size),
+ num_warps=8,
+ )
+ return residual_out, modulated_out
+
+
+__all__ = [
+ "indexed_gate",
+ "indexed_gate_rms_norm_scale_shift",
+ "indexed_scale_shift_",
+ "rms_norm_indexed_scale_shift",
+]
diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py
index 76f7a78da..5b644ca45 100644
--- a/vllm_omni/diffusion/attention/parallel/factory.py
+++ b/vllm_omni/diffusion/attention/parallel/factory.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -47,6 +47,7 @@ def build_parallel_attention_strategy(
ulysses_degree = getattr(p, "ulysses_degree", 1)
ring_degree = getattr(p, "ring_degree", 1)
allgather_degree = getattr(p, "allgather_degree", 1)
+ ulysses_a2a_permute = getattr(p, "ulysses_a2a_permute", False)
try:
sp_group = get_sp_group()
@@ -83,6 +84,7 @@ def build_parallel_attention_strategy(
scatter_idx=scatter_idx,
gather_idx=gather_idx,
use_sync=use_sync,
+ ulysses_a2a_permute=ulysses_a2a_permute,
)
# Pure Ring Attention
diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py
index fc903d050..69f6a810e 100644
--- a/vllm_omni/diffusion/attention/parallel/ulysses.py
+++ b/vllm_omni/diffusion/attention/parallel/ulysses.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
@@ -16,6 +16,11 @@ from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGr
from vllm_omni.diffusion.forward_context import get_ulysses_mode
+def _a2a_permute_enabled(enabled: bool, scatter_idx: int, gather_idx: int, world_size: int) -> bool:
+ """Whether the fused permute-free all-to-all applies (strict layout only)."""
+ return enabled and world_size > 1 and scatter_idx == 2 and gather_idx == 1
+
+
def _ceil_div(n: int, d: int) -> int:
return (n + d - 1) // d
@@ -181,12 +186,23 @@ class UlyssesParallelAttention:
scatter_idx: int,
gather_idx: int,
use_sync: bool,
+ ulysses_a2a_permute: bool = False,
) -> None:
self._sp_group = sp_group
self._ulysses_pg = sp_group.ulysses_group
self._scatter_idx = scatter_idx
self._gather_idx = gather_idx
self._use_sync = use_sync
+ self._ulysses_a2a_permute = ulysses_a2a_permute
+ if _a2a_permute_enabled(
+ ulysses_a2a_permute,
+ scatter_idx,
+ gather_idx,
+ sp_group.ulysses_world_size,
+ ):
+ from vllm_omni.diffusion.distributed.a2a_permute import ensure_a2a_permute_available
+
+ ensure_a2a_permute_available()
@property
def enabled(self) -> bool:
@@ -347,9 +363,26 @@ class UlyssesParallelAttention:
)
# (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size)
- query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync)
- key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync)
- value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync)
+ if _a2a_permute_enabled(
+ self._ulysses_a2a_permute,
+ self._scatter_idx,
+ self._gather_idx,
+ ulysses_world_size,
+ ):
+ from vllm_omni.diffusion.distributed.a2a_permute import ulysses_qkv_fwd
+
+ group_name = self._ulysses_pg.group_name
+ query = ulysses_qkv_fwd(query, group_name, ulysses_world_size)
+ key = ulysses_qkv_fwd(key, group_name, ulysses_world_size)
+ value = ulysses_qkv_fwd(value, group_name, ulysses_world_size)
+ else:
+ query = SeqAllToAll4D.apply(
+ self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync
+ )
+ key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync)
+ value = SeqAllToAll4D.apply(
+ self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync
+ )
seq_lens = []
local_seq_len = 0
orig_head_cnt = 0
@@ -456,6 +489,19 @@ class UlyssesParallelAttention:
orig_head_cnt=ctx.orig_head_cnt,
use_sync=ctx.use_sync,
)
+ elif _a2a_permute_enabled(
+ self._ulysses_a2a_permute,
+ ctx.scatter_idx,
+ ctx.gather_idx,
+ dist.get_world_size(ctx.ulysses_pg),
+ ):
+ from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev
+
+ output_img = ulysses_o_rev(
+ output_img,
+ ctx.ulysses_pg.group_name,
+ dist.get_world_size(ctx.ulysses_pg),
+ )
else:
output_img = SeqAllToAll4D.apply(
ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync
@@ -488,4 +534,17 @@ class UlyssesParallelAttention:
orig_head_cnt=ctx.orig_head_cnt,
use_sync=ctx.use_sync,
)
+ if _a2a_permute_enabled(
+ self._ulysses_a2a_permute,
+ ctx.scatter_idx,
+ ctx.gather_idx,
+ dist.get_world_size(ctx.ulysses_pg),
+ ):
+ from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev
+
+ return ulysses_o_rev(
+ attn_output,
+ ctx.ulysses_pg.group_name,
+ dist.get_world_size(ctx.ulysses_pg),
+ )
return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync)
diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py
index 5f1ca1d22..7f0242d57 100644
--- a/vllm_omni/diffusion/cache/teacache/extractors.py
+++ b/vllm_omni/diffusion/cache/teacache/extractors.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Model-specific extractors for TeaCache.
@@ -13,6 +13,7 @@ all model-specific information needed for generic caching, including preprocessi
transformer execution, and postprocessing logic.
"""
+import math
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
@@ -723,11 +724,11 @@ def extract_flux2_klein_context(
def extract_longcat_context(
module: nn.Module, # LongCatImageTransformer2DModel
hidden_states,
- timestep,
- guidance,
- encoder_hidden_states,
- txt_ids,
- img_ids,
+ encoder_hidden_states=None,
+ timestep=None,
+ img_ids=None,
+ txt_ids=None,
+ guidance=None,
**kwargs,
) -> CacheContext:
"""Extract the cache context for LongCat Image.
@@ -1259,11 +1260,10 @@ def extract_minimax_h3_context(
preprocessing via ``_embed``, cacheable ``blocks`` loop, and
``final_layer`` postprocessing with row selection and update masks.
"""
+ from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
- _BF16_DTYPE,
_FORWARD_SUPPORTED_KWARGS,
MINIMAX_H3_ADALN_MODALITY_NUM,
- _modulate_scale_shift,
_required_kwarg,
)
@@ -1304,6 +1304,9 @@ def extract_minimax_h3_context(
psp = _required_kwarg(kwargs, "packed_seq_params")
cu_seqlens = module._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32)
max_seqlen = int(module._psp_field(psp, "packed_seq_params", "max_seqlen_q"))
+ # Mirror forward()'s host-side scalar so the refiner sees the number of
+ # packed requests without reading cu_seqlens off-device.
+ num_requests = int(module._psp_optional(psp, "num_requests", 1))
refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params")
refiner_cu = module._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32)
refiner_max = int(module._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q"))
@@ -1317,8 +1320,20 @@ def extract_minimax_h3_context(
if inverse_indices.shape[0] != seq_len:
raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
device = x.device
-
- rope_freqs = module.rope(img_position_ids).to(device)
+ local_span = module._rope_local_span(seq_len)
+ local_start, local_len = local_span
+ rope_table = kwargs.get("rope_table")
+ if rope_table is None:
+ rope_table = module.prepare_rope_table(
+ img_position_ids,
+ seq_len=seq_len,
+ )
+ else:
+ module._validate_prepared_rope_table(
+ rope_table,
+ local_len=local_len,
+ device=device,
+ )
decoder_input, t_emb = module._embed(
x=x,
@@ -1330,50 +1345,103 @@ def extract_minimax_h3_context(
text_pos=text_pos.to(device),
refiner_cu_seqlens=refiner_cu.to(device),
refiner_max_seqlen=refiner_max,
+ num_requests=num_requests,
seq_len=seq_len,
device=device,
+ local_span=local_span,
)
combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)
inverse_indices = inverse_indices.to(device)
cu_seqlens = cu_seqlens.to(device)
+ # In strict SP, _embed already returns this rank's packed rows while
+ # combined_indices still describes the full packed sequence. TeaCache uses
+ # the first block's modulation as part of its rank-local cache state, so
+ # the modulation indices must be sliced to the same local layout. Keep the
+ # full indices below: local_sp_prepare owns the global-to-local conversion
+ # for transformer-block inputs.
+ state_combined_indices = combined_indices.narrow(0, local_start, local_len)
+ if local_len == seq_len:
+ state_combined_indices = combined_indices
+
shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb)
- modulated_hidden = module.blocks[0].norm1(decoder_input)
- modulated_input = _modulate_scale_shift(
- modulated_hidden,
+ modulated_input = module.blocks[0].norm1(decoder_input)
+ modulated_input = indexed_scale_shift_(
+ modulated_input,
shift_msa,
scale_msa,
- combined_indices,
- dtype=_BF16_DTYPE,
+ state_combined_indices,
)
- def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
- hidden, block_rope, block_combined = module.sp_prepare(
- decoder_input,
- rope_freqs,
- combined_indices,
+ def synchronize_cache_decision(local_should_compute: bool) -> bool:
+ if local_len == seq_len:
+ return local_should_compute
+
+ from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
+
+ decision = torch.tensor(
+ int(local_should_compute),
+ dtype=torch.int32,
+ device=device,
)
+ get_sp_group().all_reduce(decision, op=torch.distributed.ReduceOp.MAX)
+ return bool(decision.item())
+
+ def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
+ if local_len == seq_len:
+ hidden, block_rope, block_combined = module.sp_prepare(
+ decoder_input,
+ rope_table,
+ combined_indices,
+ )
+ else:
+ hidden, block_rope, block_combined = module.local_sp_prepare(
+ decoder_input,
+ rope_table,
+ combined_indices,
+ )
for block in module.blocks:
hidden = block(
hidden,
t_emb=t_emb,
combined_indices=block_combined,
- rope_freqs=block_rope,
+ rope_table=block_rope,
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
packed_total=seq_len,
video_layout=video_layout,
)
- hidden = module.sp_gather(hidden)
+ # TeaCache stores residuals as block_output - decoder_input. Strict SP
+ # must keep both tensors rank-local until postprocess projects local
+ # rows to compact logits and gathers them exactly once.
+ if local_len == seq_len:
+ hidden = module.sp_gather(hidden)
return (hidden,)
def postprocess(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
- video_logits, audio_logits = module.final_layer(
- hidden,
- t_emb=t_emb,
- inverse_indices=inverse_indices,
- )
+ if local_len == seq_len:
+ video_logits, audio_logits = module.final_layer(
+ hidden,
+ t_emb=t_emb,
+ inverse_indices=inverse_indices,
+ )
+ else:
+ local_inverse_indices = inverse_indices.narrow(
+ 0,
+ local_start,
+ local_len,
+ )
+ video_logits, audio_logits = module.final_layer(
+ hidden,
+ t_emb=t_emb,
+ inverse_indices=local_inverse_indices,
+ )
+ compact_logits = torch.cat((video_logits, audio_logits), dim=-1)
+ compact_logits = module.sp_gather(compact_logits)
+ video_width = module.arch.latents_dim * math.prod(module.arch.patch_size)
+ video_logits = compact_logits[..., :video_width]
+ audio_logits = compact_logits[..., video_width:]
video_logits = video_logits.index_select(0, infer_out_pos.to(device))
audio_logits = audio_logits.index_select(0, audio_pos.to(device))
if not skip_mask_out_condition:
@@ -1392,6 +1460,9 @@ def extract_minimax_h3_context(
temb=t_emb,
run_transformer_blocks=run_transformer_blocks,
postprocess=postprocess,
+ extra_states={
+ "synchronize_cache_decision": synchronize_cache_decision,
+ },
)
diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py
index 85a42adbb..b60e17836 100644
--- a/vllm_omni/diffusion/cache/teacache/hook.py
+++ b/vllm_omni/diffusion/cache/teacache/hook.py
@@ -138,7 +138,16 @@ class TeaCacheHook(ModelHook):
state = self.state_manager.get_state()
# Decide whether to compute or cache based on modulated input similarity
- should_compute = self._should_compute_full_transformer(state, ctx.modulated_input)
+ local_should_compute = self._should_compute_full_transformer(state, ctx.modulated_input)
+ sync_cache_decision = (ctx.extra_states or {}).get("synchronize_cache_decision")
+ if sync_cache_decision is not None:
+ should_compute = sync_cache_decision(local_should_compute)
+ # A rank that locally chose the cache path must reset its counter
+ # when another SP rank requires a full collective block execution.
+ if should_compute and not local_should_compute:
+ state.accumulated_rel_l1_distance = 0.0
+ else:
+ should_compute = local_should_compute
if not should_compute and state.previous_residual is not None:
# ============================================================================
diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py
index 4a050ab38..11fc2c113 100644
--- a/vllm_omni/diffusion/compile.py
+++ b/vllm_omni/diffusion/compile.py
@@ -56,26 +56,50 @@ def regionally_compile(
# Build all compiled callables before mutating the model. This keeps setup
# failures atomic: callers can safely continue with the uncompiled model if
# torch.compile raises synchronously for any repeated block.
- compiled_forwards: list[tuple[nn.Module, Any]] = []
+ # Keep the pre-hook callable separately. DLO and the ordinary layerwise
+ # offloader install a HookRegistry wrapper in ``module.forward``. That
+ # wrapper performs stream/event synchronization and storage rebinding, so
+ # compiling it pulls offload control flow into the graph. Compile the
+ # original block compute instead and leave the wrapper outside the graph.
+ compiled_forwards: list[tuple[nn.Module, Any | None, Any]] = []
for name, submod in model.named_modules():
if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
# Compile the block compute while keeping nn.Module.__call__ hooks
- # outside the compiled graph. NOTE: anything that wraps this
- # callable must stay signature-transparent — cache-dit's
- # BlockAdapter matches blocks by inspecting forward's parameter
- # names and return annotation, which torch.compile preserves.
+ # outside the compiled graph. If a HookRegistry is already
+ # installed, ``submod.forward`` is the hook dispatcher and the
+ # original callable is stored in ``_omni_original_forward``.
+ # NOTE: anything that wraps this callable must stay
+ # signature-transparent — cache-dit's BlockAdapter matches blocks
+ # by inspecting forward's parameter names and return annotation,
+ # which torch.compile preserves.
+ original_forward = getattr(submod, "_omni_original_forward", None)
+ forward_to_compile = original_forward if original_forward is not None else submod.forward
compiled_forwards.append(
(
submod,
- torch.compile(submod.forward, *compile_args, **compile_kwargs),
+ original_forward,
+ torch.compile(forward_to_compile, *compile_args, **compile_kwargs),
)
)
if not compiled_forwards:
logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
else:
- for submod, compiled_forward in compiled_forwards:
- submod.forward = compiled_forward
+ for submod, original_forward, compiled_forward in compiled_forwards:
+ if original_forward is None:
+ submod.forward = compiled_forward
+ continue
+
+ # Keep HookRegistry's dispatcher as ``forward`` and replace only
+ # the callable it dispatches to. Hooks such as MagCache retain a
+ # reference to the original callable as well, so update those
+ # references when they still point at the pre-compile function.
+ submod._omni_original_forward = compiled_forward
+ registry = getattr(submod, "_hook_registry", None)
+ for hook in getattr(registry, "_hooks", {}).values():
+ fn_ref = getattr(hook, "fn_ref", None)
+ if fn_ref is not None and fn_ref.original_forward is original_forward:
+ fn_ref.original_forward = compiled_forward
logger.info(
"Regional compilation applied to %d module(s) for repeated blocks %s.",
len(compiled_forwards),
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 137fb9860..4c9c37c35 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -1,13 +1,13 @@
# adapted from sglang and fastvideo
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import copy
import math
import os
import random
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field, fields
-from enum import Enum
+from enum import Enum, StrEnum
from typing import TYPE_CHECKING, Any
import diffusers
@@ -77,11 +77,13 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
# Handle "diffusion_attention_backend" shorthand: merge into
# diffusion_attention_config before field filtering.
diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
- if diffusion_attn_backend is not None:
+ fastvideo_vsa_topk = normalized.pop("fastvideo_vsa_topk", None)
+ if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
existing = normalized.get("diffusion_attention_config")
normalized["diffusion_attention_config"] = parse_attention_config(
existing,
attention_backend=diffusion_attn_backend,
+ fastvideo_vsa_topk=fastvideo_vsa_topk,
)
# Check environment variable as fallback for cache_backend.
@@ -98,12 +100,42 @@ def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]
return normalized
+def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
+ """Validate HWR policy without touching the configured storage domain.
+
+ Filesystem locality and store construction belong to the eligible loader
+ path. Keeping this check purely structural is what lets disabled and
+ AllGather DLO configurations retain zero HWR interaction.
+ """
+ if mode not in {"disabled", "preferred", "required"}:
+ raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required")
+ if mode != "disabled" and (not isinstance(root, str) or not root.strip()):
+ raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")
+
+
+def validate_dlo_host_registration_options(
+ *,
+ limit_gib: object,
+ enable_dlo: bool,
+ use_allgather: bool,
+ hwr_mode: object,
+) -> float:
+ """Validate the optional transport budget without probing CUDA or HWR."""
+ value = float(limit_gib)
+ if not math.isfinite(value) or value < 0:
+ raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
+ if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
+ raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime")
+ return value
+
+
def parse_kv_cache_skip_selector(
selector: str | list[int] | tuple[int, ...] | set[int] | None,
) -> set[int] | None:
"""Parse a non-negative index selector such as "0-9,20,25-30"."""
if selector is None:
return None
+ values: set[int]
if isinstance(selector, set):
values = selector
elif isinstance(selector, (list, tuple)):
@@ -112,7 +144,7 @@ def parse_kv_cache_skip_selector(
text = selector.strip()
if not text:
return None
- values: set[int] = set()
+ values = set()
for chunk in text.split(","):
token = chunk.strip()
if not token:
@@ -154,8 +186,8 @@ class DiffusionParallelConfig:
pipeline_parallel_size: int = 1
"""Number of pipeline parallel stages."""
- data_parallel_size: int = 1
- """Number of data parallel groups."""
+ data_parallel_size: int | None = None
+ """Number of data parallel groups. Inferred from WORLD when omitted."""
tensor_parallel_size: int = 1
"""Number of tensor parallel groups."""
@@ -191,6 +223,9 @@ class DiffusionParallelConfig:
sequence shapes across the ring group.
"""
+ ulysses_a2a_permute: bool = False
+ """Use fused permute-free all-to-all for eligible strict Ulysses exchanges."""
+
cfg_parallel_size: int = 1
"""Number of ranks used to execute guidance passes in parallel."""
@@ -243,8 +278,9 @@ class DiffusionParallelConfig:
def _validate_parallel_config(self) -> Self:
"""Validates the config relationships among the parallel strategies."""
assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
- assert self.data_parallel_size > 0, "Data parallel size must be > 0"
+ assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
+ assert self.sequence_parallel_size is not None
assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
assert self.ring_degree > 0, "Ring degree must be > 0"
@@ -283,10 +319,12 @@ class DiffusionParallelConfig:
self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
)
- # Calculate world_size from other parallelism dimensions
+ # Until the runtime WORLD size is known, an omitted DP dimension means
+ # one replica. OmniDiffusionConfig resolves it against num_gpus below.
+ data_parallel_size = self.data_parallel_size or 1
other_parallel_world_size = (
self.pipeline_parallel_size
- * self.data_parallel_size
+ * data_parallel_size
* self.tensor_parallel_size
* self.sequence_parallel_size
* self.cfg_parallel_size
@@ -297,13 +335,17 @@ class DiffusionParallelConfig:
# 1. Standalone: when other parallelism is all 1, HSDP determines world_size
# 2. Combined: HSDP overlays on top of other parallelism
if self.use_hsdp:
- if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
- raise ValueError(
- "HSDP (use_hsdp=True) cannot be used with TP or DP "
- f"(tensor_parallel_size={self.tensor_parallel_size}, "
- f"data_parallel_size={self.data_parallel_size}). "
- "Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
- )
+ incompatible = []
+ if self.tensor_parallel_size > 1:
+ incompatible.append("TP")
+ if data_parallel_size > 1:
+ incompatible.append("DP")
+ if self.pipeline_parallel_size > 1:
+ incompatible.append("PP")
+ if self.enable_expert_parallel:
+ incompatible.append("EP")
+ if incompatible:
+ raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
if self.hsdp_shard_size == -1:
# Auto-calculate: use other_parallel_world_size as shard_size
if self.hsdp_replicate_size <= 0:
@@ -338,6 +380,43 @@ class DiffusionParallelConfig:
else:
self.world_size = other_parallel_world_size
+ def resolve_data_parallel_size(self, world_size: int) -> int:
+ """Resolve or validate diffusion DP against the actual WORLD size."""
+ if world_size <= 0:
+ raise ValueError(f"WORLD size must be > 0, but got {world_size}")
+
+ if self.use_hsdp:
+ if self.data_parallel_size not in (None, 1):
+ raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
+ expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
+ if world_size != expected_world_size:
+ raise ValueError(
+ f"WORLD size ({world_size}) must equal HSDP size "
+ f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})"
+ )
+ self.data_parallel_size = 1
+ self.world_size = world_size
+ return 1
+
+ assert self.sequence_parallel_size is not None
+ non_dp_size = (
+ self.pipeline_parallel_size
+ * self.tensor_parallel_size
+ * self.sequence_parallel_size
+ * self.cfg_parallel_size
+ )
+ if world_size % non_dp_size != 0:
+ raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
+ inferred_data_parallel_size = world_size // non_dp_size
+ if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size:
+ raise ValueError(
+ f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value "
+ f"({inferred_data_parallel_size})"
+ )
+ self.data_parallel_size = inferred_data_parallel_size
+ self.world_size = world_size
+ return inferred_data_parallel_size
+
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
"""
@@ -538,7 +617,11 @@ class DiffusionCacheConfig:
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
-def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None:
+def resolve_model_class_name(
+ model: str | None,
+ diffusion_load_format: str = "default",
+ revision: str | None = None,
+) -> str | None:
"""Resolve the diffusion pipeline class name from the model config.
Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call
@@ -550,11 +633,14 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
if not model:
return None
-
is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
- # Diffusers models: read _class_name from the pipeline index.
- model_index = get_diffusion_model_index(model)
+ # Diffusers models: read _class_name from the pipeline index. Missing
+ # local paths can otherwise be interpreted as invalid Hub repo IDs.
+ try:
+ model_index = get_diffusion_model_index(model, revision=revision)
+ except Exception:
+ model_index = None
if model_index is not None:
return model_index.get("_class_name")
if diffusion_load_format == "diffusers":
@@ -562,12 +648,16 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
# Other models: map model_type / architecture from config.json.
try:
- cfg = get_hf_file_to_dict("config.json", model) or {}
+ cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {}
except Exception:
cfg = {}
model_type = cfg.get("model_type")
architectures = cfg.get("architectures") or []
+ from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
+
+ if _looks_like_hidream_o1(model, cfg):
+ return "HiDreamO1ImagePipeline"
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
return "BagelPipeline"
if (
@@ -598,6 +688,19 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
return None
+def uses_diffusers_adapter(od_config: object) -> bool:
+ """Return whether execution uses the Diffusers adapter backend.
+
+ A custom pipeline takes precedence over ``diffusion_load_format`` in the
+ worker, so adapter-specific behavior must only apply when no custom
+ pipeline override is configured.
+ """
+ return (
+ getattr(od_config, "custom_pipeline_args", None) is None
+ and getattr(od_config, "diffusion_load_format", "default") == "diffusers"
+ )
+
+
@dataclass
class OmniDiffusionConfig:
# Model and path configuration (for convenience)
@@ -651,8 +754,10 @@ class OmniDiffusionConfig:
# in both directions. See docs/features/session_state_manager.md.
enable_session_state_manager: bool = False
- # Distributed executor backend
- distributed_executor_backend: str = "mp"
+ # Distributed executor backend. ``None`` means auto: ``"uni"`` at
+ # ``num_gpus == 1``, ``"mp"`` otherwise. Explicit ``"mp"`` keeps a
+ # worker subprocess (process isolation and RPC timeouts).
+ distributed_executor_backend: str | None = None
nccl_port: int | None = None
# Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
@@ -665,6 +770,9 @@ class OmniDiffusionConfig:
# Local Diffusion KV ownership and cache-layout mode.
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+ # Maximum number of native BlockTable rows one public request can own
+ # (sequences plus independent contexts). The model adapter defines it.
+ diffusion_kv_max_rows_per_request: int | None = None
# Optional override for the diffusion model runner class (import path).
# Precedence in the worker: this override > the runner declared by the
@@ -707,6 +815,14 @@ class OmniDiffusionConfig:
dlo_use_allgather: bool = True
# Leading main-DiT blocks kept resident by distributed layerwise offload.
dlo_resident_layers: int = 0
+ # Final-layout Host Weight Runtime policy. The loader only activates this
+ # for eligible no-AllGather DLO; all other configurations preserve their
+ # existing loader/storage path.
+ host_weight_runtime_mode: str = "disabled"
+ host_weight_runtime_root: str | None = None
+ # Optional per-worker ceiling for registering final-layout HWR mappings.
+ # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
+ dlo_host_registration_limit_gib: float = 0.0
pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading
@@ -785,6 +901,7 @@ class OmniDiffusionConfig:
default_factory=lambda: {
"transformer": True,
"vae": True,
+ "text_encoder": True,
}
)
override_transformer_cls_name: str | None = None
@@ -848,6 +965,14 @@ class OmniDiffusionConfig:
# Maximum number of sequences to generate in a batch
max_num_seqs: int = 1
+ # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode.
+ # These mirror the structured stage cache/scheduler config until all
+ # diffusion runtime construction consumes VllmOmniConfig directly.
+ kv_cache_memory_bytes: int | None = None
+ gpu_memory_utilization: float = 0.9
+ max_num_batched_tokens: int | None = None
+ max_model_len: int | None = None
+
# Request-mode batch admission: wait briefly for compatible requests to
# accumulate in the scheduler waiting queue before the first schedule() of
# a wave. Improves fused forward batch sizes under bursty HTTP ingress.
@@ -933,6 +1058,23 @@ class OmniDiffusionConfig:
if not isinstance(self.diffusion_compile_dynamic, bool):
raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+ if self.diffusion_kv_max_rows_per_request is not None and (
+ type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0
+ ):
+ raise ValueError("diffusion_kv_max_rows_per_request must be a positive integer when set")
+ if (
+ self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+ and self.diffusion_kv_max_rows_per_request is None
+ ):
+ raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set")
+ if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0:
+ raise ValueError("kv_cache_memory_bytes must be non-negative")
+ if not 0.0 < self.gpu_memory_utilization <= 1.0:
+ raise ValueError("gpu_memory_utilization must be in (0, 1]")
+ if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0:
+ raise ValueError("max_num_batched_tokens must be positive")
+ if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0:
+ raise ValueError("max_model_len must be positive or -1")
if self.omni_kv_config is None:
self.omni_kv_config = {}
@@ -944,10 +1086,9 @@ class OmniDiffusionConfig:
"need_recv_cache", False
):
raise ValueError(
- "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
- "disable need_recv_cache until connector-aware admission is implemented"
+ "paged_scheduler Diffusion KV does not support imported AR KV; "
+ "disable need_recv_cache until connector-aware import is implemented"
)
-
self.master_port = self._resolve_master_port()
self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:
@@ -980,13 +1121,11 @@ class OmniDiffusionConfig:
else:
self.num_gpus = 1
- if self.num_gpus < self.parallel_config.world_size:
- raise ValueError(
- f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
- )
+ self.parallel_config.resolve_data_parallel_size(self.num_gpus)
if self.diffusion_compile_granularity == "full":
incompatible_features = []
+ assert self.parallel_config.sequence_parallel_size is not None
if self.parallel_config.use_hsdp:
incompatible_features.append("HSDP")
if self.parallel_config.sequence_parallel_size > 1:
@@ -1060,6 +1199,17 @@ class OmniDiffusionConfig:
elif self.max_cpu_loras < 1:
raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
+ validate_host_weight_runtime_options(
+ mode=self.host_weight_runtime_mode,
+ root=self.host_weight_runtime_root,
+ )
+ self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
+ limit_gib=self.dlo_host_registration_limit_gib,
+ enable_dlo=self.enable_distributed_layerwise_offload,
+ use_allgather=self.dlo_use_allgather,
+ hwr_mode=self.host_weight_runtime_mode,
+ )
+
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
raise ValueError(
"diffusers_load_kwargs and diffusers_call_kwargs are only "
@@ -1182,10 +1332,7 @@ class OmniDiffusionConfig:
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
- # Default model_class_name for diffusers adapter
- if self.model_class_name is None and self.diffusion_load_format == "diffusers":
- self.model_class_name = "DiffusersAdapterPipeline"
-
+ assert self.model is not None
try:
config_dict = get_diffusion_model_index(
self.model,
@@ -1194,6 +1341,8 @@ class OmniDiffusionConfig:
if config_dict is not None:
if self.model_class_name is None:
self.model_class_name = config_dict.get("_class_name", None)
+ if self.model_class_name is None and self.diffusion_load_format == "diffusers":
+ self.model_class_name = "DiffusersAdapterPipeline"
self.update_multimodal_support()
# Skip transformer config loading for diffusers adapter
@@ -1210,9 +1359,20 @@ class OmniDiffusionConfig:
exc,
)
else:
- tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+ from vllm_omni.model_extras import get_transformer_config_subfolder
+
+ transformer_subfolder = get_transformer_config_subfolder(
+ self.model_class_name,
+ model=self.model,
+ revision=self.revision,
+ )
+ tf_config_dict = get_hf_file_to_dict(
+ f"{transformer_subfolder}/config.json",
+ self.model,
+ revision=self.revision,
+ )
if tf_config_dict is None:
- tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model)
+ tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
if tf_config_dict is not None:
self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
else:
@@ -1223,6 +1383,8 @@ class OmniDiffusionConfig:
# Skip transformer config loading for diffusers adapter
# (non-DiT models don't have a separate transformer folder/config)
if self.diffusion_load_format == "diffusers":
+ if self.model_class_name is None:
+ self.model_class_name = "DiffusersAdapterPipeline"
self.set_tf_model_config(TransformerConfig())
logger.warning(
"Could not find a valid pipeline index per Diffusers format. "
@@ -1232,7 +1394,7 @@ class OmniDiffusionConfig:
"that require additional inputs."
)
else:
- cfg = get_hf_file_to_dict("config.json", self.model)
+ cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision)
if cfg is None:
# Lance ships its top-level config.json one directory above
# the per-checkpoint subfolders (``Lance_3B/`` or
@@ -1249,7 +1411,18 @@ class OmniDiffusionConfig:
model_type = cfg.get("model_type")
architectures = cfg.get("architectures") or []
- if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+ from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
+
+ assert self.model is not None
+ is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg)
+ if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1:
+ raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature")
+
+ if is_hidream_o1:
+ self.model_class_name = "HiDreamO1ImagePipeline"
+ self.set_tf_model_config(TransformerConfig())
+ self.update_multimodal_support()
+ elif model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
self.model_class_name = "BagelPipeline"
self.set_tf_model_config(TransformerConfig())
self.update_multimodal_support()
@@ -1298,12 +1471,19 @@ class OmniDiffusionConfig:
elif model_type == "vla":
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
+ assert self.model is not None
if _looks_like_dreamzero(self.model):
self.model_class_name = "DreamZeroPipeline"
self.set_tf_model_config(TransformerConfig())
self.update_multimodal_support()
else:
raise
+ elif cfg.get("type") == "pi0":
+ # π0 (Pi-Zero) VLA — the LeRobot config.json uses ``type: "pi0"``.
+ if self.model_class_name is None:
+ self.model_class_name = "Pi0Pipeline"
+ self.set_tf_model_config(TransformerConfig())
+ self.update_multimodal_support()
elif architectures and len(architectures) == 1:
architecture = architectures[0]
from vllm_omni.diffusion.registry import DiffusionModelRegistry
@@ -1368,6 +1548,9 @@ class DiffusionOutput:
# memory usage info
peak_memory_mb: float = 0.0
+ # KV-recv wall-clock (ms) from the runner; feeds vllm_omni:diffusion_kv_load_s.
+ kv_recv_ms: float = 0.0
+
# When True, move tensor fields to CPU at construction time. Useful when
# the output is shipped across process boundaries (e.g. step-execution
# mode) and the receiving side must not initialise a stray CUDA context.
@@ -1490,10 +1673,11 @@ class AttnQuantSpec:
for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
if v is not None and v not in self._VALID_DTYPES:
raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
- for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
- if v not in self._VALID_BLOCK_SIZES:
+ for block_name, block_value in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
+ if block_value not in self._VALID_BLOCK_SIZES:
raise ValueError(
- f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
+ f"quant.{block_name}={block_value!r} unsupported; "
+ f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
)
@property
@@ -1506,6 +1690,20 @@ class AttnQuantSpec:
BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})
+class RainFusionPrecision(StrEnum):
+ """Execution precision for block-sparse RainFusion (rf_v3) attention.
+
+ ``bf16``: no quantization, pure BF16 sparse attention (official baseline).
+ ``fp8``: BSA FP8 path - Hadamard rotation then full FP8 block quantization
+ of Q/K/V before the BSA kernel.
+ ``mix``: EagleQBSA mixed precision - Q/K per-block INT8 + V per-channel FP8.
+ """
+
+ BF16 = "bf16"
+ FP8 = "fp8"
+ MIX = "mix"
+
+
@dataclass
class BlockSparseSpec:
"""User-facing controls shared by block-sparse attention backends.
@@ -1514,10 +1712,13 @@ class BlockSparseSpec:
``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an
index selector such as "0-3,38") exempts individual DiT blocks. Those two are
the accuracy knobs to trade back quality at a fixed ``sparsity``.
+ ``precision`` selects the kernel precision mode (see ``RainFusionPrecision``).
"""
sparsity: float = 0.8
start_step: int = 0
+ end_step: int = 0
+ precision: str = RainFusionPrecision.BF16.value
skip_layers: str | list[int] | None = None
skip_layer_indices: set[int] | None = field(default=None, repr=False)
@@ -1525,6 +1726,13 @@ class BlockSparseSpec:
self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0
if self.start_step < 0:
raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.")
+ if self.end_step < 0:
+ raise ValueError(f"block_sparse.end_step must be >= 0; got {self.end_step!r}.")
+ if self.precision not in {p.value for p in RainFusionPrecision}:
+ raise ValueError(
+ f"block_sparse.precision must be one of {sorted(p.value for p in RainFusionPrecision)}; "
+ f"got {self.precision!r}."
+ )
self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)
@@ -1535,6 +1743,7 @@ class AttentionSpec:
backend: str
skip_softmax: SkipSoftmaxSpec | None = None
quant: AttnQuantSpec | None = None
+ fastvideo_vsa_topk: int | None = None
block_sparse: BlockSparseSpec | None = None
skip_calibration: dict | None = field(default=None, repr=False)
@@ -1554,6 +1763,11 @@ class AttentionSpec:
f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
f"backend={self.backend!r}. Remove quant or set a supported backend."
)
+ if self.fastvideo_vsa_topk is not None:
+ if self.backend.upper() != "FASTVIDEO_VSA":
+ raise ValueError("fastvideo_vsa_topk is only supported by the FASTVIDEO_VSA backend.")
+ if self.fastvideo_vsa_topk <= 0:
+ raise ValueError("fastvideo_vsa_topk must be positive.")
if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
# Selecting the backend is the opt-in; without an explicit block the
# defaults apply rather than silently running dense.
@@ -1595,12 +1809,17 @@ class AttentionSpec:
if q.flashinfer_backend is not None:
quant_kw["flashinfer_backend"] = q.flashinfer_backend
kw["quant"] = quant_kw
+ if self.fastvideo_vsa_topk is not None:
+ kw["topk"] = self.fastvideo_vsa_topk
if self.block_sparse is not None:
bs = self.block_sparse
kw["sparsity"] = bs.sparsity
kw["start_step"] = bs.start_step
+ kw["end_step"] = bs.end_step
+ kw["precision"] = bs.precision
if bs.skip_layer_indices:
kw["skip_layers"] = sorted(bs.skip_layer_indices)
+
return kw or None
@@ -1670,7 +1889,7 @@ class AttentionConfig:
normalized[role] = node
return
- spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
+ spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"}
node_dict = dict(node)
node_keys = set(node_dict)
if node_keys & spec_keys:
@@ -1709,6 +1928,7 @@ def parse_attention_config(
attention_config: AttentionConfig | Mapping[str, Any] | None = None,
*,
attention_backend: str | None = None,
+ fastvideo_vsa_topk: int | None = None,
) -> AttentionConfig:
"""Pure type-conversion: coerce *attention_config* to an AttentionConfig.
@@ -1736,6 +1956,14 @@ def parse_attention_config(
if attention_backend.lower() != "auto":
normalized.default = AttentionSpec(backend=attention_backend)
+ if fastvideo_vsa_topk is not None:
+ if normalized.default is None:
+ raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.")
+ if normalized.default.backend.upper() != "FASTVIDEO_VSA":
+ raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.")
+ normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk
+ normalized.default.__post_init__()
+
return normalized
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 08218122e..0929e9419 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -1,11 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import asyncio
import concurrent.futures
+import copy
import inspect
+import os
import queue
import threading
import time
@@ -17,18 +19,24 @@ from typing import TYPE_CHECKING, Any
import numpy as np
import PIL.Image
import torch
+from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.utils.import_utils import resolve_obj_by_qualname
from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm_omni.diffusion.data import (
DiffusionOutput,
DiffusionRequestAbortedError,
OmniDiffusionConfig,
+ uses_diffusers_adapter,
)
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
+from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
from vllm_omni.diffusion.io_support import (
get_dummy_run_num_frames,
+ get_dummy_run_num_image_inputs,
image_color_format,
supports_audio_output,
supports_multimodal_input,
@@ -49,13 +57,61 @@ from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput
from vllm_omni.errors import client_error_from_metadata, is_client_error_status
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.metrics.utils import (
+ diffusion_scheduler_waiting_metrics,
+ extract_diffusion_denoise_ms,
+ extract_diffusion_vae_decode_ms,
+)
if TYPE_CHECKING:
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
-_ASYNC_OUTPUT_TIMEOUT = 30.0 # seconds
+_ASYNC_OUTPUT_TIMEOUT_ENV = "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT"
+_ASYNC_OUTPUT_TIMEOUT_DEFAULT = 600.0 # seconds
+
+
+def _async_output_timeout() -> float:
+ """Seconds to wait for one step's background D2H/SHM copy.
+
+ The copy itself finishes in milliseconds, but it is queued behind the GPU
+ work for that step, so the wall-clock wait tracks step time — a single-GPU
+ box legitimately runs tens of seconds per step on large shapes. A tight
+ bound therefore does not catch a hung engine (worker death and a dead
+ result pump are surfaced by the worker monitor and ``check_health``); it
+ only aborts renders that are still making progress, throwing away the
+ denoise that already completed. The default matches
+ ``_DLO_DP_WAVE_TIMEOUT_S`` in the same subsystem.
+
+ Resolved here rather than at import so a malformed value degrades to the
+ default instead of raising on the request path: this runs inside
+ ``step_streaming``/``add_req_and_wait_for_response``, where a typo in the
+ environment must not start failing generations.
+ """
+ raw = os.environ.get(_ASYNC_OUTPUT_TIMEOUT_ENV)
+ if raw is None:
+ return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+ try:
+ timeout = float(raw)
+ except ValueError:
+ logger.warning_once(
+ "Ignoring %s=%r: not a number. Using the default %.1fs.",
+ _ASYNC_OUTPUT_TIMEOUT_ENV,
+ raw,
+ _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+ )
+ return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+ if timeout <= 0:
+ logger.warning_once(
+ "Ignoring %s=%r: must be positive. Using the default %.1fs.",
+ _ASYNC_OUTPUT_TIMEOUT_ENV,
+ raw,
+ _ASYNC_OUTPUT_TIMEOUT_DEFAULT,
+ )
+ return _ASYNC_OUTPUT_TIMEOUT_DEFAULT
+ return timeout
+
__all__ = [
"DiffusionEngine",
@@ -102,8 +158,14 @@ def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) ->
def supports_request_batch(od_config: OmniDiffusionConfig) -> bool:
model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None))
- if model_cls is None:
- model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None))
+ if model_cls is not None:
+ return bool(getattr(model_cls, "supports_request_batch", False))
+
+ if uses_diffusers_adapter(od_config):
+ model_cls = DiffusionModelRegistry._try_load_model_cls("DiffusersAdapterPipeline")
+ return bool(getattr(model_cls, "supports_request_batch", False))
+
+ model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None))
if model_cls is None:
return False
return bool(getattr(model_cls, "supports_request_batch", False))
@@ -184,12 +246,39 @@ class DiffusionEngine:
from the resolved execution mode.
"""
self.od_config = od_config
+ # Set after the paged-KV profile request has gone through model-owned
+ # preprocessing. Real requests are admitted only within this measured
+ # activation envelope: (max execution sequences, max seq_len,
+ # max target_len).
+ self._diffusion_kv_profile_limits: tuple[int, int, int] | None = None
self._init_process_hooks(od_config)
self.execution_mode = self._resolve_execution_mode(od_config)
self._init_executor(od_config)
try:
- self._init_scheduler(od_config, scheduler)
+ profile_requests = self._prepare_diffusion_kv_profile_requests()
+ kv_control_plane = initialize_diffusion_kv_control_plane(
+ self.executor,
+ od_config,
+ profile_requests=profile_requests,
+ )
+ if kv_control_plane is None:
+ self._init_scheduler(od_config, scheduler)
+ else:
+ (
+ kv_cache_config,
+ scheduler_block_size,
+ hash_block_size,
+ kv_vllm_config,
+ ) = kv_control_plane
+ self._init_scheduler(
+ od_config,
+ scheduler,
+ kv_cache_config,
+ scheduler_block_size=scheduler_block_size,
+ hash_block_size=hash_block_size,
+ kv_vllm_config=kv_vllm_config,
+ )
self._init_runtime_state()
self._init_execute_fn()
self._log_execution_mode(od_config)
@@ -242,6 +331,11 @@ class DiffusionEngine:
self,
od_config: OmniDiffusionConfig,
scheduler: BaseScheduler | None = None,
+ kv_cache_config: KVCacheConfig | None = None,
+ *,
+ scheduler_block_size: int | None = None,
+ hash_block_size: int | None = None,
+ kv_vllm_config: VllmConfig | None = None,
) -> None:
if scheduler is not None:
self.scheduler = scheduler
@@ -249,7 +343,16 @@ class DiffusionEngine:
self.scheduler = StepScheduler()
else:
self.scheduler = RequestScheduler()
- self.scheduler.initialize(od_config)
+ if kv_cache_config is None:
+ self.scheduler.initialize(od_config)
+ else:
+ self.scheduler.initialize(
+ od_config,
+ kv_cache_config=kv_cache_config,
+ scheduler_block_size=scheduler_block_size,
+ hash_block_size=hash_block_size,
+ kv_vllm_config=kv_vllm_config,
+ )
def _init_runtime_state(self) -> None:
# DP multi-concurrency: allow batching dp_size requests so each
@@ -284,6 +387,9 @@ class DiffusionEngine:
self._shutdown_complete = False
self.abort_queue: queue.Queue[str] = queue.Queue()
self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
+ # Copied onto the existing output metrics payload so queue monitoring
+ # reuses the normal diffusion result path without additional IPC.
+ self._scheduler_num_waiting_reqs = 0
def _init_execute_fn(self) -> None:
if self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
@@ -339,9 +445,28 @@ class DiffusionEngine:
# Async mode: wait for background D2H/SHM to complete.
if output.async_output_id:
fut = self.executor.wait_output_ready(output.async_output_id)
- output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=_ASYNC_OUTPUT_TIMEOUT)
+ timeout = _async_output_timeout()
+ try:
+ output = await asyncio.wait_for(asyncio.wrap_future(fut), timeout=timeout)
+ except (TimeoutError, asyncio.TimeoutError):
+ describe = getattr(self.executor, "describe_pending_state", None)
+ logger.error(
+ "Timed out after %.1fs waiting for async output; set %s to a larger value "
+ "to allow slower steps. Executor state: %s",
+ timeout,
+ _ASYNC_OUTPUT_TIMEOUT_ENV,
+ describe(output.async_output_id) if describe else "unavailable",
+ )
+ raise
postprocess_start_time = time.perf_counter()
- formatted_outputs = self.postprocess_output(request, output)
+ scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0))
+ try:
+ formatted_outputs = self.postprocess_output(request, output)
+ except Exception as exc:
+ # Preserve the latest scheduler snapshot across terminal
+ # abort/error paths, which do not produce formatted outputs.
+ setattr(exc, "diffusion_metrics", scheduler_metrics)
+ raise
postprocess_time = time.perf_counter() - postprocess_start_time
step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
logger.debug(
@@ -353,14 +478,24 @@ class DiffusionEngine:
step_total_ms,
)
for request_output in formatted_outputs:
- request_output.metrics.update(
- {
- "preprocess_time_ms": preprocess_time * 1000,
- "diffusion_engine_exec_time_ms": exec_total_time * 1000,
- "diffusion_engine_total_time_ms": step_total_ms,
- "postprocess_time_ms": postprocess_time * 1000,
- }
- )
+ metrics_update = {
+ "preprocess_time_ms": preprocess_time * 1000,
+ "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+ "postprocess_time_ms": postprocess_time * 1000,
+ **scheduler_metrics,
+ }
+ if request.scheduler_queue_wait_ms is not None:
+ metrics_update["scheduler_queue_wait_ms"] = request.scheduler_queue_wait_ms
+ vae_decode_ms = extract_diffusion_vae_decode_ms(output)
+ if vae_decode_ms is not None:
+ metrics_update["vae_decode_time_ms"] = vae_decode_ms
+ forward_ms = extract_diffusion_denoise_ms(output)
+ if forward_ms is not None:
+ metrics_update["forward_time_ms"] = forward_ms
+ kv_recv_ms = getattr(output, "kv_recv_ms", 0.0)
+ if kv_recv_ms > 0:
+ metrics_update["kv_recv_time_ms"] = kv_recv_ms
+ request_output.metrics.update(metrics_update)
yield formatted_outputs
async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
@@ -451,6 +586,7 @@ class DiffusionEngine:
self._wait_for_admission_if_needed_locked()
sched_output = self.scheduler.schedule()
+ self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0)
if sched_output.is_empty:
self._emit_finished_outputs(sched_output.finished_req_ids, None)
@@ -578,12 +714,25 @@ class DiffusionEngine:
if not task.future.done():
task.future.set_exception(exc)
+ def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None:
+ """Clear terminal Worker rows while Scheduler owns the allocations."""
+
+ od_config = getattr(self, "od_config", None)
+ if od_config is None or not is_scheduler_paged_kv_mode(
+ getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ ):
+ return
+ unique_request_ids = list(dict.fromkeys(request_ids))
+ if unique_request_ids:
+ self.executor.remove_diffusion_kv_requests(unique_request_ids)
+
def _emit_finished_outputs(
self,
finished_ids: set[str],
runner_output: BaseRunnerOutput | None = None,
missing_result_error: str = "Diffusion execution finished without a final output",
) -> None:
+ self._remove_diffusion_kv_requests(finished_ids)
for rid in finished_ids:
if runner_output is not None:
_output = runner_output.get_request_output(rid)
@@ -603,6 +752,8 @@ class DiffusionEngine:
self._emit_finished_outputs(finished_ids, runner_output)
return
+ self._remove_diffusion_kv_requests(finished_ids)
+
delivered_finished_req_ids: set[str] = set()
# finished_ids may have some requests that are not scheduler in this round.
@@ -694,9 +845,36 @@ class DiffusionEngine:
"""Run model-owned preprocessing once, before entering Engine locks."""
pre_process_func = getattr(self, "pre_process_func", None)
- if pre_process_func is None:
- return request
- return pre_process_func(request)
+ if pre_process_func is not None:
+ request = pre_process_func(request)
+ self._validate_diffusion_kv_profile_limits(request)
+ return request
+
+ def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None:
+ """Keep admitted paged-KV requests within the profiled activation shape."""
+
+ profile_limits = getattr(self, "_diffusion_kv_profile_limits", None)
+ if profile_limits is None:
+ return
+ kv_requests = request.diffusion_kv_requests
+ if not kv_requests:
+ return
+
+ request_limits = (
+ len(kv_requests),
+ max(kv_request.seq_len for kv_request in kv_requests),
+ max(kv_request.target_len for kv_request in kv_requests),
+ )
+ if any(actual > profiled for actual, profiled in zip(request_limits, profile_limits, strict=True)):
+ request_sequences, request_seq_len, request_target_len = request_limits
+ profile_sequences, profile_seq_len, profile_target_len = profile_limits
+ raise ValueError(
+ f"Diffusion KV request {request.request_id!r} exceeds the startup memory-profile envelope: "
+ f"sequences={request_sequences} (profiled={profile_sequences}), "
+ f"max_seq_len={request_seq_len} (profiled={profile_seq_len}), "
+ f"max_target_len={request_target_len} (profiled={profile_target_len}). "
+ "Reduce the request shape or extend the model's paged-KV profile recipe."
+ )
def _add_prepared_request(self, request: OmniDiffusionRequest) -> str:
"""Admit a request whose model-owned preprocessing is complete."""
@@ -769,6 +947,7 @@ class DiffusionEngine:
sched_output = self.scheduler.schedule()
if sched_output.is_empty:
if target_request_id in sched_output.finished_req_ids:
+ self._remove_diffusion_kv_requests([target_request_id])
return self._finalize_finished_request(target_request_id)
if not self.scheduler.has_requests():
raise RuntimeError("Diffusion scheduler has no runnable requests.")
@@ -798,6 +977,7 @@ class DiffusionEngine:
if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
raise ValueError("Sync func should receive one result at one time")
if target_request_id in finished_req_ids:
+ self._remove_diffusion_kv_requests([target_request_id])
req_output = runner_output.get_request_output(target_request_id)
output = self._finalize_finished_request(
target_request_id,
@@ -806,7 +986,7 @@ class DiffusionEngine:
)
if output.async_output_id:
fut = self.executor.wait_output_ready(output.async_output_id)
- output = fut.result(timeout=_ASYNC_OUTPUT_TIMEOUT)
+ output = fut.result(timeout=_async_output_timeout())
return output
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
@@ -860,47 +1040,111 @@ class DiffusionEngine:
self.close()
raise e
- def _dummy_run(self):
- """A dummy run to warm up the model."""
- num_inference_steps = 1
- height = 512
- width = 512
- prompt: OmniTextPrompt = {"prompt": "dummy run"}
+ def _make_dummy_request(
+ self,
+ *,
+ height: int,
+ width: int,
+ guidance_scale: float,
+ num_image_inputs: int = 1,
+ ) -> OmniDiffusionRequest | None:
+ """Build a one-step model request for startup profiling or warmup."""
+ prompt: OmniTextPrompt = {"prompt": "dummy run"}
supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
if supports_image_input:
- # Provide a dummy image input if the model supports it
color_format = image_color_format(self.od_config.model_class_name)
- dummy_image = PIL.Image.new(color_format, (width, height))
- prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
+ images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)]
+ prompt.setdefault("multi_modal_data", {})["image"] = images[0] if len(images) == 1 else images
if supports_audio_input:
audio_sr = 16000
- dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
- prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
+ prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32)
num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
if num_frames <= 0:
- logger.info("Skipping dummy warmup run (num_frames=0)")
- return
- req = OmniDiffusionRequest(
+ return None
+ return OmniDiffusionRequest(
prompt=prompt,
request_id=DUMMY_DIFFUSION_REQUEST_ID,
sampling_params=OmniDiffusionSamplingParams(
height=height,
width=width,
- num_inference_steps=num_inference_steps,
+ num_inference_steps=1,
num_frames=num_frames,
- # Keep warmup path minimal and robust across text encoders.
- # Some models may fail when warmup implicitly triggers
- # classifier-free guidance with an empty negative prompt.
- guidance_scale=0.0,
+ guidance_scale=guidance_scale,
num_outputs_per_prompt=1,
- # Disable CFG for warmup to avoid triggering CFG parallel
- # validation when cfg_parallel_size > 1.
extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
),
)
+
+ def _prepare_diffusion_kv_profile_requests(self) -> list[OmniDiffusionRequest] | None:
+ """Prepare the per-rank request batch used to profile paged-KV headroom.
+
+ The profile executes directly on each Worker before the Scheduler and
+ its KV manager exist. It uses the maximum number of requests that one
+ rank can execute together. DLO+DP is the exception because each rank
+ executes one request from the collective wave.
+
+ Hunyuan is currently the only model integrated with
+ ``paged_scheduler``; 1024x1024, enabled CFG, and the maximum advertised
+ reference-image count exercise its first-step activation peak.
+ Admission compares each preprocessed request's CFG count and tokenized
+ sequence/target shape with the resulting per-request profile envelope.
+ Future paged model integrations must extend this recipe for their
+ serving limits (for example, video frame count) rather than reusing it
+ silently.
+ """
+
+ if (
+ getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ is not DiffusionKVCacheMode.PAGED_SCHEDULER
+ ):
+ return None
+ request = self._make_dummy_request(
+ height=1024,
+ width=1024,
+ guidance_scale=5.0,
+ num_image_inputs=get_dummy_run_num_image_inputs(self.od_config.model_class_name),
+ )
+ if request is None:
+ raise RuntimeError("paged_scheduler requires a runnable Diffusion KV memory profile request")
+ request = self._prepare_request_for_admission(request)
+ kv_requests = request.diffusion_kv_requests
+ if not kv_requests:
+ raise RuntimeError("paged_scheduler profile preprocessing must produce Diffusion KV requests")
+ self._diffusion_kv_profile_limits = (
+ len(kv_requests),
+ max(kv_request.seq_len for kv_request in kv_requests),
+ max(kv_request.target_len for kv_request in kv_requests),
+ )
+ dlo_dp_request_mode = self.execution_mode is DiffusionExecutionMode.REQUEST_BATCH and _uses_dlo_dp_concurrency(
+ self.od_config
+ )
+ profile_batch_size = 1 if dlo_dp_request_mode else _max_num_seqs(self.od_config)
+ profile_requests: list[OmniDiffusionRequest] = []
+ for index in range(profile_batch_size):
+ profile_request = copy.copy(request)
+ profile_request.request_id = f"{DUMMY_DIFFUSION_REQUEST_ID}/kv-profile-{index}"
+ profile_request.sampling_params = copy.deepcopy(request.sampling_params)
+ # Preprocessing owns request geometry and prepared model inputs,
+ # but native KV requests remain Scheduler-only mutable state. The
+ # profile bypasses Scheduler admission and must not send them to
+ # the Worker.
+ profile_request.diffusion_kv_requests = None
+ profile_requests.append(profile_request)
+ return profile_requests
+
+ def _dummy_run(self):
+ """A dummy run to warm up the model."""
+ req = self._make_dummy_request(
+ height=512,
+ width=512,
+ guidance_scale=0.0,
+ )
+ if req is None:
+ logger.info("Skipping dummy warmup run (num_frames=0)")
+ return
logger.info("dummy run to warm up the model")
output = self.add_req_and_wait_for_response(req)
if output.error:
@@ -1085,8 +1329,11 @@ class DiffusionEngine:
def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
+ request_ids = list(dict.fromkeys(request_ids))
- for request_id in dict.fromkeys(request_ids):
+ self._remove_diffusion_kv_requests(request_ids)
+
+ for request_id in request_ids:
if self.scheduler.get_request_state(request_id) is not None:
self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
@@ -1118,4 +1365,7 @@ class DiffusionEngine:
if runner_output is not None and runner_output.async_output_id is not None:
return DiffusionOutput(async_output_id=runner_output.async_output_id)
+ if state.status == DiffusionRequestStatus.FINISHED_ERROR and state.error:
+ return DiffusionOutput(error=state.error)
+
return DiffusionOutput(error=missing_result_error)
diff --git a/vllm_omni/diffusion/diffusion_kv/config.py b/vllm_omni/diffusion/diffusion_kv/config.py
index cd69d8988..0a7d6070b 100644
--- a/vllm_omni/diffusion/diffusion_kv/config.py
+++ b/vllm_omni/diffusion/diffusion_kv/config.py
@@ -37,3 +37,8 @@ def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode:
"use 'dense_legacy' or 'paged_scheduler'"
)
return mode
+
+
+def is_scheduler_paged_kv_mode(mode: DiffusionKVCacheMode) -> bool:
+ """Return whether a parsed cache mode uses Scheduler-owned paging."""
+ return mode is DiffusionKVCacheMode.PAGED_SCHEDULER
diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py
new file mode 100644
index 000000000..31aaaa46e
--- /dev/null
+++ b/vllm_omni/diffusion/diffusion_kv/initialization.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.config import VllmConfig
+from vllm.v1.core.kv_cache_utils import (
+ generate_scheduler_kv_cache_config,
+ get_kv_cache_configs,
+ resolve_kv_cache_block_sizes,
+)
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+ from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+ from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+def build_native_kv_cache_configs(
+ vllm_config: VllmConfig,
+ worker_specs: list[dict[str, KVCacheSpec]],
+ available_memory: list[int],
+) -> tuple[list[KVCacheConfig], KVCacheConfig]:
+ """Build rank-local and Scheduler-native configs using vLLM utilities."""
+
+ from vllm.v1.core import single_type_kv_cache_manager
+
+ single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config)
+ worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory)
+ return worker_configs, generate_scheduler_kv_cache_config(worker_configs)
+
+
+def initialize_diffusion_kv_control_plane(
+ executor: DiffusionExecutor,
+ od_config: OmniDiffusionConfig,
+ *,
+ profile_requests: list[OmniDiffusionRequest] | None = None,
+ device: torch.device | None = None,
+) -> tuple[KVCacheConfig, int, int, VllmConfig] | None:
+ """Run the native Worker-spec to Scheduler-config initialization chain."""
+
+ if (
+ getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ is not DiffusionKVCacheMode.PAGED_SCHEDULER
+ ):
+ return None
+ if not profile_requests:
+ raise ValueError("paged_scheduler Diffusion KV initialization requires prepared memory profile requests")
+ if device is None:
+ device = current_omni_platform.get_torch_device(0)
+ vllm_config = create_diffusion_vllm_config(device, od_config)
+ worker_specs = executor.get_kv_cache_specs()
+ available_memory = executor.determine_available_kv_memory(profile_requests)
+ if len(worker_specs) != od_config.num_gpus or len(available_memory) != od_config.num_gpus:
+ raise ValueError(
+ "Diffusion KV initialization rank count mismatch: "
+ f"expected={od_config.num_gpus}, specs={len(worker_specs)}, memory={len(available_memory)}"
+ )
+ worker_configs, scheduler_kv_cache_config = build_native_kv_cache_configs(
+ vllm_config,
+ worker_specs,
+ available_memory,
+ )
+ scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes(
+ scheduler_kv_cache_config,
+ vllm_config,
+ )
+ # Native sizing may auto-fit ``max_model_len=-1`` against the profiled
+ # memory budget. Forward that final value so every Worker sizes its
+ # physical BlockTables from the same bound as the Scheduler.
+ resolved_max_model_len = vllm_config.model_config.max_model_len
+ if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0:
+ raise ValueError("native Diffusion KV sizing must resolve a positive max_model_len")
+ executor.set_kv_cache_configs(worker_configs, resolved_max_model_len)
+ return (
+ scheduler_kv_cache_config,
+ scheduler_block_size,
+ hash_block_size,
+ vllm_config,
+ )
diff --git a/vllm_omni/diffusion/diffusion_kv/manager.py b/vllm_omni/diffusion/diffusion_kv/manager.py
new file mode 100644
index 000000000..eee51b864
--- /dev/null
+++ b/vllm_omni/diffusion/diffusion_kv/manager.py
@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from vllm_omni.diffusion.diffusion_kv.metadata import (
+ DiffusionKVMetadata,
+ DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+
+
+class DiffusionKVAdmissionError(ValueError):
+ """A request-specific KV admission failure safe to return to the caller."""
+
+
+class DiffusionKVCacheManager:
+ """Atomic public-request facade over native vLLM KVCacheManager."""
+
+ def __init__(
+ self,
+ kv_cache_config: KVCacheConfig,
+ *,
+ max_model_len: int,
+ scheduler_block_size: int,
+ hash_block_size: int,
+ max_in_flight_tokens: int | None = None,
+ ) -> None:
+ if max_model_len <= 0:
+ raise ValueError(f"max_model_len must be positive, got {max_model_len}")
+ if kv_cache_config.num_blocks <= 0 or not kv_cache_config.kv_cache_groups:
+ raise ValueError("Diffusion KVCacheConfig must contain a positive block pool and at least one group")
+ if scheduler_block_size <= 0 or hash_block_size <= 0:
+ raise ValueError("scheduler_block_size and hash_block_size must be positive")
+ if max_in_flight_tokens is not None and max_in_flight_tokens <= 0:
+ raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}")
+ self.native_manager = KVCacheManager(
+ kv_cache_config=kv_cache_config,
+ max_model_len=max_model_len,
+ scheduler_block_size=scheduler_block_size,
+ hash_block_size=hash_block_size,
+ max_in_flight_tokens=max_in_flight_tokens,
+ enable_caching=False,
+ )
+ self.max_model_len = max_model_len
+ # Native vLLM may reserve a null block, so an idle BlockPool does not
+ # necessarily report ``kv_cache_config.num_blocks`` free blocks.
+ self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks()
+ self._requests: dict[str, tuple[DiffusionKVRequest, ...]] = {}
+ self._metadata: dict[str, DiffusionKVMetadata] = {}
+ self._internal_request_ids: set[str] = set()
+ self._next_allocation_generation = 1
+
+ def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int:
+ """Return the native full-sequence reservation for one public request.
+
+ vLLM performs this calculation inside ``allocate_slots`` when
+ ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic
+ requirement that native requests do not have: every execution sequence
+ must fit atomically. Sum the native per-sequence result so a request
+ that can never fit is rejected independently of current pool usage.
+ """
+
+ coordinator = self.native_manager.coordinator
+ empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks
+ return sum(
+ coordinator.get_num_blocks_to_allocate(
+ request_id=request.request_id,
+ num_tokens=request.num_tokens,
+ new_computed_blocks=empty_blocks,
+ num_encoder_tokens=0,
+ total_computed_tokens=0,
+ num_local_computed_tokens=0,
+ num_tokens_main_model=request.num_tokens,
+ apply_admission_cap=True,
+ )
+ for request in requests
+ )
+
+ def has_request(self, public_request_id: str) -> bool:
+ return public_request_id in self._requests
+
+ def reserve_request(
+ self,
+ public_request_id: str,
+ kv_requests: Sequence[DiffusionKVRequest],
+ ) -> DiffusionKVMetadata | None:
+ """Reserve all CFG sequences or roll back the complete public request."""
+
+ if not public_request_id:
+ raise ValueError("public_request_id must be non-empty")
+ if public_request_id in self._requests:
+ raise ValueError(f"Diffusion KV request {public_request_id!r} is already allocated")
+ requests = tuple(kv_requests)
+ if not requests:
+ raise ValueError("Diffusion KV allocation requires at least one sequence")
+
+ sequence_ids = [request.sequence_id for request in requests]
+ internal_ids = [request.request_id for request in requests]
+ expected_sequence_ids = list(range(len(requests)))
+ if sequence_ids != expected_sequence_ids:
+ raise ValueError(
+ "Diffusion KV requests must be ordered by contiguous sequence_id values: "
+ f"expected={expected_sequence_ids}, got={sequence_ids}"
+ )
+ if len(internal_ids) != len(set(internal_ids)):
+ raise ValueError(f"Diffusion KV internal request IDs must be unique, got {internal_ids}")
+ conflicts = self._internal_request_ids.intersection(internal_ids)
+ if conflicts:
+ raise ValueError(f"Diffusion KV internal request IDs are already allocated: {sorted(conflicts)}")
+ for request in requests:
+ if request.kv_contexts:
+ raise DiffusionKVAdmissionError("independent DiffusionKVContext allocation is not implemented yet")
+ if request.seq_len > self.max_model_len:
+ raise DiffusionKVAdmissionError(
+ f"Diffusion KV sequence {request.request_id!r} exceeds max_model_len: "
+ f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
+ )
+
+ required_blocks = self._get_empty_pool_required_blocks(requests)
+ if required_blocks > self._empty_pool_num_free_blocks:
+ raise DiffusionKVAdmissionError(
+ f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: "
+ f"required_blocks={required_blocks}, available_blocks={self._empty_pool_num_free_blocks}; "
+ "increase KV cache capacity or reduce the request sequence count/length"
+ )
+
+ allocated: list[DiffusionKVRequest] = []
+ sequence_metadata: list[DiffusionKVSequenceMetadata] = []
+ try:
+ for request in requests:
+ blocks = self.native_manager.allocate_slots(
+ request,
+ num_new_tokens=request.seq_len,
+ delay_cache_blocks=True,
+ full_sequence_must_fit=True,
+ )
+ if blocks is None:
+ self._rollback(allocated)
+ allocated.clear()
+ return None
+ allocated.append(request)
+ sequence_metadata.append(
+ DiffusionKVSequenceMetadata(
+ sequence_id=request.sequence_id,
+ prefix_len=request.prefix_len,
+ target_len=request.target_len,
+ seq_len=request.seq_len,
+ block_ids=blocks.get_block_ids(),
+ )
+ )
+ except Exception:
+ self._rollback(allocated)
+ raise
+
+ metadata = DiffusionKVMetadata(
+ request_id=public_request_id,
+ allocation_generation=self._next_allocation_generation,
+ sequences=tuple(sequence_metadata),
+ )
+ self._next_allocation_generation += 1
+ self._requests[public_request_id] = requests
+ self._metadata[public_request_id] = metadata
+ self._internal_request_ids.update(internal_ids)
+ return metadata
+
+ def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata:
+ try:
+ return self._metadata[public_request_id]
+ except KeyError as exc:
+ raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
+
+ def free_request(self, public_request_id: str) -> None:
+ requests = self._requests.pop(public_request_id, ())
+ self._metadata.pop(public_request_id, None)
+ for request in reversed(requests):
+ self.native_manager.free(request)
+ self._internal_request_ids.discard(request.request_id)
+
+ def close(self) -> None:
+ for public_request_id in tuple(self._requests):
+ self.free_request(public_request_id)
+
+ def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
+ for request in reversed(requests):
+ self.native_manager.free(request)
diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py
new file mode 100644
index 000000000..c5ab3fb16
--- /dev/null
+++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py
@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed import get_dcp_group
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
+from vllm.v1.worker import block_table as native_block_table
+from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
+from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+ DiffusionPagedAttentionAdapter,
+ DiffusionPagedAttentionLayerAdapter,
+ DiffusionPagedAttentionMetadata,
+ DiffusionPagedAttentionRowBinding,
+ DiffusionPagedAttentionRuntime,
+)
+from vllm_omni.platforms import current_omni_platform
+
+DiffusionKVIdentity = tuple[str, int | None, str | None]
+DiffusionKVSnapshot = tuple[object, ...]
+
+
+@dataclass(frozen=True)
+class _DiffusionKVRowInstall:
+ identity: DiffusionKVIdentity
+ token_len: int
+ block_ids: tuple[tuple[int, ...], ...]
+
+
+@dataclass(frozen=True)
+class _DiffusionKVRequestState:
+ generation: int
+ snapshot: DiffusionKVSnapshot
+ row_installs: tuple[_DiffusionKVRowInstall, ...]
+
+
+class DiffusionKVModelRunnerBackend:
+ """Native paged-KV state and BlockTable operations for a model runner."""
+
+ def __init__(
+ self,
+ *,
+ vllm_config: VllmConfig,
+ od_config: OmniDiffusionConfig,
+ device: torch.device,
+ ) -> None:
+ self.vllm_config = vllm_config
+ self.od_config = od_config
+ self.device = device
+ self.kv_cache_config: KVCacheConfig | None = None
+ self.kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+ self.block_tables: BlockTables | None = None
+ self.paged_attention_adapter: DiffusionPagedAttentionAdapter | None = None
+ self._kv_cache_layer_adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
+ self._diffusion_kv_identity_to_row: dict[DiffusionKVIdentity, int] = {}
+ self._diffusion_kv_free_rows: list[int] = []
+ self._diffusion_kv_request_states: dict[str, _DiffusionKVRequestState] = {}
+
+ def register_kv_cache_layers(
+ self,
+ layers: Mapping[str, tuple[Any, KVCacheSpec]],
+ ) -> dict[str, KVCacheSpec]:
+ """Register native adapters and return their canonical cache specs."""
+
+ if self.kv_cache_config is not None:
+ raise RuntimeError("Diffusion KV cache layers cannot be changed after physical initialization")
+
+ if layers:
+ current_omni_platform.get_diffusion_kv_block_tables_cls()
+ if not callable(current_omni_platform.build_diffusion_kv_attn_metadata):
+ raise NotImplementedError(
+ "Diffusion paged KV requires platform-native BlockTables and attention metadata hooks"
+ )
+
+ forward_context = self.vllm_config.compilation_config.static_forward_context
+ previous_adapters = self._kv_cache_layer_adapters
+ for layer_name, (layer, spec) in layers.items():
+ if not isinstance(spec, AttentionSpec):
+ raise TypeError(
+ f"Diffusion KV layer {layer_name!r} produced unsupported spec {type(spec).__name__}; "
+ "only native attention specs are supported"
+ )
+ if not layer.attn_backend.supports_paged_kv:
+ raise NotImplementedError(
+ f"Diffusion paged KV layer {layer_name!r} requires an Omni backend with paged support; "
+ f"selected {layer.attn_backend.get_name()}"
+ )
+ existing = forward_context.get(layer_name)
+ if existing is not None and existing is not layer and existing is not previous_adapters.get(layer_name):
+ raise RuntimeError(f"Diffusion KV layer name {layer_name!r} collides with the native forward context")
+
+ previous_use_non_causal = self.vllm_config.attention_config.use_non_causal
+ self.vllm_config.attention_config.use_non_causal = any(
+ bool(getattr(spec, "non_causal", False)) for _, spec in layers.values()
+ )
+ adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
+ try:
+ with set_current_vllm_config(self.vllm_config):
+ for layer_name, (layer, spec) in layers.items():
+ adapters[layer_name] = DiffusionPagedAttentionLayerAdapter(
+ layer_name=layer_name,
+ layer=layer,
+ spec=spec,
+ vllm_config=self.vllm_config,
+ device=self.device,
+ ulysses_degree=self._get_paged_attention_ulysses_degree(layer),
+ )
+ except Exception:
+ self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+ raise
+
+ geometries = {(adapter.num_heads, adapter.num_kv_heads, adapter.head_size) for adapter in adapters.values()}
+ if len(geometries) > 1:
+ self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+ raise ValueError(
+ "Native paged attention metadata builders require one rank-local attention geometry; "
+ f"got {sorted(geometries)!r}"
+ )
+ attention_geometry = next(iter(geometries)) if geometries else None
+ set_attention_geometry = (
+ getattr(self.vllm_config.model_config, "set_attention_geometry", None)
+ if attention_geometry is not None
+ else None
+ )
+ if attention_geometry is not None and not callable(set_attention_geometry):
+ self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+ raise RuntimeError(
+ "Diffusion native paged attention requires a model config that exposes set_attention_geometry"
+ )
+
+ updated_forward_context = dict(forward_context)
+ for layer_name in set(layers) | set(previous_adapters):
+ entry = updated_forward_context.get(layer_name)
+ layer_entry = layers.get(layer_name)
+ if entry is previous_adapters.get(layer_name) or (layer_entry is not None and entry is layer_entry[0]):
+ updated_forward_context.pop(layer_name, None)
+ updated_forward_context.update(adapters)
+
+ if attention_geometry is not None:
+ num_heads, num_kv_heads, head_size = attention_geometry
+ try:
+ set_attention_geometry(
+ num_heads=num_heads,
+ num_kv_heads=num_kv_heads,
+ head_size=head_size,
+ )
+ except Exception:
+ self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+ raise
+
+ forward_context.clear()
+ forward_context.update(updated_forward_context)
+ self._kv_cache_layer_adapters = adapters
+ return {layer_name: adapter.spec for layer_name, adapter in adapters.items()}
+
+ def _get_paged_attention_ulysses_degree(self, layer: Any) -> int:
+ if bool(getattr(layer, "skip_sequence_parallel", False)):
+ return 1
+
+ parallel_config = getattr(self.od_config, "parallel_config", None)
+ ulysses_degree = int(getattr(parallel_config, "ulysses_degree", 1))
+ ring_degree = int(getattr(parallel_config, "ring_degree", 1))
+ allgather_degree = int(getattr(parallel_config, "allgather_degree", 1))
+ ulysses_mode = str(getattr(parallel_config, "ulysses_mode", "strict"))
+ if ring_degree > 1:
+ raise NotImplementedError(
+ "Diffusion paged attention does not support Ring or hybrid Ulysses+Ring attention"
+ )
+ if allgather_degree > 1:
+ raise NotImplementedError("Diffusion paged attention does not support AllGather-KV sequence parallelism")
+ if ulysses_degree > 1 and ulysses_mode != "strict":
+ raise NotImplementedError(
+ f"Diffusion paged attention currently supports only strict Ulysses; got ulysses_mode={ulysses_mode!r}"
+ )
+ return ulysses_degree
+
+ def _get_max_rows_per_request(self) -> int:
+ max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None)
+ if type(max_rows) is not int or max_rows <= 0:
+ raise ValueError(
+ "paged_scheduler requires a positive integer diffusion_kv_max_rows_per_request adapter limit"
+ )
+ return max_rows
+
+ def _get_max_model_len(self) -> int:
+ max_len = getattr(self.vllm_config.model_config, "max_model_len", None)
+ if type(max_len) is not int or max_len <= 0:
+ raise ValueError("paged_scheduler requires a positive model_config.max_model_len")
+ return max_len
+
+ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+ """Allocate and bind native attention KV tensors for this rank."""
+ if self.kv_cache_config is not None:
+ raise RuntimeError("Diffusion native KV cache is already initialized")
+ if not self._kv_cache_layer_adapters:
+ raise RuntimeError("Diffusion KV cache layers must be registered before physical initialization")
+
+ configured_layers = {
+ layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
+ }
+ registered_layers = set(self._kv_cache_layer_adapters)
+ if configured_layers != registered_layers:
+ raise ValueError(
+ "Rank-local Diffusion KVCacheConfig layer mismatch: "
+ f"expected={sorted(registered_layers)}, configured={sorted(configured_layers)}"
+ )
+
+ kv_cache_config = copy.deepcopy(kv_cache_config)
+ max_len = self._get_max_model_len()
+ max_rows_per_request = self._get_max_rows_per_request()
+ scheduler_config = self.vllm_config.scheduler_config
+ max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
+ if type(max_num_seqs) is not int or max_num_seqs <= 0:
+ raise ValueError("paged_scheduler requires a positive scheduler_config.max_num_seqs")
+
+ kv_cache_groups = kv_cache_config.kv_cache_groups
+ if not kv_cache_groups:
+ raise ValueError("paged_scheduler requires at least one native KV cache group")
+ block_sizes = [group.kv_cache_spec.block_size for group in kv_cache_groups]
+ if any(type(block_size) is not int or block_size <= 0 for block_size in block_sizes):
+ raise ValueError("native diffusion KV cache groups must have positive integer block sizes")
+ unaligned_max_num_blocks_per_group = [
+ group.kv_cache_spec.max_num_blocks_per_req(self.vllm_config, max_len) for group in kv_cache_groups
+ ]
+ if any(type(capacity) is not int or capacity <= 0 for capacity in unaligned_max_num_blocks_per_group):
+ raise ValueError("native diffusion KV cache groups must have positive row capacities")
+ max_num_blocks_per_group = [
+ native_block_table.get_block_table_width(capacity, block_size)
+ for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True)
+ ]
+
+ max_num_reqs = max_num_seqs * max_rows_per_request
+ max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None)
+ if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0:
+ raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer")
+ parallel_config = self.vllm_config.parallel_config
+ cp_size = parallel_config.decode_context_parallel_size
+ cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0
+ cp_interleave = parallel_config.cp_kv_cache_interleave_size
+
+ # Native metadata builders size per-row buffers from max_num_seqs.
+ # A diffusion public request can occupy several sequence/context rows.
+ scheduler_config.max_num_seqs = max_num_reqs
+ kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+ previous_adapter_caches = {
+ layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items()
+ }
+ try:
+ with set_current_vllm_config(self.vllm_config):
+ attn_groups, _, kernel_block_sizes = init_attn_backend(
+ kv_cache_config,
+ self.vllm_config,
+ self.device,
+ )
+ block_tables_cls = current_omni_platform.get_diffusion_kv_block_tables_cls()
+ block_tables = block_tables_cls(
+ block_sizes=block_sizes,
+ max_num_reqs=max_num_reqs,
+ max_num_batched_tokens=max_num_batched_tokens,
+ max_num_blocks_per_group=max_num_blocks_per_group,
+ device=self.device,
+ kernel_block_sizes=kernel_block_sizes,
+ cp_size=cp_size,
+ cp_rank=cp_rank,
+ cp_interleave=cp_interleave,
+ )
+ paged_attention_adapter = DiffusionPagedAttentionAdapter(
+ vllm_config=self.vllm_config,
+ device=self.device,
+ kv_cache_config=kv_cache_config,
+ block_tables=block_tables,
+ attn_groups=attn_groups,
+ layers=self._kv_cache_layer_adapters,
+ resolve_row=self._resolve_paged_attention_row,
+ )
+ init_kv_cache(
+ kv_caches,
+ self.vllm_config.compilation_config.static_forward_context,
+ kv_cache_config,
+ attn_groups,
+ self.device,
+ self.vllm_config.cache_config.cache_dtype,
+ kernel_block_sizes,
+ self.vllm_config,
+ )
+ except Exception:
+ scheduler_config.max_num_seqs = max_num_seqs
+ for layer_name, previous_cache in previous_adapter_caches.items():
+ self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache
+ raise
+
+ self.kv_caches = kv_caches
+ self.block_tables = block_tables
+ self.paged_attention_adapter = paged_attention_adapter
+ self._diffusion_kv_identity_to_row.clear()
+ self._diffusion_kv_free_rows = list(range(max_num_reqs - 1, -1, -1))
+ self._diffusion_kv_request_states.clear()
+ self.kv_cache_config = kv_cache_config
+
+ def _validate_row(
+ self,
+ *,
+ identity: DiffusionKVIdentity,
+ token_len: int,
+ block_ids: tuple[list[int], ...],
+ seen_identities: set[DiffusionKVIdentity],
+ ) -> _DiffusionKVRowInstall:
+ if identity in seen_identities:
+ raise ValueError(f"Duplicate diffusion KV row identity: {identity!r}")
+ seen_identities.add(identity)
+ if type(token_len) is not int or token_len < 0:
+ raise ValueError(f"Diffusion KV row {identity!r} has invalid token length {token_len!r}")
+
+ assert self.kv_cache_config is not None
+ groups = self.kv_cache_config.kv_cache_groups
+ if not isinstance(block_ids, tuple) or len(block_ids) != len(groups):
+ actual_groups = len(block_ids) if isinstance(block_ids, (tuple, list)) else "invalid"
+ raise ValueError(f"Diffusion KV row {identity!r} has {actual_groups} block groups; expected {len(groups)}")
+
+ max_len = self._get_max_model_len()
+ normalized_groups: list[tuple[int, ...]] = []
+ for group_index, (group_block_ids, kv_cache_group) in enumerate(zip(block_ids, groups, strict=True)):
+ if not isinstance(group_block_ids, list):
+ raise ValueError(f"Diffusion KV row {identity!r} group {group_index} block IDs must be a list")
+ spec = kv_cache_group.kv_cache_spec
+ expected_count = spec.max_num_blocks_per_req(self.vllm_config, token_len)
+ if len(group_block_ids) != expected_count:
+ raise ValueError(
+ f"Diffusion KV row {identity!r} group {group_index} has {len(group_block_ids)} blocks; "
+ f"expected {expected_count} for {token_len} tokens"
+ )
+ row_capacity = spec.max_num_blocks_per_req(self.vllm_config, max_len)
+ if len(group_block_ids) > row_capacity:
+ raise ValueError(
+ f"Diffusion KV row {identity!r} group {group_index} exceeds row capacity {row_capacity}"
+ )
+ normalized_ids: list[int] = []
+ for block_id in group_block_ids:
+ if type(block_id) is not int:
+ raise ValueError(
+ f"Diffusion KV row {identity!r} group {group_index} has non-integer block ID {block_id!r}"
+ )
+ if not 0 <= block_id < self.kv_cache_config.num_blocks:
+ raise ValueError(
+ f"Diffusion KV row {identity!r} group {group_index} block ID {block_id} is outside "
+ f"[0, {self.kv_cache_config.num_blocks})"
+ )
+ normalized_ids.append(block_id)
+ normalized_groups.append(tuple(normalized_ids))
+
+ return _DiffusionKVRowInstall(
+ identity=identity,
+ token_len=token_len,
+ block_ids=tuple(normalized_groups),
+ )
+
+ def _prepare_install(
+ self,
+ metadata: DiffusionKVMetadata,
+ ) -> tuple[list[_DiffusionKVRowInstall], DiffusionKVSnapshot]:
+ if self.block_tables is None or self.kv_cache_config is None:
+ raise RuntimeError("paged_scheduler native KV cache is not initialized")
+ if type(metadata.request_id) is not str or not metadata.request_id:
+ raise ValueError("Diffusion KV metadata request_id must be a non-empty string")
+ if type(metadata.allocation_generation) is not int or metadata.allocation_generation < 0:
+ raise ValueError("Diffusion KV allocation_generation must be a non-negative integer")
+ if type(self.kv_cache_config.num_blocks) is not int or self.kv_cache_config.num_blocks <= 0:
+ raise ValueError("native diffusion KV cache config must provide a positive num_blocks")
+ if not metadata.sequences:
+ raise ValueError(f"Diffusion KV request {metadata.request_id!r} must contain at least one sequence")
+
+ seen_identities: set[DiffusionKVIdentity] = set()
+ installs: list[_DiffusionKVRowInstall] = []
+ sequence_snapshots: list[object] = []
+ for sequence in metadata.sequences:
+ if type(sequence.sequence_id) is not int or sequence.sequence_id < 0:
+ raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}")
+ sequence_identity = (metadata.request_id, sequence.sequence_id, None)
+ sequence_install = self._validate_row(
+ identity=sequence_identity,
+ token_len=sequence.seq_len,
+ block_ids=sequence.block_ids,
+ seen_identities=seen_identities,
+ )
+ installs.append(sequence_install)
+ sequence_snapshots.append(
+ (
+ sequence.sequence_id,
+ sequence.prefix_len,
+ sequence.target_len,
+ sequence.seq_len,
+ sequence_install.block_ids,
+ sequence.context_ids,
+ )
+ )
+
+ context_snapshots: list[object] = []
+ for context in metadata.contexts:
+ if type(context.context_id) is not str or not context.context_id:
+ raise ValueError("Diffusion KV context_id must be a non-empty string")
+ context_identity = (metadata.request_id, None, context.context_id)
+ context_install = self._validate_row(
+ identity=context_identity,
+ token_len=context.num_tokens,
+ block_ids=context.block_ids,
+ seen_identities=seen_identities,
+ )
+ installs.append(context_install)
+ context_snapshots.append(
+ (
+ context.context_id,
+ context.cache_role,
+ context.num_tokens,
+ context_install.block_ids,
+ )
+ )
+
+ max_rows_per_request = self._get_max_rows_per_request()
+ if len(installs) > max_rows_per_request:
+ raise ValueError(
+ f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows; "
+ f"adapter limit is {max_rows_per_request}"
+ )
+ snapshot: DiffusionKVSnapshot = (
+ metadata.request_id,
+ tuple(sequence_snapshots),
+ tuple(context_snapshots),
+ )
+ return installs, snapshot
+
+ def _rollback_staged_writes(self, rows: list[int], previous_num_blocks: object | None) -> None:
+ block_tables = self.block_tables
+ if block_tables is None:
+ return
+ for group_table in getattr(block_tables, "block_tables", ()):
+ clear_staged_writes = getattr(group_table, "clear_staged_writes", None)
+ if clear_staged_writes is not None:
+ clear_staged_writes()
+ if previous_num_blocks is None:
+ return
+ native_num_blocks = getattr(block_tables, "num_blocks", None)
+ if native_num_blocks is None or not hasattr(native_num_blocks, "np"):
+ return
+ native_num_blocks.np[:, rows] = previous_num_blocks
+ copy_to_uva = getattr(native_num_blocks, "copy_to_uva", None)
+ if copy_to_uva is not None:
+ copy_to_uva()
+
+ def _apply_rows(self, rows: list[int], installs: list[_DiffusionKVRowInstall]) -> None:
+ assert self.block_tables is not None
+ if self.paged_attention_adapter is not None:
+ self.paged_attention_adapter.invalidate_prepared_batches()
+ native_num_blocks = getattr(self.block_tables, "num_blocks", None)
+ previous_num_blocks = None
+ if native_num_blocks is not None and hasattr(native_num_blocks, "np"):
+ previous_num_blocks = native_num_blocks.np[:, rows].copy()
+ try:
+ for row, install in zip(rows, installs, strict=True):
+ self.block_tables.append_block_ids(
+ row,
+ tuple(list(group_ids) for group_ids in install.block_ids),
+ overwrite=True,
+ )
+ self.block_tables.apply_staged_writes()
+ except Exception:
+ self._rollback_staged_writes(rows, previous_num_blocks)
+ raise
+
+ def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
+ """Install one Scheduler allocation snapshot into native Worker rows."""
+ if not is_scheduler_paged_kv_mode(
+ getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ ):
+ raise ValueError("Dense diffusion execution must not install Diffusion KV metadata")
+
+ installs, snapshot = self._prepare_install(metadata)
+ current_state = self._diffusion_kv_request_states.get(metadata.request_id)
+ if current_state is not None:
+ if metadata.allocation_generation < current_state.generation:
+ raise ValueError(
+ f"Stale Diffusion KV allocation generation {metadata.allocation_generation} for request "
+ f"{metadata.request_id!r}; installed generation is {current_state.generation}"
+ )
+ if snapshot != current_state.snapshot:
+ raise ValueError(
+ f"Conflicting Diffusion KV allocation snapshot for active request {metadata.request_id!r}; "
+ "remove the request before replacing its rows"
+ )
+ if metadata.allocation_generation == current_state.generation:
+ return False
+ self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
+ generation=metadata.allocation_generation,
+ snapshot=snapshot,
+ row_installs=current_state.row_installs,
+ )
+ return False
+
+ existing_identities = [
+ install.identity for install in installs if install.identity in self._diffusion_kv_identity_to_row
+ ]
+ if existing_identities:
+ raise RuntimeError(f"Diffusion KV row registry contains orphaned identities: {existing_identities!r}")
+ if len(installs) > len(self._diffusion_kv_free_rows):
+ raise ValueError(
+ f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows, but only "
+ f"{len(self._diffusion_kv_free_rows)} rows are free"
+ )
+
+ rows = list(reversed(self._diffusion_kv_free_rows[-len(installs) :])) if installs else []
+ if rows:
+ self._apply_rows(rows, installs)
+
+ if installs:
+ del self._diffusion_kv_free_rows[-len(installs) :]
+ self._diffusion_kv_identity_to_row.update(
+ (install.identity, row) for install, row in zip(installs, rows, strict=True)
+ )
+ self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
+ generation=metadata.allocation_generation,
+ snapshot=snapshot,
+ row_installs=tuple(installs),
+ )
+ return True
+
+ def get_diffusion_kv_row(
+ self,
+ request_id: str,
+ sequence_id: int | None,
+ context_id: str | None = None,
+ ) -> int:
+ """Resolve a Scheduler allocation identity to its native table row."""
+ if (
+ not is_scheduler_paged_kv_mode(
+ getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ )
+ or self.block_tables is None
+ ):
+ raise RuntimeError("paged_scheduler native BlockTables are not initialized")
+ identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
+ try:
+ return self._diffusion_kv_identity_to_row[identity]
+ except KeyError as exc:
+ raise KeyError(f"No native diffusion KV row is installed for {identity!r}") from exc
+
+ def get_paged_attention_adapter(self) -> DiffusionPagedAttentionAdapter:
+ if self.paged_attention_adapter is None:
+ raise RuntimeError("paged_scheduler native attention adapter is not initialized")
+ return self.paged_attention_adapter
+
+ def activate_paged_attention_metadata(self, metadata: DiffusionPagedAttentionMetadata):
+ """Activate Runner-owned request metadata for an internal denoise loop."""
+
+ runtime = DiffusionPagedAttentionRuntime(self.get_paged_attention_adapter(), metadata)
+ return runtime, runtime.activate()
+
+ def _resolve_paged_attention_row(
+ self,
+ request_id: str,
+ sequence_id: int | None,
+ context_id: str | None,
+ ) -> DiffusionPagedAttentionRowBinding:
+ row_index = self.get_diffusion_kv_row(request_id, sequence_id, context_id)
+ identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
+ request_state = self._diffusion_kv_request_states.get(request_id)
+ if request_state is not None:
+ for install in request_state.row_installs:
+ if install.identity == identity:
+ return DiffusionPagedAttentionRowBinding(
+ row_index=row_index,
+ max_seq_len=install.token_len,
+ block_ids=install.block_ids,
+ )
+ raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
+
+ def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+ """Retire Worker rows without logically freeing Scheduler-owned blocks."""
+ if (
+ not is_scheduler_paged_kv_mode(
+ getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ )
+ or self.block_tables is None
+ ):
+ return 0
+ request_id_set = set(request_ids)
+ identities_and_rows = [
+ (identity, row)
+ for identity, row in self._diffusion_kv_identity_to_row.items()
+ if identity[0] in request_id_set
+ ]
+ rows = [row for _, row in identities_and_rows]
+ if rows:
+ assert self.kv_cache_config is not None
+ num_groups = len(self.kv_cache_config.kv_cache_groups)
+ installs = [
+ _DiffusionKVRowInstall(
+ identity=identity,
+ token_len=0,
+ block_ids=tuple(tuple() for _ in range(num_groups)),
+ )
+ for identity, _ in identities_and_rows
+ ]
+ self._apply_rows(rows, installs)
+
+ for identity, _ in identities_and_rows:
+ del self._diffusion_kv_identity_to_row[identity]
+ for request_id in request_id_set:
+ self._diffusion_kv_request_states.pop(request_id, None)
+ self._diffusion_kv_free_rows.extend(rows)
+ self._diffusion_kv_free_rows.sort(reverse=True)
+ return len(rows)
+
+ def refresh_block_table_layout(self) -> None:
+ """Refresh native pointer tensors after a CuMem KV-cache wake-up."""
+ if (
+ not is_scheduler_paged_kv_mode(
+ getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+ )
+ or self.block_tables is None
+ ):
+ return
+ if self.paged_attention_adapter is not None:
+ self.paged_attention_adapter.invalidate_prepared_batches()
+ self.block_tables.init_block_table_layout_tensors()
diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py
new file mode 100644
index 000000000..6a7862669
--- /dev/null
+++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py
@@ -0,0 +1,946 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
+ PagedPiecewisePlan,
+ build_paged_piecewise_plan,
+)
+from vllm_omni.diffusion.forward_context import override_paged_kv_adapter
+from vllm_omni.platforms import current_omni_platform
+
+
+@dataclass(frozen=True)
+class DiffusionPagedAttentionRowBinding:
+ """Worker row and logical length installed for one allocation identity."""
+
+ row_index: int
+ max_seq_len: int
+ block_ids: tuple[tuple[int, ...], ...] = ()
+
+
+DiffusionKVRowResolver = Callable[
+ [str, int | None, str | None],
+ DiffusionPagedAttentionRowBinding,
+]
+
+
+class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
+ """Register a diffusion layer with vLLM's native cache machinery.
+
+ This object deliberately does *not* subclass ``vllm.Attention``. The
+ latter owns a second execution path and would bypass Omni's sequence
+ parallel pre/post hooks. The wrapper only supplies the small
+ ``AttentionLayerBase`` contract needed by ``init_attn_backend`` and keeps
+ the platform-native attention implementation/cache view available to the
+ diffusion adapter.
+ """
+
+ def __init__(
+ self,
+ *,
+ layer_name: str,
+ layer: Any,
+ spec: AttentionSpec,
+ vllm_config: VllmConfig,
+ device: torch.device,
+ ulysses_degree: int = 1,
+ ) -> None:
+ if type(ulysses_degree) is not int or ulysses_degree <= 0:
+ raise ValueError(f"ulysses_degree must be a positive integer, got {ulysses_degree!r}")
+ num_heads = int(layer.num_heads)
+ num_kv_heads = int(spec.num_kv_heads)
+ if num_heads <= 0 or num_kv_heads <= 0 or num_heads % num_kv_heads != 0:
+ raise ValueError(
+ "Paged attention requires positive Q/KV heads with num_heads divisible by num_kv_heads: "
+ f"num_heads={num_heads}, num_kv_heads={num_kv_heads}"
+ )
+ if num_heads % ulysses_degree != 0 or num_kv_heads % ulysses_degree != 0:
+ raise ValueError(
+ "Paged attention requires Q/KV heads divisible by ulysses_degree: "
+ f"num_heads={num_heads}, num_kv_heads={num_kv_heads}, ulysses_degree={ulysses_degree}"
+ )
+ num_heads //= ulysses_degree
+ num_kv_heads //= ulysses_degree
+
+ attention_config = vllm_config.attention_config
+ previous_backend = attention_config.backend
+ previous_backend_per_kind = attention_config.backend_per_kind
+ try:
+ # This is a portable vLLM backend request. The active platform
+ # resolves it to its native implementation, such as FlashAttention
+ # on CUDA or AscendAttentionBackend on NPU.
+ attention_config.backend = AttentionBackendEnum.FLASH_ATTN
+ attention_config.backend_per_kind = {}
+ with set_current_vllm_config(vllm_config):
+ attn_backend = get_attn_backend(
+ head_size=spec.head_size,
+ dtype=vllm_config.model_config.dtype,
+ kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+ attn_type=AttentionType.DECODER,
+ num_heads=num_heads,
+ )
+ finally:
+ attention_config.backend = previous_backend
+ attention_config.backend_per_kind = previous_backend_per_kind
+ attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
+ attn_backend,
+ ulysses_degree=ulysses_degree,
+ )
+ canonical_spec = replace(
+ spec,
+ num_kv_heads=num_kv_heads,
+ indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
+ )
+ self.layer_name = layer_name
+ self.spec = canonical_spec
+ self.num_heads = num_heads
+ self.num_kv_heads = num_kv_heads
+ self.head_size = int(canonical_spec.head_size)
+ self.head_size_v = int(getattr(canonical_spec, "head_size_v", canonical_spec.head_size))
+ if self.head_size_v != self.head_size:
+ raise ValueError(
+ "Diffusion native paged attention requires equal key/value head sizes; "
+ f"got head_size={self.head_size}, head_size_v={self.head_size_v}"
+ )
+ self.softmax_scale = float(layer.softmax_scale)
+ self.attn_backend = attn_backend
+ self.kv_cache: torch.Tensor | None = None
+ # Native attention backends consume either device tensors or host
+ # scalar copies of these scales. The paged diffusion path is currently
+ # unquantized, so initialize both contracts to identity.
+ self._q_scale = torch.ones(1, device=device, dtype=torch.float32)
+ self._k_scale = torch.ones(1, device=device, dtype=torch.float32)
+ self._v_scale = torch.ones(1, device=device, dtype=torch.float32)
+ self._q_scale_float = 1.0
+ self._k_scale_float = 1.0
+ self._v_scale_float = 1.0
+ with set_current_vllm_config(vllm_config):
+ self.impl = self._create_native_impl(vllm_config)
+
+ def get_attn_backend(self):
+ return self.attn_backend
+
+ def _create_native_impl(self, vllm_config: VllmConfig):
+ impl_cls = self.attn_backend.get_impl_cls()
+ impl = impl_cls(
+ num_heads=self.num_heads,
+ head_size=self.head_size,
+ scale=self.softmax_scale,
+ num_kv_heads=self.num_kv_heads,
+ alibi_slopes=None,
+ sliding_window=None,
+ kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+ logits_soft_cap=None,
+ attn_type=AttentionType.DECODER,
+ kv_sharing_target_layer_name=None,
+ )
+ if not callable(getattr(impl, "forward", None)):
+ raise RuntimeError(
+ f"Native paged-attention backend {self.attn_backend.get_name()!r} has no forward implementation"
+ )
+ needs_explicit_cache_update = (
+ not self.attn_backend.forward_includes_kv_cache_update
+ or current_omni_platform.requires_diffusion_paged_kv_prewrite()
+ )
+ if needs_explicit_cache_update and not callable(getattr(impl, "do_kv_cache_update", None)):
+ raise RuntimeError(
+ f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache"
+ )
+ if hasattr(impl, "vllm_flash_attn_version") and impl.vllm_flash_attn_version is None:
+ raise RuntimeError(f"Native paged-attention kernel is unavailable for diffusion layer {self.layer_name!r}")
+ return impl
+
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> AttentionSpec:
+ del vllm_config
+ return self.spec
+
+
+@dataclass(frozen=True)
+class DiffusionPagedAttentionRow:
+ """One logical BlockTable row participating in a paged attention call."""
+
+ request_id: str
+ query_len: int
+ seq_len: int
+ kv_start_pos: int = 0
+ sequence_id: int | None = None
+ context_id: str | None = None
+
+ def __post_init__(self) -> None:
+ if type(self.request_id) is not str or not self.request_id:
+ raise ValueError("Paged attention request_id must be a non-empty string")
+ if (self.sequence_id is None) == (self.context_id is None):
+ raise ValueError("Paged attention row requires exactly one of sequence_id or context_id")
+ if self.sequence_id is not None and (type(self.sequence_id) is not int or self.sequence_id < 0):
+ raise ValueError("Paged attention sequence_id must be a non-negative integer")
+ if self.context_id is not None and (type(self.context_id) is not str or not self.context_id):
+ raise ValueError("Paged attention context_id must be a non-empty string")
+ if type(self.query_len) is not int or self.query_len <= 0:
+ raise ValueError("Paged attention query_len must be a positive integer")
+ if type(self.seq_len) is not int or self.seq_len <= 0:
+ raise ValueError("Paged attention seq_len must be a positive integer")
+ if type(self.kv_start_pos) is not int or self.kv_start_pos < 0:
+ raise ValueError("Paged attention kv_start_pos must be a non-negative integer")
+ if self.kv_start_pos + self.query_len > self.seq_len:
+ raise ValueError(
+ "Paged attention write span exceeds seq_len: "
+ f"start={self.kv_start_pos}, query_len={self.query_len}, seq_len={self.seq_len}"
+ )
+
+ @property
+ def identity(self) -> tuple[str, int | None, str | None]:
+ return (self.request_id, self.sequence_id, self.context_id)
+
+
+@dataclass(frozen=True)
+class DiffusionPagedAttentionMetadata:
+ """Runner-owned row layouts for one request-level denoise loop."""
+
+ prefill_rows: tuple[DiffusionPagedAttentionRow, ...]
+ denoise_rows: tuple[DiffusionPagedAttentionRow, ...]
+
+ def __post_init__(self) -> None:
+ if not self.prefill_rows or not self.denoise_rows:
+ raise ValueError("Paged attention metadata requires both prefill and denoise rows")
+ prefill_identities = tuple(row.identity for row in self.prefill_rows)
+ denoise_identities = tuple(row.identity for row in self.denoise_rows)
+ if prefill_identities != denoise_identities:
+ raise ValueError(
+ "Paged attention prefill and denoise rows must use the same ordered identities: "
+ f"prefill={prefill_identities!r}, denoise={denoise_identities!r}"
+ )
+
+ def rows_for_step(self, step_idx: int | None) -> tuple[DiffusionPagedAttentionRow, ...]:
+ if type(step_idx) is not int or step_idx < 0:
+ raise ValueError(f"Paged attention requires a non-negative denoise step index, got {step_idx!r}")
+ return self.prefill_rows if step_idx == 0 else self.denoise_rows
+
+
+@dataclass(frozen=True)
+class PreparedDiffusionPagedAttentionBatch:
+ """Native metadata shared by all paged attention layers in one forward."""
+
+ rows: tuple[DiffusionPagedAttentionRow, ...]
+ row_indices: torch.Tensor
+ query_start_loc: torch.Tensor
+ seq_lens: torch.Tensor
+ positions: torch.Tensor
+ block_tables: tuple[torch.Tensor, ...]
+ slot_mappings: torch.Tensor
+ attn_metadata: dict[str, Any]
+ slot_mappings_by_layer: dict[str, torch.Tensor]
+ num_tokens: int
+ _owner: object = field(repr=False, compare=False)
+ _generation: int = field(repr=False, compare=False)
+
+
+@dataclass(frozen=True, slots=True)
+class DiffusionPagedAttentionContext:
+ """One layer's native inputs for an Omni paged-backend invocation."""
+
+ layer: DiffusionPagedAttentionLayerAdapter
+ query: torch.Tensor
+ key_write: torch.Tensor
+ value_write: torch.Tensor
+ slot_mapping: torch.Tensor
+ native_metadata: Any
+ piecewise_plan: PagedPiecewisePlan | None
+ piecewise_native_metadata: tuple[Any, ...]
+ query_token_shape: tuple[int, ...]
+ query_has_head_dims: bool
+ output_scatter_indices: torch.Tensor | None = None
+
+ def restore_output(self, output: torch.Tensor) -> torch.Tensor:
+ if self.output_scatter_indices is not None:
+ flat_output = output.reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+ restored = output.new_zeros(
+ (*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
+ ).reshape(-1, self.layer.num_heads, self.layer.head_size_v)
+ restored.index_copy_(0, self.output_scatter_indices, flat_output)
+ if self.query_has_head_dims:
+ return restored.reshape(
+ *self.query_token_shape,
+ self.layer.num_heads,
+ self.layer.head_size_v,
+ )
+ return restored.reshape(
+ *self.query_token_shape,
+ self.layer.num_heads * self.layer.head_size_v,
+ )
+ if self.query_has_head_dims:
+ return output.reshape(*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
+ return output.reshape(*self.query_token_shape, self.layer.num_heads * self.layer.head_size_v)
+
+
+class DiffusionPagedAttentionRuntime:
+ """Activate Runner-prepared rows lazily as the denoise loop advances."""
+
+ def __init__(
+ self,
+ adapter: DiffusionPagedAttentionAdapter,
+ metadata: DiffusionPagedAttentionMetadata,
+ ) -> None:
+ self.adapter = adapter
+ self.metadata = metadata
+ self._active_phase: bool | None = None
+ self._activation: Any | None = None
+ self._closed = False
+
+ def ensure_active(self, step_idx: int | None) -> DiffusionPagedAttentionAdapter:
+ if self._closed:
+ raise RuntimeError("Paged attention runtime is already closed")
+ is_prefill = step_idx == 0
+ rows = self.metadata.rows_for_step(step_idx)
+ if self._active_phase == is_prefill:
+ return self.adapter
+
+ self._deactivate()
+ batch = self.adapter.prepare_batch(rows)
+ activation = self.adapter.activate(batch)
+ activation.__enter__()
+ self._activation = activation
+ self._active_phase = is_prefill
+ return self.adapter
+
+ def _deactivate(self) -> None:
+ activation = self._activation
+ self._activation = None
+ self._active_phase = None
+ if activation is not None:
+ activation.__exit__(None, None, None)
+
+ def close(self) -> None:
+ if self._closed:
+ return
+ self._deactivate()
+ self._closed = True
+
+ @contextmanager
+ def activate(self) -> Iterator[DiffusionPagedAttentionRuntime]:
+ try:
+ yield self
+ finally:
+ self.close()
+
+
+class DiffusionPagedAttentionAdapter:
+ """Translate diffusion rows into vLLM metadata for an Omni paged backend."""
+
+ def __init__(
+ self,
+ *,
+ vllm_config: VllmConfig,
+ device: torch.device,
+ kv_cache_config: KVCacheConfig,
+ block_tables: BlockTables,
+ attn_groups: list[list[Any]],
+ layers: Mapping[str, DiffusionPagedAttentionLayerAdapter],
+ resolve_row: DiffusionKVRowResolver,
+ ) -> None:
+ if not layers:
+ raise ValueError("Paged attention requires at least one native attention layer")
+ if len(attn_groups) != len(kv_cache_config.kv_cache_groups):
+ raise ValueError(
+ "Paged attention group mismatch: "
+ f"builders={len(attn_groups)}, cache_groups={len(kv_cache_config.kv_cache_groups)}"
+ )
+ self.vllm_config = vllm_config
+ self.device = torch.device(device)
+ self.kv_cache_config = kv_cache_config
+ self.block_tables = block_tables
+ self.attn_groups = attn_groups
+ self.layers = dict(layers)
+ self.resolve_row = resolve_row
+ self._owner = object()
+ self._prepare_generation = 0
+ self._active_batch: PreparedDiffusionPagedAttentionBatch | None = None
+ self._active_piecewise_plan: PagedPiecewisePlan | None = None
+ self._active_piecewise_native_metadata: tuple[dict[str, Any], ...] | None = None
+ self._causal_by_group = self._resolve_group_causality()
+ self._reorder_batch_threshold = self._resolve_reorder_batch_threshold()
+
+ def _resolve_group_causality(self) -> dict[int, bool]:
+ causal_by_group: dict[int, bool] = {}
+ for group_index, cache_group in enumerate(self.kv_cache_config.kv_cache_groups):
+ layer_causality = {
+ not bool(getattr(self.layers[layer_name].spec, "non_causal", False))
+ for layer_name in cache_group.layer_names
+ }
+ if len(layer_causality) != 1:
+ raise ValueError(f"Paged attention cache group {group_index} mixes causal and non-causal layers")
+ causal_by_group[group_index] = layer_causality.pop()
+ return causal_by_group
+
+ def _resolve_reorder_batch_threshold(self) -> int | None:
+ thresholds = [
+ group.get_metadata_builder(0).reorder_batch_threshold
+ for group_index, groups in enumerate(self.attn_groups)
+ if self._causal_by_group[group_index]
+ for group in groups
+ ]
+ concrete_thresholds = [threshold for threshold in thresholds if threshold is not None]
+ return min(concrete_thresholds, default=None)
+
+ def _validate_native_batch_order(self, rows: tuple[DiffusionPagedAttentionRow, ...]) -> None:
+ threshold = self._reorder_batch_threshold
+ if threshold is None or not any(self._causal_by_group.values()):
+ return
+
+ found_long_query = False
+ for row in rows:
+ if row.query_len > threshold:
+ found_long_query = True
+ elif found_long_query:
+ raise ValueError(
+ "Causal paged attention rows must place native decode/short-query rows before "
+ f"prefill/long-query rows (decode threshold={threshold})"
+ )
+
+ def _validate_row_capacity(
+ self,
+ rows: tuple[DiffusionPagedAttentionRow, ...],
+ row_indices: list[int],
+ ) -> None:
+ native_num_blocks = self.block_tables.num_blocks.np
+ blocks_per_kv_block = self.block_tables.blocks_per_kv_block
+ for row, row_index in zip(rows, row_indices, strict=True):
+ if type(row_index) is not int or not 0 <= row_index < self.block_tables.max_num_reqs:
+ raise ValueError(f"Paged attention row {row.identity!r} resolved to invalid Worker row {row_index!r}")
+ for group_index, (block_size, block_multiplier) in enumerate(
+ zip(
+ self.block_tables.block_sizes,
+ blocks_per_kv_block,
+ strict=True,
+ )
+ ):
+ required_blocks = cdiv(
+ row.seq_len,
+ block_size * self.block_tables.cp_size,
+ )
+ required_kernel_blocks = required_blocks * block_multiplier
+ installed_kernel_blocks = int(native_num_blocks[group_index, row_index])
+ if required_kernel_blocks > installed_kernel_blocks:
+ raise ValueError(
+ f"Paged attention row {row.identity!r} requires {required_blocks} blocks in "
+ f"cache group {group_index} for seq_len={row.seq_len}, but its installed Worker "
+ f"row contains only {installed_kernel_blocks // block_multiplier} blocks"
+ )
+
+ def _build_native_metadata(
+ self,
+ *,
+ query_lens: Sequence[int],
+ query_start_loc: torch.Tensor,
+ query_start_loc_cpu: torch.Tensor,
+ seq_lens: torch.Tensor,
+ seq_lens_cpu: torch.Tensor,
+ positions: torch.Tensor,
+ block_tables: Sequence[torch.Tensor],
+ slot_mappings: torch.Tensor,
+ causal: bool | Mapping[int, bool],
+ ) -> dict[str, Any]:
+ dcp_local_seq_lens = None
+ if self.block_tables.cp_size > 1:
+ dcp_local_seq_lens = get_dcp_local_seq_lens(
+ seq_lens,
+ dcp_size=self.block_tables.cp_size,
+ dcp_rank=self.block_tables.cp_rank,
+ cp_kv_cache_interleave_size=self.block_tables.cp_interleave,
+ )
+ with set_current_vllm_config(self.vllm_config):
+ return current_omni_platform.build_diffusion_kv_attn_metadata(
+ attn_groups=self.attn_groups,
+ num_reqs=len(query_lens),
+ num_tokens=sum(query_lens),
+ query_start_loc_gpu=query_start_loc,
+ query_start_loc_cpu=query_start_loc_cpu,
+ max_query_len=max(query_lens),
+ seq_lens=seq_lens,
+ max_seq_len=max(int(seq_len) for seq_len in seq_lens_cpu),
+ block_tables=block_tables,
+ slot_mappings=slot_mappings,
+ kv_cache_config=self.kv_cache_config,
+ seq_lens_cpu=seq_lens_cpu,
+ seq_lens_cpu_upper_bound=seq_lens_cpu,
+ dcp_local_seq_lens=dcp_local_seq_lens,
+ positions=positions,
+ causal=causal,
+ )
+
+ def prepare_batch(
+ self,
+ rows: Sequence[DiffusionPagedAttentionRow],
+ ) -> PreparedDiffusionPagedAttentionBatch:
+ rows = tuple(rows)
+ if not rows:
+ raise ValueError("Paged attention batch must contain at least one row")
+ if self._active_batch is not None:
+ raise RuntimeError("Cannot prepare paged attention metadata during an active forward")
+ if len(rows) > self.block_tables.max_num_reqs:
+ raise ValueError(
+ f"Paged attention batch has {len(rows)} rows; Worker capacity is {self.block_tables.max_num_reqs}"
+ )
+
+ identities = [row.identity for row in rows]
+ if len(set(identities)) != len(identities):
+ raise ValueError("Paged attention batch contains duplicate row identities")
+ row_bindings = [self.resolve_row(row.request_id, row.sequence_id, row.context_id) for row in rows]
+ row_indices_list = [binding.row_index for binding in row_bindings]
+ if len(set(row_indices_list)) != len(row_indices_list):
+ raise ValueError("Paged attention batch resolves multiple inputs to the same Worker row")
+ for row, binding in zip(rows, row_bindings, strict=True):
+ if type(binding.max_seq_len) is not int or binding.max_seq_len < 0:
+ raise ValueError(
+ f"Paged attention row {row.identity!r} resolved to invalid logical capacity {binding.max_seq_len!r}"
+ )
+ if row.seq_len > binding.max_seq_len:
+ raise ValueError(
+ f"Paged attention row {row.identity!r} uses seq_len={row.seq_len}, but its installed "
+ f"allocation has logical length {binding.max_seq_len}"
+ )
+ self._validate_native_batch_order(rows)
+ self._validate_row_capacity(rows, row_indices_list)
+
+ query_lens = [row.query_len for row in rows]
+ num_tokens = sum(query_lens)
+ if num_tokens > self.block_tables.max_num_batched_tokens:
+ raise ValueError(
+ f"Paged attention batch has {num_tokens} tokens; Worker capacity is "
+ f"{self.block_tables.max_num_batched_tokens}"
+ )
+ query_offsets = [0]
+ for query_len in query_lens:
+ query_offsets.append(query_offsets[-1] + query_len)
+
+ # Native BlockTables reuses persistent gather/slot buffers. Starting a
+ # new preparation invalidates every batch prepared before it, even if a
+ # later metadata builder raises.
+ self._prepare_generation += 1
+ generation = self._prepare_generation
+ row_indices = torch.tensor(row_indices_list, dtype=torch.int32, device=self.device)
+ query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
+ query_start_loc = query_start_loc_cpu.to(self.device)
+ seq_lens_cpu = torch.tensor([row.seq_len for row in rows], dtype=torch.int32)
+ seq_lens = seq_lens_cpu.to(self.device)
+ positions = torch.cat(
+ [
+ torch.arange(
+ row.kv_start_pos,
+ row.kv_start_pos + row.query_len,
+ dtype=torch.int64,
+ device=self.device,
+ )
+ for row in rows
+ ]
+ )
+
+ block_tables = self.block_tables.gather_block_tables(
+ row_indices,
+ num_reqs_padded=len(rows),
+ )
+ slot_mappings = self.block_tables.compute_slot_mappings(
+ row_indices,
+ query_start_loc,
+ positions,
+ num_tokens_padded=num_tokens,
+ )
+ causal: bool | Mapping[int, bool]
+ causal_values = set(self._causal_by_group.values())
+ causal = causal_values.pop() if len(causal_values) == 1 else self._causal_by_group
+ attn_metadata = self._build_native_metadata(
+ query_lens=query_lens,
+ query_start_loc=query_start_loc,
+ query_start_loc_cpu=query_start_loc_cpu,
+ seq_lens=seq_lens,
+ seq_lens_cpu=seq_lens_cpu,
+ positions=positions,
+ block_tables=block_tables,
+ slot_mappings=slot_mappings,
+ causal=causal,
+ )
+ slot_mappings_by_layer = build_slot_mappings_by_layer(
+ slot_mappings,
+ self.kv_cache_config,
+ )
+ return PreparedDiffusionPagedAttentionBatch(
+ rows=rows,
+ row_indices=row_indices,
+ query_start_loc=query_start_loc,
+ seq_lens=seq_lens,
+ positions=positions,
+ block_tables=tuple(block_tables),
+ slot_mappings=slot_mappings,
+ attn_metadata=attn_metadata,
+ slot_mappings_by_layer=slot_mappings_by_layer,
+ num_tokens=num_tokens,
+ _owner=self._owner,
+ _generation=generation,
+ )
+
+ def invalidate_prepared_batches(self) -> None:
+ """Invalidate native buffer views after BlockTable state changes."""
+
+ if self._active_batch is not None:
+ raise RuntimeError("Cannot change paged attention BlockTables during an active forward")
+ self._prepare_generation += 1
+
+ @contextmanager
+ def activate(
+ self,
+ batch: PreparedDiffusionPagedAttentionBatch,
+ ) -> Iterator[DiffusionPagedAttentionAdapter]:
+ if batch._owner is not self._owner:
+ raise ValueError("Prepared paged attention batch belongs to a different adapter")
+ if batch._generation != self._prepare_generation:
+ raise ValueError("Prepared paged attention batch is stale after a newer batch preparation")
+ if self._active_batch is not None:
+ raise RuntimeError("Paged attention adapter already has an active forward")
+ self._active_batch = batch
+ self._active_piecewise_plan = None
+ self._active_piecewise_native_metadata = None
+ try:
+ with override_paged_kv_adapter(self), set_current_vllm_config(self.vllm_config):
+ yield self
+ finally:
+ self._active_piecewise_plan = None
+ self._active_piecewise_native_metadata = None
+ self._active_batch = None
+
+ @staticmethod
+ def _flatten_tensor(
+ tensor: torch.Tensor,
+ *,
+ num_heads: int,
+ head_size: int,
+ name: str,
+ ) -> tuple[torch.Tensor, tuple[int, ...], bool]:
+ if tensor.ndim >= 3 and tensor.shape[-2:] == (num_heads, head_size):
+ return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-2]), True
+ hidden_size = num_heads * head_size
+ if tensor.ndim >= 1 and tensor.shape[-1] == hidden_size:
+ return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-1]), False
+ raise ValueError(
+ f"Paged attention {name} must end in ({num_heads}, {head_size}) or ({hidden_size},); "
+ f"got shape={tuple(tensor.shape)}"
+ )
+
+ @staticmethod
+ def _token_packing_indices(
+ token_shape: tuple[int, ...],
+ batch: PreparedDiffusionPagedAttentionBatch,
+ *,
+ name: str,
+ device: torch.device,
+ ) -> torch.Tensor | None:
+ if len(token_shape) == 1:
+ if token_shape[0] == batch.num_tokens:
+ return None
+ raise ValueError(
+ f"Paged attention {name} token count must match the prepared write batch: "
+ f"tokens={token_shape[0]}, prepared={batch.num_tokens}"
+ )
+ if len(token_shape) == 2:
+ batch_size, tokens_per_row = token_shape
+ query_lens = tuple(row.query_len for row in batch.rows)
+ if batch_size != len(batch.rows) or any(query_len > tokens_per_row for query_len in query_lens):
+ raise ValueError(
+ f"Paged attention batched {name} layout must match prepared rows for the current write: "
+ f"shape={token_shape}, row_query_lens={query_lens}"
+ )
+ if all(query_len == tokens_per_row for query_len in query_lens):
+ return None
+ return torch.cat(
+ [
+ torch.arange(
+ row_index * tokens_per_row,
+ row_index * tokens_per_row + query_len,
+ dtype=torch.long,
+ device=device,
+ )
+ for row_index, query_len in enumerate(query_lens)
+ ]
+ )
+ raise ValueError(
+ f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; "
+ f"got token shape={token_shape}"
+ )
+
+ def _get_piecewise_plan(
+ self,
+ full_attn_spans: list[list[tuple[int, int]]],
+ ) -> PagedPiecewisePlan:
+ batch = self._active_batch
+ if batch is None:
+ raise RuntimeError("Piecewise paged attention requires an active prepared batch")
+ if any(row.kv_start_pos + row.query_len != row.seq_len for row in batch.rows):
+ invalid_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len]
+ raise ValueError(
+ "Paged piecewise attention requires each query/write span to end at seq_len; "
+ f"invalid rows={invalid_rows!r}"
+ )
+
+ cached_plan = self._active_piecewise_plan
+ if cached_plan is not None:
+ frozen_spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans)
+ if cached_plan.spans != frozen_spans:
+ raise ValueError("Paged piecewise attention metadata changed between layers in one active batch")
+ return cached_plan
+
+ plan = build_paged_piecewise_plan(
+ full_attn_spans,
+ query_offsets=[row.kv_start_pos for row in batch.rows],
+ query_lens=[row.query_len for row in batch.rows],
+ seq_lens=[row.seq_len for row in batch.rows],
+ device=self.device,
+ )
+ self._active_piecewise_plan = plan
+ return plan
+
+ def _get_piecewise_native_metadata(
+ self,
+ plan: PagedPiecewisePlan,
+ ) -> tuple[dict[str, Any], ...]:
+ """Build one native metadata object per segment and cache it by batch."""
+
+ batch = self._active_batch
+ if batch is None:
+ raise RuntimeError("Piecewise paged attention requires an active prepared batch")
+ cached_metadata = self._active_piecewise_native_metadata
+ if cached_metadata is not None:
+ return cached_metadata
+
+ metadata_by_segment: list[dict[str, Any]] = []
+ for packed_segment in plan.segments:
+ row_segments = packed_segment.row_segments
+ query_lens = [segment.q_end - segment.q_start for segment in row_segments]
+ seq_lens_values = [segment.kv_end for segment in row_segments]
+ query_offsets = [0]
+ for query_len in query_lens:
+ query_offsets.append(query_offsets[-1] + query_len)
+ query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
+ seq_lens_cpu = torch.tensor(seq_lens_values, dtype=torch.int32)
+ if packed_segment.query_range is None:
+ positions = batch.positions.index_select(0, packed_segment.query_indices)
+ slot_mappings = batch.slot_mappings.index_select(-1, packed_segment.query_indices)
+ else:
+ start, end = packed_segment.query_range
+ positions = batch.positions[start:end]
+ slot_mappings = batch.slot_mappings[..., start:end]
+ segment_metadata = self._build_native_metadata(
+ query_lens=query_lens,
+ query_start_loc=query_start_loc_cpu.to(self.device),
+ query_start_loc_cpu=query_start_loc_cpu,
+ seq_lens=seq_lens_cpu.to(self.device),
+ seq_lens_cpu=seq_lens_cpu,
+ positions=positions,
+ block_tables=batch.block_tables,
+ slot_mappings=slot_mappings,
+ causal=(row_segments[0].mode == "causal"),
+ )
+ # FA3's full-CUDA-graph metadata builder reuses one persistent
+ # scheduler buffer across builds. Piecewise attention prepares all
+ # segments before executing any of them, so each segment needs its
+ # own snapshot rather than a view overwritten by the next build.
+ snapshotted = set()
+ for native_metadata in segment_metadata.values():
+ metadata_id = id(native_metadata)
+ if metadata_id in snapshotted:
+ continue
+ snapshotted.add(metadata_id)
+ for field_name in ("scheduler_metadata", "prefix_scheduler_metadata"):
+ scheduler_metadata = getattr(native_metadata, field_name, None)
+ if isinstance(scheduler_metadata, torch.Tensor):
+ setattr(native_metadata, field_name, scheduler_metadata.clone())
+ metadata_by_segment.append(segment_metadata)
+ self._active_piecewise_native_metadata = tuple(metadata_by_segment)
+ return self._active_piecewise_native_metadata
+
+ def prepare_layer_context(
+ self,
+ layer_name: str,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ *,
+ omni_attn_metadata: Any | None = None,
+ ) -> DiffusionPagedAttentionContext:
+ if self._active_batch is None:
+ raise RuntimeError("Paged attention forward must run inside adapter.activate(batch)")
+ try:
+ layer = self.layers[layer_name]
+ except KeyError as exc:
+ raise KeyError(f"Unknown diffusion paged attention layer {layer_name!r}") from exc
+ batch = self._active_batch
+ full_attn_spans = self._validate_omni_attn_metadata(omni_attn_metadata)
+
+ query_flat, query_token_shape, query_has_head_dims = self._flatten_tensor(
+ query,
+ num_heads=layer.num_heads,
+ head_size=layer.head_size,
+ name="query",
+ )
+ key_flat, key_token_shape, _ = self._flatten_tensor(
+ key,
+ num_heads=layer.num_kv_heads,
+ head_size=layer.head_size,
+ name="key",
+ )
+ value_flat, value_token_shape, _ = self._flatten_tensor(
+ value,
+ num_heads=layer.num_kv_heads,
+ head_size=layer.head_size_v,
+ name="value",
+ )
+ if key_token_shape != query_token_shape:
+ if len(key_token_shape) == 1:
+ raise ValueError(
+ "Paged attention key token count must match the query layout: "
+ f"query={query_token_shape}, key={key_token_shape}"
+ )
+ raise ValueError(
+ "Paged attention batched key layout must match prepared rows and the query layout: "
+ f"query={query_token_shape}, key={key_token_shape}"
+ )
+ if value_token_shape != query_token_shape:
+ if len(value_token_shape) == 1:
+ raise ValueError(
+ "Paged attention value token count must match the query layout: "
+ f"query={query_token_shape}, value={value_token_shape}"
+ )
+ raise ValueError(
+ "Paged attention batched value layout must match prepared rows and the query layout: "
+ f"query={query_token_shape}, value={value_token_shape}"
+ )
+ output_scatter_indices = self._token_packing_indices(
+ query_token_shape,
+ batch,
+ name="Q/K/V",
+ device=query.device,
+ )
+ if output_scatter_indices is not None:
+ query_flat = query_flat.index_select(0, output_scatter_indices)
+ key_flat = key_flat.index_select(0, output_scatter_indices)
+ value_flat = value_flat.index_select(0, output_scatter_indices)
+ if query.device != self.device or key.device != self.device or value.device != self.device:
+ raise ValueError(
+ f"Paged attention Q/K/V must be on {self.device}; "
+ f"got query={query.device}, key={key.device}, value={value.device}"
+ )
+ if query.dtype != key.dtype or query.dtype != value.dtype:
+ raise ValueError(f"Paged attention Q/K/V dtypes must match; got {query.dtype}, {key.dtype}, {value.dtype}")
+ expected_dtype = self.vllm_config.model_config.dtype
+ if query.dtype != expected_dtype:
+ raise ValueError(
+ f"Paged attention Q/K/V dtype must match model activation dtype {expected_dtype}; got {query.dtype}"
+ )
+ if not bool(getattr(layer.spec, "non_causal", False)):
+ non_suffix_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len]
+ if non_suffix_rows:
+ raise ValueError(
+ "Causal paged attention requires each query/write span to end at seq_len; "
+ f"invalid rows={non_suffix_rows!r}"
+ )
+
+ slot_mapping = batch.slot_mappings_by_layer.get(layer_name)
+ if slot_mapping is None:
+ raise KeyError(f"No native slot mapping was built for diffusion layer {layer_name!r}")
+ slot_mapping = slot_mapping.reshape(-1)[: batch.num_tokens]
+ if slot_mapping.numel() != batch.num_tokens:
+ raise ValueError(
+ f"Paged attention slot mapping for {layer_name!r} has {slot_mapping.numel()} entries; "
+ f"expected {batch.num_tokens}"
+ )
+
+ # vLLM's CUDA writer consumes the source token strides directly, so
+ # preserve strided QKV-projection views instead of materializing an
+ # extra K/V copy here. Platform writers with stricter contracts (for
+ # example, Ascend PA scatter) make their inputs contiguous locally.
+ # The native CUDA writer still assumes heads and head dimensions are
+ # packed within each token; materialize unusual inner layouts only.
+ if key_flat.stride(-1) != 1 or (key_flat.shape[-2] > 1 and key_flat.stride(-2) != key_flat.shape[-1]):
+ key_flat = key_flat.contiguous()
+ if value_flat.stride(-1) != 1 or (value_flat.shape[-2] > 1 and value_flat.stride(-2) != value_flat.shape[-1]):
+ value_flat = value_flat.contiguous()
+ # The writer also consumes slot_mapping directly, so a span that starts
+ # inside a physical block naturally writes only the requested suffix.
+ try:
+ native_metadata = batch.attn_metadata[layer_name]
+ except KeyError as exc:
+ raise KeyError(f"No native attention metadata was built for diffusion layer {layer_name!r}") from exc
+ piecewise_plan: PagedPiecewisePlan | None = None
+ piecewise_native_metadata: tuple[Any, ...] = ()
+ if full_attn_spans is not None:
+ piecewise_plan = self._get_piecewise_plan(full_attn_spans)
+ native_metadata_by_segment = self._get_piecewise_native_metadata(piecewise_plan)
+ try:
+ piecewise_native_metadata = tuple(
+ segment_metadata[layer_name] for segment_metadata in native_metadata_by_segment
+ )
+ except KeyError as exc:
+ raise KeyError(
+ f"No piecewise native attention metadata was built for diffusion layer {layer_name!r}"
+ ) from exc
+ return DiffusionPagedAttentionContext(
+ layer=layer,
+ query=query_flat.contiguous(),
+ key_write=key_flat,
+ value_write=value_flat,
+ slot_mapping=slot_mapping,
+ native_metadata=native_metadata,
+ piecewise_plan=piecewise_plan,
+ piecewise_native_metadata=piecewise_native_metadata,
+ query_token_shape=query_token_shape,
+ query_has_head_dims=query_has_head_dims,
+ output_scatter_indices=output_scatter_indices,
+ )
+
+ @staticmethod
+ def _validate_omni_attn_metadata(
+ metadata: Any | None,
+ ) -> list[list[tuple[int, int]]] | None:
+ if metadata is None:
+ return None
+ full_attn_spans = getattr(metadata, "full_attn_spans", None)
+ attn_mask = getattr(metadata, "attn_mask", None)
+ unsupported_fields = [
+ field_name
+ for field_name in (
+ "joint_attn_mask",
+ "query_ranges",
+ "video_layout",
+ )
+ if getattr(metadata, field_name, None) is not None
+ ]
+ if attn_mask is not None:
+ if full_attn_spans is None:
+ unsupported_fields.append("attn_mask")
+ elif not isinstance(attn_mask, torch.Tensor) or attn_mask.ndim != 4:
+ unsupported_fields.append("attn_mask (piecewise paging requires a 4D tensor)")
+ extra = getattr(metadata, "extra", None)
+ if extra:
+ unsupported_fields.append(f"extra={sorted(extra)}")
+ if unsupported_fields:
+ raise NotImplementedError(
+ "Diffusion paged attention cannot translate Omni attention metadata fields "
+ f"{unsupported_fields!r} to native FlashAttention metadata"
+ )
+ return full_attn_spans
diff --git a/vllm_omni/diffusion/diffusion_kv/request.py b/vllm_omni/diffusion/diffusion_kv/request.py
index 4db9701a1..fd4740099 100644
--- a/vllm_omni/diffusion/diffusion_kv/request.py
+++ b/vllm_omni/diffusion/diffusion_kv/request.py
@@ -105,7 +105,7 @@ class DiffusionKVRequest:
self.status = RequestStatus.WAITING
self.num_preemptions = 0
self.num_in_flight_tokens = 0
- # vLLM 0.26 uses this optional boundary when publishing cached blocks;
+ # vLLM 0.26+ uses this optional boundary when publishing cached blocks;
# zero means that no sparse-retention boundary is active.
self.shared_prefix_boundary = 0
diff --git a/vllm_omni/diffusion/distributed/__init__.py b/vllm_omni/diffusion/distributed/__init__.py
index 668ace06d..2a1159026 100644
--- a/vllm_omni/diffusion/distributed/__init__.py
+++ b/vllm_omni/diffusion/distributed/__init__.py
@@ -6,11 +6,6 @@ from vllm_omni.diffusion.distributed.hsdp import (
HSDPInferenceConfig,
apply_hsdp_to_model,
)
-from vllm_omni.diffusion.distributed.parallel_state import (
- get_fs_group,
- get_fully_shard_rank,
- get_fully_shard_world_size,
-)
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelConfig,
SequenceParallelInput,
@@ -46,8 +41,4 @@ __all__ = [
"get_sharding_validator",
"HSDPInferenceConfig",
"apply_hsdp_to_model",
- # FS (Fully Shard) utilities
- "get_fs_group",
- "get_fully_shard_rank",
- "get_fully_shard_world_size",
]
diff --git a/vllm_omni/diffusion/distributed/a2a_permute.py b/vllm_omni/diffusion/distributed/a2a_permute.py
new file mode 100644
index 000000000..d6ea05218
--- /dev/null
+++ b/vllm_omni/diffusion/distributed/a2a_permute.py
@@ -0,0 +1,215 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused permute-free Ulysses all-to-all over NCCL symmetric memory.
+
+JIT-compiles the CUDA kernel from pytorch/pytorch#178230 (``all_to_all_permute``)
+and exposes it as two *functional* custom ops:
+
+ ulysses_qkv_fwd(x, group_name, world_size) # (B, S/p, H, D) -> (B, S, H/p, D)
+ ulysses_o_rev(y, group_name, world_size) # (B, S, H/p, D) -> (B, S/p, H, D)
+
+These replace the synchronous ``all_to_all_4D`` (permute + NCCL all_to_all_single)
+used by Ulysses SP. All symmetric-memory bookkeeping (a shape-keyed, rendezvoused
+buffer cache + the copy-in) lives *inside* the ops, so from torch.compile's point
+of view each op is an opaque ``input -> output`` function: no graph break, no
+visible buffer mutation, no Python control flow in the traced graph. A
+``register_fake`` provides output metadata so Dynamo keeps the op in-graph.
+"""
+
+from __future__ import annotations
+
+import glob
+import math
+import os
+import sysconfig
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+from torch.distributed.distributed_c10d import _resolve_process_group
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_BUILD_LOCK = threading.Lock()
+_BUILT = False
+_WORKSPACE_LOCK = threading.Lock()
+
+
+@dataclass(slots=True)
+class _SymmWorkspace:
+ allocation: torch.Tensor
+ handle: Any
+ capacity_bytes: int
+ stream_id: int
+
+
+# A single grow-only byte workspace per device/process-group. Shape changes
+# reuse a typed view into the workspace instead of retaining one allocation per
+# historical request shape.
+_SYMM_WORKSPACES: dict[tuple[torch.device, str], _SymmWorkspace] = {}
+
+
+def _ensure_built() -> None:
+ """JIT-compile + load the kernel once (process-wide). Cached by torch."""
+ global _BUILT
+ if _BUILT:
+ return
+ with _BUILD_LOCK:
+ if _BUILT:
+ return
+ from torch.utils.cpp_extension import load
+
+ here = os.path.dirname(__file__)
+ src = os.path.join(here, "csrc", "a2a_permute.cu")
+ site_packages = sysconfig.get_paths()["purelib"]
+ nvidia_root = os.path.join(site_packages, "nvidia")
+ include_paths = glob.glob(os.path.join(nvidia_root, "*", "include"))
+ nccl_libs = glob.glob(os.path.join(nvidia_root, "nccl", "lib", "libnccl.so*"))
+ if not nccl_libs:
+ raise RuntimeError("a2a_permute: could not locate nvidia-nccl libnccl.so")
+ if not any(os.path.isfile(os.path.join(path, "nccl.h")) for path in include_paths):
+ raise RuntimeError("a2a_permute: could not locate nvidia-nccl nccl.h")
+ load(
+ name="vllm_omni_a2a_permute",
+ sources=[src],
+ # PyTorch wheel-only installations keep headers such as cusparse.h
+ # under nvidia/{cu13,cusparse}/include rather than CUDA_HOME.
+ extra_include_paths=list(include_paths),
+ extra_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3"],
+ extra_cuda_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3", "--expt-relaxed-constexpr"],
+ extra_ldflags=[nccl_libs[0]],
+ is_python_module=False,
+ verbose=False,
+ )
+ symm_mem.set_backend("NCCL")
+ logger.info("[a2a_permute] JIT kernel built and loaded; symm-mem backend=NCCL")
+ _BUILT = True
+
+
+def ensure_a2a_permute_available() -> None:
+ """Build the extension during worker/model initialization, not a request."""
+ _ensure_built()
+
+
+def _required_nbytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
+ return math.prod(shape) * torch.empty((), dtype=dtype).element_size()
+
+
+def _get_symm_buffer(
+ symm_shape: tuple[int, ...], dtype: torch.dtype, device: torch.device, group_name: str
+) -> torch.Tensor:
+ """Get a typed view from a grow-only symmetric-memory workspace.
+
+ Allocation/growth and rendezvous are collective. Every rank reaches them in
+ lockstep because sequence parallel execution is symmetric. Steady-state
+ requests at or below the peak capacity take no synchronization path.
+ """
+ device = torch.device(device)
+ key = (device, group_name)
+ required_bytes = _required_nbytes(symm_shape, dtype)
+ stream_id = torch.cuda.current_stream(device).cuda_stream
+ with _WORKSPACE_LOCK:
+ workspace = _SYMM_WORKSPACES.get(key)
+ if workspace is not None and workspace.stream_id != stream_id:
+ raise RuntimeError("a2a_permute workspace must be used from a single CUDA stream")
+ if workspace is None or workspace.capacity_bytes < required_bytes:
+ if torch.cuda.is_current_stream_capturing():
+ capacity = 0 if workspace is None else workspace.capacity_bytes
+ raise RuntimeError(
+ "a2a_permute: cannot grow the symmetric-memory workspace "
+ f"from {capacity} to {required_bytes} bytes during CUDA graph capture; "
+ "warm up the maximum request shape before capture"
+ )
+ if workspace is None:
+ # NCCL symmetric-memory rendezvous requires the communicator to
+ # exist before the first allocation.
+ pg = _resolve_process_group(group_name)
+ warm = torch.ones(1, device=device)
+ dist.all_reduce(warm, group=pg)
+ # Growth is rare and must not release the previous allocation while
+ # an earlier kernel is still consuming it.
+ torch.accelerator.synchronize(device)
+ allocation = symm_mem.empty(required_bytes, dtype=torch.uint8, device=device)
+ handle = symm_mem.rendezvous(allocation, group_name)
+ workspace = _SymmWorkspace(
+ allocation=allocation,
+ handle=handle,
+ capacity_bytes=required_bytes,
+ stream_id=stream_id,
+ )
+ _SYMM_WORKSPACES[key] = workspace
+ return workspace.allocation[:required_bytes].view(dtype).view(symm_shape)
+
+
+def clear_a2a_permute_workspaces() -> None:
+ """Release cached symmetric-memory workspaces during worker shutdown."""
+ with _WORKSPACE_LOCK:
+ if not _SYMM_WORKSPACES:
+ return
+ devices = {workspace.allocation.device for workspace in _SYMM_WORKSPACES.values()}
+ for device in devices:
+ torch.accelerator.synchronize(device)
+ count = len(_SYMM_WORKSPACES)
+ _SYMM_WORKSPACES.clear()
+ logger.info("[a2a_permute] Released %d symmetric-memory workspace(s)", count)
+
+
+# ---------------------------------------------------------------------------
+# Forward: (B, S_local, H, D) -> (B, S_global, H_local, D)
+# scatter heads (dim 2), gather sequence (dim 1). Matches all_to_all_4D(x,2,1).
+# ---------------------------------------------------------------------------
+@torch.library.custom_op("vllm_omni_a2a::ulysses_qkv_fwd", mutates_args=(), device_types="cuda")
+def ulysses_qkv_fwd(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+ _ensure_built()
+ p = world_size
+ B, s_local, H, D = x.shape
+ Hl = H // p
+ lc = Hl * D
+ rows = B * s_local
+ symm_in = _get_symm_buffer((rows, p, lc), x.dtype, x.device, group_name)
+ # (B, S_local, H, D) -> (rows, p, lc); H is row-major so column block r = heads [r*Hl:(r+1)*Hl]
+ # Fused QKV projections produce row-strided views. Stage complete rows with
+ # the copy engine instead of materializing the view with TensorIterator.
+ torch.ops.a2ap.copy_rows(x.view(rows, p, lc), symm_in)
+ out = torch.empty(p, rows, lc, device=x.device, dtype=x.dtype)
+ torch.ops.a2ap.all_to_all_permute(symm_in, out, 1, 0, group_name)
+ # (p, rows, lc) -> (B, S_global, H_local, D), sequence ordered rank-major
+ return out.reshape(p, B, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(B, p * s_local, Hl, D).contiguous()
+
+
+@ulysses_qkv_fwd.register_fake
+def _(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+ B, s_local, H, D = x.shape
+ return x.new_empty(B, s_local * world_size, H // world_size, D)
+
+
+# ---------------------------------------------------------------------------
+# Reverse: (B, S_global, H_local, D) -> (B, S_local, H, D)
+# scatter sequence (dim 1), gather heads (dim 2). Matches all_to_all_4D(y,1,2).
+# ---------------------------------------------------------------------------
+@torch.library.custom_op("vllm_omni_a2a::ulysses_o_rev", mutates_args=(), device_types="cuda")
+def ulysses_o_rev(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+ _ensure_built()
+ p = world_size
+ B, s_global, Hl, D = y.shape
+ s_local = s_global // p
+ H = Hl * p
+ cols = Hl * D
+ rows = B * s_local
+ symm_in = _get_symm_buffer((p, rows, cols), y.dtype, y.device, group_name)
+ # (B, p*S_local, Hl, D) -> (p, B*S_local, cols): block r = sequence shard destined to rank r
+ symm_in.copy_(y.reshape(B, p, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(p, rows, cols))
+ out = torch.empty(rows, p, cols, device=y.device, dtype=y.dtype)
+ torch.ops.a2ap.all_to_all_permute(symm_in, out, 0, 1, group_name)
+ # (rows, p, cols) -> (B, S_local, H, D), heads ordered rank-major
+ return out.reshape(B, s_local, p, Hl, D).reshape(B, s_local, H, D).contiguous()
+
+
+@ulysses_o_rev.register_fake
+def _(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+ B, s_global, Hl, D = y.shape
+ return y.new_empty(B, s_global // world_size, Hl * world_size, D)
diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
index d804740f9..e8ae6bf32 100644
--- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
+++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
@@ -9,7 +9,7 @@ import torch
import torch.distributed as dist
from vllm.logger import init_logger
-from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
logger = init_logger(__name__)
@@ -45,7 +45,8 @@ class DistributedVaeExecutor:
"""
def __init__(self):
- self.group = get_dit_group()
+ # Use a dedicated process group spanning the complete worker WORLD.
+ self.group = get_world_group().device_group
self.world_size = dist.get_world_size(self.group)
self.rank = dist.get_rank(self.group)
self.parallel_size = 1
@@ -189,7 +190,7 @@ class DistributedVaeMixin:
if self.distributed_executor.parallel_size > pp_size:
logger.warning(
f"vae_patch_parallel_size={self.distributed_executor.parallel_size} "
- f"is greater than dit_group={world_size};"
- f" using dit_group size={world_size}"
+ f"is greater than WORLD={world_size};"
+ f" using WORLD size={world_size}"
)
return True
diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py
index 55130e0cf..4c20f1587 100644
--- a/vllm_omni/diffusion/distributed/cfg_parallel.py
+++ b/vllm_omni/diffusion/distributed/cfg_parallel.py
@@ -235,8 +235,8 @@ class CFGParallelMixin(metaclass=ABCMeta):
Predict noise with N-branch CFG dispatch across M GPUs.
This is the multi-branch counterpart of predict_noise_maybe_with_cfg().
- Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel,
- DreamID). Existing 2-branch models should continue using
+ Use this for models with 3 or more CFG branches (e.g., OmniGen2, Bagel).
+ Existing 2-branch models should continue using
predict_noise_maybe_with_cfg().
Args:
diff --git a/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu
new file mode 100644
index 000000000..1c218a1cf
--- /dev/null
+++ b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu
@@ -0,0 +1,268 @@
+// JIT build of pytorch/pytorch#178230 all_to_all_permute (Ulysses-style fused
+// permute-free all-to-all over NCCL symmetric memory).
+//
+// Kernel is verbatim from the PR; the host entry point is adapted to the
+// NCCLDevCommManager API shipped in torch 2.11 (we obtain the host ncclComm_t
+// via the manager and create our own ncclDevComm with enough LSA barriers,
+// cached per group — the PR's keyed-devcomm manager API is not in 2.11).
+
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+#include
+
+#include |