[Model] Optimize MiniMax-H3 DLO component lifecycle

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/cuda/test-ready.yml

@@ -125,32 +125,8 @@ steps:
           - vllm_omni/outputs/
         commands:
           - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              # vLLM 0.27 added a pre-flight capacity check in
-              # distributed/device_communicators/shm_broadcast.py. Without this
-              # the container gets Docker's 64 MiB default and every test in
-              # this step dies at engine start with
-              #   RuntimeError: Insufficient space in /dev/shm:
-              #                 240 MiB required, 64 MiB available
-              # (9 failed in 8.7s on build #2949). vLLM 0.26 had no such check
-              # and allocated regardless, so this step passed on main while
-              # being under-provisioned the whole time — 0.27 turned a silent
-              # risk into an explicit error.
-              #
-              # This is the only inline docker plugin block in .buildkite/cuda/;
-              # every other step inherits common/ci_mirror_hardwares.yml, which
-              # already sets 8gb. Matching that rather than inventing a value.
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
+        mirror_hardwares: l4_1
+
 
       - label: "RLHF · VeRL-Omni E2E Test"
         source_file_dependencies:

.claude/skills/add-tts-model/SKILL.md

@@ -345,8 +345,8 @@ Heavier scenarios in the same file use **`advanced_model` only** (streaming, ext
 @pytest.mark.tts
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None:
-    openai_client.send_audio_speech_request({...})
+def test_voice_clone_en_non_streaming_001(omni_server, online_client) -> None:
+    online_client.send_audio_speech_request({...})
 ```
 
 **L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)).
@@ -363,8 +363,8 @@ def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None:
 
 | Path | Fixture | Call |
 |------|---------|------|
-| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` |
-| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` |
+| Online `/v1/*` | `online_client` | `online_client.send_*_request(request_config)` |
+| Offline inference | `offline_client` | `offline_client.send_*_request(request_config)` |
 
 1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc.
 2. **No matching helper** → add `send_<feature>_request` (or `send_<route>_http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test.
@@ -396,7 +396,7 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers
   `/health` return non-200 until you're actually ready.
 - **Mark tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model`
 - **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file
-- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests
+- **Online and offline e2e go through `runtime.py`** — `online_client.send_audio_speech_request` (online); `offline_client.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests
 
 ## Phase 4: Async Chunk (Streaming)
 

.claude/skills/vllm-omni-test/SKILL.md

@@ -94,8 +94,8 @@ Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-styl
 
 | Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level |
 |----------|------------------|---------------------------|---------------------------|
-| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
-| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
+| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `offline_client`. **Function (isolation only):** `omni_runner_function` + `offline_client_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
+| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `online_client`. **Function (isolation only):** `omni_server_function` + `online_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
 | **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/cuda/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). |
 | **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: JSON `mark` with `full_model` + `omni`/`tts`/`diffusion`; wire `test-nightly.yml` Perf steps |
 | **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) |
@@ -110,9 +110,9 @@ After choosing offline/online/docs/perf, classify the **product under test** and
 |-----------|------------------------------|----------------------------|----------------------------------------|
 | **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) |
 | **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo |
-| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` with deploy YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` |
-| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) |
-| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` |
+| **Offline runner** | `omni_runner` + `offline_client`, `generate_multimodal` | `Omni(...).generate` with deploy YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` |
+| **Online client** | `online_client.send_omni_request` (chat completions, modalities) | `online_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) |
+| **Typical assertions** | Stage outputs, text/audio keywords via `OfflineOmniClient` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` |
 | **Deploy YAML** | Per-model deploy configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT |
 | **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) |
 | **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group |
@@ -175,9 +175,9 @@ See **L1 unit test constraints (mocking)** below for the full do/don't list.
 @pytest.mark.omni
 @hardware_test(...)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_<scenario>(omni_runner, omni_runner_handler) -> None:
+def test_<scenario>(omni_runner, offline_client) -> None:
     request_config = {"prompts": ..., "modalities": [...]}  # optional: images, videos, audios
-    omni_runner_handler.send_omni_request(request_config)
+    offline_client.send_omni_request(request_config)
 ```
 
 *Offline generative e2e — **Diffusion** (representative):*
@@ -187,8 +187,8 @@ def test_<scenario>(omni_runner, omni_runner_handler) -> None:
 @pytest.mark.diffusion
 @hardware_test(...)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
-def test_text_to_image_001(omni_runner_handler) -> None:
-    omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}})
+def test_text_to_image_001(offline_client) -> None:
+    offline_client.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}})
 ```
 
 *Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):*
@@ -198,8 +198,8 @@ def test_text_to_image_001(omni_runner_handler) -> None:
 @pytest.mark.tts
 @hardware_test(...)
 @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
-def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
-    omni_runner_handler.send_audio_speech_request({
+def test_text_to_audio_001(omni_runner, offline_client) -> None:
+    offline_client.send_audio_speech_request({
         "input": "...",
         "task_type": "Base",
         "ref_audio": REF_AUDIO_URL,
@@ -214,8 +214,8 @@ def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
 @pytest.mark.tts
 @hardware_test(...)
 @pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
-def test_voice_clone_001(omni_runner_handler) -> None:
-    omni_runner_handler.send_single_stage_tts_request({
+def test_voice_clone_001(offline_client) -> None:
+    offline_client.send_single_stage_tts_request({
         "input": "...",
         "language": "en",
         "prompt_audio_path": REF_AUDIO_PATH,
@@ -231,9 +231,9 @@ def test_voice_clone_001(omni_runner_handler) -> None:
 @pytest.mark.omni
 @hardware_test(...)
 @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:
     request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]}
-    openai_client.send_omni_request(request_config)
+    online_client.send_omni_request(request_config)
 ```
 
 *Online serving e2e — **TTS** (representative):*
@@ -244,9 +244,9 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
 @pytest.mark.tts
 @hardware_test(...)
 @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:
     request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...}
-    openai_client.send_audio_speech_request(request_config)
+    online_client.send_audio_speech_request(request_config)
 ```
 
 *Online serving e2e — **Diffusion** X2I (representative):*
@@ -255,8 +255,8 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
-def test_text_to_image_001(omni_server, openai_client) -> None:
-    openai_client.send_diffusion_request({...})  # chat completions + extra_body
+def test_text_to_image_001(omni_server, online_client) -> None:
+    online_client.send_diffusion_request({...})  # chat completions + extra_body
 ```
 
 *Online serving e2e — **Diffusion** X2V (representative):*
@@ -264,8 +264,8 @@ def test_text_to_image_001(omni_server, openai_client) -> None:
 ```python
 @pytest.mark.core_model
 @pytest.mark.diffusion
-def test_text_to_video_001(omni_server, openai_client) -> None:
-    openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}})  # /v1/videos
+def test_text_to_video_001(omni_server, online_client) -> None:
+    online_client.send_video_diffusion_request({"model": ..., "form_data": {...}})  # /v1/videos
 ```
 
 *Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts).
@@ -324,21 +324,21 @@ vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offl
 
 | Scope | Online fixtures | Offline fixtures | When to use |
 |-------|-----------------|------------------|-------------|
-| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. |
-| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. |
+| **Module** (default) | `omni_server` → `online_client` | `omni_runner` → `offline_client` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. |
+| **Function** | `omni_server_function` → `online_client_function` | `omni_runner_function` → `offline_client_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. |
 
 **Rules:**
 
 1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function.
-2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa).
+2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `online_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `online_client_function`. Do not mix module fixture with function client (or vice versa).
 3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ.
 4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`).
 
 ```python
 # Default — module-scoped server (L4 expansion)
 @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True)
-def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
-    openai_client.send_images_generations_request({...})
+def test_foo_images_generations_default_1024(omni_server, online_client) -> None:
+    online_client.send_images_generations_request({...})
 
 # Function-scoped — reliability / per-test clean state only
 @pytest.mark.parametrize(
@@ -346,9 +346,9 @@ def test_foo_images_generations_default_1024(omni_server, openai_client) -> None
     [pytest.param(OmniServerParams(model=MODEL), marks=H100)],
     indirect=True,
 )
-def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None:
-    openai_client_function.send_omni_sleep_http_request({...})
-    openai_client_function.send_omni_wakeup_http_request({...})
+def test_foo_sleep_wakeup_cycle(omni_server_function, online_client_function) -> None:
+    online_client_function.send_omni_sleep_http_request({...})
+    online_client_function.send_omni_wakeup_http_request({...})
 ```
 
 #### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`)
@@ -357,19 +357,19 @@ def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) ->
 
 | Principle | Action |
 |-----------|--------|
-| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. |
+| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OnlineOmniClient` (online) or `OfflineOmniClient` (offline) that matches the endpoint / pipeline shape. |
 | **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. |
 | **Add when missing** | No suitable helper → add **`send_<feature>_request`** (high-level: call + general `assert_*`) or **`send_<route>_<verb>_http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. |
 | **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. |
 
-**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`.
+**Online** (`online_client` from `omni_server`): `OnlineOmniClient.send_*_request`.
 
-**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`.
+**Offline** (`offline_client` from `omni_runner`): `OfflineOmniClient.send_*_request`.
 
 | Do | Don't |
 |----|-------|
-| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test |
-| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` |
+| Online: `online_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test |
+| Offline: `offline_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` |
 | Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module |
 | Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint |
 
@@ -392,7 +392,7 @@ def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) ->
 
 Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`.
 
-**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):**
+**`send_*` ↔ `assert_*` pairing (online `OnlineOmniClient`):**
 
 | High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling |
 |----------------------------------------|--------------------------------------------|-----------------------------|
@@ -405,7 +405,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not**
 
 **Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation.
 
-**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding):
+**Common `OnlineOmniClient` entry points** (non-exhaustive — grep `runtime.py` before adding):
 
 | Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) |
 |------|---------------------------|-----------------------------------|
@@ -416,7 +416,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not**
 | TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … |
 | Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … |
 
-**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding):
+**Common `OfflineOmniClient` entry points** (offline — grep `runtime.py` before adding):
 
 | Area | High-level (`send_*_request` + assert) | Notes |
 |------|----------------------------------------|-------|
@@ -426,7 +426,7 @@ Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not**
 | Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests |
 | New model family | **Add** `send_<family>_request` here first | Then call from `tests/e2e/offline_inference/test_*.py` |
 
-L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers.
+L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OnlineOmniClient` / `OfflineOmniClient`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers.
 
 #### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`)
 
@@ -452,7 +452,7 @@ When the user asks for **invalid parameter validation**, **invalid request bodie
 | **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.<omni or tts or diffusion>]` |
 | **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` |
 | **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) |
-| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) |
+| **HTTP client** | **Low-level** `online_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) |
 | **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test_<route>_malformed_json` when not parametrized |
 | **Body helpers** | `_minimal_<endpoint>_json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` |
 | **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) |
@@ -485,13 +485,13 @@ When the user asks for **invalid parameter validation**, **invalid request bodie
 @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: dict[str, object],
     err_message: str | tuple[str, ...],
 ) -> None:
     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}
     )
 ```
@@ -542,18 +542,18 @@ def send_images_generations_request(self, request_config: dict[str, Any], ...) -
 
 # tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching)
 @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
-def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
-    openai_client.send_images_generations_request({"json": body, "timeout": 300})
+def test_foo_images_generations_default_1024(omni_server, online_client) -> None:
+    online_client.send_images_generations_request({"json": body, "timeout": 300})
 
 @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
-def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None:
+def test_foo_images_generations_sizes_256_512_1024(omni_server, online_client) -> None:
     for size in ("256x256", "512x512", "1024x1024"):
-        openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300})
+        online_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300})
 
 # chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response
 @pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
-def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None:
-    openai_client.send_diffusion_request({...})
+def test_foo_chat_completions_t2i_fallback(omni_server, online_client) -> None:
+    online_client.send_diffusion_request({...})
 ```
 
 L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e).

.claude/skills/vllm-omni-test/references/test-routing.md

@@ -279,7 +279,7 @@ When adding or modifying tests, do not stop at “where the file lives” — al
 4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e.
 5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`.
 6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios.
-7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`).
+7. **Fixture scope**: default **`omni_server` + `online_client`** / **`omni_runner` + `offline_client`** (module). Use **`omni_server_function` + `online_client_function`** / **`omni_runner_function` + `offline_client_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`).
 8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module.
 9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3.
 10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test_<slug>_vllm_omni.json` + **Perf Test · &lt;Model&gt;** step (artifact upload); not part of “L4 functional cases” by default.

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -1,8 +1,8 @@
-**Scope**
+#### Scope
 
 For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
 
-**Test Case Design**
+#### Test Case Design
 
 For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows:
 
@@ -34,12 +34,12 @@ For a *normal priority* model, further reduce the number of test cases.
 
 Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`.
 
-**Code Style**
+#### Code Style
 
-- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
+- Validation: test that the multimodal output files of your model have the correct shapes. `OnlineOmniClient.send_diffusion_request` should have taken care of this.
 - Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/#markers-for-tests).
 - To maximize code reuse, you may refer to
-    - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
+    - `tests/conftest.py` for `omni_server` (running server in subprocess) and `online_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
     - `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
     - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
 - Doc: add a concise docstring for each test function.

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -20,7 +20,7 @@ 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. |
@@ -169,10 +169,192 @@ This mode means:
   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 implementation. The representation-independent identity,
+> producer, and restoration contracts landed in
+> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer
+> wires the final-layout BF16 path into eligible no-AllGather DLO.
+
+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 <node-local-root>`. Registration of shared mappings
+and direct asynchronous H2D remain a separate PR3 transport change.
+
+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 today). It supports ordinary-loader final layouts for TP1
+and TP2 rank identities plus SP layout identities; 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.
+
+#### 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 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. The backend uses the existing two bounded rank-local staging
+slots for the immutable HWR tensors, and its transactional `enable()` cleanup
+removes partial hooks and closes the lease before reporting a setup failure.
+Preferred mode then uses the runner's fresh canonical retry; required mode
+propagates the failure.
+
+#### PR2 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.
+- 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. |
@@ -185,7 +367,8 @@ This mode means:
   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.
+  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.
@@ -211,10 +394,11 @@ reconstructs those tensors with their recorded layouts. Other online methods
 must use `--dlo-no-use-allgather` or disable online quantization until their
 runtime layouts are validated.
 
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+The Host Weight Runtime representation and publication contracts are merged;
+see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and
+[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned
+no-AllGather consumer above is the PR2 implementation; registration and direct
+asynchronous H2D remain a separate transport follow-up.
 
 ## Validation coverage
 
@@ -231,6 +415,8 @@ Current source-level validation includes:
   transforms without parameter-side flags;
 - 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.
 
@@ -243,7 +429,7 @@ 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 |
@@ -275,7 +461,7 @@ 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%)** |

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,

docs/models/supported_models.md

@@ -52,6 +52,8 @@ th {
 | `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) |

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) |
 

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -63,6 +63,8 @@ omni = Omni(
 | `--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 |
 
 ## Host-weight loading
 
@@ -104,6 +106,88 @@ 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:
+
+```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, then DLO streams them
+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.
+
+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:
@@ -135,7 +219,9 @@ must enter each collective.
 - 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,
-  but it provides no shared-mmap host-memory benefit or guarantee.
+  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
@@ -147,10 +233,10 @@ must enter each collective.
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
+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 direct-checkpoint path.
+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.

docs/user_guide/diffusion_features.md

@@ -162,6 +162,7 @@ The following tables show which models support each feature:
 | **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**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -182,8 +183,8 @@ The following tables show which models support each feature:
 
 **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, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni, SANA-WM
+- **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
 

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -4,8 +4,8 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
 
 
 This example demonstrates how to generate videos from images using Wan2.2,
-LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
-offline inference API.
+LTX-2, HunyuanVideo-1.5, and SANA-Video Image-to-Video models with
+vLLM-Omni's offline inference API.
 
 ## Local CLI Usage
 
@@ -66,9 +66,49 @@ 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
+
+SANA checkpoints identify their upstream T2V pipeline in `model_index.json`.
+Select the native I2V implementation explicitly:
+
+```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
+```
+
+For the 720p checkpoint, use
+`Efficient-Large-Model/SANA-Video_2B_720p_diffusers` with
+`--height 704 --width 1280`. The native I2V path supports both checkpoint
+variants: the 480p checkpoint uses the Wan VAE, while the 720p checkpoint uses
+the LTX-2 Video VAE.
+
+For SANA-Video, 81 frames at 16 FPS is the standard checkpoint request
+(approximately five seconds), not minute-scale long-video generation.
+Minute-scale SANA generation requires the separate LongSANA/LongLive
+autoregressive workflow, which this pipeline does not implement.
+
+See the [SANA-Video recipe](../../../../recipes/NVIDIA/SANA-Video-2B.md) for
+online serving, backend boundaries, and the tested hardware profile.
+
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
+- `--model`: Model ID or local path (for example Wan I2V/TI2V, LTX-2, or
+  SANA-Video).
+- `--model-class-name`: Explicit pipeline override. SANA-Video I2V requires
+  `SanaImageToVideoPipeline`.
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
 - `--height/--width`: Output resolution (auto-calculated from image if not set).

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 <path>`.
+
 ## Key Arguments
 
 ### Common

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.
@@ -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

examples/offline_inference/image_to_video/image_to_video.py

@@ -3,13 +3,14 @@
 
 """
 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." \
@@ -107,13 +115,16 @@ 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 +388,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 +446,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
 

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 <path>`.
+
 ### Cosmos3
 
 ```bash

examples/offline_inference/text_to_video/text_to_video.py

@@ -114,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:

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

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}"

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"

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"

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
 

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"

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"

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"

pyproject.toml

@@ -73,6 +73,8 @@ dev = [
     "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)
 ]
 
@@ -266,6 +268,7 @@ markers = [
     # 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",
+    "H800: [hardware-resource] Tests that require H800 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",

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: <https://github.com/NVlabs/Sana>
+- 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: <https://huggingface.co/docs/diffusers/api/pipelines/sana_video>
+- 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.

recipes/README.md

@@ -57,6 +57,8 @@ recipes/
 | [`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) |

tests/config/test_config_factory.py

@@ -2753,22 +2753,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")

tests/config/test_omni_config.py

@@ -244,6 +244,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"),
     [
@@ -443,10 +473,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
@@ -567,6 +599,7 @@ def test_sub_config_fields_match_structured_scopes():
         # legacy engine-args path.
         "model_subdir",
         "tokenizer_subdir",
+        "requires_full_payload_input",
     }
     vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
     assert issubclass(OmniStageLoadConfig, VllmLoadConfig)

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,12 +22,11 @@ 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.media",
     "tests.helpers.fixtures.run_args",
     "tests.helpers.fixtures.runtime",
-    "tests.helpers.fixtures.speaker_cache",
     "tests.helpers.fixtures.xdist",
 )
 
@@ -32,46 +34,3 @@ pytest_plugins = (
 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},
-}

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."""
 
@@ -64,6 +64,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}
@@ -75,11 +78,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()
@@ -105,6 +111,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)
 
@@ -160,6 +167,41 @@ 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.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,
+        confirmed_num_computed_tokens=26,
+    )
+
+
 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.
 

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -4,10 +4,8 @@ from types import SimpleNamespace
 import pytest
 from vllm.v1.engine import FinishReason
 
-from vllm_omni.core.sched.omni_scheduler_mixin import (
-    DEFAULT_INPUT_WAIT_TIMEOUT_S,
-    OmniSchedulerMixin,
-)
+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
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -17,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()
@@ -58,7 +127,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
         ("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", DEFAULT_INPUT_WAIT_TIMEOUT_S),
+        ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
         ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
         ("restore-full", scheduler.waiting),

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."""
 

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(

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]
 
@@ -95,20 +97,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,

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,
@@ -242,13 +243,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 +262,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}
     )
 
@@ -328,12 +329,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 +360,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 +376,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,
@@ -618,14 +619,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 +639,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 +671,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 +688,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 +894,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 +910,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 +963,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,
@@ -1021,7 +1022,7 @@ _MINIMAX_MUSIC3_MAX_FRAMES = 9000
 @pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
 def test_minimax_music3_speech_invalid_field_values(
     omni_server: OmniServer,
-    openai_client: OpenAIClientHandler,
+    online_client: OnlineOmniClient,
     overrides: dict[str, object],
     err_message: str | tuple[str, ...],
 ) -> None:
@@ -1038,6 +1039,6 @@ def test_minimax_music3_speech_invalid_field_values(
             body.pop(key, None)
         else:
             body[key] = value
-    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}
     )

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)

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}
     )

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]
@@ -140,7 +140,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
 @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 +148,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 +174,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 +183,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 +243,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 +263,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,

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,

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,

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,

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))
@@ -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,

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,

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}")

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)
 

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)
 

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)
 

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)
 

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)
 

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": {

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"]

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"
         },

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"
         },

tests/dfx/stability/tests/test_voxcpm2.json

@@ -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"
-            }
-        ]
-    }
-]

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}
+                ]
+            }
+        ]
     }
 ]

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,
         },
     )
 

tests/diffusion/cache/test_cache_dit.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 tests for CacheDiT enablement.
@@ -13,6 +13,7 @@ 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
@@ -61,6 +62,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
 
@@ -110,7 +123,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)
 

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
@@ -10,7 +10,7 @@ from diffusers import DiffusionPipeline  # pyright: ignore[reportPrivateImportUs
 from PIL import Image
 
 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,
@@ -651,7 +651,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 +667,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)

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.
 
@@ -18,8 +18,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

tests/diffusion/model_loader/test_diffusers_loader.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
 
 """
 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
@@ -23,6 +25,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import (
     TensorBinding,
 )
 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]
@@ -73,6 +76,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,
@@ -92,6 +162,81 @@ 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"
+
+    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()
+
+
+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 _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
@@ -103,6 +248,64 @@ def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
     )
 
 
+@pytest.mark.parametrize(
+    ("dist_offload", "use_allgather"),
+    [
+        (False, False),
+        (True, True),
+    ],
+)
+def test_hwr_disabled_for_noneligible_dlo_paths_without_store_interaction(
+    monkeypatch,
+    tmp_path,
+    dist_offload,
+    use_allgather,
+):
+    """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))
+    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

tests/diffusion/model_loader/test_final_layout_host_weights.py

@@ -726,6 +726,36 @@ def test_identity_uses_resolved_revision_and_exact_semantics(tmp_path: Path) ->
     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"),
     [

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -774,6 +774,12 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
     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():

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)

tests/diffusion/models/minimax_h3/test_minimax_h3_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
 """MiniMax H3 step-wise execution (continuous batching) contract tests.
 
 These run on CPU against a stand-in DiT, so they cover the packing, the
@@ -31,6 +31,8 @@ class _SegmentMeanModel:
         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"]
 
@@ -158,9 +160,9 @@ def test_batched_step_execution_matches_independent_requests():
 
     model = _SegmentMeanModel()
     specs = [
-        dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
         # 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))]
@@ -178,7 +180,7 @@ def test_batched_step_execution_matches_independent_requests():
     states = []
     for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
         branch, video_rows, audio_rows = _make_branch(**spec)
-        if index == 1:
+        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))
 

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]))

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)

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 "

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -1,36 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+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": 1500, "rocm": 1500},
 }
 
-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,7 +60,7 @@ 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)
@@ -140,6 +145,7 @@ 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

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -38,6 +38,7 @@ 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.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -695,6 +696,67 @@ 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):
+        self.closed = False
+        self.provenance = SimpleNamespace(resolution_id=resolution_id)
+
+    def close(self):
+        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
+
+
 class _MultiBlockModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
@@ -838,6 +900,47 @@ class TestMmapWeightLoading:
         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_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):
@@ -1504,6 +1607,48 @@ class TestMmapValidation:
 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
 
@@ -1538,6 +1683,53 @@ class TestConfigValidation:
                 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

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()

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

tests/diffusion/test_diffusion_model_runner.py

@@ -825,9 +825,6 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
             del kwargs
             return SimpleNamespace(transformer=torch.nn.Identity())
 
-        def take_host_weight_plan(self):
-            return None
-
     class _DummyMemoryProfiler:
         consumed_memory = 0
 
@@ -869,8 +866,8 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
     monkeypatch.setattr(
         model_runner_module,
-        "get_offload_backend",
-        lambda od_config, device, host_weight_plan: None,
+        "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
@@ -883,8 +880,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

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
@@ -879,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"):

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:
@@ -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"):

tests/distributed/omni_connectors/test_chunk_transfer_adapter.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 threading
 from collections import deque
@@ -9,9 +9,10 @@ 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
@@ -337,6 +338,121 @@ 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.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,
+    )
+
+    assert len(adapter._pending_save_reqs) == 1
+    boundary_task = adapter._pending_save_reqs.popleft()
+    assert "ext-stream" not in adapter.requests_num_chunks_sent
+
+    request.num_computed_tokens = 3
+    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])
+    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.

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()

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,
     )

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}"

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()

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

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})
 
 

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
+<output-root>/<variant>/<task>/
+```
+
+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/<variant>/<task>/
+```
+
+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.

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()

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,
+    )

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

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()

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

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
 
@@ -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

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,7 +15,7 @@ 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
@@ -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(

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,7 +16,7 @@ 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
 
@@ -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

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,7 +12,7 @@ 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.slow]
 
@@ -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)

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,
     )

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()
     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()

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)

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,7 +11,7 @@ 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
 
@@ -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)

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

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)

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)