fix pre commit

Signed-off-by: princepride <wangzhipeng628@gmail.com>

.buildkite/test-nightly.yml

@@ -459,52 +459,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: VoxCPM2 · Perf Test"
-        key: nightly-voxcpm2-performance
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -670,7 +624,7 @@ steps:
                   - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                     resources:
                       limits:
-                        nvidia.com/gpu: 2
+                        nvidia.com/gpu: 3
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -889,8 +843,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
         timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -898,7 +852,7 @@ steps:
           - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -1033,6 +987,60 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
+        key: nightly-hunyuan-image3-performance
+        soft_fail: true
+        timeout_in_minutes: 120
+        if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.3.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
+            EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
+            EXIT2=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
+            EXIT3=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$((EXIT1 | EXIT2 | EXIT3))
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
   # Diffusion x2v only (Wan, HunyuanVideo, …). x2i/x2a/x2t live in the X2I group above, not here.
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
@@ -1280,7 +1288,7 @@ steps:
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
       - nightly-diffusion-x2iat-performance-qwen-image-layered
       - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
@@ -1297,7 +1305,7 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance

.gitignore

@@ -270,6 +270,3 @@ vllm_omni/_version.py
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/
-
-# Local Lance precision-matching dump scripts (not part of the repo)
-lance_compare/

benchmarks/diffusion/bench_sage_attn3_matrix.sh

@@ -1,108 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
-# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
-#
-#     bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
-#
-# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
-
-set -u
-
-STEPS=${STEPS:-50}
-SEED=${SEED:-42}
-PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
-OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
-mkdir -p "$OUT_DIR"
-
-declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
-
-run() {
-    local name="$1"; shift
-    local model="$1"; shift
-    local h="$1"; shift
-    local w="$1"; shift
-    local f="$1"; shift
-    local g="$1"; shift
-    local log="$OUT_DIR/${name}.log"
-    echo "[$(date +%T)] Running $name ..."
-    env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
-        --model "$model" \
-        --prompt "$PROMPT" \
-        --height "$h" --width "$w" --num-frames "$f" \
-        --num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
-        --enable-diffusion-pipeline-profiler \
-        --output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
-            echo "  FAILED — see $log"
-            TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
-            TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
-            return
-        }
-
-    local wait_ms
-    wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
-    local total_s
-    total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
-    local text_enc_s
-    text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local vae_s
-    vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local pipeline_s
-    pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
-    local transformer_s
-    transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
-
-    TOTAL[$name]="${total_s:-0}"
-    PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
-    TEXT_ENC[$name]="$text_enc_s"
-    TRANSFORMER[$name]="$transformer_s"
-    VAE[$name]="$vae_s"
-
-    echo "  total=${TOTAL[$name]}s  per-step=${PER_STEP[$name]}s/it  transformer=${TRANSFORMER[$name]}s  vae=${VAE[$name]}s"
-}
-
-# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
-HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
-run hv15_sdpa_eager     "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
-run hv15_sdpa_compiled  "$HV_MODEL" 480 832 33 6.0
-run hv15_sage3_eager    "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
-WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
-run wan22_sdpa_eager     "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
-run wan22_sdpa_compiled  "$WAN_MODEL" 704 1280 33 5.0
-run wan22_sage3_eager    "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-print_table() {
-    local title="$1"; shift
-    echo
-    echo "$title (${STEPS} steps, seed=${SEED})"
-    echo
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
-    for row in "$@"; do
-        case "$row" in
-            *_sdpa_eager)     label="SDPA + Eager" ;;
-            *_sdpa_compiled)  label="SDPA + Compiled" ;;
-            *_sage3_eager)    label="SAGE_ATTN_3 + Eager" ;;
-            *_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
-        esac
-        printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-            "$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
-    done
-}
-
-print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
-    hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
-
-print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
-    wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
-
-echo
-echo "Logs: $OUT_DIR/*.log"

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -79,6 +79,7 @@ Usage:
 import argparse
 import ast
 import asyncio
+import base64
 import glob
 import json
 import logging
@@ -1013,6 +1014,73 @@ def calculate_metrics(
     return metrics
 
 
+def _save_generated_outputs(
+    outputs: list[RequestFuncOutput],
+    requests_list: list[RequestFuncInput],
+    save_dir: str,
+) -> None:
+    """Decode and save base64 images/videos from successful responses."""
+    os.makedirs(save_dir, exist_ok=True)
+    saved = 0
+    failed = 0
+
+    for idx, (req, out) in enumerate(zip(requests_list, outputs)):
+        if not out.success or not out.response_body:
+            continue
+
+        media_urls: list[str] = []
+
+        # Chat-completions style: choices[*].message.content[*].image_url.url
+        choices = out.response_body.get("choices", [])
+        if isinstance(choices, list):
+            for choice in choices:
+                content = (choice or {}).get("message", {}).get("content")
+                if not isinstance(content, list):
+                    continue
+                for item in content:
+                    if not isinstance(item, dict) or item.get("type") != "image_url":
+                        continue
+                    url = (item.get("image_url") or {}).get("url", "")
+                    if isinstance(url, str) and url.startswith("data:"):
+                        media_urls.append(url)
+
+        # Images endpoint style: data[*].b64_json
+        data_items = out.response_body.get("data", [])
+        if isinstance(data_items, list):
+            for data_item in data_items:
+                if not isinstance(data_item, dict):
+                    continue
+                b64_json = data_item.get("b64_json", "")
+                if isinstance(b64_json, str) and b64_json:
+                    media_urls.append(f"data:image/png;base64,{b64_json}")
+
+        for img_idx, url in enumerate(media_urls):
+            if "," not in url:
+                continue
+
+            try:
+                header, b64_data = url.split(",", 1)
+                ext = "png"
+                if "image/jpeg" in header:
+                    ext = "jpg"
+                elif "image/webp" in header:
+                    ext = "webp"
+                elif "video/mp4" in header:
+                    ext = "mp4"
+
+                img_bytes = base64.b64decode(b64_data)
+                fname = f"req_{idx:04d}_{img_idx}.{ext}"
+                fpath = os.path.join(save_dir, fname)
+                with open(fpath, "wb") as f:
+                    f.write(img_bytes)
+                saved += 1
+            except Exception as e:
+                failed += 1
+                logger.warning(f"Failed to save image for request {idx}: {e}", exc_info=True)
+
+    logger.info(f"Saved {saved} generated image(s) to {save_dir}. Failed to save {failed} image(s).")
+
+
 def wait_for_service(base_url: str, timeout: int = 120) -> None:
     print(f"Waiting for service at {base_url}...")
     start_time = time.time()
@@ -1202,6 +1270,9 @@ async def benchmark(args):
 
     print("\n" + "=" * 60)
 
+    if args.save_dir:
+        _save_generated_outputs(outputs, requests_list, args.save_dir)
+
     if args.output_file:
         with open(args.output_file, "w") as f:
             json.dump(metrics, f, indent=2)
@@ -1278,11 +1349,18 @@ if __name__ == "__main__":
         default=1,
         help="Number of warmup requests to run before measurement.",
     )
+    # NOTE Changed default from 1 to 2 because some models (e.g., Bagel) run
+    # `num_timesteps - 1` denoising iterations. A default of 1 results in 0 steps,
+    # which causes errors.
+    # TODO If this slightly longer warmup causes regression issues for other
+    # diffusion pipelines in the future, consider implementing model-specific
+    # overrides instead of a global default.
     parser.add_argument(
         "--warmup-num-inference-steps",
         type=int,
-        default=1,
-        help="num_inference_steps used for warmup requests.",
+        default=2,
+        help="Number of inference steps used for warmup requests. "
+        "Default is 2 to ensure at least one denoising step is executed.",
     )
     parser.add_argument(
         "--warmup-concurrency",
@@ -1323,6 +1401,13 @@ if __name__ == "__main__":
         default=3.0,
         help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
     )
+    parser.add_argument(
+        "--save-dir",
+        type=str,
+        default=None,
+        help="Directory to save generated images/outputs for visual inspection. "
+        "If not set, generated outputs are discarded after metric collection.",
+    )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
     parser.add_argument(
         "--enable-negative-prompt",

benchmarks/tts/bench_tts.py

@@ -130,9 +130,6 @@ def build_bench_args(
     if task_extra_body:
         cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
 
-    if model_cfg.get("trust_remote_code"):
-        cmd.append("--trust-remote-code")
-
     if concurrency is not None:
         cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
 

benchmarks/tts/model_configs.yaml

@@ -32,11 +32,8 @@ models:
         task_type: Base
 
   openbmb/VoxCPM2:
-    supported_tasks: [voice_clone, default_voice]
-    trust_remote_code: true
+    supported_tasks: [voice_clone]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
     task_extra_body:
       voice_clone: {}
-      default_voice:
-        voice: default

docs/models/supported_models.md

@@ -47,7 +47,7 @@ th {
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion/attention_backends.md

@@ -19,8 +19,6 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
-| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
-
 
 ## Configuration
 
@@ -187,29 +185,6 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
-## SageAttention3 Installation
-
-vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
-
-Build from source:
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention/sageattention3_blackwell
-python setup.py install
-```
-
-Quick check:
-
-```bash
-python -c "import sageattn3; print(sageattn3.__file__)"
-```
-
-Notes:
-
-- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
-- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
-
 ## Usage Examples
 
 ### Default (auto-route)
@@ -260,20 +235,6 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --output outputs/wan22_sage.mp4
 ```
 
-### Enable SageAttention3
-
-Example:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
-    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 480 --width 832 --num-frames 33 \
-    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
-    --tensor-parallel-size 2 \
-    --output outputs/hv15_sage3.mp4
-```
-
 ### Mixed backends across roles
 
 Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:

docs/user_guide/diffusion/mot_config.md

@@ -0,0 +1,50 @@
+This directory contains auto-tuned Triton kernel configurations for the
+MoT (Mixture-of-Tokens) GEMM and RMSNorm operators used by BAGEL and other
+MoT-architecture diffusion models.
+
+File naming convention:
+    device_name=<GPU_NAME>,dtype=<DTYPE>.json
+
+For example:
+    device_name=NVIDIA_A100-SXM4-80GB,dtype=w16a16.json
+
+Each JSON file maps (K, N) matrix shapes to a dictionary of batch sizes (M)
+and their optimal Triton tile configurations:
+
+    {
+        "3584_4608": {              // K=3584, N=4608 (QKV projection)
+            "1024": {
+                "BLOCK_SIZE_M": 64,
+                "BLOCK_SIZE_N": 128,
+                "BLOCK_SIZE_K": 64,
+                "GROUP_SIZE_M": 8,
+                "num_warps": 4,
+                "num_stages": 3
+            },
+            ...
+        },
+        "3584_3584": { ... }        // K=3584, N=3584 (output projection)
+    }
+
+Config loading order (3-tier, see ops/mot_gemm.py):
+    1. $VLLM_TUNED_CONFIG_FOLDER/<filename>  (env override)
+    2. This directory: vllm_omni/diffusion/layers/mot/configs/<filename>
+    3. Conservative default config (compiles everywhere, sub-optimal perf)
+
+If no config file matches the current device, a warning is printed with
+instructions to run the auto-tuning benchmark.
+
+To generate configs for your hardware:
+
+    python benchmarks/kernels/mot_linear_benchmarks.py \
+        --model ByteDance-Seed/BAGEL-7B-MoT \
+        --tp-size 1 --dtype w16a16 --tune \
+        --save-dir vllm_omni/diffusion/layers/mot/configs/
+
+For multi-GPU tuning (uses Ray for parallel search):
+
+    python benchmarks/kernels/mot_linear_benchmarks.py \
+        --model ByteDance-Seed/BAGEL-7B-MoT \
+        --tp-size 2 --tune
+
+See benchmarks/kernels/mot_linear_benchmarks.py for full options.

docs/user_guide/diffusion_features.md

@@ -129,7 +129,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
 | **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -147,7 +147,7 @@ The following tables show which models support each feature:
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 **Frame Interpolation Support**

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -329,16 +329,6 @@ Single-GPU serves now default to the uniproc executor (lower IPC overhead, the B
 
 To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
 
-### Tuning stage 1 `max_num_seqs` per task type
-The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
-
-CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
-
-```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
-    --stage-overrides '{"1": {"max_num_seqs": 1}}'
-```
-
 ### Sending requests
 ```bash
 # CustomVoice with a predefined speaker

examples/offline_inference/hunyuan_image3/end2end.py

@@ -61,8 +61,8 @@ def parse_args():
     parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
     parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--height", type=int, default=None, help="Output image height.")
-    parser.add_argument("--width", type=int, default=None, help="Output image width.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
+    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
     parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
     parser.add_argument(
         "--bot-task",
@@ -229,21 +229,7 @@ def main():
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    if (args.height is None) != (args.width is None):
-        raise ValueError("--height and --width must both be specified or both omitted.")
-    user_specified_size = args.height is not None and args.width is not None
-    if args.modality in ("img2text", "text2text"):
-        ar_image_size = "auto"
-    elif user_specified_size:
-        ar_image_size = f"{args.width}x{args.height}"
-    else:
-        ar_image_size = None
-    ar_stop_token_ids = resolve_stop_token_ids(
-        task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size
-    )
-    print(
-        f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}"
-    )
+    ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
     for sp in params_list:
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps

examples/offline_inference/lance/README.md

@@ -1,52 +0,0 @@
-# Lance (ByteDance) — offline inference
-
-[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
-autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
-**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
-checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
-implements it by reusing the BAGEL transformer core and specializing only the
-ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
-
-This example covers all six Lance modalities from the upstream HF model card:
-`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
-`x2t_video` (video understanding).
-
-## Hardware
-
-Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
-
-## Run
-
-```bash
-# Text-to-image
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance \
-    --prompts "a corgi astronaut on the moon, cinematic" \
-    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
-    --height 1024 --width 1024 \
-    --output ./out
-
-# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
-# choices for all six task variants)
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
-    --num-frames 25 --video-height 480 --video-width 768 \
-    --prompts "a cat playing piano, cinematic" \
-    --steps 30 --fps 8 --output ./out
-```
-
-`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
-the 3-D `latent_pos_embed` table is loaded; the other paths can point at
-the top-level `bytedance-research/Lance` repo and resolve the right
-sub-checkpoint automatically.
-
-The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
-`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
-
-## Defaults
-
-Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
-text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`).  For
-the understanding paths (`img2text` / `video2text`), sampling is enabled by
-default at `--text-temperature 0.8` because Lance's greedy decoder emits an
-immediate EOS for many prompts.

examples/offline_inference/lance/end2end.py

@@ -1,333 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance-3B offline inference (single-stage).
-
-Lance is BAGEL-lineage; this example covers the single-stage paths wired in
-this PR:
-
-    text2img  — image generation from text
-    img2text  — image understanding (caption / VQA)
-
-Examples:
-
-    python examples/offline_inference/lance/end2end.py \
-        --model bytedance-research/Lance \
-        --prompts "a corgi astronaut on the moon, cinematic" \
-        --output ./out
-
-    python examples/offline_inference/lance/end2end.py \
-        --model bytedance-research/Lance --modality img2text \
-        --image-path /path/to/photo.jpg \
-        --prompts "Describe this image in detail."
-"""
-
-import argparse
-import os
-
-
-def parse_args():
-    p = argparse.ArgumentParser()
-    p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
-    p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
-    p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
-    p.add_argument(
-        "--modality",
-        default="text2img",
-        choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
-        help="Lance single-stage modality.",
-    )
-    p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
-    p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
-    p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
-    p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
-    p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
-    p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
-    p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
-    p.add_argument(
-        "--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
-    )
-    p.add_argument("--output", type=str, default=".", help="Output directory.")
-    p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
-    p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
-    p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
-    p.add_argument("--negative-prompt", type=str, default=None)
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
-    p.add_argument(
-        "--do-sample",
-        action="store_true",
-        default=True,
-        help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
-    )
-    p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
-    p.add_argument(
-        "--text-temperature",
-        type=float,
-        default=0.8,
-        help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
-    )
-    p.add_argument(
-        "--system-prompt",
-        type=str,
-        default=None,
-        help=(
-            "Override the task system prompt (used by x2t_image / x2t_video for "
-            'per-example QA instructions like "Look at the image carefully and '
-            'answer the question."; without it x2t falls back to a '
-            "caption-style default and the model describes instead of answering)."
-        ),
-    )
-    # Lance is single-stage diffusion — no deploy YAML needed.  Required
-    # engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
-    # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
-    # and ``create_default_diffusion`` materializes the stage config.
-
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(p)
-    return p.parse_args()
-
-
-from vllm_omni.diffusion.models.lance.prompts import (  # noqa: E402
-    VIDEO_PAD,
-    VISION_END,
-    VISION_START,
-    render_lance_prompt,
-)
-
-# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
-# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
-_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
-
-
-def _format_text2img_prompts(prompts):
-    return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
-
-
-def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
-
-
-def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
-    """Image-to-Video (no first-frame pin): image + text → long video.
-
-    Passes the input image via ``multi_modal_data.first_frame``.  The pipeline
-    treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
-    multi-frame video at the requested ``num_frames`` × ``video_h`` ×
-    ``video_w`` shape.
-    """
-    import os as _os
-
-    if not image_path or not _os.path.exists(image_path):
-        raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
-    from PIL import Image as _Image
-
-    img = _Image.open(image_path).convert("RGB")
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [
-        {
-            "prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"first_frame": img},
-            "modalities": ["video"],
-            "extra_args": extra,
-        }
-        for p in prompts
-    ]
-
-
-def _format_img2img_prompts(prompts, image_path):
-    import os as _os
-
-    if not image_path or not _os.path.exists(image_path):
-        raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
-    from PIL import Image as _Image
-
-    img = _Image.open(image_path).convert("RGB")
-    return [
-        {
-            "prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"img2img": img},
-            "modalities": ["image"],
-        }
-        for p in prompts
-    ]
-
-
-def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
-    import os as _os
-
-    if not video_path or not _os.path.exists(video_path):
-        raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
-
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [
-        {
-            "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"video": video_path},
-            "modalities": ["video"],
-            "extra_args": extra,
-        }
-        for p in prompts
-    ]
-
-
-def _format_video2text_prompts(prompts, video_path, system_prompt=None):
-    import os as _os
-
-    if not video_path or not _os.path.exists(video_path):
-        raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
-    import imageio.v3 as iio
-
-    video = iio.imread(video_path)
-    return [
-        {
-            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
-            "multi_modal_data": {"video": video},
-            "modalities": ["text"],
-        }
-        for p in prompts
-    ]
-
-
-def _format_img2text_prompts(prompts, image_path, system_prompt=None):
-    from PIL import Image
-
-    if not image_path or not os.path.exists(image_path):
-        raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
-    img = Image.open(image_path).convert("RGB")
-    return [
-        {
-            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
-            "multi_modal_data": {"image": img},
-            "modalities": ["text"],
-        }
-        for p in prompts
-    ]
-
-
-def main():
-    args = parse_args()
-    os.makedirs(args.output, exist_ok=True)
-
-    if args.txt_prompts:
-        with open(args.txt_prompts, encoding="utf-8") as f:
-            prompts = [ln.strip() for ln in f if ln.strip()]
-    elif args.modality == "img2text":
-        prompts = args.prompts or ["Describe this image in detail."]
-    else:
-        prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
-
-    from vllm_omni.entrypoints.omni import Omni
-
-    # Video modalities require the Lance_3B_Video checkpoint (3D
-    # latent_pos_embed shaped (126976, 2048)).  Without it, vllm-omni silently
-    # loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
-    # immediately indexes out of bounds.  Rewrite the model path so
-    # LancePipeline._select_video_variant picks the video subfolder.
-    resolved_model = args.model
-    if args.modality in ("text2video", "video2video", "video2text", "image2video"):
-        if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
-            candidate = os.path.join(resolved_model, "Lance_3B_Video")
-            if os.path.isdir(candidate):
-                resolved_model = candidate
-
-    omni_kwargs = vars(args).copy()
-    omni_kwargs["model"] = resolved_model
-    # Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
-    omni_kwargs.setdefault("pipeline", "lance")
-    omni_kwargs.setdefault("max_num_batched_tokens", 32768)
-    omni_kwargs.setdefault("max_num_seqs", 1)
-    omni_kwargs.setdefault("enforce_eager", True)
-    omni_kwargs.setdefault("trust_remote_code", True)
-    omni_kwargs.setdefault("enable_prefix_caching", False)
-    omni_kwargs.setdefault("async_chunk", False)
-    omni = Omni(**omni_kwargs)
-
-    if args.modality == "img2text":
-        formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
-    elif args.modality == "text2video":
-        formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
-    elif args.modality == "video2text":
-        formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
-    elif args.modality == "img2img":
-        formatted = _format_img2img_prompts(prompts, args.image_path)
-    elif args.modality == "video2video":
-        formatted = _format_video2video_prompts(
-            prompts, args.video_path, args.num_frames, args.video_height, args.video_width
-        )
-    elif args.modality == "image2video":
-        formatted = _format_image2video_prompts(
-            prompts, args.image_path, args.num_frames, args.video_height, args.video_width
-        )
-    else:
-        formatted = _format_text2img_prompts(prompts)
-
-    params_list = omni.default_sampling_params_list
-    diffusion_params = params_list[0]  # single-stage: one param set
-    diffusion_params.num_inference_steps = args.steps  # type: ignore
-    if args.seed is not None:
-        diffusion_params.seed = args.seed  # type: ignore
-    if args.height is not None:
-        diffusion_params.height = args.height  # type: ignore
-    if args.width is not None:
-        diffusion_params.width = args.width  # type: ignore
-    extra = getattr(diffusion_params, "extra_args", {}) or {}
-    extra["cfg_text_scale"] = args.cfg_text_scale
-    extra["timestep_shift"] = args.timestep_shift
-    if args.modality in ("img2text", "video2text"):
-        extra["max_think_tokens"] = args.max_text_tokens
-        extra["do_sample"] = args.do_sample
-        extra["text_temperature"] = args.text_temperature
-    if args.negative_prompt is not None:
-        extra["negative_prompt"] = args.negative_prompt
-    diffusion_params.extra_args = extra  # type: ignore
-
-    outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
-
-    if args.modality in ("img2text", "video2text"):
-        for i, req_output in enumerate(outputs):
-            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
-            if text is None and hasattr(req_output, "custom_output"):
-                text = (req_output.custom_output or {}).get("text_output")
-            print(f"[Output {i}] {text}")
-        return
-
-    if args.modality in ("text2video", "video2video", "image2video"):
-        for i, req_output in enumerate(outputs):
-            custom = getattr(req_output, "custom_output", None) or {}
-            frames = custom.get("video_frames")
-            if not frames:
-                print(f"[Output {i}] no video frames returned")
-                continue
-            mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
-            try:
-                import imageio.v3 as iio
-                import numpy as np
-
-                arr = np.stack([np.asarray(f) for f in frames], axis=0)
-                iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
-                print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
-            except Exception as exc:
-                print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
-                frame_dir = os.path.join(args.output, f"lance_{i}_frames")
-                os.makedirs(frame_dir, exist_ok=True)
-                for j, frame in enumerate(frames):
-                    frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
-                print(f"  → {frame_dir} ({len(frames)} frames)")
-        return
-
-    idx = 0
-    for req_output in outputs:
-        images = getattr(req_output, "images", None)
-        if not images:
-            continue
-        for j, img in enumerate(images):
-            path = os.path.join(args.output, f"lance_{idx}_{j}.png")
-            img.save(path)
-            print(f"[Output] Saved image to {path}")
-        idx += 1
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/ming_flash_omni/README.md

@@ -8,7 +8,7 @@ vLLM-Omni supports three deployment modes:
 |------|--------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
-| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.yaml` | Image (online-serving only at the moment) |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 

examples/offline_inference/sensenova_u1/README.md

@@ -37,10 +37,6 @@ python end2end.py --modality text2text \
 # Custom resolution (image generation)
 python end2end.py --prompt "A futuristic cityscape at sunset" \
                   --width 2048 --height 1024 --think
-
-# Text-to-image with Cache-DiT acceleration
-python end2end.py --prompt "A cute cat" \
-                  --cache-backend cache_dit \
 ```
 
 > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
@@ -139,8 +135,6 @@ python end2end.py \
 | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
 | `--enforce-eager` | False | Disable torch.compile |
 | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
-| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
-| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
 
 ## Reducing GPU Memory Usage
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -155,18 +155,6 @@ def parse_args():
             "of extra CPU<->GPU transfers."
         ),
     )
-    parser.add_argument(
-        "--cache-backend",
-        type=str,
-        default=None,
-        choices=["cache_dit"],
-        help="Cache backend to use for image generation acceleration.",
-    )
-    parser.add_argument(
-        "--enable-cache-dit-summary",
-        action="store_true",
-        help="Enable cache-dit summary logging after diffusion forward passes.",
-    )
 
     nullify_stage_engine_defaults(parser)
     return parser.parse_args()
@@ -194,8 +182,6 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
-        cache_backend=args.cache_backend,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
     )
 
     extra_args = {
@@ -239,7 +225,6 @@ def main():
     print(f"  Seed           : {args.seed}")
     print(f"  Think mode     : {args.think}")
     print(f"  TP size        : {args.tensor_parallel_size}")
-    print(f"  Cache backend  : {args.cache_backend or 'none'}")
     print(f"{'=' * 60}\n")
 
     # Build prompt dict

examples/offline_inference/text_to_speech/README.md

@@ -271,14 +271,6 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --lang zh
 ```
 
-### Seed for Reproducibility
-```bash
-python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
-    --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --seed 42
-```
-
 ### Notes
 - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
 - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -80,12 +80,6 @@ def run_e2e():
         default=600,
         help="Stage initialization timeout in seconds",
     )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for generation",
-    )
     nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
 
@@ -128,7 +122,7 @@ def run_e2e():
     if mm_processor_kwargs:
         prompts["mm_processor_kwargs"] = mm_processor_kwargs
 
-    sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
+    sampling_params_list = [OmniDiffusionSamplingParams()]
 
     print(f"Generating speech for: {args.text}")
 

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -43,8 +43,8 @@ def _estimate_prompt_len(
     """
     try:
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-        from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-            Qwen3TTSPromptEmbedsBuilder,
+        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+            Qwen3TTSTalkerForConditionalGeneration,
         )
 
         if model_name not in _cache:
@@ -124,7 +124,7 @@ def _estimate_prompt_len(
             except Exception:
                 return None
 
-        return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],

examples/online_serving/lance/README.md

@@ -1,40 +0,0 @@
-# Lance — online serving
-
-OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
-
-## Start the server
-
-```bash
-bash examples/online_serving/lance/run_server.sh
-# or, with overrides
-MODEL=bytedance-research/Lance \
-DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
-PORT=8091 \
-    bash examples/online_serving/lance/run_server.sh
-```
-
-The default deploy config is the single-stage `lance.yaml`; for `video_edit`
-or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint
-config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to
-`run_server.sh`).
-
-## Send a request
-
-```bash
-# text-to-image
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "A cute corgi astronaut on the moon, cinematic" \
-    --modality text2img \
-    --output corgi.png
-
-# image edit
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "Convert this into a vibrant cartoon-style illustration" \
-    --modality img2img \
-    --image-url path/to/photo.png \
-    --output edited.png
-```
-
-The client is shared with BAGEL — same OpenAI message format, same
-``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` /
-``width`` knobs.

examples/online_serving/lance/openai_chat_client.py

@@ -1,184 +0,0 @@
-#!/usr/bin/env python3
-"""
-Lance OpenAI-compatible chat client for image generation and multimodal tasks.
-
-Usage:
-    python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png
-    python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
-"""
-
-import argparse
-import base64
-from pathlib import Path
-
-import requests
-
-
-def generate_image(
-    prompt: str,
-    server_url: str = "http://localhost:8091",
-    image_url: str | None = None,
-    height: int | None = None,
-    width: int | None = None,
-    steps: int | None = None,
-    seed: int | None = None,
-    negative_prompt: str | None = None,
-    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
-) -> bytes | str | None:
-    """Generate an image or text using the chat completions API.
-
-    Args:
-        prompt: Text description or prompt
-        server_url: Server URL
-        image_url: URL or path to input image (for img2img/img2text)
-        height: Image height in pixels
-        width: Image width in pixels
-        steps: Number of inference steps
-        seed: Random seed
-        negative_prompt: Negative prompt
-        modality: Task modality hint
-
-    Returns:
-        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
-    """
-
-    # Construct Message Content
-    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
-
-    if image_url:
-        # Check if local file
-        if Path(image_url).exists():
-            with open(image_url, "rb") as f:
-                b64_data = base64.b64encode(f.read()).decode("utf-8")
-                final_image_url = f"data:image/jpeg;base64,{b64_data}"
-        else:
-            final_image_url = image_url
-
-        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
-
-    messages = [{"role": "user", "content": content}]
-
-    # Build request payload with all parameters at top level
-    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
-    payload = {"messages": messages}
-
-    # Set output modalities at top level
-    if modality == "text2img" or modality == "img2img":
-        payload["modalities"] = ["image"]
-    elif modality == "img2text" or modality == "text2text":
-        payload["modalities"] = ["text"]
-
-    # Add generation parameters directly to payload
-    if height is not None:
-        payload["height"] = height
-    if width is not None:
-        payload["width"] = width
-    if steps is not None:
-        payload["num_inference_steps"] = steps
-    if seed is not None:
-        payload["seed"] = seed
-    if negative_prompt:
-        payload["negative_prompt"] = negative_prompt
-
-    # Send request
-    try:
-        print(f"Sending request to {server_url} with modality {modality}...")
-        response = requests.post(
-            f"{server_url}/v1/chat/completions",
-            headers={"Content-Type": "application/json"},
-            json=payload,
-            timeout=300,
-        )
-        response.raise_for_status()
-        data = response.json()
-
-        # Extract content - check ALL choices since server may return multiple
-        # (e.g., text in choices[0], image in choices[1])
-        choices = data.get("choices", [])
-
-        # First pass: look for image output in any choice
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-
-            # Handle Image Output
-            if isinstance(choice_content, list) and len(choice_content) > 0:
-                first_item = choice_content[0]
-                if isinstance(first_item, dict) and "image_url" in first_item:
-                    img_url_str = first_item["image_url"].get("url", "")
-                    if img_url_str.startswith("data:image"):
-                        _, b64_data = img_url_str.split(",", 1)
-                        return base64.b64decode(b64_data)
-
-        # Second pass: look for text output if no image found
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-            if isinstance(choice_content, str) and choice_content:
-                return choice_content
-
-        print(f"Unexpected response format: {choices}")
-        return None
-
-    except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
-    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
-    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
-    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
-
-    # Modality Control
-    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
-    parser.add_argument(
-        "--modality",
-        "-m",
-        default="text2img",
-        choices=["text2img", "img2img", "img2text", "text2text"],
-        help="Task modality",
-    )
-
-    # Generation Params
-    parser.add_argument("--height", type=int, default=512, help="Image height")
-    parser.add_argument("--width", type=int, default=512, help="Image width")
-    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed")
-    parser.add_argument("--negative", help="Negative prompt")
-
-    args = parser.parse_args()
-
-    print(f"Mode: {args.modality}")
-    if args.image_url:
-        print(f"Input Image: {args.image_url}")
-
-    result = generate_image(
-        prompt=args.prompt,
-        server_url=args.server,
-        image_url=args.image_url,
-        height=args.height,
-        width=args.width,
-        steps=args.steps,
-        seed=args.seed,
-        negative_prompt=args.negative,
-        modality=args.modality,
-    )
-
-    if result:
-        if isinstance(result, bytes):
-            # It's an image
-            output_path = Path(args.output)
-            output_path.write_bytes(result)
-            print(f"Image saved to: {output_path}")
-            print(f"Size: {len(result) / 1024:.1f} KB")
-        elif isinstance(result, str):
-            # It's text
-            print("Response:")
-            print(result)
-    else:
-        print("Failed to generate response")
-        exit(1)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/lance/run_server.sh

@@ -1,15 +0,0 @@
-#!/bin/bash
-# Lance online serving startup script
-
-MODEL="${MODEL:-bytedance-research/Lance}"
-PORT="${PORT:-8091}"
-DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}"
-
-echo "Starting Lance server..."
-echo "Model:         $MODEL"
-echo "Deploy config: $DEPLOY_CONFIG"
-echo "Port:          $PORT"
-
-vllm serve "$MODEL" --omni \
-    --deploy-config "$DEPLOY_CONFIG" \
-    --port "$PORT"

examples/online_serving/ming_flash_omni/README.md

@@ -10,7 +10,7 @@ Please refer to [README.md](../../../README.md)
 |------|---------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
-| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml` | Image |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -63,27 +63,34 @@ bash run_curl_multimodal_generation.sh use_image_gen
 ## Image generation (text-to-image)
 
 Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
-Launch with the image deploy YAML, which adds an image-gen stage behind
-the thinker:
+Launch with the dual-stage YAML, which adds an image-gen stage behind the
+thinker:
 
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
     --stage-init-timeout 1800 \
-    --init-timeout 1800 \
-    --port 8091
+    --port 8188
 ```
 
 Then request image output by passing `"modalities": ["image"]`:
 
 ```bash
-curl -s http://127.0.0.1:8091/v1/chat/completions \
+curl http://127.0.0.1:8188/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
+      "messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
       "modalities": ["image"]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
+    }' -o /tmp/ming_response.json
+
+python - <<'PY'
+import base64, json
+r = json.load(open('/tmp/ming_response.json'))
+url = r['choices'][0]['message']['content'][0]['image_url']['url']
+open('/tmp/ming_gen.png', 'wb').write(base64.b64decode(url.split(',', 1)[1]))
+print('saved /tmp/ming_gen.png')
+PY
 ```
 
 ### Optional knobs
@@ -102,7 +109,7 @@ Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_
 Example with all the knobs:
 
 ```bash
-curl http://127.0.0.1:8091/v1/chat/completions \
+curl http://127.0.0.1:8188/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "Jonathan1909/Ming-flash-omni-2.0",
@@ -115,7 +122,7 @@ curl http://127.0.0.1:8091/v1/chat/completions \
             "byte5_text":["理解与生成统一"]}}}
       ],
       "messages": [{"role": "user", "content": "Draw a poster."}]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
+    }' -o /tmp/ming_response.json
 ```
 
 ### img2img (reference image)
@@ -135,9 +142,9 @@ into the diffusion stage as `extra[reference_image]`:
 
 ### GPU layout
 
-The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
+The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
 (TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
-`devices` per stage to relocate; with fewer GPUs available, drop
+`runtime.devices` per stage to relocate; with fewer GPUs available, drop
 the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
 warmup takes roughly an extra 30–60 s on top of the thinker — set
 `--stage-init-timeout 1800` if the default 300 s is too tight.
@@ -149,7 +156,7 @@ warmup takes roughly an extra 30–60 s on top of the thinker — set
 | `["text"]` or omitted | Thinker only | Text |
 | `["audio"]` | Thinker + Talker | Audio (speech) |
 | `["text", "audio"]` | Thinker + Talker | Text + Audio |
-| `["image"]` | Thinker + Imagegen (image deploy YAML) | Image (PNG, base64 in `choices[0].message.content`) |
+| `["image"]` | Thinker + Imagegen (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
 
 For ready-to-copy curl examples (text / audio / multimodal input, SSE
 streaming, reasoning mode), see the recipe at

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -103,7 +103,7 @@ echo ""
 
 if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
     # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
-    # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
+    # The shipped ming_flash_omni_dual.yaml runs the thinker on cards 0-3
     # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
     # section in README.md for knobs and img2img.
     out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"

examples/online_serving/qwen3_omni/README.md

@@ -18,6 +18,10 @@ Asynchronous chunk streaming operates as **enabled by default** within this bund
 Additionally, NPU, ROCm, and XPU per-platform configuration deltas are deterministically merged from the
 `platforms`: section of the corresponding YAML.
 
+**Note:** The OpenAI-style **`/v1/realtime`** WebSocket interface (facilitating streaming PCM audio input alongside audio and transcription output)
+is currently **unsupported** while the `async_chunk` configuration attribute is enabled.
+It is requisite to instantiate the default omni architecture or utilize a deployment configuration specifying `async_chunk: false` to facilitate real-time streaming sessions.
+
 To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
@@ -285,7 +289,7 @@ python openai_realtime_client.py \
 | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
 | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
 
-Ensure the server is running, for example:
+Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091

examples/online_serving/sensenova_u1/README.md

@@ -6,21 +6,11 @@
 vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
 ```
 
-With Cache-DiT acceleration:
-
-```bash
-vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \
-    --cache-backend cache_dit
-```
-
 Or use the convenience script:
 
 ```bash
 cd examples/online_serving/sensenova_u1
 bash run_server.sh
-
-# Cache-DiT
-CACHE_BACKEND=cache_dit bash run_server.sh
 ```
 
 ### Tensor Parallelism (TP)

examples/online_serving/sensenova_u1/run_server.sh

@@ -3,18 +3,10 @@
 
 MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}"
 PORT="${PORT:-8091}"
-CACHE_BACKEND="${CACHE_BACKEND:-none}"
 
 echo "Starting SenseNova-U1 server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
-echo "Cache backend: $CACHE_BACKEND"
-
-EXTRA_ARGS=()
-if [ "$CACHE_BACKEND" != "none" ]; then
-    EXTRA_ARGS+=(--cache-backend "$CACHE_BACKEND")
-fi
 
 vllm serve "$MODEL" --omni \
-    --port "$PORT" \
-    "${EXTRA_ARGS[@]}"
+    --port "$PORT"

examples/online_serving/text_to_speech/README.md

@@ -18,7 +18,7 @@ For the full list of supported architectures across all modalities, see
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
-| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
+| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
 | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
@@ -321,29 +321,16 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 ### CLI client
 ```bash
 cd examples/online_serving/text_to_speech/omnivoice
-# Text-only (auto voice)
 python speech_client.py --text "Hello, how are you?"
-
-# Language hint
 python speech_client.py --text "Bonjour, comment allez-vous?" --language French
-# Voice cloning (reference audio + optional ref_text)
-python speech_client.py \
---text "Bonjour, comment allez-vous?" \
---ref-audio /path/to/ref_audio.wav \
---ref-text "Bonjour, comment allez-vous?"
-
-# Style instruction (voice design-style control)
-python speech_client.py \
---text "Bonjour, comment allez-vous?" \
---language French \
---instructions "loud voice"
-
-# Deterministic output with seed parameter
-python speech_client.py --text "Hello, how are you?" --seed 42
 ```
 
-The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.
+The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
 
+### Notes
+- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](../../offline_inference/text_to_speech/README.md#omnivoice).
+
+---
 
 ## Qwen3-TTS
 

examples/online_serving/text_to_speech/omnivoice/speech_client.py

@@ -12,8 +12,6 @@ Examples:
 """
 
 import argparse
-import base64
-import os
 
 import httpx
 
@@ -21,24 +19,6 @@ DEFAULT_API_BASE = "http://localhost:8091"
 DEFAULT_API_KEY = "EMPTY"
 
 
-def encode_audio_to_base64(audio_path: str) -> str:
-    """Encode a local audio file to a base64 data URL."""
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    ext = audio_path.lower().rsplit(".", 1)[-1]
-    mime = {
-        "wav": "audio/wav",
-        "mp3": "audio/mpeg",
-        "flac": "audio/flac",
-        "ogg": "audio/ogg",
-    }.get(ext, "audio/wav")
-
-    with open(audio_path, "rb") as f:
-        b64 = base64.b64encode(f.read()).decode("utf-8")
-    return f"data:{mime};base64,{b64}"
-
-
 def run_tts(args) -> None:
     """Generate speech via /v1/audio/speech API."""
     payload = {
@@ -46,36 +26,16 @@ def run_tts(args) -> None:
         "input": args.text,
         "response_format": args.response_format,
     }
-    if args.seed is not None:
-        payload["extra_params"] = {}
-        payload["extra_params"]["seed"] = args.seed
 
     if args.voice:
         payload["voice"] = args.voice
     if args.language:
         payload["language"] = args.language
 
-    if args.ref_audio:
-        ref = args.ref_audio
-        if ref.startswith(("http://", "https://", "data:")):
-            payload["ref_audio"] = ref
-        else:
-            payload["ref_audio"] = encode_audio_to_base64(ref)
-
-    if args.ref_text:
-        payload["ref_text"] = args.ref_text
-
-    if args.instructions:
-        payload["instructions"] = args.instructions
-
     print(f"Model: {args.model}")
     print(f"Text: {args.text}")
-    if args.seed:
-        print(f"Seed: {args.seed}")
-
     if args.voice:
         print(f"Voice: {args.voice}")
-
     if args.language:
         print(f"Language: {args.language}")
     print("Generating audio...")
@@ -85,6 +45,7 @@ def run_tts(args) -> None:
         "Content-Type": "application/json",
         "Authorization": f"Bearer {args.api_key}",
     }
+
     with httpx.Client(timeout=300.0) as client:
         response = client.post(api_url, json=payload, headers=headers)
 
@@ -125,30 +86,6 @@ def main():
         choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
         help="Audio format (default: wav)",
     )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Reference audio for voice cloning (local path, URL, or data: URI)",
-    )
-    parser.add_argument(
-        "--ref-text",
-        type=str,
-        default=None,
-        help="Reference text for voice cloning",
-    )
-    parser.add_argument(
-        "--instructions",
-        type=str,
-        default=None,
-        help="Voice style/emotion instructions",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for generation, default: None for stochastic output)",
-    )
     parser.add_argument("--output", "-o", default=None, help="Output file path")
     args = parser.parse_args()
     run_tts(args)

recipes/ByteDance/Lance.md

@@ -1,177 +0,0 @@
-# Lance
-
-> Unified autoregressive + diffusion multimodal (text / image / video)
-
-## Summary
-
-- Vendor: ByteDance
-- Model: [`bytedance-research/Lance`](https://huggingface.co/bytedance-research/Lance) (`Lance_3B`, `Lance_3B_Video`)
-- Task: text2img, text2video, img2img (image edit), video2video (video edit),
-  img2text (image understanding), video2text (video understanding)
-- Mode: Offline inference, Online serving (OpenAI-compatible API)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe to run Lance — a 3B unified AR + diffusion multimodal model
-on a Qwen2.5-VL backbone — via vLLM-Omni. Lance is **BAGEL-lineage**: the
-released checkpoint uses the same `*_moe_gen` Mixture-of-Transformers weight
-layout as BAGEL, so vLLM-Omni reuses BAGEL's transformer core and specializes
-only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
-
-All six single-stage modalities are supported on the `Lance_3B` checkpoint;
-`text2video` and `video2video` additionally need `Lance_3B_Video` for the 3-D
-`latent_pos_embed` table. The two-stage AR + DiT topology is deferred and
-needs `LanceConfig` / `LanceProcessor` registered in the `vllm` package (a
-separate upstream PR).
-
-## References
-
-- Offline example:
-  [`examples/offline_inference/lance/end2end.py`](../../examples/offline_inference/lance/end2end.py)
-- Online serving:
-  [`examples/online_serving/lance/`](../../examples/online_serving/lance/)
-- Deploy config:
-  [`vllm_omni/deploy/lance.yaml`](../../vllm_omni/deploy/lance.yaml)
-- E2E tests:
-  [`tests/e2e/online_serving/test_lance.py`](../../tests/e2e/online_serving/test_lance.py)
-- HuggingFace model page:
-  [bytedance-research/Lance](https://huggingface.co/bytedance-research/Lance)
-- PR: [vllm-project/vllm-omni#3710](https://github.com/vllm-project/vllm-omni/pull/3710)
-  (closes [#3697](https://github.com/vllm-project/vllm-omni/issues/3697))
-
-## Hardware Support
-
-## GPU
-
-### 1x NVIDIA B300 / A100 80GB
-
-#### Environment
-
-- OS: Linux
-- Python: 3.12
-- Driver / runtime: CUDA ≥ 12.4
-- vLLM-Omni version: 0.18.x.dev
-
-#### Command — text-to-image (default)
-
-```bash
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance \
-    --prompts "a corgi astronaut on the moon, cinematic" \
-    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
-    --height 1024 --width 1024 \
-    --seed 42 --output ./out
-```
-
-Defaults match upstream `inference_lance.sh`: 30 denoising steps,
-`timestep-shift 3.5`, text CFG 4.0, seed 42, 1024×1024.
-
-#### Command — image edit (img2img)
-
-```bash
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance --modality img2img \
-    --image-path /path/to/input.png \
-    --prompts "Convert this into a vibrant cartoon-style illustration" \
-    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
-    --output ./out
-```
-
-The Lance-native VAE prefill scatters Wan2.2 latents into the LLM query
-sequence; no separate image encoder is needed.
-
-#### Command — text-to-video
-
-```bash
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
-    --num-frames 25 --video-height 480 --video-width 768 \
-    --prompts "a cat playing piano, cinematic" \
-    --steps 30 --fps 8 --output ./out
-```
-
-Use the `Lance_3B_Video` subfolder for any video path so the 3-D
-`latent_pos_embed` table is loaded; image / understanding paths can point at
-the top-level repo and resolve the right sub-checkpoint automatically.
-
-#### Command — image / video understanding
-
-```bash
-# Image → text (caption / VQA)
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance --modality img2text \
-    --image-path /path/to/photo.jpg \
-    --prompts "Describe this image in detail." \
-    --do-sample --text-temperature 0.8
-
-# Video → text
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance --modality video2text \
-    --video-path /path/to/clip.mp4 \
-    --prompts "What is happening in this video?"
-```
-
-Sampling is enabled by default at `--text-temperature 0.8` for the
-understanding paths because Lance's greedy decoder emits an immediate EOS for
-many prompts.
-
-#### Verification
-
-```bash
-pytest -s -v tests/e2e/online_serving/test_lance.py
-```
-
-#### Notes
-
-- BF16 footprint: ~7 GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE; comfortably fits
-  on a single 16 GB+ GPU for `Lance_3B`.
-- `rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
-  through `BagelRotaryEmbedding`. `text2img` uses scalar position ids
-  (BAGEL-equivalent); `img2text` / `video` / edit paths thread per-axis 3-D
-  position ids.
-- `video_edit` quality is more abstract than `text2video` at the same
-  resolution (known position-id offset between VAE-prefill and gen-latent
-  blocks). Functionally correct end-to-end.
-
-## Online Serving
-
-Lance supports all single-stage modalities via the OpenAI-compatible
-`/v1/chat/completions` API.
-
-### Launch
-
-```bash
-bash examples/online_serving/lance/run_server.sh
-# or, with overrides
-MODEL=bytedance-research/Lance \
-DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
-PORT=8091 \
-    bash examples/online_serving/lance/run_server.sh
-```
-
-For `text2video` / `video2video`, set
-`MODEL=bytedance-research/Lance/Lance_3B_Video`.
-
-### Send requests
-
-```bash
-# Text-to-image
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "A cute corgi astronaut on the moon, cinematic" \
-    --modality text2img --output corgi.png
-
-# Image edit
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "Convert this into a vibrant cartoon-style illustration" \
-    --modality img2img --image-url path/to/photo.png \
-    --output edited.png
-
-# Image understanding
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "Describe this image" \
-    --modality img2text --image-url photo.jpg
-```
-
-The client is shared with BAGEL — same OpenAI message format, same
-`modalities` / `num_inference_steps` / `seed` / `height` / `width` knobs.

recipes/GLM/GLM-Image.md

@@ -1,112 +0,0 @@
-# GLM-Image for text-to-image and image editing
-
-## Summary
-
-- Vendor: Z.ai
-- Model: `GLM/GLM-Image`
-- Task: Text-to-image (T2I) and image-to-image
-- Mode: Online serving with the OpenAI-compatible API
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for serving
-`GLM/GLM-Image` with vLLM-Omni on **two 80 GB NVIDIA A800** GPUs (Ampere-class,
-same default layout as the upstream **2×A100 80GB** example: Stage 0 AR on GPU 0,
-Stage 1 diffusion on GPU 1) and validate the deployment with the existing
-`examples/online_serving/glm_image` clients.
-
-## References
-
-- Upstream or canonical docs:
-  [`docs/user_guide/examples/online_serving/glm_image.md`](../../docs/user_guide/examples/online_serving/glm_image.md)
-- Related issue or discussion:
-  [#2888](https://github.com/vllm-project/vllm-omni/pull/2888)
-
-## Hardware Support
-
-This recipe documents **dual-GPU** CUDA layouts on A800 80 GB
-for the same software stack. Add more platforms (for example ROCm / NPU) as
-community validation lands.
-
-## GPU
-
-### 2× A800 80GB
-
-#### Environment
-
-These versions were taken from a working **editable** install: activate `vllm-omni/.venv` (or your equivalent), then align `pip` / Git with the rows below when reproducing this recipe.
-
-- OS: Linux
-- Python: 3.12
-- Driver / runtime: NVIDIA CUDA stack with **two** A800 80 GB GPUs visible (set `CUDA_VISIBLE_DEVICES` on your host if needed)
-- vLLM: **0.19.0**
-- vLLM-Omni: **0.19.0rc2.dev138+g38d5f2d53** (editable install from this repo; Git **`38d5f2d5`**, `git describe` ≈ **`v0.19.0rc1-138-g38d5f2d5`**)
-- Transformers: **5.5.4** (same `.venv` as above; required so `glm_image` configs load for Stage 0)
-
-#### Command
-
-Start the server from the repository root:
-
-```bash
-vllm serve zai-org/GLM-Image --omni --port 8091
-```
-
-To use the bundled stage config explicitly (same default as above):
-
-```bash
-vllm serve zai-org/GLM-Image \
-  --omni \
-  --port 8091 \
-  --deploy-config vllm_omni/deploy/glm_image.yaml
-```
-
-#### Verification
-
-Run one of the existing example clients after the server is ready:
-
-```bash
-curl -s http://172.18.69.133:8000/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [
-      {"role": "user", "content": "A beautiful landscape painting"}
-    ],
-    "extra_body": {
-      "height": 1920,
-      "width": 1920,
-      "num_inference_steps": 50,
-      "true_cfg_scale": 1.5,
-      "seed": 42
-    }
-  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > land.png
-```
-After the command finishes, check for the output files:
-
-```bash
-ls output.png
-```
-
-#### Sample end-to-end metrics
-
-One representative **offline** GLM-Image E2E run on this recipe’s **2× A800 80GB**.
-Overall summary from the run’s metrics. Rough wall-time split: **Stage 0 (AR)** ~**25 s**,
-**Stage 1 (diffusion)** ~**34 s** (see `e2e_stage_*_wall_time_ms` below).
-
-| Field | Value |
-| --- | ---: |
-| e2e_requests | 1 |
-| e2e_wall_time_ms | 61,148.679 |
-| e2e_total_tokens | 1,300 |
-| e2e_avg_time_per_request_ms | 61,148.679 |
-| e2e_avg_tokens_per_s | 21.260 |
-| e2e_stage_0_wall_time_ms | 24,708.760 |
-| e2e_stage_1_wall_time_ms | 33,787.442 |
-
-#### Notes
-
-- Memory usage: Roughly **~38 GiB + KV** on Stage 0 (AR) and **~20 GiB** on Stage 1 (DiT+VAE) per the user guide; two 80 GB cards match the default split.
-- Key flags: `--omni` is required; `--stage-configs-path` is optional unless you use a custom YAML (for example single-GPU).
-- Keep **Transformers ≥ 5.5.1** (this recipe used **5.5.4**) so `glm_image` configs resolve; otherwise Stage 0 can fail at `ModelConfig` validation.
-- Known limitations: This starter recipe follows the dual-GPU online path documented under `examples/online_serving/glm_image`. The first request may be slower due to warmup.
-- Generation time: about **61 s** wall time end-to-end for the sample above (50 inference steps, 1024×1024).

recipes/Helios/Helios.md

@@ -1,109 +0,0 @@
-# Helios for video generation
-
-## Summary
-
-- Vendor: Helios
-- Model: `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled`
-- Task: Text-to-video, image-to-video, and video-to-video generation
-- Mode: Offline inference with the bundled Helios example
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for running Helios
-video generation with vLLM-Omni. The concrete command below focuses on
-`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU. The
-same example script also supports Helios-Mid, Helios-Distilled, image-to-video,
-and video-to-video; see the linked example docs for those model-specific flags.
-
-## References
-
-- Upstream repository: <https://github.com/PKU-YuanGroup/Helios>
-- Model weights:
-  <https://huggingface.co/BestWishYsh/Helios-Base>
-- Upstream or canonical docs:
-  [`docs/user_guide/examples/offline_inference/helios.md`](../../docs/user_guide/examples/offline_inference/helios.md)
-- Related example under `examples/`:
-  [`examples/offline_inference/helios/README.md`](../../examples/offline_inference/helios/README.md)
-
-## Hardware Support
-
-## GPU
-
-### 1x NVIDIA H20
-
-#### Environment
-
-- OS: Ubuntu Linux x86_64
-- Python: 3.12.12
-- Driver / runtime: NVIDIA driver `580.126.20`, CUDA `13.0`
-- Hardware: 1x NVIDIA H20 GPU from an 8x H20 host
-- vLLM version: `0.19.0`
-- vLLM-Omni version or commit: `a3903810`
-
-#### Command
-
-Run the baseline Helios-Base text-to-video example from the repository root:
-
-```bash
-cd examples/offline_inference/helios
-
-python end2end.py \
-  --model BestWishYsh/Helios-Base \
-  --sample-type t2v \
-  --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-  --guidance-scale 5.0 \
-  --output helios_t2v_base.mp4
-```
-
-To use cache-dit acceleration, run on a vLLM-Omni checkout that includes Helios cache-dit support and enable the cache backend:
-
-```bash
-cd examples/offline_inference/helios
-
-python end2end.py \
-  --cache-backend cache_dit \
-  --enable-cache-dit-summary \
-  --model BestWishYsh/Helios-Base \
-  --sample-type t2v \
-  --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
-  --guidance-scale 5.0 \
-  --output helios_t2v_base.mp4
-```
-
-#### Verification
-
-The script should print the resolved generation configuration, total generation
-time, and output path. A successful run ends with output similar to:
-
-```text
-Total generation time: <seconds> seconds (<milliseconds> ms)
-Saved generated video to helios_t2v_base.mp4
-```
-
-#### Important flags and stage configs
-
-- `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by
-  this recipe.
-- `--sample-type t2v` selects text-to-video generation. Use `i2v` with
-  `--image-path` for image-to-video, or `v2v` with `--video-path` for
-  video-to-video.
-- `--guidance-scale 5.0` matches the Helios-Base recommendation. Use
-  `--guidance-scale 1.0` for Helios-Distilled.
-- `--cache-backend cache_dit` enables the cache-dit acceleration path.
-- `--enable-cache-dit-summary` prints cache-dit summary information after
-  diffusion forward passes.
-- No separate stage config is required for this offline recipe; the bundled
-  Helios example configures the pipeline through `end2end.py` arguments.
-- For Helios-Mid, add `--is-enable-stage2`,
-  `--pyramid-num-inference-steps-list 20 20 20`, `--use-cfg-zero-star`,
-  `--use-zero-init`, and `--zero-steps 1`.
-- For Helios-Distilled, add `--is-enable-stage2`,
-  `--pyramid-num-inference-steps-list 2 2 2`, and
-  `--is-amplify-first-chunk`.
-
-#### Known limitations
-
-- Helios generates video in 33-frame chunks. For best performance, set
-  `--num-frames` to a multiple of `33`; non-multiple values are rounded up to
-  the nearest multiple of `33`.

recipes/README.md

@@ -29,13 +29,10 @@ recipes/
 |--------|------|----------|
 | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
-| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
-| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
-| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
@@ -47,7 +44,6 @@ recipes/
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
-| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like

recipes/mistralai/Voxtral-TTS.md

@@ -1,107 +0,0 @@
-# Voxtral TTS for text-to-speech
-
-## Summary
-
-- Vendor: Mistral AI
-- Model: `mistralai/Voxtral-4B-TTS-2603`
-- Task: Text-to-speech with model-provided voice presets or voice cloning
-- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for serving Voxtral TTS
-models with vLLM-Omni and validate the deployment with the existing TTS client
-examples in this repository.
-
-## References
-
-- Upstream model card: <https://huggingface.co/mistralai/Voxtral-4B-TTS-2603>
-- Related examples under `examples/`:
-  [`examples/offline_inference/text_to_speech/voxtral_tts/end2end.py`](../../examples/offline_inference/text_to_speech/voxtral_tts/end2end.py)
-  [`examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py`](../../examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py)
-- Related issue or discussion:
-  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
-
-## Hardware Support
-
-## GPU
-
-### 1x RTX 4090 24GB
-
-#### Environment
-
-- OS: Ubuntu 22.04.5 LTS
-- Python: 3.12.11
-- PyTorch: 2.11.0+cu130
-- Driver / runtime: CUDA 13.0
-- GPU: NVIDIA GeForce RTX 4090, 24 GB
-- vLLM version: 0.20.0
-- vLLM-Omni version or commit: 0.20.0rc2.dev99+g857356d5b
-- Required package: mistral_common >= 1.10.0
-
-#### Command
-
-Online:
-
-```bash
-vllm serve mistralai/Voxtral-4B-TTS-2603 \
-    --omni \
-    --port 8091
-```
-
-The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`
-
-#### Verification
-
-**Quick smoke test with curl:**
-
-```bash
-curl -X POST http://127.0.0.1:8091/v1/audio/speech \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "mistralai/Voxtral-4B-TTS-2603",
-    "input": "Hello, this is Voxtral TTS running with vLLM-Omni.",
-    "voice": "casual_female",
-    "language": "English",
-    "response_format": "wav"
-  }' \
-  --output voxtral.wav
-```
-
-**Streaming audio:**
-
-```bash
-curl -X POST http://127.0.0.1:8091/v1/audio/speech \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "mistralai/Voxtral-4B-TTS-2603",
-    "input": "Hello, this is Voxtral TTS streaming PCM.",
-    "voice": "casual_female",
-    "language": "English",
-    "stream": true,
-    "response_format": "pcm"
-  }' \
-  --output voxtral_stream.pcm
-```
-Convert the raw PCM output to WAV
-```bash
-ffmpeg -f s16le -ar 24000 -ac 1 -i voxtral_stream.pcm voxtral_stream.wav -y
-ffprobe -hide_banner voxtral_stream.wav
-```
-
-#### Notes
-
-- Memory usage: Stage 0 (`audio_generation`) ~18.95 GiB, Stage 1
-  (`audio_tokenizer`) ~1.55 GiB. The observed server-startup peak was
-  `21006 MiB / 24564 MiB` (`20.51 GiB / 23.99 GiB`, about 85.5%).
-  Both stages share GPU 0 via the deploy config (`gpu_memory_utilization: 0.8`
-  for Stage 0, `0.1` for Stage 1).
-- Key flags: `--omni` is required. `--port 8091` is optional and only selects
-  the HTTP port.
-- Known limitation: Voxtral TTS outputs mono audio at 24 kHz. And with the public
-mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
-voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
-) or adding a `ref_audio` request parameter will fail
-because the public checkpoint does not provide the encoder weights needed to
-convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -1,92 +0,0 @@
-"""Unit tests for Omni AR streaming-session async placeholder handling."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-
-# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
-# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them.
-# isort: off
-import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
-from vllm.sampling_params import SamplingParams
-from vllm.v1.request import Request, RequestStatus, StreamingUpdate
-from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
-
-# isort: on
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
-    sched = OmniARScheduler.__new__(OmniARScheduler)
-    sched._new_prompt_len_snapshot = {}
-    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
-    sched.num_waiting_for_streaming_input = 0
-    sched.log_stats = False
-    sched.chunk_transfer_adapter = None
-    return sched
-
-
-def _make_request() -> Request:
-    return Request(
-        request_id="req-ar-streaming-test",
-        prompt_token_ids=[1, 2, 3],
-        sampling_params=SamplingParams(max_tokens=8),
-        pooling_params=None,
-        arrival_time=100.0,
-        block_hasher=None,
-    )
-
-
-def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
-    return StreamingUpdate(
-        mm_features=None,
-        prompt_token_ids=[10, 20] if prompt_token_ids is None else prompt_token_ids,
-        max_tokens=32,
-        arrival_time=200.0,
-        sampling_params=SamplingParams(max_tokens=16),
-    )
-
-
-def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
-    sched = _make_scheduler(stage_id=0)
-    session = _make_request()
-    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    session.append_output_token_ids([7, 8, 9])
-    session.num_computed_tokens = 6
-    session.num_output_placeholders = 1
-    session.spec_token_ids = [-1]
-
-    sched._update_request_as_session(session, _make_update([10, 20]))
-
-    assert session.discard_latest_async_tokens is True
-    assert session.num_output_placeholders == 0
-    assert session.spec_token_ids == []
-    # The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
-    # carried into the next streaming prompt before the new chunk tokens.
-    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 10, 20]
-    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 10, 20]
-    assert session._output_token_ids == []
-    assert session.num_prompt_tokens == 7
-    assert sched._new_prompt_len_snapshot[session.request_id] == 2
-
-
-def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() -> None:
-    sched = _make_scheduler(stage_id=0)
-    session = _make_request()
-    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-    session.append_output_token_ids([7, 8, 9])
-    session.num_computed_tokens = 6
-    session.num_output_placeholders = 0
-
-    sched._update_request_as_session(session, _make_update([10, 20]))
-
-    assert session.discard_latest_async_tokens is False
-    assert session.num_output_placeholders == 0
-    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
-    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
-    assert session._output_token_ids == []
-    assert session.num_prompt_tokens == 8
-    assert sched._new_prompt_len_snapshot[session.request_id] == 2

tests/core/sched/test_omni_scheduler_mixin.py

@@ -1,6 +1,6 @@
-"""Unit tests for generation streaming session replacement.
+"""Unit tests for OmniSchedulerMixin streaming session replacement.
 
-These tests pin the behavior of `_update_request_as_session` against
+These tests pin the behavior of `_replace_session_with_streaming_update` against
 current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading
 vLLM, failures here should highlight incompatible changes to request state or
 update payloads early.
@@ -19,34 +19,19 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreEventType
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
-from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
 # isort: on
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class _SkippedWaitingStub:
-    def __contains__(self, request: Request) -> bool:
-        return False
-
-
-class _ChunkTransferAdapterStub:
-    def __init__(self) -> None:
-        self.segment_finished_requests: set[str] = set()
-
-
-class _SchedulerStub(OmniGenerationScheduler):
-    """Minimal scheduler surface required by OmniGenerationScheduler."""
+class _SchedulerStub(OmniSchedulerMixin):
+    """Minimal scheduler surface required by OmniSchedulerMixin."""
 
     def __init__(self, *, log_stats: bool = False) -> None:
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
-        self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
-        self.skipped_waiting = _SkippedWaitingStub()
-
-    def _enqueue_waiting_request(self, session: Request) -> None:
-        raise AssertionError("unexpected enqueue for skipped_waiting miss")
 
 
 def _make_request(**kwargs) -> Request:
@@ -86,7 +71,7 @@ class TestReplaceSessionWithStreamingUpdate:
 
         update = _make_update(prompt_token_ids=[40, 41, 42])
         sched.num_waiting_for_streaming_input = 3
-        sched._update_request_as_session(session, update)
+        sched._replace_session_with_streaming_update(session, update)
 
         assert session._output_token_ids == []
         assert list(session._all_token_ids) == [40, 41, 42]
@@ -103,7 +88,7 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         update = _make_update(prompt_token_ids=None)
-        sched._update_request_as_session(session, update)
+        sched._replace_session_with_streaming_update(session, update)
 
         assert session.prompt_token_ids == ()
         assert list(session._all_token_ids) == []
@@ -123,7 +108,7 @@ class TestReplaceSessionWithStreamingUpdate:
             pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)")
         update = replace(base, additional_information=None)
 
-        sched._update_request_as_session(session, update)
+        sched._replace_session_with_streaming_update(session, update)
         assert session.additional_information is None
 
     def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None:
@@ -131,14 +116,14 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         sched.num_waiting_for_streaming_input = 5
-        sched._update_request_as_session(session, _make_update())
+        sched._replace_session_with_streaming_update(session, _make_update())
         assert sched.num_waiting_for_streaming_input == 5
 
     def test_records_queued_event_when_log_stats_enabled(self) -> None:
         sched = _SchedulerStub(log_stats=True)
         session = _make_request()
         session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-        sched._update_request_as_session(session, _make_update())
+        sched._replace_session_with_streaming_update(session, _make_update())
 
         assert session.events
         assert session.events[-1].type == EngineCoreEventType.QUEUED

tests/dfx/perf/scripts/run_benchmark.py

@@ -191,16 +191,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         raise ValueError("The number of prompts does not match the QPS or max_concurrency")
 
     args = ["--host", host, "--port", str(port)]
-    exclude_keys = {
-        "request_rate",
-        "baseline",
-        "num_prompts",
-        "max_concurrency",
-        "task",
-        "enabled",
-        "eval_phase",
-        "trust_remote_code",
-    }
+    exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
 
     for key, value in params.items():
         if key in exclude_keys or value is None:
@@ -216,14 +207,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         elif not isinstance(value, bool):
             args.extend([arg_name, str(value)])
 
-    for config in BENCHMARK_CONFIGS:
-        if config.get("test_name") != test_name:
-            continue
-        server_params = config.get("server_params") or {}
-        if server_params.get("trust_remote_code") or params.get("trust_remote_code"):
-            args.append("--trust-remote-code")
-        break
-
     # QPS / request-rate sweep
     for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]

tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json

@@ -1,33 +1,35 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_cfgp2",
-        "description": "TP=2 CfgP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_fp8_cfgp2",
+        "description": "TP=2 Quantization=fp8 CfgP=2 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
-                "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 2,
                 "cfg-parallel-size": 2,
+                "quantization": "fp8",
+                "distributed-executor-backend": "mp",
+                "enforce-eager": true,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps8",
+                "name": "1024x1024_steps50",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 8,
+                "num-inference-steps": 50,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_p99": 4.7469,
-                    "peak_memory_mb_max": 101100,
-                    "peak_memory_mb_mean": 101100
+                    "throughput_qps": 0.1035,
+                    "latency_p99": 9.9057,
+                    "peak_memory_mb_max": 66470,
+                    "peak_memory_mb_mean": 66470
                 }
             }
         ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json

@@ -1,33 +1,35 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_sp2",
-        "description": "TP=2 SP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_fp8_sp2",
+        "description": "TP=2 Quantization=fp8 SP=2 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 2,
-                "ulysses-degree": 2,
+                "usp": 2,
+                "quantization": "fp8",
+                "distributed-executor-backend": "mp",
+                "enforce-eager": true,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps8",
+                "name": "1024x1024_steps50",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 8,
+                "num-inference-steps": 50,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.20,
-                    "latency_p99": 5.1025,
-                    "peak_memory_mb_max": 97402,
-                    "peak_memory_mb_mean": 97402
+                    "throughput_qps": 0.08,
+                    "latency_p99": 12.0731,
+                    "peak_memory_mb_max": 66314,
+                    "peak_memory_mb_mean": 66314
                 }
             }
         ]

tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json

@@ -1,32 +1,34 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp4",
-        "description": "TP=4 baseline",
+        "test_name": "test_hunyuan_image_tp4_fp8",
+        "description": "TP=4 Quantization=fp8 baseline",
         "server_type": "vllm-omni",
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
                 "tensor-parallel-size": 4,
+                "quantization": "fp8",
+                "distributed-executor-backend": "mp",
+                "enforce-eager": true,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps8",
+                "name": "1024x1024_steps50",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 8,
+                "num-inference-steps": 50,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_p99": 4.9235,
-                    "peak_memory_mb_max": 57576,
-                    "peak_memory_mb_mean": 57576
+                    "throughput_qps": 0.08,
+                    "latency_p99": 13.1227,
+                    "peak_memory_mb_max": 46838,
+                    "peak_memory_mb_mean": 46838
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -63,6 +63,24 @@
             }
         },
         "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.15,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 55000,
+                    "peak_memory_mb_mean": 55000
+                }
+            },
             {
                 "name": "1536x1536_steps35_i2i_2img",
                 "dataset": "random",

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -61,6 +61,23 @@
             }
         },
         "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.21,
+                    "latency_mean": 10,
+                    "peak_memory_mb_max": 56000,
+                    "peak_memory_mb_mean": 56000
+                }
+            },
             {
                 "name": "1536x1536_steps35_i2i",
                 "dataset": "random",

tests/dfx/perf/tests/test_runner_metadata.py

@@ -17,16 +17,7 @@ def test_task_excluded_from_cli_args():
         "percentile-metrics": "audio_rtf,audio_ttfp",
         "baseline": {"mean_audio_rtf": [0.5]},
     }
-    exclude_keys = {
-        "request_rate",
-        "baseline",
-        "num_prompts",
-        "max_concurrency",
-        "task",
-        "enabled",
-        "eval_phase",
-        "trust_remote_code",
-    }
+    exclude_keys = {"request_rate", "baseline", "num_prompts", "max_concurrency", "task", "enabled", "eval_phase"}
     args = []
     for key, value in params.items():
         if key in exclude_keys or value is None:

tests/dfx/perf/tests/test_voxcpm2.json

@@ -1,198 +0,0 @@
-[
-  {
-    "test_name": "test_voxcpm2",
-    "server_params": {
-      "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true
-    },
-    "benchmark_params": [
-      {
-        "task": "voice_clone",
-        "eval_phase": "latency",
-        "dataset_name": "seed-tts",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
-          20
-        ],
-        "max_concurrency": [
-          1
-        ],
-        "seed_tts_locale": "en",
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "median_audio_ttfp_ms": [
-            600
-          ],
-          "median_audio_rtf": [
-            0.2
-          ]
-        }
-      },
-      {
-        "task": "voice_clone",
-        "eval_phase": "throughput",
-        "dataset_name": "seed-tts",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
-          80
-        ],
-        "max_concurrency": [
-          8
-        ],
-        "seed_tts_locale": "en",
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "median_audio_ttfp_ms": [
-            3500
-          ],
-          "median_audio_rtf": [
-            0.85
-          ],
-          "audio_throughput": [
-            8
-          ]
-        }
-      },
-      {
-        "task": "voice_clone",
-        "eval_phase": "quality",
-        "enabled": false,
-        "dataset_name": "seed-tts",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
-          200
-        ],
-        "max_concurrency": [
-          4
-        ],
-        "seed_tts_locale": "en",
-        "seed_tts_wer_eval": true,
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_audio_rtf": [
-            0.45
-          ]
-        }
-      },
-      {
-        "task": "default_voice",
-        "eval_phase": "latency",
-        "dataset_name": "seed-tts-text",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
-          20
-        ],
-        "max_concurrency": [
-          1
-        ],
-        "seed_tts_locale": "en",
-        "extra_body": {
-          "voice": "default"
-        },
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "median_audio_ttfp_ms": [
-            350
-          ],
-          "median_audio_rtf": [
-            0.15
-          ]
-        }
-      },
-      {
-        "task": "default_voice",
-        "eval_phase": "throughput",
-        "dataset_name": "seed-tts-text",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
-          80
-        ],
-        "max_concurrency": [
-          8
-        ],
-        "seed_tts_locale": "en",
-        "extra_body": {
-          "voice": "default"
-        },
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "median_audio_ttfp_ms": [
-            3000
-          ],
-          "median_audio_rtf": [
-            0.5
-          ],
-          "audio_throughput": [
-            10
-          ]
-        }
-      },
-      {
-        "task": "default_voice",
-        "eval_phase": "stress",
-        "dataset_name": "seed-tts-text",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
-          100
-        ],
-        "request_rate": [
-          2.0
-        ],
-        "seed_tts_locale": "en",
-        "extra_body": {
-          "voice": "default"
-        },
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "median_audio_ttfp_ms": [
-            500
-          ],
-          "median_audio_rtf": [
-            0.3
-          ],
-          "audio_throughput": [
-            5
-          ]
-        }
-      },
-      {
-        "task": "default_voice",
-        "eval_phase": "quality",
-        "enabled": false,
-        "dataset_name": "seed-tts-text",
-        "backend": "openai-audio-speech",
-        "endpoint": "/v1/audio/speech",
-        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
-          200
-        ],
-        "max_concurrency": [
-          4
-        ],
-        "seed_tts_locale": "en",
-        "extra_body": {
-          "voice": "default"
-        },
-        "seed_tts_wer_eval": true,
-        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_audio_rtf": [
-            0.35
-          ]
-        }
-      }
-    ]
-  }
-]

tests/dfx/reliability/helpers.py

@@ -4,16 +4,10 @@ This module keeps fault injection callable from tests directly:
 - GPU OOM (CUDA sidecar memory hog)
 - process kill by pattern and signal
 - post-ready hooks via ``fault_injector`` / ``omni_server_after_fault`` fixtures
-
-Worker PID classification (``worker_pids`` in snapshots) is still used for logging /
-markers; post-fault **process** cleanup assertions use the full captured
-``tree_pids`` (serve root + descendants at fault time) — see
-:func:`assert_no_server_tree_process_residual_and_gpu_release`.
 """
 
 from __future__ import annotations
 
-import concurrent.futures
 import http.client
 import json
 import logging
@@ -33,16 +27,6 @@ import pytest
 
 logger = logging.getLogger(__name__)
 
-# Substrings matched against ``psutil.Process.name`` + argv text for classifying GPU
-# workers under the test server tree. Extend per-suite via
-# ``server.reliability_worker_markers_extra`` or replace via
-# ``server.reliability_worker_markers`` (see ``_resolve_runtime_worker_markers``).
-DEFAULT_RUNTIME_WORKER_MARKERS: tuple[str, ...] = (
-    "multiprocessing.spawn",
-    # Covers ``VLLM::Worker``, ``VLLM::StageEngineCoreProc_*`` (Omni stage engines), etc.
-    "VLLM::",
-)
-
 
 @dataclass
 class OomHandle:
@@ -468,61 +452,6 @@ def _safe_proc_info(pid: int) -> tuple[str, str]:
         return "<unknown>", "<unavailable>"
 
 
-def _normalize_worker_marker_substrings(seq: Sequence[Any], *, context: str) -> tuple[str, ...]:
-    """Return deduped non-empty marker substrings; raises if any entry is blank."""
-    out: list[str] = []
-    for item in seq:
-        s = str(item).strip()
-        if not s:
-            raise ValueError(f"{context}: worker marker must be a non-empty string, got {item!r}")
-        out.append(s)
-    return tuple(dict.fromkeys(out))
-
-
-def _resolve_runtime_worker_markers(
-    server: Any,
-    injector_markers_extra: Sequence[str] | None,
-) -> tuple[str, ...]:
-    """Resolve which substrings classify worker PIDs for fault snapshots.
-
-    If ``server.reliability_worker_markers`` is set to a sequence, **only** those
-    markers are used (no defaults). Otherwise defaults from
-    :data:`DEFAULT_RUNTIME_WORKER_MARKERS` are merged with:
-
-    - ``server.reliability_worker_markers_extra`` when present, and
-    - ``injector_markers_extra`` from :func:`make_server_root_kill_fault_injector` /
-      :func:`make_server_tree_kill_fault_injector`.
-    """
-    replace = getattr(server, "reliability_worker_markers", None)
-    if replace is not None:
-        return _normalize_worker_marker_substrings(replace, context="server.reliability_worker_markers")
-
-    merged: list[str] = []
-    merged.extend(DEFAULT_RUNTIME_WORKER_MARKERS)
-    extra_attr = getattr(server, "reliability_worker_markers_extra", None)
-    if extra_attr is not None:
-        merged.extend(
-            _normalize_worker_marker_substrings(extra_attr, context="server.reliability_worker_markers_extra")
-        )
-    if injector_markers_extra is not None:
-        merged.extend(
-            _normalize_worker_marker_substrings(
-                injector_markers_extra,
-                context="injector worker_markers_extra",
-            )
-        )
-    return tuple(dict.fromkeys(merged))
-
-
-def _pid_looks_like_runtime_worker(pid: int, markers: Sequence[str]) -> bool:
-    """True when process name or argv contains any resolved worker marker substring."""
-    if not markers:
-        return False
-    name, cmdline = _safe_proc_info(pid)
-    text = f"{name} {cmdline}"
-    return any(marker in text for marker in markers)
-
-
 def _list_server_process_tree(server: Any) -> list[int]:
     """Return [root, descendants...] PIDs for the current test server instance."""
     root_proc = getattr(server, "proc", None)
@@ -539,29 +468,6 @@ def _list_server_process_tree(server: Any) -> list[int]:
     return [root_pid, *descendants]
 
 
-def _pids_in_server_tree_substring_match(server: Any, pattern: str, *, limit: int) -> list[int]:
-    """Return up to ``limit`` PIDs in the server tree whose name+cmdline contains ``pattern`` literally.
-
-    ``pgrep -f`` matches the process command line used by procps and does not consult the
-    short process name; vLLM-Omni often exposes titles like ``VLLM::Worker`` or
-    ``VLLM::StageEngineCoreProc_*`` via prctl/setproctitle in ``ps``/``name()`` only.
-    This helper keeps kill injection aligned with :func:`_safe_proc_info` / fault snapshots.
-    """
-    if limit <= 0:
-        return []
-    tree = _list_server_process_tree(server)
-    if not tree or not pattern:
-        return []
-    hits: list[int] = []
-    for pid in tree:
-        if len(hits) >= limit:
-            break
-        name, cmdline = _safe_proc_info(int(pid))
-        if pattern in f"{name} {cmdline}":
-            hits.append(int(pid))
-    return hits
-
-
 def _log_server_process_tree(server: Any) -> None:
     """Print server process tree for debugging fault injection targets."""
     pids = _list_server_process_tree(server)
@@ -580,315 +486,6 @@ FaultInjector = Callable[[Any], None]
 """Callable invoked with the live ``OmniServer`` after it is ready (see ``omni_server_after_fault``)."""
 
 
-def run_fault_injection_with_rate_load(
-    *,
-    submit_request: Callable[[], Any],
-    inject_fault: Callable[[], None],
-    num_requests: int = 10,
-    request_rate: float = 0.3,
-    submit_interval_sec: float | None = None,
-    completion_timeout_sec: float = 120.0,
-) -> dict[str, Any]:
-    """Submit requests with a fixed rate, then inject fault once an in-flight request is observed."""
-    if num_requests <= 0:
-        raise ValueError("num_requests must be > 0")
-    if request_rate <= 0:
-        raise ValueError("request_rate must be > 0")
-
-    interval_sec = submit_interval_sec if submit_interval_sec is not None else (1.0 / request_rate)
-    if interval_sec < 0:
-        raise ValueError("submit_interval_sec must be >= 0")
-
-    max_workers = min(max(2, num_requests), 8)
-    injected = False
-    futures: list[concurrent.futures.Future[Any]] = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-        for idx in range(num_requests):
-            futures.append(executor.submit(submit_request))
-            in_flight = any(not future.done() for future in futures)
-            if in_flight and not injected:
-                inject_fault()
-                injected = True
-            if idx != num_requests - 1 and interval_sec > 0:
-                time.sleep(interval_sec)
-
-        if not injected:
-            in_flight = any(not future.done() for future in futures)
-            if in_flight:
-                inject_fault()
-                injected = True
-
-        if not injected:
-            pytest.skip(
-                "no in-flight request observed before fault injection; "
-                "increase request cost or request rate for this environment"
-            )
-
-        done, pending = concurrent.futures.wait(
-            futures,
-            timeout=completion_timeout_sec,
-            return_when=concurrent.futures.ALL_COMPLETED,
-        )
-        assert not pending, (
-            "some in-flight-load requests did not converge after fault injection: "
-            f"pending={len(pending)} done={len(done)}"
-        )
-    failure_observed = False
-    completed = 0
-    success = 0
-    exceptions = 0
-    for future in futures:
-        completed += 1
-        try:
-            future.result()
-            success += 1
-        except Exception:  # noqa: BLE001
-            exceptions += 1
-            failure_observed = True
-    return {
-        "num_submitted": num_requests,
-        "completed": completed,
-        "success": success,
-        "exceptions": exceptions,
-        "failure_observed": failure_observed,
-        "inflight_observed": True,
-    }
-
-
-def list_alive_pids(pids: Sequence[int]) -> list[int]:
-    """Return PIDs from ``pids`` that still exist in the kernel and are **not** zombies.
-
-    After SIGTERM/SIGKILL the serve child may exit before the test harness calls
-    ``Popen.wait()``; until reaped, ``psutil.pid_exists`` stays true for the defunct
-    slot. Those PIDs must not count as "residual server processes" for leak checks.
-    """
-    out: list[int] = []
-    for pid in pids:
-        pid_i = int(pid)
-        if not psutil.pid_exists(pid_i):
-            continue
-        try:
-            if psutil.Process(pid_i).status() == psutil.STATUS_ZOMBIE:
-                continue
-        except psutil.Error:
-            continue
-        out.append(pid_i)
-    return out
-
-
-def query_gpu_compute_pid_used_memory_mb() -> dict[int, int] | None:
-    """Query NVIDIA compute-process memory map (pid -> used MB).
-
-    Returns ``None`` when ``nvidia-smi`` is unavailable/unreadable on current host.
-    """
-    out = subprocess.run(
-        ["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
-        check=False,
-        capture_output=True,
-        text=True,
-    )
-    if out.returncode != 0:
-        logger.warning("[reliability][gpu] nvidia-smi query failed: %s", out.stderr.strip())
-        return None
-
-    pid_to_mem_mb: dict[int, int] = {}
-    for line in out.stdout.splitlines():
-        row = line.strip()
-        if not row:
-            continue
-        pieces = [part.strip() for part in row.split(",", maxsplit=1)]
-        if len(pieces) != 2:
-            continue
-        pid_raw, mem_raw = pieces
-        if not pid_raw.isdigit():
-            continue
-        mem_digits = "".join(ch for ch in mem_raw if ch.isdigit())
-        if not mem_digits:
-            continue
-        pid_to_mem_mb[int(pid_raw)] = int(mem_digits)
-    return pid_to_mem_mb
-
-
-def worker_residual_timeout_after_kill_signal(signal_name: str) -> float:
-    """Wall-clock budget for fault-snapshot PIDs / GPU to clear after ``signal_name``."""
-    if signal_name == "SIGKILL":
-        return 30.0
-    if signal_name == "SIGINT":
-        return 120.0
-    if signal_name == "SIGTERM":
-        return 90.0
-    return 30.0
-
-
-def assert_no_server_tree_process_residual_and_gpu_release(
-    server: Any,
-    *,
-    scenario: str,
-    timeout_sec: float = 30.0,
-    poll_interval_sec: float = 0.5,
-) -> None:
-    """Assert no live PIDs remain from the fault-time server tree and no GPU use by those PIDs.
-
-    Uses ``reliability_fault_snapshot["tree_pids"]`` (root + all descendants captured at
-    injection time), not the marker-filtered ``worker_pids`` subset — so helpers like
-    ``multiprocessing.resource_tracker`` are included in the residual check.
-
-    **Zombie** PIDs (exited children not yet reaped by the test ``Popen``) are excluded
-    from the "alive" list; ``psutil.pid_exists`` alone would false-positive on them.
-
-    Note: if a child is reparented to PID 1 while staying alive, its PID is unchanged
-    and remains detected; if it **exits and a new unrelated process reuses the same
-    numeric PID**, this check may false-positive (rare on short windows).
-    """
-    snapshot = getattr(server, "reliability_fault_snapshot", None)
-    if not isinstance(snapshot, dict):
-        pytest.fail(f"[{scenario}] missing reliability fault snapshot on server")
-    tree_pids = [int(pid) for pid in snapshot.get("tree_pids", [])]
-    if not tree_pids:
-        pytest.skip(f"[{scenario}] no server process tree PIDs captured for this run")
-
-    tree_pid_set = set(tree_pids)
-    deadline = time.monotonic() + timeout_sec
-    last_alive: list[int] = []
-    last_gpu_leaks: dict[int, int] = {}
-    while time.monotonic() < deadline:
-        last_alive = list_alive_pids(tree_pids)
-        gpu_map = query_gpu_compute_pid_used_memory_mb()
-        if gpu_map is None:
-            pytest.skip(f"[{scenario}] nvidia-smi unavailable; skip GPU release assertion")
-        last_gpu_leaks = {pid: mem_mb for pid, mem_mb in gpu_map.items() if pid in tree_pid_set and mem_mb > 0}
-        if not last_alive and not last_gpu_leaks:
-            return
-        time.sleep(poll_interval_sec)
-
-    assert not last_alive, f"[{scenario}] residual server-tree processes remain alive: {last_alive}"
-    assert not last_gpu_leaks, f"[{scenario}] server-tree PID GPU memory not released: {last_gpu_leaks}"
-
-
-def assert_no_worker_residual_and_gpu_release(
-    server: Any,
-    *,
-    scenario: str,
-    timeout_sec: float = 30.0,
-    poll_interval_sec: float = 0.5,
-) -> None:
-    """Deprecated: use :func:`assert_no_server_tree_process_residual_and_gpu_release`."""
-    assert_no_server_tree_process_residual_and_gpu_release(
-        server,
-        scenario=scenario,
-        timeout_sec=timeout_sec,
-        poll_interval_sec=poll_interval_sec,
-    )
-
-
-def _capture_server_fault_snapshot(
-    server: Any,
-    *,
-    worker_markers_extra: Sequence[str] | None = None,
-) -> dict[str, list[int]]:
-    """Capture current server process snapshot for post-fault assertions.
-
-    ``tree_pids`` is ``[root, ...descendants]`` and is used by
-    :func:`assert_no_server_tree_process_residual_and_gpu_release`.
-    ``worker_pids`` is the marker-filtered subset for debugging only.
-    """
-    markers = _resolve_runtime_worker_markers(server, worker_markers_extra)
-    tree_pids = _list_server_process_tree(server)
-    root_pid = tree_pids[0] if tree_pids else None
-    worker_pids: list[int] = []
-    for pid in tree_pids:
-        if root_pid is not None and int(pid) == int(root_pid):
-            continue
-        if _pid_looks_like_runtime_worker(pid, markers):
-            worker_pids.append(pid)
-    snapshot = {
-        "root_pids": [root_pid] if root_pid is not None else [],
-        "tree_pids": tree_pids,
-        "worker_pids": worker_pids,
-        "worker_markers": list(markers),
-    }
-    setattr(server, "reliability_fault_snapshot", snapshot)
-    return snapshot
-
-
-def make_server_root_kill_fault_injector(
-    *,
-    signal_name: str = "SIGKILL",
-    post_kill_wait_seconds: float = 0.0,
-    worker_markers_extra: Sequence[str] | None = None,
-) -> FaultInjector:
-    """Kill only the current server root process (simulate first strike on PID1-like process).
-
-    ``worker_markers_extra`` is merged into worker PID detection for the fault snapshot
-    (after :data:`DEFAULT_RUNTIME_WORKER_MARKERS` unless
-    ``server.reliability_worker_markers`` replaces the set entirely).
-    """
-
-    def _inject(server: Any) -> None:
-        _log_server_process_tree(server)
-        snapshot = _capture_server_fault_snapshot(server, worker_markers_extra=worker_markers_extra)
-        root_pids = snapshot["root_pids"]
-        if not root_pids:
-            pytest.skip("no server root process found for root-kill injection")
-        root_pid = int(root_pids[0])
-        sig = getattr(signal, signal_name, None)
-        if sig is None:
-            raise ValueError(f"Unsupported signal_name: {signal_name}")
-        name, cmdline = _safe_proc_info(root_pid)
-        print(
-            f"[reliability][process-kill] root-kill pid={root_pid} name={name} signal={signal_name} cmdline={cmdline}",
-            flush=True,
-        )
-        os.kill(root_pid, sig)
-        if post_kill_wait_seconds > 0:
-            time.sleep(post_kill_wait_seconds)
-
-    return _inject
-
-
-def make_server_tree_kill_fault_injector(
-    *,
-    signal_name: str = "SIGKILL",
-    post_kill_wait_seconds: float = 0.0,
-    inter_kill_wait_seconds: float = 0.2,
-    worker_markers_extra: Sequence[str] | None = None,
-) -> FaultInjector:
-    """Kill current server root first, then kill all remaining descendants.
-
-    See :func:`make_server_root_kill_fault_injector` for ``worker_markers_extra``.
-    """
-
-    def _inject(server: Any) -> None:
-        _log_server_process_tree(server)
-        snapshot = _capture_server_fault_snapshot(server, worker_markers_extra=worker_markers_extra)
-        tree_pids = [int(pid) for pid in snapshot["tree_pids"]]
-        if not tree_pids:
-            pytest.skip("no server process tree found for tree-kill injection")
-
-        sig = getattr(signal, signal_name, None)
-        if sig is None:
-            raise ValueError(f"Unsupported signal_name: {signal_name}")
-
-        root_pid = tree_pids[0]
-        kill_order = [root_pid, *[pid for pid in tree_pids if pid != root_pid]]
-        for pid in kill_order:
-            if not psutil.pid_exists(pid):
-                continue
-            name, cmdline = _safe_proc_info(pid)
-            print(
-                f"[reliability][process-kill] tree-kill pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
-                flush=True,
-            )
-            os.kill(pid, sig)
-            if inter_kill_wait_seconds > 0:
-                time.sleep(inter_kill_wait_seconds)
-
-        if post_kill_wait_seconds > 0:
-            time.sleep(post_kill_wait_seconds)
-
-    return _inject
-
-
 def make_process_kill_fault_injector(
     *,
     grep_patterns: str | Sequence[str],
@@ -896,18 +493,11 @@ def make_process_kill_fault_injector(
     limit: int = 1,
     post_kill_wait_seconds: float = 0.0,
 ) -> FaultInjector:
-    """Build a post-ready injector that kills processes matched by ordered ``grep_patterns``.
+    """Build a post-ready injector that kills processes matched by ``pgrep -f``.
 
-    Matching is tried in two phases per pattern (first hit wins):
-
-    1. **Server process tree** (``psutil``): literal substring match against
-       ``Process.name()`` plus argv text. This matches how :func:`_safe_proc_info`
-       logs processes and catches vLLM-style titles (``VLLM::...``) that often do not
-       appear in ``pgrep -f``\'s command-line view.
-    2. **``pgrep -f``** (legacy): scoped to the server PID tree when it is known;
-       uses procps regular-expression rules for the pattern.
-
-    If neither phase finds a target, the returned callable issues ``pytest.skip``.
+    Tries each pattern in order until at least one PID is killed. If none match,
+    the returned callable issues ``pytest.skip`` (same behavior as the previous
+    inline reliability test).
 
     Args:
         grep_patterns: One pattern or an ordered list of patterns.
@@ -924,45 +514,9 @@ def make_process_kill_fault_injector(
             logger.warning(
                 "[reliability][process-kill] no server process tree found; fallback to global pgrep matching"
             )
-
-        def _kill_and_wait(filtered: list[int], pattern: str, *, source: str) -> None:
-            sig = getattr(signal, signal_name, None)
-            if sig is None:
-                raise ValueError(f"Unsupported signal_name: {signal_name}")
-            for pid in filtered:
-                name, cmdline = _safe_proc_info(pid)
-                print(
-                    f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} "
-                    f"source={source} cmdline={cmdline}",
-                    flush=True,
-                )
-                os.kill(pid, sig)
-            print(
-                f"[reliability][process-kill] matched pattern={pattern!r} source={source} "
-                f"killed_pids={filtered} killed_count={len(filtered)}",
-                flush=True,
-            )
-            if post_kill_wait_seconds > 0:
-                print(
-                    f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
-                    flush=True,
-                )
-                time.sleep(post_kill_wait_seconds)
-
         for pattern in patterns:
             print(
-                f"[reliability][process-kill] trying server_tree substring pattern={pattern!r} "
-                f"signal={signal_name} limit={limit}",
-                flush=True,
-            )
-            tree_hits = _pids_in_server_tree_substring_match(server, pattern, limit=limit)
-            if tree_hits:
-                _kill_and_wait(tree_hits, pattern, source="server_tree")
-                return
-
-        for pattern in patterns:
-            print(
-                f"[reliability][process-kill] trying pgrep -f pattern={pattern!r} signal={signal_name} limit={limit}",
+                f"[reliability][process-kill] trying pattern={pattern} signal={signal_name} limit={limit}",
                 flush=True,
             )
             pids = inject_process_kill(
@@ -981,7 +535,26 @@ def make_process_kill_fault_injector(
                 )
                 continue
             if filtered:
-                _kill_and_wait(filtered, pattern, source="pgrep")
+                sig = getattr(signal, signal_name, None)
+                if sig is None:
+                    raise ValueError(f"Unsupported signal_name: {signal_name}")
+                for pid in filtered:
+                    name, cmdline = _safe_proc_info(pid)
+                    print(
+                        f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
+                        flush=True,
+                    )
+                    os.kill(pid, sig)
+                print(
+                    f"[reliability][process-kill] matched pattern={pattern} killed_pids={filtered} killed_count={len(filtered)}",
+                    flush=True,
+                )
+                if post_kill_wait_seconds > 0:
+                    print(
+                        f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
+                        flush=True,
+                    )
+                    time.sleep(post_kill_wait_seconds)
                 return
         logger.warning(
             "[reliability][process-kill] no process matched patterns=%s signal=%s limit=%s",

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -21,19 +21,13 @@ from tests.dfx.conftest import (
     resolve_oom_device_spec,
 )
 from tests.dfx.reliability.helpers import (
-    FaultInjector,
-    assert_no_server_tree_process_residual_and_gpu_release,
     extract_openai_error_contract_from_bytes,
     get_health_raw,
     inject_gpu_oom,
     make_process_kill_fault_injector,
-    make_server_root_kill_fault_injector,
-    make_server_tree_kill_fault_injector,
     post_chat_completions_raw,
     post_json_raw,
-    run_fault_injection_with_rate_load,
     stop_gpu_oom_hogs,
-    worker_residual_timeout_after_kill_signal,
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -103,73 +97,7 @@ FAULT_ERROR_KEYWORDS = (
     "500",
     "503",
 )
-PROCESS_KILL_ERROR_KEYWORDS = (
-    "timeout",
-    "did not complete within",
-    "connection",
-    "engine",
-    "orchestrator",
-    "dead",
-    "internal",
-    "500",
-    "503",
-)
-# Prefix match for vLLM-titled children (``VLLM::Worker``, ``VLLM::StageEngineCoreProc_*``, …).
-# Must be a literal substring of ``name``/``cmdline`` — no trailing whitespace (would break match).
-RUNTIME_WORKER_PATTERN = "VLLM::"
-# RFC#2366 signal x target matrix:
-# - worker: SIGTERM, SIGKILL
-# - serve-root no-load: SIGINT; SIGTERM/SIGKILL skipped (vllm issue#43060)
-# - serve-root with-load: SIGINT; SIGTERM skipped (issue#3683); SIGKILL skipped (issue#43060)
-# - serve-tree no-load: SIGTERM/SIGKILL
-# - serve-tree with-load: SIGKILL; SIGTERM skipped (issue#3683)
-_SERVE_ROOT_NON_SIGINT_SKIP = pytest.mark.skip(reason="vllm issue#43060")
-_WITH_LOAD_SIGTERM_SKIP = pytest.mark.skip(reason="issue#3683")
-SERVE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill", marks=_SERVE_ROOT_NON_SIGINT_SKIP),
-]
-SERVE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm", marks=_WITH_LOAD_SIGTERM_SKIP),
-    pytest.param("SIGINT", id="sigint", marks=_WITH_LOAD_SIGTERM_SKIP),
-    pytest.param("SIGKILL", id="sigkill", marks=_WITH_LOAD_SIGTERM_SKIP),
-]
-TREE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm", marks=_WITH_LOAD_SIGTERM_SKIP),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-WORKER_SIGNAL_FAULT_PARAMS = [
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGTERM",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigterm",
-    ),
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns=RUNTIME_WORKER_PATTERN,
-            signal_name="SIGKILL",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigkill",
-    ),
-]
-INFLIGHT_INJECTION_REQUEST_RATE = 0.3
-INFLIGHT_INJECTION_REQUEST_COUNT = 10
-
-# Post-fault chat probe: ``MAX_ELAPSED`` aligns with HTTP read timeout; assert allows +1s jitter
-# when the client raises at timeout (e.g. ~30.03s).
-_PROCESS_KILL_WORKER_CHAT_FF_HTTP_TIMEOUT_SEC = 30
-_PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC = 30
+RUNTIME_WORKER_PATTERN = "VLLM::Worker"
 
 
 class _HasServeArgs(Protocol):
@@ -266,55 +194,6 @@ def _looks_like_server_unreachable(exc: BaseException) -> bool:
     return "connection refused" in msg or "actively refused" in msg
 
 
-def _chat_request_config(omni_server: _HasServeArgs) -> dict[str, Any]:
-    messages = dummy_messages_from_mix_data(
-        system_prompt=_get_system_prompt(),
-        content_text="What is the capital of China? Answer in one short sentence.",
-    )
-    return {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": False,
-        "modalities": ["text"],
-    }
-
-
-def _assert_post_fault_chat_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
-    payload = {
-        "model": model,
-        "messages": [{"role": "user", "content": "post-fault fast-fail check"}],
-        "stream": False,
-        "modalities": ["text"],
-    }
-    start = time.monotonic()
-    try:
-        status, body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] /v1/chat/completions did not fail fast: {elapsed:.2f}s"
-        assert status >= 500, (
-            f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
-        )
-    except Exception:  # noqa: BLE001
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
-
-
-def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            if status == 503:
-                return
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-            return
-        time.sleep(0.5)
-    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
-
-
 QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
 
 
@@ -475,7 +354,17 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
@@ -496,7 +385,7 @@ def test_reliability_fault_process_kill_request_failure(
     try:
         openai_client_function.send_omni_request(request_config, request_num=1)
     except Exception as exc:
-        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+        assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
     else:
         pytest.fail("expected request failure after process-kill injection")
 
@@ -505,22 +394,48 @@ def test_reliability_fault_process_kill_request_failure(
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     omni_server_after_fault_function,
 ) -> None:
-    """Black-box: after worker process kill, chat fails fast, concurrent chat does not hang; /health→503 last.
-
-    ``SIGTERM`` on a stage engine often exceeds a 15s client-visible path; ``SIGKILL`` is
-    typically quicker. See module constants ``_PROCESS_KILL_WORKER_CHAT_FF_*``.
-    """
+    """Black-box: after worker SIGKILL, /health→503, chat fails fast, concurrent chat does not hang."""
     host = omni_server_after_fault_function.host
     port = omni_server_after_fault_function.port
     model = omni_server_after_fault_function.model
 
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    saw_503 = False
+    health_final_status: int | None = None
+    health_final_body = b""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            health_final_status, health_final_body = status, body
+            if status == 503:
+                saw_503 = True
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    assert saw_503, (
+        f"[process_kill health] expected /health to become 503 after fault injection, got {last_observation}"
+    )
+
     payload = {
         "model": model,
         "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
@@ -529,28 +444,20 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     }
     ff_status: int | None = None
     ff_body = b""
+    ff_exc: BaseException | None = None
     start = time.monotonic()
     try:
-        ff_status, ff_body = post_json_raw(
-            host,
-            port,
-            "/v1/chat/completions",
-            payload,
-            timeout_sec=_PROCESS_KILL_WORKER_CHAT_FF_HTTP_TIMEOUT_SEC,
-        )
+        ff_status, ff_body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
         elapsed = time.monotonic() - start
-        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 elapsed < 15, f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
         assert ff_status >= 500, (
             f"[process_kill fast_fail] expected server-side failure after fault, "
             f"got status={ff_status}, body={ff_body[:200]!r}"
         )
-    except Exception:
+    except Exception as exc:
+        ff_exc = exc
         elapsed = time.monotonic() - start
-        assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
-            f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
-        )
+        assert elapsed < 15, f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
 
     payload_json = json.dumps(
         {
@@ -583,195 +490,32 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     assert elapsed < 30, f"[process_kill concurrent] fault-time request convergence is too slow: {elapsed:.2f}s"
 
     fault_observed = False
+    conc_debug: list[Any] = []
     for future in done:
         try:
             status, body = future.result()
+            conc_debug.append((status, body[:200]))
             if status >= 500:
                 fault_observed = True
         except Exception as exc:
-            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+            conc_debug.append(repr(exc))
             fault_observed = True
+    # DEBUG: remove before merge
+    print(
+        health_final_status,
+        health_final_body[:200],
+        ff_status,
+        ff_body[:200],
+        ff_exc,
+        conc_debug,
+    )
     assert fault_observed, (
         "[process_kill concurrent] expected at least one request to fail after process-kill fault injection"
     )
 
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    final_health_status: int | None = None
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            final_health_status = status
-            if status == 503:
-                break
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-        time.sleep(0.5)
-    if final_health_status != 503:
-        pytest.skip("issue#3050")
-    assert final_health_status == 503, (
-        "[process_kill health] expected /health 503 after fatal fault, "
-        f"got status={final_health_status}, last_observation={last_observation}"
-    )
-
 
 @pytest.mark.slow
-@pytest.mark.skip(reason="issue#3683")
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_worker_with_load_request_failure(
-    omni_server_function,
-    openai_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),
-        inject_fault=lambda: fault_injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-
-
-@pytest.mark.slow
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    """Black-box: kill serve root without load; verify fast-fail/health/cleanup."""
-    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    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),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=300.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    """Black-box: kill server tree without load; verify fast-fail/health/cleanup."""
-    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    request_config = _chat_request_config(omni_server_function)
-    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: openai_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,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_chat_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
+@pytest.mark.skip(reason="issue#2327")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(

tests/dfx/reliability/test_reliability_wan22.py

@@ -20,25 +20,28 @@ from tests.dfx.conftest import (
     supports_video_generation,
 )
 from tests.dfx.reliability.helpers import (
-    FaultInjector,
-    assert_no_server_tree_process_residual_and_gpu_release,
     get_health_raw,
     inject_gpu_oom,
     make_process_kill_fault_injector,
-    make_server_root_kill_fault_injector,
-    make_server_tree_kill_fault_injector,
-    run_fault_injection_with_rate_load,
     stop_gpu_oom_hogs,
-    worker_residual_timeout_after_kill_signal,
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_image
 
 RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
     {
-        "test_name": "wan22_i2v_reliability_default",
+        "test_name": "wan22_t2v_reliability_default",
         "server_params": {
-            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "server_args": [
+                "--num-gpus",
+                "1",
+                "--boundary-ratio",
+                "0.875",
+                "--flow-shift",
+                "5.0",
+                "--disable-log-stats",
+            ],
         },
     }
 ]
@@ -70,114 +73,9 @@ PROCESS_KILL_ERROR_KEYWORDS = (
     "500",
     "503",
 )
-# RFC#2366 signal x target matrix (aligned with test_reliability_qwen3_omni.py):
-# - worker: SIGTERM, SIGKILL
-# - serve-root: SIGTERM, SIGINT; SIGKILL skipped (issue#3725)
-# - serve-tree no-load: SIGTERM, SIGKILL
-# - serve-tree with-load: SIGTERM, SIGKILL
-_SERVE_ROOT_SIGKILL_SKIP = pytest.mark.skip(reason="issue#3725")
-SERVE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGINT", id="sigint"),
-    pytest.param("SIGKILL", id="sigkill", marks=_SERVE_ROOT_SIGKILL_SKIP),
-]
-TREE_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-TREE_WITH_LOAD_SIGNAL_PARAMS = [
-    pytest.param("SIGTERM", id="sigterm"),
-    pytest.param("SIGKILL", id="sigkill"),
-]
-WORKER_SIGNAL_FAULT_PARAMS = [
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns="vLLM-Omni::",
-            signal_name="SIGTERM",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigterm",
-    ),
-    pytest.param(
-        make_process_kill_fault_injector(
-            grep_patterns="vLLM-Omni::",
-            signal_name="SIGKILL",
-            limit=1,
-            post_kill_wait_seconds=2.0,
-        ),
-        id="runtime_process_chain_sigkill",
-    ),
-]
 
 WAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, E2E_STAGE_CONFIGS_DIR)
 DIFFUSION_VIDEO_PARAMS = [param for param in WAN_PARAMS if supports_video_generation(param.model)]
-INFLIGHT_INJECTION_REQUEST_RATE = 0.3
-INFLIGHT_INJECTION_REQUEST_COUNT = 10
-
-
-def _video_request_config() -> dict[str, Any]:
-    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
-    return {
-        "form_data": {
-            "prompt": "Generate a realistic road-driving video with camera motion.",
-            "width": 512,
-            "height": 512,
-            "fps": 8,
-            "num_frames": 8,
-            "guidance_scale": 1.0,
-            "flow_shift": 5.0,
-            "num_inference_steps": 4,
-            "seed": 42,
-        },
-        "image_reference": image_data_url,
-        "stream": False,
-    }
-
-
-def _assert_post_fault_video_fast_fail(host: str, port: int, *, scenario: str) -> None:
-    url = f"http://{host}:{port}/v1/videos"
-    payload = {
-        "prompt": "post-fault fast-fail check",
-        "width": "512",
-        "height": "512",
-        "fps": "8",
-        "num_frames": "8",
-        "num_inference_steps": "4",
-    }
-    start = time.monotonic()
-    try:
-        response = requests.post(
-            url,
-            data=payload,
-            headers={"Accept": "application/json"},
-            timeout=20,
-        )
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] /v1/videos did not fail fast: {elapsed:.2f}s"
-        assert response.status_code >= 500, (
-            f"[{scenario} fast_fail] expected server-side error after fault, "
-            f"got status={response.status_code}, body={response.text[:200]!r}"
-        )
-    except Exception:  # noqa: BLE001
-        elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
-
-
-def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
-    deadline = time.monotonic() + 20.0
-    last_observation = ""
-    while time.monotonic() < deadline:
-        try:
-            status, body = get_health_raw(host, port, timeout_sec=5)
-            last_observation = f"http={status}, body={body[:200]!r}"
-            if status == 503:
-                return
-        except Exception as exc:  # noqa: BLE001
-            last_observation = f"exception={exc!r}"
-            return
-        time.sleep(0.5)
-    pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
 
 
 @pytest.mark.slow
@@ -224,15 +122,40 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_request_failure(
+def test_reliability_fault_process_kill_video_request_failure(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
-    request_config = _video_request_config()
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    request_config = {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
     try:
         openai_client_function.send_video_diffusion_request(request_config, request_num=1)
     except Exception as exc:
@@ -245,11 +168,21 @@ def test_reliability_fault_process_kill_request_failure(
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
+def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
@@ -357,37 +290,6 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     )
 
 
-@pytest.mark.slow
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize(
-    "fault_injector",
-    WORKER_SIGNAL_FAULT_PARAMS,
-    indirect=True,
-)
-@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_worker_with_load_request_failure(
-    omni_server_function,
-    openai_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),
-        inject_fault=lambda: fault_injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-
-
 @pytest.mark.slow
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.skip(reason="issue#2327")
@@ -507,127 +409,3 @@ def test_reliability_video_oom_recovers_after_fault_removed(
             "unhealthy terminal state should fail fast on /v1/videos, "
             f"got health={terminal_health}, request_status={recovery_status}, body={recovery_text_prefix!r}"
         )
-
-
-@pytest.mark.slow
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", 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,
-    signal_name: str,
-) -> None:
-    """Black-box: during in-flight video requests, kill serve root and verify post-fault behavior/cleanup."""
-    request_config = _video_request_config()
-    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),
-        inject_fault=lambda: injector(omni_server_function),
-        num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
-        request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    """Black-box: kill serve root without load; verify fast-fail/health/cleanup."""
-    scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
-    injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
-    omni_server_function,
-    signal_name: str,
-) -> None:
-    """Black-box: kill server tree without load; verify fast-fail/health/cleanup."""
-    scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    injector(omni_server_function)
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
-@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
-@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
-    omni_server_function,
-    openai_client_function,
-    signal_name: str,
-) -> None:
-    """Black-box: during in-flight video requests, kill server process tree and verify post-fault behavior/cleanup."""
-    request_config = _video_request_config()
-    scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
-    injector = make_server_tree_kill_fault_injector(
-        signal_name=signal_name,
-        post_kill_wait_seconds=2.0,
-        inter_kill_wait_seconds=0.1,
-    )
-    load_result = run_fault_injection_with_rate_load(
-        submit_request=lambda: openai_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,
-        completion_timeout_sec=120.0,
-    )
-    assert load_result["failure_observed"], (
-        f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
-    )
-    host = omni_server_function.host
-    port = omni_server_function.port
-    _assert_post_fault_video_fast_fail(host, port, scenario=scenario)
-    _assert_post_fault_health_terminal(host, port, scenario=scenario)
-    assert_no_server_tree_process_residual_and_gpu_release(
-        omni_server_function,
-        scenario=scenario,
-        timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
-    )

tests/diffusion/attention/test_sage_attn3.py

@@ -1,131 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import importlib
-import sys
-import types
-
-import pytest
-import torch
-from vllm.platforms import current_platform
-
-SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
-
-
-def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
-    fake_module = types.ModuleType("sageattn3")
-    fake_module.sageattn3_blackwell = kernel_impl
-    monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
-    sys.modules.pop(SAGE_ATTN3_MODULE, None)
-    return importlib.import_module(SAGE_ATTN3_MODULE)
-
-
-def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatch):
-    calls = {}
-
-    def fake_kernel(query, key, value, is_causal=False):
-        calls["query_shape"] = query.shape
-        calls["is_causal"] = is_causal
-        return query + key + value
-
-    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
-    impl = backend_module.SageAttention3Impl(
-        num_heads=4,
-        head_size=64,
-        softmax_scale=1.0 / 8.0,
-        causal=False,
-    )
-
-    query = torch.randn(2, 8, 4, 64)
-    key = torch.randn(2, 8, 4, 64)
-    value = torch.randn(2, 8, 4, 64)
-
-    output = impl.forward_cuda(query, key, value)
-
-    assert calls["query_shape"] == (2, 4, 8, 64)
-    assert calls["is_causal"] is False
-    expected = (query.transpose(1, 2) + key.transpose(1, 2) + value.transpose(1, 2)).transpose(1, 2)
-    assert torch.allclose(output, expected)
-
-
-def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
-    def fake_kernel(*args, **kwargs):
-        raise AssertionError("sageattn3_blackwell should not be used for GQA")
-
-    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
-    sdpa_calls = {}
-
-    def fake_sdpa(query, key, value, **kwargs):
-        sdpa_calls["query_shape"] = query.shape
-        sdpa_calls["key_shape"] = key.shape
-        sdpa_calls["enable_gqa"] = kwargs["enable_gqa"]
-        return query + 1
-
-    monkeypatch.setattr(backend_module.F, "scaled_dot_product_attention", fake_sdpa)
-
-    impl = backend_module.SageAttention3Impl(
-        num_heads=4,
-        head_size=64,
-        softmax_scale=1.0 / 8.0,
-        causal=False,
-    )
-
-    query = torch.randn(2, 8, 4, 64)
-    key = torch.randn(2, 8, 2, 64)
-    value = torch.randn(2, 8, 2, 64)
-
-    output = impl.forward_cuda(query, key, value)
-
-    assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
-    assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
-    assert sdpa_calls["enable_gqa"] is True
-    expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
-    assert torch.allclose(output, expected)
-
-
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
-def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch):
-    from vllm.platforms.interface import DeviceCapability
-
-    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
-    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
-    from vllm_omni.platforms.cuda import platform as cuda_platform_module
-    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
-
-    original_import_module = importlib.import_module
-
-    monkeypatch.setattr(
-        CudaOmniPlatform,
-        "get_device_capability",
-        classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
-    )
-    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
-    monkeypatch.setattr(
-        cuda_platform_module.importlib,
-        "import_module",
-        lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name),
-    )
-
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
-
-    assert backend_path == DiffusionAttentionBackendEnum.SAGE_ATTN_3.get_path()
-
-
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
-def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
-    from vllm.platforms.interface import DeviceCapability
-
-    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
-    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
-    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
-
-    monkeypatch.setattr(
-        CudaOmniPlatform,
-        "get_device_capability",
-        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
-    )
-    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
-
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
-
-    assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()

tests/diffusion/cache/test_cache_backends.py

@@ -19,7 +19,6 @@ from vllm_omni.diffusion.cache.cache_dit_backend import (
     CUSTOM_DIT_ENABLERS,
     CacheDiTBackend,
 )
-from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
 from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
@@ -314,150 +313,3 @@ class TestCacheSelector:
         """Test getting invalid backend raises error."""
         with pytest.raises(ValueError, match="Unsupported cache backend"):
             get_cache_backend("invalid_backend", {})
-
-
-class TestMagCacheBackend:
-    """Test MagCacheBackend implementation."""
-
-    from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
-
-    def test_init(self):
-        """Test initialization."""
-        config = DiffusionCacheConfig(mag_threshold=0.1, mag_max_skip_steps=2, mag_calibrate=True)
-        backend = MagCacheBackend(config)
-        assert backend.config.mag_threshold == 0.1
-        assert backend.config.mag_max_skip_steps == 2
-        assert backend.enabled is False
-
-    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
-    def test_enable(self, mock_apply_hook):
-        """Test enabling MagCache on pipeline."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "FluxPipeline"
-        mock_transformer = Mock()
-        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
-        mock_pipeline.transformer = mock_transformer
-
-        mock_ratios = [1.0] * 28
-        config = DiffusionCacheConfig(
-            mag_ratios=mock_ratios,
-        )
-        backend = MagCacheBackend(config)
-        backend.enable(mock_pipeline)
-
-        assert backend.enabled is True
-        mock_apply_hook.assert_called_once()
-
-        call_args = mock_apply_hook.call_args
-        assert call_args[0][0] == mock_transformer
-
-    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
-    def test_enable_with_calibration(self, mock_apply_hook):
-        """Test enabling MagCache in calibration mode."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "FluxPipeline"
-        mock_transformer = Mock()
-        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
-        mock_pipeline.transformer = mock_transformer
-
-        config = DiffusionCacheConfig(
-            mag_calibrate=True,
-        )
-        backend = MagCacheBackend(config)
-        backend.enable(mock_pipeline)
-
-        assert backend.enabled is True
-        mock_apply_hook.assert_called_once()
-
-    def test_refresh(self):
-        """Test refreshing MagCache state calls enable when not registered."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "FluxPipeline"
-        mock_transformer = Mock()
-        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
-        mock_pipeline.transformer = mock_transformer
-
-        mock_transformer.named_children = Mock(return_value=[])
-
-        mock_ratios = [1.0] * 28
-        config = DiffusionCacheConfig(
-            mag_ratios=mock_ratios,
-        )
-        backend = MagCacheBackend(config)
-
-        assert backend._registered is False
-        backend.refresh(mock_pipeline, num_inference_steps=50)
-        assert backend._registered is True
-
-    def test_is_enabled(self):
-        """Test is_enabled method."""
-        mock_ratios = [1.0] * 28
-        config = DiffusionCacheConfig(mag_ratios=mock_ratios)
-        backend = MagCacheBackend(config)
-        assert backend.is_enabled() is False
-
-    def test_get_mag_cache_backend(self):
-        """Test getting MagCache backend via selector."""
-        mock_ratios = [1.0] * 28
-        config_dict = {
-            "mag_ratios": mock_ratios,
-            "num_inference_steps": 28,
-            "threshold": 0.06,
-            "max_skip_steps": 3,
-            "retention_ratio": 0.2,
-        }
-        backend = get_cache_backend("mag_cache", config_dict)
-        assert backend is not None
-        assert isinstance(backend, MagCacheBackend)
-        assert backend.config.threshold == 0.06
-
-    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
-    def test_enable_single_block(self, mock_apply_hook):
-        """Test enabling MagCache on single transformer block."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "FluxPipeline"
-
-        mock_block = Mock()
-        mock_block.__class__.__name__ = "FluxTransformer2DModel"
-        mock_blocks = [mock_block]
-
-        mock_transformer = Mock()
-        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
-        mock_transformer.blocks = mock_blocks
-        mock_pipeline.transformer = mock_transformer
-
-        mock_ratios = [1.0] * 28
-        config = DiffusionCacheConfig(
-            mag_ratios=mock_ratios,
-        )
-        backend = MagCacheBackend(config)
-        backend.enable(mock_pipeline)
-
-        assert backend.enabled is True
-        mock_apply_hook.assert_called_once()
-
-        call_args = mock_apply_hook.call_args
-        assert call_args[0][0] == mock_transformer
-
-    @patch("vllm_omni.diffusion.cache.magcache.backend.apply_mag_cache_hook")
-    def test_enable_multi_block(self, mock_apply_hook):
-        """Test enabling MagCache on multiple transformer blocks."""
-        mock_pipeline = Mock()
-        mock_pipeline.__class__.__name__ = "FluxPipeline"
-
-        mock_blocks = [Mock() for _ in range(24)]
-
-        mock_transformer = Mock()
-        mock_transformer.__class__.__name__ = "FluxTransformer2DModel"
-        mock_transformer.blocks = mock_blocks
-        mock_pipeline.transformer = mock_transformer
-
-        mock_ratios = [1.0] * 28
-        config = DiffusionCacheConfig(
-            mag_ratios=mock_ratios,
-        )
-        backend = MagCacheBackend(config)
-        backend.enable(mock_pipeline)
-
-        assert backend.enabled is True
-        mock_apply_hook.assert_called_once()

tests/diffusion/cache/test_prompt_embed_cache.py

@@ -1,529 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Unit tests for prompt-embedding cache for diffusion pipelines.
-
-Covers the module :mod:`vllm_omni.diffusion.cache.prompt_embed_cache`:
-
-* :class:`PromptEmbedCache` LRU semantics, stats and thread-safety.
-* :func:`_hashable` safe-input classification (scalars, torch dtype/device,
-  numpy scalars, nested containers, tensors/PIL bypass).
-* :func:`install_prompt_embed_cache` / :func:`uninstall_prompt_embed_cache`
-  behaviour: caching, bypass on tensors, bypass on precomputed embeds,
-  idempotent install, restore on uninstall, detachment of cached tensors.
-* :func:`resolve_prompt_embed_cache_config` env-var overrides.
-"""
-
-from __future__ import annotations
-
-import threading
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.cache.prompt_embed_cache import (
-    _CACHE_MISS,
-    _NOT_HASHABLE,
-    PromptEmbedCache,
-    _build_key,
-    _detach_output,
-    _hashable,
-    install_prompt_embed_cache,
-    resolve_prompt_embed_cache_config,
-    uninstall_prompt_embed_cache,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-# ---------------------------------------------------------------------------
-# _hashable
-# ---------------------------------------------------------------------------
-
-
-class TestHashable:
-    """Tests for the ``_hashable`` helper."""
-
-    @pytest.mark.parametrize(
-        "value",
-        [None, "prompt", "", 0, 1, -5, 3.14, True, False, b"bytes"],
-    )
-    def test_scalars_are_returned_unchanged(self, value):
-        assert _hashable(value) == value
-
-    def test_torch_device_is_normalized(self):
-        out = _hashable(torch.device("cpu"))
-        assert out == ("__torch_device__", "cpu")
-        assert hash(out) is not None  # must actually be hashable
-
-    def test_torch_dtype_is_normalized(self):
-        out = _hashable(torch.float32)
-        assert out[0] == "__torch_dtype__"
-        assert "float32" in out[1]
-        assert hash(out) is not None
-
-    def test_numpy_scalar_is_unwrapped(self):
-        np = pytest.importorskip("numpy")
-        assert _hashable(np.int64(7)) == 7
-        assert _hashable(np.float32(1.5)) == pytest.approx(1.5)
-
-    def test_numpy_ndarray_is_not_hashable(self):
-        np = pytest.importorskip("numpy")
-        assert _hashable(np.array([1, 2, 3])) is _NOT_HASHABLE
-
-    def test_nested_list_and_tuple(self):
-        out = _hashable(["a", ("b", 1), [2, 3]])
-        assert out[0] == "__seq__"
-        # fully hashable
-        assert hash(out) is not None
-
-    def test_dict_is_normalized_by_sorted_keys(self):
-        a = _hashable({"b": 1, "a": 2})
-        b = _hashable({"a": 2, "b": 1})
-        assert a == b
-        assert hash(a) is not None
-
-    def test_tensor_is_not_hashable(self):
-        assert _hashable(torch.zeros(2)) is _NOT_HASHABLE
-
-    def test_list_containing_tensor_is_not_hashable(self):
-        assert _hashable(["a", torch.zeros(1)]) is _NOT_HASHABLE
-
-    def test_dict_containing_tensor_is_not_hashable(self):
-        assert _hashable({"x": torch.zeros(1)}) is _NOT_HASHABLE
-
-    def test_custom_object_is_not_hashable(self):
-        class Obj:
-            pass
-
-        assert _hashable(Obj()) is _NOT_HASHABLE
-
-
-# ---------------------------------------------------------------------------
-# _detach_output
-# ---------------------------------------------------------------------------
-
-
-class TestDetachOutput:
-    def test_detaches_tensor(self):
-        t = torch.zeros(2, requires_grad=True)
-        out = _detach_output(t)
-        assert isinstance(out, torch.Tensor)
-        assert not out.requires_grad
-
-    def test_detaches_inside_tuple(self):
-        t = torch.ones(1, requires_grad=True)
-        out = _detach_output((t, None))
-        assert isinstance(out, tuple)
-        assert not out[0].requires_grad
-        assert out[1] is None
-
-    def test_detaches_inside_list_and_dict(self):
-        t = torch.ones(1, requires_grad=True)
-        out = _detach_output({"a": [t]})
-        assert not out["a"][0].requires_grad
-
-    def test_preserves_namedtuple_type(self):
-        from collections import namedtuple
-
-        NT = namedtuple("NT", ["x", "y"])
-        t = torch.ones(1, requires_grad=True)
-        out = _detach_output(NT(x=t, y=1))
-        assert isinstance(out, NT)
-        assert not out.x.requires_grad
-        assert out.y == 1
-
-    def test_passthrough_for_non_tensor(self):
-        assert _detach_output("hi") == "hi"
-        assert _detach_output(None) is None
-        assert _detach_output(5) == 5
-
-
-# ---------------------------------------------------------------------------
-# PromptEmbedCache
-# ---------------------------------------------------------------------------
-
-
-class TestPromptEmbedCache:
-    def test_rejects_non_positive_max_size(self):
-        with pytest.raises(ValueError):
-            PromptEmbedCache(max_size=0)
-        with pytest.raises(ValueError):
-            PromptEmbedCache(max_size=-1)
-
-    def test_miss_returns_sentinel_and_bumps_miss_counter(self):
-        cache = PromptEmbedCache(max_size=2)
-        assert cache.get("k") is _CACHE_MISS
-        assert cache.stats()["misses"] == 1
-        assert cache.stats()["hits"] == 0
-
-    def test_put_then_get_returns_value_and_bumps_hit_counter(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("k", 42)
-        assert cache.get("k") == 42
-        assert cache.stats()["hits"] == 1
-
-    def test_none_is_a_cacheable_value(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("k", None)
-        # Distinguishable from a miss.
-        assert cache.get("k") is None
-        assert cache.stats()["hits"] == 1
-
-    def test_lru_eviction_keeps_recently_used(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("a", 1)
-        cache.put("b", 2)
-        # Access "a" to mark it MRU.
-        assert cache.get("a") == 1
-        cache.put("c", 3)  # should evict "b"
-        assert cache.get("b") is _CACHE_MISS
-        assert cache.get("a") == 1
-        assert cache.get("c") == 3
-        assert cache.stats()["size"] == 2
-
-    def test_get_refreshes_recency(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("a", 1)
-        cache.put("b", 2)
-        # Hit on "a" → "b" becomes LRU.
-        cache.get("a")
-        cache.put("c", 3)
-        assert cache.get("b") is _CACHE_MISS
-        assert cache.get("a") == 1
-
-    def test_put_overwrites_existing_key(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("k", 1)
-        cache.put("k", 2)
-        assert cache.get("k") == 2
-        assert cache.stats()["size"] == 1
-
-    def test_clear_empties_store_but_keeps_counters(self):
-        cache = PromptEmbedCache(max_size=2)
-        cache.put("k", 1)
-        cache.get("k")
-        cache.clear()
-        assert cache.stats()["size"] == 0
-        # clear only drops entries, counters reflect prior activity.
-        assert cache.stats()["hits"] == 1
-
-    def test_stats_reports_expected_fields(self):
-        cache = PromptEmbedCache(max_size=4)
-        stats = cache.stats()
-        assert set(stats.keys()) == {"size", "max_size", "hits", "misses", "bypassed"}
-        assert stats["max_size"] == 4
-        assert stats["size"] == 0
-
-    def test_thread_safety_under_concurrent_writes(self):
-        cache = PromptEmbedCache(max_size=1024)
-
-        def worker(start: int) -> None:
-            for i in range(start, start + 200):
-                cache.put(f"k{i}", i)
-                cache.get(f"k{i}")
-
-        threads = [threading.Thread(target=worker, args=(i * 200,)) for i in range(4)]
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
-
-        stats = cache.stats()
-        # 4 * 200 = 800 writes, each followed by a successful read.
-        assert stats["hits"] == 800
-        assert stats["size"] == 800
-
-
-# ---------------------------------------------------------------------------
-# _build_key
-# ---------------------------------------------------------------------------
-
-
-def _make_sig(fn):
-    import inspect
-
-    return inspect.signature(fn)
-
-
-class TestBuildKey:
-    def test_positional_and_keyword_collide(self):
-        def encode(prompt, device=None):
-            return None
-
-        sig = _make_sig(encode)
-        k1 = _build_key(sig, "m", ("hello",), {"device": "cpu"})
-        k2 = _build_key(sig, "m", (), {"prompt": "hello", "device": "cpu"})
-        assert k1 == k2
-        assert k1 is not None
-
-    def test_different_prompts_produce_different_keys(self):
-        def encode(prompt):
-            return None
-
-        sig = _make_sig(encode)
-        assert _build_key(sig, "m", ("a",), {}) != _build_key(sig, "m", ("b",), {})
-
-    def test_different_model_tags_produce_different_keys(self):
-        def encode(prompt):
-            return None
-
-        sig = _make_sig(encode)
-        assert _build_key(sig, "m1", ("a",), {}) != _build_key(sig, "m2", ("a",), {})
-
-    def test_tensor_argument_bypasses(self):
-        def encode(prompt, extra=None):
-            return None
-
-        sig = _make_sig(encode)
-        assert _build_key(sig, "m", ("hi",), {"extra": torch.zeros(1)}) is None
-
-    def test_precomputed_prompt_embeds_bypasses(self):
-        def encode(prompt=None, prompt_embeds=None):
-            return None
-
-        sig = _make_sig(encode)
-        # None → not a bypass.
-        assert _build_key(sig, "m", (), {"prompt": "hi"}) is not None
-        # A non-None precomputed embed triggers bypass even if it would
-        # otherwise be "hashable".
-        assert _build_key(sig, "m", (), {"prompt": "hi", "prompt_embeds": "x"}) is None
-
-    def test_invalid_binding_returns_none(self):
-        def encode(prompt):
-            return None
-
-        sig = _make_sig(encode)
-        # Unknown kwarg → TypeError on bind → bypass.
-        assert _build_key(sig, "m", (), {"nope": 1}) is None
-
-
-# ---------------------------------------------------------------------------
-# install_prompt_embed_cache / uninstall_prompt_embed_cache
-# ---------------------------------------------------------------------------
-
-
-class _FakePipeline:
-    """Minimal pipeline with a signature-bearing ``encode_prompt``."""
-
-    def __init__(self):
-        self.call_count = 0
-
-    def encode_prompt(
-        self,
-        prompt: str,
-        device=None,
-        num_images_per_prompt: int = 1,
-        negative_prompt: str | None = None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-    ):
-        self.call_count += 1
-        # Return a fresh tensor so identity changes between real calls.
-        return torch.tensor([float(self.call_count)])
-
-
-class TestInstallAndUninstall:
-    def test_returns_none_when_pipeline_has_no_encode_prompt(self):
-        class Empty:
-            pass
-
-        assert install_prompt_embed_cache(Empty()) is None
-
-    def test_cache_hits_on_identical_args(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe, max_size=4)
-        assert cache is not None
-
-        out1 = pipe.encode_prompt("cat", device="cpu")
-        out2 = pipe.encode_prompt("cat", device="cpu")
-        # Second call must be served from cache (wrapped fn not invoked).
-        assert pipe.call_count == 1
-        # Cache returns the same tensor object.
-        assert torch.equal(out1, out2)
-        stats = cache.stats()
-        assert stats["hits"] == 1
-        assert stats["misses"] == 1
-        assert stats["bypassed"] == 0
-
-    def test_positional_and_keyword_share_cache_slot(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe)
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt(prompt="cat")
-        assert pipe.call_count == 1
-        assert cache.stats()["hits"] == 1
-
-    def test_different_prompts_miss(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe)
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt("dog")
-        assert pipe.call_count == 2
-        assert cache.stats()["misses"] == 2
-
-    def test_tensor_argument_bypasses_and_increments_counter(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe)
-        # Put a tensor into a non-precomputed-embed slot to trigger bypass.
-        pipe.encode_prompt("cat", device=torch.device("cpu"), num_images_per_prompt=1)
-
-        # Now pass a fake tensor via a harmless positional-like call: use
-        # the negative_prompt slot with an unhashable object.
-        class Unhashable:
-            pass
-
-        pipe.encode_prompt("cat", negative_prompt=Unhashable())  # type: ignore[arg-type]
-        pipe.encode_prompt("cat", negative_prompt=Unhashable())  # type: ignore[arg-type]
-        assert cache.stats()["bypassed"] == 2
-        # All three calls actually executed the underlying function.
-        assert pipe.call_count == 3
-
-    def test_precomputed_embeds_bypass_cache(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe)
-        pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
-        pipe.encode_prompt("cat", prompt_embeds=torch.zeros(1))
-        assert cache.stats()["bypassed"] == 2
-        assert cache.stats()["hits"] == 0
-        assert pipe.call_count == 2
-
-    def test_install_is_idempotent(self):
-        pipe = _FakePipeline()
-        c1 = install_prompt_embed_cache(pipe, max_size=4)
-        c2 = install_prompt_embed_cache(pipe, max_size=16)
-        assert c1 is c2
-        # The second call must not re-wrap (max_size stays from first install).
-        assert c1.max_size == 4
-
-    def test_disabled_cache_passes_through(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe, enabled=False)
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt("cat")
-        assert pipe.call_count == 2
-        assert cache.stats()["hits"] == 0
-        assert cache.stats()["misses"] == 0
-
-    def test_toggle_enabled_flag_takes_effect(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe, enabled=True)
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt("cat")
-        assert cache.stats()["hits"] == 1
-
-        cache.enabled = False
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt("cat")
-        # Cache-bypassed calls go through; hit count unchanged.
-        assert cache.stats()["hits"] == 1
-
-    def test_cached_tensor_is_detached(self):
-        class GradPipeline:
-            def encode_prompt(self, prompt):
-                return torch.ones(1, requires_grad=True)
-
-        pipe = GradPipeline()
-        install_prompt_embed_cache(pipe)
-        # First call populates the cache.
-        pipe.encode_prompt("cat")
-        # Second call returns the cached tensor which must be detached.
-        cached = pipe.encode_prompt("cat")
-        assert isinstance(cached, torch.Tensor)
-        assert not cached.requires_grad
-
-    def test_lru_eviction_through_wrapper(self):
-        pipe = _FakePipeline()
-        cache = install_prompt_embed_cache(pipe, max_size=2)
-        pipe.encode_prompt("a")
-        pipe.encode_prompt("b")
-        pipe.encode_prompt("c")  # evicts "a"
-        pipe.encode_prompt("a")  # miss again → re-run
-        assert pipe.call_count == 4
-        assert cache.stats()["size"] == 2
-
-    def test_uninstall_restores_original(self):
-        pipe = _FakePipeline()
-        original = pipe.encode_prompt
-        install_prompt_embed_cache(pipe)
-        assert pipe.encode_prompt is not original
-        uninstall_prompt_embed_cache(pipe)
-        # The bound method reference will differ, but identity of underlying
-        # function matches and wrapper attributes are gone.
-        assert not hasattr(pipe, "_prompt_embed_cache") or pipe._prompt_embed_cache is None
-        # After uninstall calls go straight to the real function.
-        pipe.encode_prompt("cat")
-        pipe.encode_prompt("cat")
-        assert pipe.call_count == 2
-
-    def test_uninstall_on_unwrapped_pipeline_is_noop(self):
-        class Empty:
-            def encode_prompt(self, prompt):
-                return None
-
-        pipe = Empty()
-        uninstall_prompt_embed_cache(pipe)  # must not raise
-
-
-# ---------------------------------------------------------------------------
-# resolve_prompt_embed_cache_config
-# ---------------------------------------------------------------------------
-
-
-@pytest.fixture
-def clean_env(monkeypatch):
-    monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raising=False)
-    monkeypatch.delenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", raising=False)
-
-
-class TestResolveConfig:
-    def test_defaults_when_nothing_set(self, clean_env):
-        enable, size = resolve_prompt_embed_cache_config()
-        assert enable is False
-        assert size == 32
-
-    def test_explicit_args_used(self, clean_env):
-        enable, size = resolve_prompt_embed_cache_config(enable=True, max_size=8)
-        assert enable is True
-        assert size == 8
-
-    @pytest.mark.parametrize("raw", ["1", "true", "TRUE", "yes", "on"])
-    def test_env_enable_truthy(self, clean_env, monkeypatch, raw):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
-        enable, _ = resolve_prompt_embed_cache_config(enable=False)
-        assert enable is True
-
-    @pytest.mark.parametrize("raw", ["0", "false", "FALSE", "no", "off"])
-    def test_env_enable_falsy_overrides_true(self, clean_env, monkeypatch, raw):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", raw)
-        enable, _ = resolve_prompt_embed_cache_config(enable=True)
-        assert enable is False
-
-    def test_env_enable_invalid_leaves_arg_untouched(self, clean_env, monkeypatch):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE", "garbage")
-        enable, _ = resolve_prompt_embed_cache_config(enable=True)
-        assert enable is True
-
-    def test_env_size_overrides_arg(self, clean_env, monkeypatch):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "7")
-        _, size = resolve_prompt_embed_cache_config(max_size=99)
-        assert size == 7
-
-    def test_env_size_zero_or_negative_ignored(self, clean_env, monkeypatch):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "0")
-        _, size = resolve_prompt_embed_cache_config(max_size=12)
-        assert size == 12
-
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "-5")
-        _, size = resolve_prompt_embed_cache_config(max_size=12)
-        assert size == 12
-
-    def test_env_size_non_integer_ignored(self, clean_env, monkeypatch):
-        monkeypatch.setenv("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE", "notanint")
-        _, size = resolve_prompt_embed_cache_config(max_size=12)
-        assert size == 12
-
-    def test_max_size_none_falls_back_to_default(self, clean_env):
-        _, size = resolve_prompt_embed_cache_config(enable=True, max_size=None)
-        assert size == 32

tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py

@@ -7,9 +7,6 @@ import torch
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
     DistributedAutoencoderKLHunyuan,
 )
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_15 import (
-    DistributedAutoencoderKLHunyuanVideo15,
-)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -39,33 +36,6 @@ class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
         return b
 
 
-class _DummyDistributedAutoencoderKLHunyuanVideo15(DistributedAutoencoderKLHunyuanVideo15):
-    def __init__(self):
-        torch.nn.Module.__init__(self)
-        self.tile_sample_min_height = 8
-        self.tile_sample_min_width = 12
-        self.tile_latent_min_height = 8
-        self.tile_latent_min_width = 12
-        self.tile_overlap_factor = 0.25
-        self.use_tiling = False
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return torch.float32
-
-    def decoder(self, x: torch.Tensor) -> torch.Tensor:
-        return x + 10
-
-    def encoder(self, x: torch.Tensor) -> torch.Tensor:
-        return x + 20
-
-    def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
-        return b
-
-    def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
-        return b
-
-
 def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
     # Verify use_tiling property maps to use_spatial_tiling.
     vae = _DummyDistributedAutoencoderKLHunyuan()
@@ -107,43 +77,3 @@ def test_hunyuan_vae_encode_tiles_round_trip():
     assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
     assert len(tile_tasks) == 4
     assert torch.equal(output, x + 20)
-
-
-def test_hunyuan_video15_vae_decode_tiles_round_trip():
-    vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
-    z = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
-
-    tile_tasks, grid_spec = vae.tile_split(z)
-    decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
-    output = vae.tile_merge(decoded_tiles, grid_spec)
-
-    assert grid_spec.split_dims == (3, 4)
-    assert grid_spec.grid_shape == (2, 2)
-    assert grid_spec.tile_spec == {
-        "blend_height": 2,
-        "blend_width": 3,
-        "row_limit_height": 6,
-        "row_limit_width": 9,
-    }
-    assert len(tile_tasks) == 4
-    assert torch.equal(output, z + 10)
-
-
-def test_hunyuan_video15_vae_encode_tiles_round_trip():
-    vae = _DummyDistributedAutoencoderKLHunyuanVideo15()
-    x = torch.arange(192, dtype=torch.float32).reshape(1, 1, 1, 12, 16)
-
-    tile_tasks, grid_spec = vae.encode_tile_split(x)
-    encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
-    output = vae.tile_merge(encoded_tiles, grid_spec)
-
-    assert grid_spec.split_dims == (3, 4)
-    assert grid_spec.grid_shape == (2, 2)
-    assert grid_spec.tile_spec == {
-        "blend_height": 2,
-        "blend_width": 3,
-        "row_limit_height": 6,
-        "row_limit_width": 9,
-    }
-    assert len(tile_tasks) == 4
-    assert torch.equal(output, x + 20)

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -0,0 +1,459 @@
+# ruff: noqa: N803, E741
+"""Layer-level correctness & performance test for MoT parallel linear layers.
+
+Compares two equivalent computation paths:
+  - Reference: 2x standard vLLM parallel linear layers + PyTorch index
+    scatter/gather (text_linear(x[text_idx]) + vae_linear(x[vae_idx]))
+  - Target:    1x MoT fused parallel linear layer
+    (mot_linear(x, text_indices, vae_indices))
+
+The reference path uses cuBLAS GEMM (always auto-tuned by cuBLAS).
+The MoT path uses a fused Triton kernel whose tile config is loaded
+from a JSON file matched by ``device + dtype``.  If no tuned
+config is found for the current GPU, the kernel falls back to a conservative
+default and a warning is printed — the correctness test still passes
+but the performance comparison is NOT representative.
+
+Usage::
+    pytest tests/diffusion/kernels/mot/test_mot_linear.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass
+
+import pytest
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+    model_parallel_is_initialized,
+)
+from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import (
+    MoTQKVParallelLinear,
+)
+from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import (
+    MoTRowParallelLinear,
+)
+from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+    get_best_mot_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+# BAGEL-7B-MoT architecture parameters
+_BAGEL_HEAD_SIZE = 128
+_BAGEL_TOTAL_NUM_HEADS = 28
+_BAGEL_TOTAL_NUM_KV_HEADS = 4
+_VAE_CHUNK_SIZE = 1024  # the token number of one image
+_IMAGE_NUM = [1, 2, 4, 8]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _init_single_rank_tp_env():
+    """Initialize single-rank distributed/TP env for vLLM linear params."""
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29501")
+
+    if not torch.distributed.is_initialized():
+        init_distributed_environment(world_size=1, rank=0, local_rank=0)
+
+    if not model_parallel_is_initialized():
+        initialize_model_parallel(
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            sequence_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+        )
+
+    yield
+    destroy_distributed_env()
+
+
+# ---------------------------------------------------------------------------
+# DType configuration — extensible for future quantized types
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DTypeConfig:
+    """Parsed dtype configuration for a test case."""
+
+    torch_dtype: torch.dtype
+    use_fp8_w8a8: bool = False
+    use_int8_w8a16: bool = False
+    use_int4_w4a16: bool = False
+
+
+def _parse_dtype(dtype_str: str) -> DTypeConfig:
+    """Parse a dtype string into quantization flags and torch dtype.
+
+    Supported now:
+        "w16a16_bf16"  — BF16 weights & activations (no quantization)
+        "w16a16_fp16"  — FP16 weights & activations (no quantization)
+    Reserved for future:
+        "fp8_w8a8"     — FP8 W8A8 quantization
+        "int8_w8a16"   — INT8 weight-only quantization
+        "int4_w4a16"   — INT4 weight-only quantization
+    """
+    supported: dict[str, DTypeConfig] = {
+        "w16a16_bf16": DTypeConfig(torch_dtype=torch.bfloat16),
+        "w16a16_fp16": DTypeConfig(torch_dtype=torch.float16),
+    }
+    if dtype_str in supported:
+        return supported[dtype_str]
+    pytest.skip(f"Quantized dtype '{dtype_str}' not yet implemented in layer test")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _report_mot_config(K: int, N: int, M: int):
+    """Print which Triton tile config the MoT kernel will use."""
+    loaded_m_key, config = get_best_mot_config(M, N, K, None)
+    if loaded_m_key == -1:
+        print(
+            "  [config] WARNING: No tuned config found — "
+            "using conservative default. "
+            "Performance numbers are NOT representative. "
+            "Run mot_linear_benchmarks.py --tune to generate configs."
+        )
+    else:
+        print(f"  [config] Tuned config loaded (actual M={M}, loaded M={loaded_m_key}) config = {config})")
+
+
+def _make_indices(image_num: int, vae_chunk_size: int, device: str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Simulate exact Bagel-MoT distributions for image generation:
+    Pattern per image like: [1 Text] + [4096 VAE] + [1 Text]
+    Returns text_indices, vae_indices, and the exact total M.
+    """
+    text_idx_list = []
+    vae_idx_list = []
+
+    current_idx = 0
+    for _ in range(image_num):
+        text_idx_list.append(current_idx)
+        current_idx += 1
+
+        vae_idx_list.extend(range(current_idx, current_idx + vae_chunk_size))
+        current_idx += vae_chunk_size
+
+        text_idx_list.append(current_idx)
+        current_idx += 1
+
+    text_indices = torch.tensor(text_idx_list, dtype=torch.long, device=device)
+    vae_indices = torch.tensor(vae_idx_list, dtype=torch.long, device=device)
+
+    exact_M = current_idx  # exact_M = image_num * (vae_chunk_size + 2)
+
+    return text_indices, vae_indices, exact_M
+
+
+def _benchmark(fn, warmup: int = 20, iters: int = 100) -> float:
+    """Return mean latency in milliseconds."""
+    cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda")
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+
+    t_flush_start = time.perf_counter()
+    for _ in range(iters):
+        cache_flusher.zero_()
+    torch.accelerator.synchronize()
+    flush_time_total = time.perf_counter() - t_flush_start
+
+    #  3. Measure the total time of "flush cache + operator execution"
+    torch.accelerator.synchronize()
+    t_total_start = time.perf_counter()
+    for _ in range(iters):
+        cache_flusher.zero_()
+        fn()
+    torch.accelerator.synchronize()
+    total_time = time.perf_counter() - t_total_start
+
+    # 4. Asynchronous subtraction separation
+    # The E2E time of the pure operator = total time - flush time
+    pure_fn_time_total = total_time - flush_time_total
+    avg_ms = (pure_fn_time_total / iters) * 1000.0
+
+    del cache_flusher
+    return max(avg_ms, 0.001)
+
+
+def _sync_weights(ref_text, ref_vae, mot_layer):
+    """Assign the same random weights to reference layers and MoT layer."""
+    with torch.no_grad():
+        W_text = torch.randn_like(ref_text.weight) * 0.02
+        W_vae = torch.randn_like(ref_vae.weight) * 0.02
+        ref_text.weight.copy_(W_text)
+        ref_vae.weight.copy_(W_vae)
+        mot_layer.weight.copy_(W_text)
+        mot_layer.gen_exp.weight.copy_(W_vae)
+        if ref_text.bias is not None and mot_layer.bias is not None:
+            b_text = torch.randn_like(ref_text.bias) * 0.02
+            ref_text.bias.copy_(b_text)
+            mot_layer.bias.copy_(b_text)
+        if ref_vae.bias is not None and mot_layer.gen_exp.bias is not None:
+            b_vae = torch.randn_like(ref_vae.bias) * 0.02
+            ref_vae.bias.copy_(b_vae)
+            mot_layer.gen_exp.bias.copy_(b_vae)
+
+
+def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
+    """Reference path: index-gather → 2x standard linear → index-scatter."""
+    M = x.size(0)
+    out_text = text_linear(x[text_indices])
+    out_vae = vae_linear(x[vae_indices])
+    if isinstance(out_text, tuple):
+        out_text = out_text[0]
+    if isinstance(out_vae, tuple):
+        out_vae = out_vae[0]
+    N = out_text.size(-1)
+    output = torch.empty(M, N, dtype=x.dtype, device=x.device)
+    output[text_indices] = out_text
+    output[vae_indices] = out_vae
+    return output
+
+
+def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
+    """Compare outputs, print metrics, and assert correctness.
+
+    Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16).
+    We upcast to fp32 solely for computing error metrics with higher
+    arithmetic precision — the actual layer outputs remain bf16.
+    """
+    # Upcast for numerically stable error computation only
+    ref_hp = ref.float()
+    mot_hp = mot.float()
+
+    abs_err = (ref_hp - mot_hp).abs()
+    max_abs = abs_err.max().item()
+
+    # Mixed metric: relative error where |ref| >= 1, absolute error otherwise
+    denom = ref_hp.abs().clamp(min=1.0)
+    max_rel = (abs_err / denom).max().item()
+
+    cos_sim = (
+        torch.nn.functional.cosine_similarity(
+            ref_hp,
+            mot_hp,
+            dim=-1,
+        )
+        .min()
+        .item()
+    )
+
+    print(f"\n  [{tag}]  max_abs={max_abs:.4e}  max_rel={max_rel:.4e}  min_cos_sim={cos_sim:.6f}")
+
+    # Cosine similarity is the primary correctness gate: robust to scale
+    # and accumulation-order differences between cuBLAS and Triton.
+    # For bf16 GEMM with K up to ~20k, cos_sim > 0.99 is easily achieved.
+    assert cos_sim > 0.98, f"Cosine similarity too low: {cos_sim:.6f}"
+    # Supplementary per-element check (generous to avoid flaky failures
+    # on extreme K dimensions like 18944)
+    assert max_rel < 0.1, f"Max relative error too large: {max_rel:.4e}"
+
+
+def _run_timing(
+    ref_fn,
+    mot_fn,
+    tag: str,
+    warmup: int = 20,
+    iters: int = 100,
+):
+    """Benchmark both paths and print timing comparison."""
+    ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters)
+    mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters)
+    speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf")
+    print(f"  [{tag}]  Ref(2x linear): {ref_ms:.3f} ms  |  MoT(fused): {mot_ms:.3f} ms  |  Speedup: {speedup:.2f}x")
+
+
+# =========================================================================
+#  Test: qkv proj
+# =========================================================================
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "image_num, K, N, dtype",
+    [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM],
+    ids=[f"img{num}_K3584_N4608" for num in _IMAGE_NUM],
+)
+def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool):
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTQKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            vae_bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        assert text_linear.output_size_per_partition == N, (
+            f"Expected output_size_per_partition={N}, "
+            f"got {text_linear.output_size_per_partition}. "
+            f"Check head parameters."
+        )
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        tag = f"QKVParallel M={M} K={K} N={N}"
+        _report_mot_config(K, N, M)
+
+        with torch.no_grad():
+            ref = _reference_forward(
+                x,
+                text_idx,
+                vae_idx,
+                text_linear,
+                vae_linear,
+            )
+            mot_out, _ = mot_linear(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)
+
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    text_linear,
+                    vae_linear,
+                ),
+                lambda: mot_linear(x, text_idx, vae_idx),
+                tag,
+            )
+
+
+# =========================================================================
+#  Test: o proj
+# =========================================================================
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "image_num, K, N, dtype",
+    [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM],
+    ids=[f"img{num}_K3584_N3584" for num in _IMAGE_NUM],
+)
+def test_mot_o_proj(
+    image_num: int,
+    K: int,
+    N: int,
+    dtype: str,
+    bias: bool,
+):
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            vae_bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        tag = f"O Proj M={M} K={K} N={N}"
+        _report_mot_config(K, N, M)
+
+        # Correctness (also warms up Triton JIT compilation)
+        with torch.no_grad():
+            ref = _reference_forward(
+                x,
+                text_idx,
+                vae_idx,
+                text_linear,
+                vae_linear,
+            )
+            mot_out, _ = mot_linear(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)
+
+        # Performance
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    text_linear,
+                    vae_linear,
+                ),
+                lambda: mot_linear(x, text_idx, vae_idx),
+                tag,
+            )

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -0,0 +1,278 @@
+# ruff: noqa: N803, E741
+"""Layer-level correctness & performance test for MoTRMSNorm.
+
+Compares two equivalent computation paths:
+  - Reference: 2x vLLM IR RMSNorm op + PyTorch index scatter/gather
+    (ir.ops.rms_norm(x[text_idx], text_w) + ir.ops.rms_norm(x[vae_idx], vae_w))
+  - Target:    1x MoTRMSNorm fused Triton kernel
+    (mot_norm(x, text_indices, vae_indices))
+
+Usage::
+    pytest tests/diffusion/kernels/mot/test_mot_norm.py -v -s
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+import torch
+from vllm import ir
+from vllm.config import VllmConfig, set_current_vllm_config
+
+from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+_EPS = 1e-6
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_indices(
+    M: int,
+    text_ratio: float,
+    device: str = "cuda",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    M_text = max(1, int(M * text_ratio))
+    perm = torch.randperm(M, device=device)
+    text_indices = perm[:M_text].sort().values
+    vae_indices = perm[M_text:].sort().values
+    return text_indices, vae_indices
+
+
+def _benchmark(fn, warmup: int = 20, iters: int = 200) -> float:
+    """Return mean latency in milliseconds."""
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.accelerator.synchronize()
+    return (time.perf_counter() - t0) / iters * 1000.0
+
+
+def _reference_forward(
+    x: torch.Tensor,
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    text_weight: torch.Tensor,
+    vae_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Reference: index-gather -> 2x vLLM IR RMSNorm op -> scatter."""
+    output = torch.empty_like(x)
+    output[text_indices] = ir.ops.rms_norm(x[text_indices], text_weight, eps)
+    output[vae_indices] = ir.ops.rms_norm(x[vae_indices], vae_weight, eps)
+    return output
+
+
+def _reference_forward_head_norm(
+    x: torch.Tensor,
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    text_weight: torch.Tensor,
+    vae_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Reference for head_norm=True path.
+
+    The MoT routing happens on token dimension while RMSNorm is applied
+    independently on each head's last dimension.
+    """
+    output = torch.empty_like(x)
+    hidden_size = x.shape[-1]
+    output[text_indices] = ir.ops.rms_norm(
+        x[text_indices].reshape(-1, hidden_size),
+        text_weight,
+        eps,
+    ).reshape_as(x[text_indices])
+    output[vae_indices] = ir.ops.rms_norm(
+        x[vae_indices].reshape(-1, hidden_size),
+        vae_weight,
+        eps,
+    ).reshape_as(x[vae_indices])
+    return output
+
+
+def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
+    """Compare outputs, print metrics, and assert correctness.
+
+    Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16).
+    We upcast to fp32 solely for computing error metrics with higher
+    arithmetic precision — the actual layer outputs remain bf16.
+    """
+    ref_hp = ref.float()
+    mot_hp = mot.float()
+
+    abs_err = (ref_hp - mot_hp).abs()
+    max_abs = abs_err.max().item()
+
+    denom = ref_hp.abs().clamp(min=1.0)
+    max_rel = (abs_err / denom).max().item()
+
+    cos_sim = (
+        torch.nn.functional.cosine_similarity(
+            ref_hp,
+            mot_hp,
+            dim=-1,
+        )
+        .min()
+        .item()
+    )
+
+    print(f"\n  [{tag}]  max_abs={max_abs:.4e}  max_rel={max_rel:.4e}  min_cos_sim={cos_sim:.6f}")
+
+    # RMSNorm is element-wise (no cross-element accumulation like GEMM),
+    # so the error between two fp32-accumulating implementations should
+    # be very small — well within bf16 rounding.
+    assert cos_sim > 0.99, f"Cosine similarity too low: {cos_sim:.6f}"
+    assert max_rel < 0.05, f"Max relative error too large: {max_rel:.4e}"
+
+
+def _run_timing(
+    ref_fn,
+    mot_fn,
+    tag: str,
+    warmup: int = 20,
+    iters: int = 200,
+):
+    """Benchmark both paths and print timing comparison."""
+    ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters)
+    mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters)
+    speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf")
+    print(f"  [{tag}]  Ref(2x rms_norm): {ref_ms:.3f} ms  |  MoT(fused): {mot_ms:.3f} ms  |  Speedup: {speedup:.2f}x")
+
+
+# =========================================================================
+#  Test: MoTRMSNorm vs 2x vLLM rms_norm
+# =========================================================================
+
+
+@pytest.mark.parametrize(
+    "M, hidden_size, text_ratio",
+    [
+        (2048, 3584, 0.01),
+        (8192, 3584, 0.01),
+        (2048, 128, 0.01),
+        (8192, 128, 0.01),
+    ],
+    ids=[
+        "M2048_H3584_layernorm",
+        "M8192_H3584_layernorm",
+        "M2048_H128_qknorm",
+        "M8192_H128_qknorm",
+    ],
+)
+def test_mot_rms_norm(M: int, hidden_size: int, text_ratio: float):
+    torch.manual_seed(42)
+    dtype = torch.bfloat16
+
+    # --- Build MoT layer ---
+    mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        W_vae = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+        mot_norm.gen_weight.data.copy_(W_vae)
+
+    # --- Build reference weights (same data) ---
+    ref_text_weight = W_text.clone()
+    ref_vae_weight = W_vae.clone()
+
+    # --- Inputs ---
+    x = torch.randn(M, hidden_size, dtype=dtype, device="cuda")
+    text_idx, vae_idx = _make_indices(M, text_ratio)
+
+    tag = f"RMSNorm M={M} H={hidden_size}"
+
+    # vLLM's IR RMSNorm op may inspect global config; wrap in VllmConfig
+    # context for safety.
+    with set_current_vllm_config(VllmConfig()):
+        # --- Correctness (also warms up Triton JIT) ---
+        with torch.no_grad():
+            ref = _reference_forward(
+                x,
+                text_idx,
+                vae_idx,
+                ref_text_weight,
+                ref_vae_weight,
+                _EPS,
+            )
+            mot_out = mot_norm(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)
+
+        # --- Performance ---
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    ref_text_weight,
+                    ref_vae_weight,
+                    _EPS,
+                ),
+                lambda: mot_norm(x, text_idx, vae_idx),
+                tag,
+            )
+
+
+@pytest.mark.parametrize(
+    "M, num_heads, head_dim, text_ratio",
+    [
+        (2048, 28, 128, 0.01),
+    ],
+    ids=[
+        "M2048_NH28_HD128_qknorm_head_norm",
+    ],
+)
+def test_mot_rms_norm_head_norm(
+    M: int,
+    num_heads: int,
+    head_dim: int,
+    text_ratio: float,
+):
+    torch.manual_seed(42)
+    dtype = torch.bfloat16
+
+    # --- Build MoT layer (head_norm path) ---
+    mot_norm = MoTRMSNorm(head_dim, head_norm=True, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(head_dim, dtype=dtype, device="cuda")
+        W_vae = torch.randn(head_dim, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+        mot_norm.gen_weight.data.copy_(W_vae)
+
+    # --- Build reference weights (same data) ---
+    ref_text_weight = W_text.clone()
+    ref_vae_weight = W_vae.clone()
+
+    # --- Inputs ---
+    x = torch.randn(M, num_heads, head_dim, dtype=dtype, device="cuda")
+    text_idx, vae_idx = _make_indices(M, text_ratio)
+
+    tag = f"RMSNorm(head_norm=True) M={M} NH={num_heads} HD={head_dim}"
+
+    with set_current_vllm_config(VllmConfig()):
+        # --- Correctness (also warms up Triton JIT) ---
+        with torch.no_grad():
+            ref = _reference_forward_head_norm(
+                x,
+                text_idx,
+                vae_idx,
+                ref_text_weight,
+                ref_vae_weight,
+                _EPS,
+            )
+            mot_out = mot_norm(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -56,7 +56,10 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
         packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
         packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
         packed_position_ids=torch.arange(seq_len, dtype=torch.long),
+        packed_indexes=torch.arange(seq_len, dtype=torch.long),
         past_key_values=NaiveCache(1),
+        key_values_lens=torch.tensor([0], dtype=torch.int),
+        packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         num_timesteps=NUM_TIMESTEPS,
         timestep_shift=1.0,
         cfg_text_scale=1.0,
@@ -65,8 +68,11 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
     if cfg:
         base |= dict(
             cfg_text_scale=4.0,
+            cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
             cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
             cfg_text_past_key_values=NaiveCache(1),
+            cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
+            cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
         )
     return base
 

tests/diffusion/models/hunyuan_image3/test_prompt_utils.py

@@ -124,58 +124,6 @@ def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range():
     assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id]
     assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id]
 
-    # image_size="auto" is equivalent to None (need_ratio=True).
-    for bot in ("think", "recaption", "think_recaption", "vanilla"):
-        assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected
-        assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok, image_size="auto") == expected
-
-
-def test_resolve_stop_token_ids_explicit_image_size_stops_on_terminator():
-    """When image_size is an explicit "WxH" string, need_ratio=False: AR
-    already knows the size and stops at the cot terminator instead of the
-    ratio range.
-
-    Mirrors upstream logic where ``need_ratio`` is false when the user
-    supplies a target size, so ``final_stop_tokens`` is set to the
-    recaption/think terminator (not the ratio tokens).
-    """
-    tok = FakeTokenizer()
-
-    end_of_think_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
-    end_of_recaption_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
-
-    # think_recaption / recaption -> [end_of_recaption]
-    for bot in ("think_recaption", "recaption"):
-        for task in ("t2i", "it2i"):
-            result = resolve_stop_token_ids(task=task, bot_task=bot, tokenizer=tok, image_size="1024x1024")
-            assert result == [end_of_recaption_id], (
-                f"task={task}, bot_task={bot}, image_size='1024x1024': expected [{end_of_recaption_id}], got {result}"
-            )
-
-    # think -> [end_of_think, end_of_recaption]
-    for task in ("t2i", "it2i"):
-        result = resolve_stop_token_ids(task=task, bot_task="think", tokenizer=tok, image_size="1024x1024")
-        assert result == [end_of_think_id, end_of_recaption_id], (
-            f"task={task}, bot_task='think', image_size='1024x1024': "
-            f"expected [{end_of_think_id}, {end_of_recaption_id}], got {result}"
-        )
-
-    # vanilla -> no early return, falls through to ratio range
-    start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
-    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
-    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
-    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
-    expected_ratio = list(range(start, end + 1)) + list(range(other_start, other_end + 1))
-    for task in ("t2i", "it2i"):
-        result = resolve_stop_token_ids(task=task, bot_task="vanilla", tokenizer=tok, image_size="1024x1024")
-        assert result == expected_ratio, (
-            f"task={task}, bot_task='vanilla', image_size='1024x1024': expected ratio range, got {result}"
-        )
-
-    # Text tasks are unaffected by image_size.
-    answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
-    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok, image_size="1024x1024") == [answer_id]
-
 
 @pytest.mark.parametrize(
     "task,bot_task",

tests/diffusion/quantization/test_quantization_quality.py

@@ -71,7 +71,7 @@ class QualityTestConfig:
     prompt: str  # generation prompt
     max_lpips: float  # fail threshold — higher = more lenient
     model: str | None = None  # HF model name
-    quantization: str | dict[str, object] | None = None  # quantization method/config, e.g. "fp8"
+    quantization: str | None = None  # quantization method, e.g. "fp8"
     baseline_model: str | None = None  # explicit BF16/local baseline path
     quantized_model: str | None = None  # explicit quantized/local model path
     height: int = 1024
@@ -118,57 +118,9 @@ QUALITY_CONFIGS = [
     QualityTestConfig(
         id="fp8_z_image",
         model="Tongyi-MAI/Z-Image-Turbo",
-        quantization={
-            "method": "fp8",
-            "ignored_layers": [
-                "img_mlp",
-                *[
-                    f"layers.{layer_id}.{suffix}"
-                    for layer_id in range(15, 30)
-                    for suffix in (
-                        "attention.to_qkv",
-                        "attention.to_out.0",
-                        "feed_forward.w13",
-                        "feed_forward.w2",
-                    )
-                ],
-                *[
-                    f"model.layers.{layer_id}.{suffix}"
-                    for layer_id in range(28, 36)
-                    for suffix in (
-                        "self_attn.q_proj",
-                        "self_attn.k_proj",
-                        "self_attn.v_proj",
-                        "self_attn.o_proj",
-                        "mlp.gate_proj",
-                        "mlp.up_proj",
-                        "mlp.down_proj",
-                    )
-                ],
-            ],
-        },
+        quantization="fp8",
         task="t2i",
-        prompt=(
-            "A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an "
-            "endless nebula-drenched sky, where waterfalls of liquid starlight cascade from the edges of each "
-            "island into the cosmic abyss below. The largest island hosts an ancient, overgrown observatory "
-            "crafted from pearlescent white stone and living silverwood trees whose bioluminescent leaves pulse "
-            "with soft cerulean and violet light. At the observatory's center stands a solitary figure—a young "
-            "astronomer in flowing robes woven from woven moonlight and deep-space fabric, their face illuminated "
-            "by the glow of a holographic star chart hovering above an altar of polished obsidian. The sky swirls "
-            "with the birth of a new galaxy: ribbons of magenta and gold gas clouds twist around nascent stars, "
-            "while distant supernovae bloom like cosmic flowers in the far reaches of the void. Below, the abyss "
-            "reveals faint echoes of forgotten civilizations—ghostly silhouettes of submerged cities shimmering "
-            "in layers of atmospheric haze. Cinematic volumetric lighting cuts through the scene as twin moons—one "
-            "copper-hued, one opalescent—rise on opposite horizons, casting long, intersecting shadows across "
-            "moss-covered ruins and crystalline flora that refract light into prismatic halos. Hyper-detailed, "
-            "photorealistic rendering with the atmospheric depth of Roger Deakins' cinematography, the "
-            "architectural grandeur of Zaha Hadid, and the cosmic wonder of James Jean's illustrations. Shot on "
-            "a mythical 150mm lens with shallow depth of field, 8K resolution, Unreal Engine 5 realism, subsurface "
-            "scattering on organic elements, and ray-traced reflections dancing across every water droplet in the "
-            "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation "
-            "at the edge of existence."
-        ),
+        prompt="a cup of coffee on a wooden table, morning light",
         max_lpips=0.15,
         num_inference_steps=20,
     ),
@@ -439,6 +391,12 @@ _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUA
 
 def _quality_param(c: QualityTestConfig):
     marks = list(_marks)
+    if c.id == "fp8_z_image":
+        marks.append(
+            pytest.mark.skip(
+                reason="Z-Image FP8 quality gate temporarily disabled: https://github.com/vllm-project/vllm-omni/issues/3531"
+            )
+        )
     if c.id == "fp8_qwen_image":
         marks.append(
             pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -1,57 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from vllm_omni.diffusion import diffusion_engine
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
-    class JointAudioVideoModel:
-        dummy_run_num_frames = 2
-
-    monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
-        "_try_load_model_cls",
-        lambda model_class_name: JointAudioVideoModel,
-    )
-
-    assert diffusion_engine.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2
-
-
-def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None:
-    class AudioOutputModel:
-        support_audio_output = True
-
-    monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
-        "_try_load_model_cls",
-        lambda model_class_name: AudioOutputModel,
-    )
-
-    assert diffusion_engine.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2
-
-
-def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None:
-    class VideoOnlyModel:
-        pass
-
-    monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
-        "_try_load_model_cls",
-        lambda model_class_name: VideoOnlyModel,
-    )
-
-    assert diffusion_engine.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1
-
-
-def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(
-        diffusion_engine.DiffusionModelRegistry,
-        "_try_load_model_cls",
-        lambda model_class_name: None,
-    )
-
-    assert diffusion_engine.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2

tests/diffusion/test_diffusion_model_runner.py

@@ -52,7 +52,6 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
     runner.cache_backend = cache_backend
     runner.offload_backend = None
-    runner.prompt_embed_cache = None
     runner.od_config = SimpleNamespace(
         cache_backend=cache_backend_name,
         enable_cache_dit_summary=enable_cache_dit_summary,

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

@@ -269,7 +269,6 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -258,7 +258,6 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):
     """Test Bagel img2img with shared memory connector."""
@@ -276,7 +275,6 @@ def test_bagel_img2img_shared_memory_connector(run_level):
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_text2img_shared_memory_connector(run_level):
     """Test Bagel text2img with shared memory connector."""

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -230,11 +230,9 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
-    )
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
 
-    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
+    assert cleanup_calls == []  # no terminal cleanup; meta.finished is unobservable
 
 
 def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
@@ -250,50 +248,12 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
 
-    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
-    )
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
 
     assert connector.put.called
     sent_payload = connector.put.call_args.kwargs["data"]
     assert isinstance(sent_payload, OmniPayloadStruct)
-    assert sent_payload.meta.finished.item() is False
-    assert sent_payload.meta.is_segment_finished.item() is False
-
-
-def test_send_single_request_struct_preserves_segment_finished(build_adapter, monkeypatch):
-    adapter, connector = build_adapter(stage_id=1)
-    request = _req("req-segment", RequestStatus.WAITING, external_req_id="ext-segment")
-
-    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
-    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
-
-    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
-    )
-
-    sent_payload = connector.put.call_args.kwargs["data"]
-    assert sent_payload.meta.finished.item() is False
-    assert sent_payload.meta.is_segment_finished.item() is True
-
-
-def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
-    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"] = 111
-
-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
-
-    assert len(adapter._pending_save_reqs) == 0
-    assert adapter.requests_num_chunks_sent["ext-stream"] == 111
-
-    adapter.requests_num_chunks_sent.pop("ext-stream")
-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
-
-    assert len(adapter._pending_save_reqs) == 1
-    assert adapter.requests_num_chunks_sent["ext-stream"] == 0
+    assert sent_payload.meta is None  # confirms it's the empty struct on the wire
 
 
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
@@ -306,9 +266,7 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
-    )
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
 
     assert len(cleanup_calls) == 1
     args, _ = cleanup_calls[0]

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -253,6 +253,37 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
     )
 
 
+def _vllm_omni_output_multiple_image(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "vllm_omni_multiple.png"
+    with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
+        output = _run_vllm_omni_image_edit(
+            omni_server=server,
+            prompt=PROMPT_MULTIPLE_IMAGE,
+            input_images=[qwen_bear_image, rabbit_image],
+            output_path=output_path,
+        )
+    return output
+
+
+def _diffusers_output_multiple_image(
+    accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
+) -> Image.Image:
+    output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
+    output_path = output_dir / "diffusers_multiple.png"
+    return _run_diffusers_image_edit(
+        model=MULTIPLE_MODEL,
+        pipeline_class=QwenImageEditPlusPipeline,
+        prompt=PROMPT_MULTIPLE_IMAGE,
+        input_images=[qwen_bear_image, rabbit_image],
+        output_path=output_path,
+    )
+
+
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_single_matches_diffusers(
@@ -278,6 +309,37 @@ def test_qwen_image_edit_single_matches_diffusers(
     )
 
 
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(
+    reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
+)
+def test_qwen_image_edit_multiple_matches_diffusers(
+    accuracy_artifact_root: Path,
+    qwen_bear_image: Image.Image,
+    rabbit_image: Image.Image,
+) -> None:
+    vllm_image = _vllm_omni_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    diffusers_image = _diffusers_output_multiple_image(
+        accuracy_artifact_root=accuracy_artifact_root,
+        qwen_bear_image=qwen_bear_image,
+        rabbit_image=rabbit_image,
+    )
+    assert_similarity(
+        model_name=MULTIPLE_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+
+
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:

tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py

@@ -21,12 +21,10 @@ Usage:
 
 from __future__ import annotations
 
-import asyncio
 import uuid
 from contextlib import ExitStack
 
 import pytest
-import torch
 
 from tests.helpers.mark import hardware_test
 from vllm_omni.entrypoints.async_omni import AsyncOmni
@@ -152,101 +150,3 @@ async def test_generate_after_list_loras_inline_mode():
         assert last_output is not None
         assert isinstance(last_output, OmniRequestOutput)
         assert last_output.images, "Expected at least one generated image"
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.asyncio
-async def test_sleep_memory_reclaimed_custom_pipeline():
-    """sleep(level=1) must physically reclaim CuMemAllocator-tracked memory for
-    custom_pipeline.
-
-    Regression test for: custom pipelines constructed under ``with target_device:``
-    (CUDA default-device context) caused safetensors >=0.20.0 to use a
-    direct-to-GPU fast path (cudaMalloc via the driver API) that bypasses
-    CuMemAllocator, leaving weights invisible to sleep() and pinned in GPU
-    memory after the call.
-
-    The fix moves custom_pipeline init outside the CUDA context so all weights
-    go through the caching allocator and are therefore fully reclaimed by
-    sleep(level=1).  A non-zero ``CuMemAllocator.get_current_usage()`` after
-    sleep is the direct signal that the bypass is still occurring.
-    """
-    with ExitStack() as after:
-        engine = AsyncOmni(
-            model=MODEL,
-            custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
-            worker_extension_cls=WORKER_EXTENSION_CLASS,
-            enforce_eager=True,
-            enable_sleep_mode=True,
-        )
-        after.callback(engine.shutdown)
-
-        assert not await engine.is_sleeping(), "Engine should be awake after creation"
-
-        # Measure global VRAM before sleep (driver view; includes inline worker
-        # thread since inline mode runs in the same process).
-        torch.accelerator.synchronize()
-        free_before, total = torch.cuda.mem_get_info()
-        used_before_gib = (total - free_before) / 1024**3
-
-        # Measure CuMemAllocator-tracked usage before sleep.  In inline mode
-        # the worker runs in a thread pool inside this process, so the allocator
-        # singleton is shared and can be read directly.
-        allocator = None
-        tracked_before = 0
-        try:
-            from vllm.device_allocator.cumem import CuMemAllocator
-
-            allocator = CuMemAllocator.get_instance()
-            tracked_before = allocator.get_current_usage()
-        except Exception:
-            pass
-
-        # Put the engine to sleep; all weights should be offloaded via the pool.
-        acks = await engine.sleep(level=1)
-        await asyncio.sleep(0.5)  # allow the CUDA driver to settle
-        torch.accelerator.synchronize()
-
-        # Measure after sleep.
-        free_after, _ = torch.cuda.mem_get_info()
-        used_after_gib = (total - free_after) / 1024**3
-        drop_gib = used_before_gib - used_after_gib
-
-        # --- Primary assertion: allocator reports zero tracked memory. ---
-        # If this fails it means weights were allocated outside the CuMem pool
-        # (safetensors direct-to-GPU bypass) — the exact regression this test
-        # is designed to catch.
-        if allocator is not None:
-            tracked_after = allocator.get_current_usage()
-            assert tracked_after == 0, (
-                f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB "
-                f"after sleep(level=1) on custom_pipeline path "
-                f"(was {tracked_before / 1024**3:.3f} GiB before sleep). "
-                "Weights were allocated outside the CuMem pool via the "
-                "safetensors direct-to-GPU fast path — loader-context fix "
-                "may not be applied."
-            )
-
-        # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms
-        # reclamation at the driver level.
-        total_freed_bytes = sum(
-            (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
-            for ack in acks
-            if ack is not None
-        )
-        freed_gib = total_freed_bytes / 1024**3
-        assert freed_gib > 0 or drop_gib > 0, (
-            f"Expected GPU memory to be reclaimed after sleep(level=1) on "
-            f"custom_pipeline + enable_sleep_mode=True. "
-            f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
-            f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
-        )
-
-        # Engine must report it is sleeping.
-        assert await engine.is_sleeping()
-
-        # Wake up and confirm the engine is functional again.
-        await engine.wake_up()
-        assert not await engine.is_sleeping()

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -7,9 +7,6 @@ from contextlib import ExitStack
 
 import pytest
 
-from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
-    QwenImagePipelineWithLogProbForTest,
-)
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
@@ -17,7 +14,9 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
+CUSTOM_PIPELINE_CLASS = (
+    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
+)
 WORKER_EXTENSION_CLASS = (
     "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
 )

tests/e2e/offline_inference/test_sensenova_u1_img2img.py

@@ -167,18 +167,3 @@ def test_sensenova_u1_img2img(run_level):
         generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image)
         if run_level == "advanced_model":
             _validate_pixels(generated_image)
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.cache
-@hardware_test(res={"cuda": "H100"})
-def test_sensenova_u1_img2img_cache_dit():
-    """Test SenseNova-U1 img2img with Cache-DiT enabled."""
-    input_image = _load_input_image()
-    with OmniRunner(
-        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
-        cache_backend="cache_dit",
-    ) as runner:
-        _generate_sensenova_u1_img2img(runner.omni, input_image)

tests/e2e/offline_inference/test_sensenova_u1_text2img.py

@@ -156,17 +156,3 @@ def test_sensenova_u1_text2img(run_level):
         generated_image = _generate_sensenova_u1_image(runner.omni)
         if run_level == "advanced_model":
             _validate_pixels(generated_image)
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.cache
-@hardware_test(res={"cuda": "H100"})
-def test_sensenova_u1_text2img_cache_dit():
-    """Test SenseNova-U1 text2img with Cache-DiT enabled."""
-    with OmniRunner(
-        "SenseNova/SenseNova-U1-8B-MoT",
-        stage_configs_path=None,
-        cache_backend="cache_dit",
-    ) as runner:
-        _generate_sensenova_u1_image(runner.omni)

tests/e2e/online_serving/test_bagel_expansion.py

@@ -18,7 +18,6 @@ import pytest
 
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
@@ -26,18 +25,7 @@ PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed
 NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
 
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
-PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-HSDP_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-
-BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
-BAGEL_PARALLEL_2_DEPLOY = modify_stage_config(
-    BAGEL_CI_DEPLOY,
-    updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}},
-)
-BAGEL_HSDP_2_DEPLOY = modify_stage_config(
-    BAGEL_CI_DEPLOY,
-    updates={"stages": {0: {"devices": "0"}, 1: {"devices": "0,1"}}},
-)
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 
 def _get_diffusion_feature_cases(model: str):
@@ -75,7 +63,6 @@ def _get_diffusion_feature_cases(model: str):
         pytest.param(
             OmniServerParams(
                 model=model,
-                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--cache-backend",
                     "tea_cache",
@@ -84,33 +71,31 @@ def _get_diffusion_feature_cases(model: str):
                 ],
             ),
             id="parallel_cfg_2",
-            marks=PARALLEL_2_FEATURE_MARKS,
+            marks=PARALLEL_FEATURE_MARKS,
         ),
         # Ulysses-SP degree=2 (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
-                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--usp",
                     "2",
                 ],
             ),
             id="sp_ulysses_2",
-            marks=PARALLEL_2_FEATURE_MARKS,
+            marks=PARALLEL_FEATURE_MARKS,
         ),
         # Ring-Attention degree=2 (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
-                stage_config_path=BAGEL_PARALLEL_2_DEPLOY,
                 server_args=[
                     "--ring",
                     "2",
                 ],
             ),
             id="sp_ring_2",
-            marks=PARALLEL_2_FEATURE_MARKS,
+            marks=PARALLEL_FEATURE_MARKS,
         ),
         # Layerwise Offloading (single-card)
         pytest.param(
@@ -121,19 +106,18 @@ def _get_diffusion_feature_cases(model: str):
             id="single_card_layerwise_offload",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),
-        # Hybrid Sharded Data Parallel (2 GPUs)
+        # Hybrid Sharded Data Parallel (4 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
-                stage_config_path=BAGEL_HSDP_2_DEPLOY,
                 server_args=[
                     "--use-hsdp",
                     "--hsdp-shard-size",
-                    "2",
+                    "4",
                 ],
             ),
-            id="parallel_hsdp_2",
-            marks=HSDP_2_FEATURE_MARKS,
+            id="parallel_hsdp_4",
+            marks=PARALLEL_FEATURE_MARKS,
         ),
     ]
 
@@ -156,7 +140,7 @@ def test_bagel(
     - Ulysses-SP (degree=2)
     - Ring-Attention (degree=2)
     - Layerwise Offloading
-    - Hybrid Sharded Data Parallel (size=2)
+    - Hybrid Sharded Data Parallel (size=4)
 
     Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.
@@ -174,7 +158,6 @@ def test_bagel(
             # Enable CFG for models that use classifier-free guidance
             "negative_prompt": NEGATIVE_PROMPT,
             "true_cfg_scale": 4.0,
-            "cfg_img_scale": 1.0,
             "seed": 42,
         },
     }

tests/e2e/online_serving/test_lance.py

@@ -1,144 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end online serving test for Lance (single-stage).
-
-Verifies that the Lance pipeline can serve text-to-image and image-edit
-requests via the OpenAI-compatible chat completions API exposed by
-``vllm-omni serve``.
-
-Equivalent to running:
-
-    vllm-omni serve "bytedance-research/Lance" --omni \\
-        --pipeline lance --enforce-eager --trust-remote-code --port 8091
-
-    # text2img
-    python3 examples/online_serving/lance/openai_chat_client.py \\
-        --prompt "A cute corgi astronaut" --modality text2img
-"""
-
-import base64
-import os
-from io import BytesIO
-
-import pytest
-from vllm.assets.image import ImageAsset
-
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServerParams
-from vllm_omni.diffusion.models.lance.prompts import (
-    VIDEO_PAD,
-    VISION_END,
-    VISION_START,
-    render_lance_prompt,
-)
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-MODEL = "bytedance-research/Lance"
-
-TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic"
-IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
-
-# Lance vision block: <|vision_start|><|video_pad|><|vision_end|>
-_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
-
-# Lance is single-stage diffusion — no deploy YAML.  Pass the engine
-# knobs that used to live in ``vllm_omni/deploy/lance.yaml`` via CLI
-# flags so ``create_default_diffusion`` builds the stage config.
-_LANCE_SERVE_ARGS = [
-    "--pipeline",
-    "lance",
-    "--max-num-batched-tokens",
-    "32768",
-    "--max-num-seqs",
-    "1",
-    "--enforce-eager",
-    "--trust-remote-code",
-    "--no-enable-prefix-caching",
-    "--no-async-chunk",
-]
-
-test_params = [
-    OmniServerParams(
-        model=MODEL,
-        server_args=_LANCE_SERVE_ARGS,
-        stage_init_timeout=300,
-    ),
-]
-
-
-def _build_text2img_messages(prompt: str) -> list[dict]:
-    """Build OpenAI-format messages for text2img generation."""
-    rendered = render_lance_prompt("t2i", prompt)
-    return [
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": rendered}],
-        }
-    ]
-
-
-def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
-    """Build OpenAI-format messages for img2img generation."""
-    rendered = render_lance_prompt("image_edit", prompt, vision_token=_VISION_BLOCK)
-    return [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": rendered},
-                {
-                    "type": "image_url",
-                    "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
-                },
-            ],
-        }
-    ]
-
-
-@pytest.mark.core_model
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_lance_text2img_online(omni_server, openai_client) -> None:
-    """Lance text2img via the OpenAI-compatible chat completions API."""
-    request_config = {
-        "model": omni_server.model,
-        "messages": _build_text2img_messages(TEXT2IMG_PROMPT),
-        "modalities": ["image"],
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 2,
-            "guidance_scale": 0.0,
-            "seed": 42,
-        },
-    }
-
-    openai_client.send_diffusion_request(request_config)
-
-
-@pytest.mark.core_model
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_lance_img2img_online(omni_server, openai_client) -> None:
-    """Lance image_edit via the OpenAI-compatible chat completions API."""
-    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
-    buffer = BytesIO()
-    input_image.save(buffer, format="JPEG")
-    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-    request_config = {
-        "model": omni_server.model,
-        "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
-        "modalities": ["image"],
-        "extra_body": {
-            "num_inference_steps": 2,
-            "guidance_scale": 0.0,
-            "seed": 42,
-        },
-    }
-
-    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_omnivoice_expansion.py

@@ -74,41 +74,6 @@ class TestOmniVoiceTTS:
         openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
-class TestOmniVoiceSeed:
-    """E2E tests for OmniVoice Seed Params."""
-
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speech_auto_voice_seed_deterministic(self, omni_server, openai_client) -> None:
-        cfg = {
-            "model": omni_server.model,
-            "input": get_prompt("text"),
-            "response_format": "wav",
-            "seed": 42,
-            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
-        }
-
-        r1 = openai_client.send_audio_speech_request(cfg)[0]
-        r2 = openai_client.send_audio_speech_request(cfg)[0]
-        assert r1.audio_bytes == r2.audio_bytes
-
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speech_auto_voice_seed_non_deterministic(self, omni_server, openai_client) -> None:
-        cfg = {
-            "model": omni_server.model,
-            "input": get_prompt("text"),
-            "response_format": "wav",
-            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
-        }
-
-        cfg1 = {**cfg, "seed": 42}
-        cfg2 = {**cfg, "seed": 43}
-
-        r1 = openai_client.send_audio_speech_request(cfg1)[0]
-        r2 = openai_client.send_audio_speech_request(cfg2)[0]
-        assert r1.audio_bytes != r2.audio_bytes
-
-
 @pytest.mark.skipif(not _HAS_VOICE_CLONE, reason="Voice cloning requires transformers>=5.3.0")
 @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
 class TestOmniVoiceVoiceCloning:

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -61,11 +61,7 @@ test_params = [
             model=model,
             stage_config_path=default_path,
             use_stage_cli=True,
-            server_args=[
-                "--no-async-chunk",
-                "--stage-overrides",
-                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
-            ],
+            server_args=["--no-async-chunk"],
         ),
         id="default",
     ),

tests/e2e/online_serving/test_qwen_image_edit.py

@@ -5,8 +5,8 @@
 Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
 
 - ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
-- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
-- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
+- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2509`` — two reference images, fixed 512×512.
+- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2509`` only, ``advanced_model`` — mixed input
   resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
   scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
 
@@ -53,6 +53,7 @@ def _get_diffusion_feature_cases(model: str):
 
 
 @pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
@@ -82,16 +83,15 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
     openai_client.send_diffusion_request(request_config)
 
 
-@pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
     indirect=True,
 )
 def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
+    """Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2509``."""
 
     image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
 
@@ -118,7 +118,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "omni_server",
-    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
+    _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
     indirect=True,
 )
 def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -2,7 +2,7 @@
 Comprehensive tests of diffusion features that are available in online serving mode
 and are supported by the following models:
 - Qwen-Image-Edit: single image input
-- Qwen-Image-Edit-2511: single image input and two image inputs
+- Qwen-Image-Edit-2509: single image input and two image inputs
 """
 
 import pytest
@@ -14,7 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
-SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
+SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
 MULTI_EDIT_PROMPT = (
     "Transform the first image into a Dadaism collage art. "
     "Transform the second image into a Vincent van Gogh style painting. "
@@ -34,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "cache_dit",
+                    "tea_cache",  # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
                     "--enable-layerwise-offload",
                 ],
             ),
@@ -59,7 +59,7 @@ def _get_diffusion_feature_cases(model: str):
                 model=model,
                 server_args=[
                     "--cache-backend",
-                    "tea_cache",
+                    "cache_dit",
                     "--ring",
                     "2",
                 ],
@@ -145,7 +145,7 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
     indirect=True,
 )
 def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2511 with a single image input.
+    """Test Qwen-Image-Edit-2509 with a single image input.
 
     Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
     postprocess closure used the doubled temb (shape 2*batch) without halving
@@ -156,7 +156,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
     """
     image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 
-    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
 
     request_config = {
         "model": omni_server.model,
@@ -180,7 +180,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
     indirect=True,
 )
 def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    """Test Qwen-Image-Edit-2511 with two image inputs."""
+    """Test Qwen-Image-Edit-2509 with two image inputs."""
     image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
     image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
 

tests/engine/test_orchestrator.py

@@ -964,7 +964,6 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
 
     pool = RecordingPool()
     orchestrator = object.__new__(Orchestrator)
-    orchestrator.async_chunk = False
     orchestrator.request_states = {
         "req-stream": OrchestratorRequestState(
             request_id="req-stream",

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -44,6 +44,7 @@ realtime_server_params = [
             model=MODEL,
             stage_config_path=default_stage_config,
             use_stage_cli=True,
+            server_args=["--no-async-chunk"],
         ),
         id="default",
     ),

tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py

@@ -28,7 +28,7 @@ MODEL_BASE = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 MODEL_BASE_1_7B = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
 
 # A synthetic 1024-dim speaker embedding (all 0.1 — not a real voice, but
-# exercises the full code path through Qwen3TTSPromptEmbedsBuilder.build_prompt_embeds).
+# exercises the full code path through the talker's _build_prompt_embeds).
 DUMMY_EMBEDDING_1024 = [0.1] * 1024
 DUMMY_EMBEDDING_2048 = [0.1] * 2048
 

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -363,104 +363,3 @@ def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat):
         f"custom system_prompt content must reach the rendered prompt; "
         f"marker {marker!r} not found in prompt of length {len(out['prompt'])}"
     )
-
-
-def test_build_multistage_generation_inputs_sets_ar_stop_token_ids_with_explicit_size(serving_chat):
-    """When height+width are provided with bot_task, the AR (llm) stage
-    must receive stop_token_ids from resolve_stop_token_ids so AR stops
-    at the correct terminator instead of generating until max_tokens.
-
-    Without this, AR would generate past the cot boundary and produce
-    garbage that the DiT bridge cannot parse.
-    """
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    class FakeTokenizer:
-        SPECIAL = {
-            "<|startoftext|>": 1,
-            "<img>": 2,
-            "<recaption>": 4,
-        }
-
-        def convert_tokens_to_ids(self, tok: str) -> int:
-            return self.SPECIAL.get(tok, 0)
-
-        def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
-            return list(range(100, 100 + len(text)))
-
-    engine = SimpleNamespace(
-        stage_configs=[
-            SimpleNamespace(stage_type="llm", is_comprehension=True),
-            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
-        ],
-        default_sampling_params_list=[
-            SamplingParams(temperature=0.0),
-            OmniDiffusionSamplingParams(),
-        ],
-    )
-    images = [Image.new("RGB", (32, 32), color="red")]
-
-    # With height+width and bot_task=think, ar_image_size="1024x768" ->
-    # need_ratio=False -> stop_token_ids=[end_of_think, end_of_recaption]
-    _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs(
-        serving_chat,
-        engine=engine,
-        prompt="draw a cat",
-        extra_body={"bot_task": "think"},
-        reference_images=images,
-        gen_params=OmniDiffusionSamplingParams(height=768, width=1024),
-        tokenizer=FakeTokenizer(),
-    )
-
-    # AR stage (index 0) must have stop_token_ids set.
-    ar_params = sampling_params_list[0]
-    assert ar_params.stop_token_ids is not None, (
-        "AR stage must have stop_token_ids set when height+width and bot_task are provided"
-    )
-    assert len(ar_params.stop_token_ids) > 0, "stop_token_ids must be non-empty"
-
-    # Diffusion stage (index 1) must NOT have stop_token_ids set.
-    diff_params = sampling_params_list[1]
-    assert getattr(diff_params, "stop_token_ids", None) is None, "Diffusion stage must not have stop_token_ids set"
-
-
-def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(serving_chat):
-    """Without height+width, ar_image_size=None -> need_ratio=True ->
-    stop_token_ids is set to the ratio range (AR predicts ratio).
-    This still assigns stop_token_ids, but the set is the ratio range,
-    not the terminator.
-
-    Without bot_task at all, ar_stop_token_ids stays None and
-    stop_token_ids is not set on any stage.
-    """
-    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
-
-    engine = SimpleNamespace(
-        stage_configs=[
-            SimpleNamespace(stage_type="llm", is_comprehension=True),
-            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
-        ],
-        default_sampling_params_list=[
-            SamplingParams(temperature=0.0),
-            OmniDiffusionSamplingParams(),
-        ],
-    )
-    images = [Image.new("RGB", (32, 32), color="red")]
-
-    # No height/width, no bot_task -> ar_stop_token_ids=None ->
-    # stop_token_ids not set on any stage.
-    _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs(
-        serving_chat,
-        engine=engine,
-        prompt="draw a cat",
-        extra_body={},
-        reference_images=images,
-        gen_params=OmniDiffusionSamplingParams(),
-    )
-
-    # SamplingParams defaults stop_token_ids=[], not None.
-    # The key contract: it was NOT set by resolve_stop_token_ids,
-    # so it stays as the SamplingParams default (empty list).
-    assert sampling_params_list[0].stop_token_ids == [], (
-        "Without bot_task, AR stage stop_token_ids must be the default empty list"
-    )

tests/entrypoints/test_async_omni.py

@@ -1,19 +1,12 @@
 import asyncio
-import re
 from types import SimpleNamespace
 
 import pytest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
-from tests.helpers.mark import hardware_test
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-pytestmark = [pytest.mark.core_model]
-
-DIFFUSION_MODEL = "riverclouds/qwen_image_random"
-OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
-OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 async def _noop(**kw):
@@ -66,89 +59,6 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     return omni
 
 
-@pytest.mark.cpu
-def test_generate_submits_randomized_id_to_engine():
-    """Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""
-
-    async def run():
-        submitted_ids = []
-        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))
-
-        req_id = "my-req-1"
-        async for _ in omni.generate(
-            prompt={"prompt": "test"},
-            request_id=req_id,
-            sampling_params_list=[SimpleNamespace()],
-            output_modalities=["text"],
-        ):
-            pass
-
-        assert len(submitted_ids) == 1
-        assert submitted_ids[0] != req_id
-        assert submitted_ids[0].startswith(f"{req_id}-")
-
-    asyncio.run(run())
-
-
-@pytest.mark.cpu
-@pytest.mark.parametrize(
-    "req_ids,cancel_prefix,expected_cancel_count",
-    [
-        (["cancel-me"], "cancel-me", 1),
-        (["cancel-me", "cancel-me"], "cancel-me", 2),
-        (["cancel-hello", "cancel-hello-world"], "cancel-hello", 1),
-    ],
-)
-def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefix: str, expected_cancel_count: int):
-    """Ensure that abort() with the user-visible ID resolves correctly.
-
-    NOTE: In the case of concurrent / colliding request(s), all requests matching the
-    user provided request ID will be aborted."""
-
-    async def run():
-        aborted_batches = []
-        omni = get_async_omni_instance(
-            fake_abort_request=get_fake_abort(aborted_batches),
-        )
-
-        async def exhaust(agen):
-            async for _ in agen:
-                pass
-
-        tasks = []
-        for user_request_id in req_ids:
-            t = asyncio.create_task(
-                exhaust(
-                    omni.generate(
-                        prompt={"prompt": "test"},
-                        request_id=user_request_id,
-                        sampling_params_list=[SimpleNamespace()],
-                    )
-                )
-            )
-            await asyncio.sleep(0)
-            tasks.append(t)
-
-        assert len(omni.request_states) == len(req_ids)
-        await omni.abort(cancel_prefix)
-
-        assert len(aborted_batches) == 1
-        aborted_ids = aborted_batches[0]
-        # Aborted requests will have fmt {ext_id}-{UUID} to avoid collisions
-        for rid in aborted_ids:
-            assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
-        assert len(aborted_ids) == expected_cancel_count
-        assert len(set(aborted_ids)) == expected_cancel_count
-        assert len(omni.request_states) == len(req_ids) - expected_cancel_count
-
-        for t in tasks:
-            t.cancel()
-        await asyncio.gather(*tasks, return_exceptions=True)
-
-    asyncio.run(run())
-
-
-@pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():
         submitted_request_ids = []
@@ -181,23 +91,22 @@ def test_generate_accepts_request_after_repeated_cancellations():
                 await task
 
         assert len(await collect_outputs("after-cancel")) == 1
-
-        # Check prefixes instead of equality since generate will add
-        # a random numeric suffix to ensure the request ID is unique
-        expected_prefixes = ["baseline-", "cancel-0-", "cancel-1-", "cancel-2-", "after-cancel-"]
-        assert len(submitted_request_ids) == len(expected_prefixes)
-        for submitted, prefix in zip(submitted_request_ids, expected_prefixes):
-            assert submitted.startswith(prefix)
-
-        assert len(aborted_request_batches) == 3
-        for batch, prefix in zip(aborted_request_batches, ["cancel-0-", "cancel-1-", "cancel-2-"]):
-            assert len(batch) == 1
-            assert batch[0].startswith(prefix)
+        assert submitted_request_ids == [
+            "baseline",
+            "cancel-0",
+            "cancel-1",
+            "cancel-2",
+            "after-cancel",
+        ]
+        assert aborted_request_batches == [
+            ["cancel-0"],
+            ["cancel-1"],
+            ["cancel-2"],
+        ]
 
     asyncio.run(run_test())
 
 
-@pytest.mark.cpu
 @pytest.mark.parametrize(
     "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
 )
@@ -223,79 +132,3 @@ def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
 
     asyncio.run(run())
     assert captured_params[0].output_kind == output_kind
-
-
-# End to end tests for ensuring internal manipulation of request ID
-# in diffusion / Omni models don't leak back to the user.
-#
-# NOTE: It seems like we currently need the shutdowns here, otherwise
-# running the tests sequentially seems to leave a zombie process in diffusion
-# that can OOM the Omni tests.
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.omni
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "req_id",
-    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
-    ids=["plain", "prefixed-img", "prefixed-chat"],
-)
-async def test_diffusion_generate_preserves_request_id(req_id):
-    """Ensure diffusion model requests don't leak internal UUID-suffixed req id back to user."""
-    engine = AsyncOmni(model=DIFFUSION_MODEL)
-    try:
-        async for output in engine.generate("a white cat", request_id=req_id):
-            assert output.request_id == req_id
-    finally:
-        engine.shutdown()
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.omni
-@pytest.mark.asyncio
-@pytest.mark.parametrize(
-    "req_id",
-    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
-    ids=["plain", "prefixed-img", "prefixed-chat"],
-)
-async def test_omni_generate_preserves_request_id(req_id):
-    """Ensure omni model requests don't leak internal UUID-suffixed req id back to user."""
-    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
-    try:
-        async for output in engine.generate(
-            "Say hello in one word.",
-            request_id=req_id,
-            output_modalities=["text"],
-        ):
-            assert output.request_id == req_id
-    finally:
-        engine.shutdown()
-
-
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-@pytest.mark.omni
-@pytest.mark.asyncio
-async def test_diffusion_generate_empty_request_id():
-    """Empty request_id should get a generated internal ID, not stay empty (diffusion)"""
-    engine = AsyncOmni(model=DIFFUSION_MODEL)
-    try:
-        async for output in engine.generate("a white cat", request_id=""):
-            assert output.request_id != ""
-    finally:
-        engine.shutdown()
-
-
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.omni
-@pytest.mark.asyncio
-async def test_omni_generate_empty_request_id():
-    """Empty request_id should get a generated internal ID, not stay empty (omni)"""
-    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
-    try:
-        async for output in engine.generate(
-            "Say hello in one word.",
-            request_id="",
-            output_modalities=["text"],
-        ):
-            assert output.request_id != ""
-    finally:
-        engine.shutdown()

tests/entrypoints/test_omni_entrypoints.py

@@ -16,7 +16,6 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
 from vllm_omni.outputs import OmniRequestOutput
@@ -166,14 +165,6 @@ class FakeAsyncOmniEngine:
 def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None:
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine)
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
-    # Don't add random UUIDs to requests calling .generate since we usually
-    # just want to check for present requests anyway, and would need to just
-    # strip the UUID. Explicit checks against the mapping are in tests for
-    # AsyncOmni, or explicitly set the client req states.
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.async_omni.AsyncOmni._get_unique_request_id",
-        staticmethod(lambda request_id: request_id),
-    )
 
 
 def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
@@ -632,23 +623,16 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
 async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
     engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
     _patch_engine(monkeypatch, engine)
-    app = AsyncOmni("dummy-model")
 
-    # Requests internally have a random UUID appended to the
-    # external ID to avoid collisions, so this also tests mapping
-    external_req_id = "req-1"
-    req_id = "req-1-12345678"
+    app = AsyncOmni("dummy-model")
     try:
-        app.request_states[req_id] = ClientRequestState(
-            request_id=req_id,
-            external_request_id=external_req_id,
-        )
-        await app.abort(external_req_id)
+        app.request_states["req-1"] = object()
+        await app.abort("req-1")
     finally:
         app.shutdown()
 
-    assert engine.aborted == [[req_id]]
-    assert external_req_id not in app.request_states
+    assert engine.aborted == [["req-1"]]
+    assert "req-1" not in app.request_states
 
 
 @pytest.mark.asyncio

tests/examples/offline_inference/test_qwen3_tts_estimator.py

@@ -57,11 +57,11 @@ def wav_file(tmp_path: Path) -> Path:
 def estimator_cache(monkeypatch):
     """
     Bypass model loading by pre-populating the estimator cache, and shortcut
-    Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information
-    so that the test directly observes the real _estimate_ref_code_len closure.
+    Qwen3TTSTalker.estimate_prompt_len_from_additional_information so that
+    the test directly observes the real _estimate_ref_code_len closure.
     """
-    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-        Qwen3TTSPromptEmbedsBuilder,
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+        Qwen3TTSTalkerForConditionalGeneration,
     )
 
     def _shortcut(**kwargs):
@@ -69,7 +69,7 @@ def estimator_cache(monkeypatch):
         return kwargs["estimate_ref_code_len"](ref_audio)
 
     monkeypatch.setattr(
-        Qwen3TTSPromptEmbedsBuilder,
+        Qwen3TTSTalkerForConditionalGeneration,
         "estimate_prompt_len_from_additional_information",
         staticmethod(_shortcut),
     )

tests/helpers/runtime.py

@@ -531,6 +531,10 @@ class OmniServerStageCli(OmniServer):
         if self.env_dict is not None:
             env.update(self.env_dict)
 
+        devices = self.stage_runtime_devices.get(stage_id)
+        if devices:
+            self._set_stage_device_env(stage_id, env, devices, replica_id=replica_id)
+
         cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id)
         print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}")
         # Capture each subprocess's stdout+stderr to a per-stage log file so
@@ -1667,7 +1671,7 @@ class OpenAIClientHandler:
         # Qwen3-TTS custom fields, forwarded via extra_body.
         extra_body: dict[str, Any] = {}
         # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
-        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"):
+        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens"):
             if key in request_config:
                 extra_body[key] = request_config[key]
 
@@ -2143,8 +2147,8 @@ class OmniRunner:
         _cache = self._prompt_len_estimate_cache
         try:
             from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-            from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-                Qwen3TTSPromptEmbedsBuilder,
+            from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+                Qwen3TTSTalkerForConditionalGeneration,
             )
 
             if model_name not in _cache:
@@ -2156,7 +2160,7 @@ class OmniRunner:
 
             tok, tcfg = _cache[model_name]
             task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
                 additional_information=additional_information,
                 task_type=task_type,
                 tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -478,29 +478,19 @@ def test_compute_capture_sizes(kwargs, expected_in, not_expected):
     "batch,channels,seq_len",
     [(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)],
 )
-@pytest.mark.parametrize(
-    "dtype,atol,rtol",
-    [
-        (torch.float32, 1e-5, 1e-5),
-        (torch.bfloat16, 1e-2, 1e-2),
-        (torch.float16, 1e-3, 1e-3),
-    ],
-    ids=["fp32", "bf16", "fp16"],
-)
-def test_snakebeta_triton_vs_eager(batch, channels, seq_len, dtype, atol, rtol):
-    """Fused Triton SnakeBeta kernel must match eager PyTorch output across dtypes."""
+def test_snakebeta_triton_vs_eager(batch, channels, seq_len):
+    """Fused Triton SnakeBeta kernel must match eager PyTorch output."""
     from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
 
     if not SnakeBeta._init_triton():
         pytest.skip("Triton not available")
 
     torch.manual_seed(42)
-    snake = SnakeBeta(in_features=channels).to(DEVICE).to(dtype).eval()
-    x = torch.randn(batch, channels, seq_len, device=DEVICE, dtype=dtype)
+    snake = SnakeBeta(in_features=channels).to(DEVICE).eval()
+    x = torch.randn(batch, channels, seq_len, device=DEVICE)
 
     with torch.no_grad():
         eager_out = snake._eager_forward(x)
         triton_out = snake._triton_forward(x)
 
-    assert triton_out.dtype == x.dtype
-    torch.testing.assert_close(triton_out, eager_out, atol=atol, rtol=rtol)
+    torch.testing.assert_close(triton_out, eager_out, atol=1e-5, rtol=1e-5)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -4,105 +4,34 @@ from types import SimpleNamespace
 import numpy as np
 import torch
 
-from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-    NORMALIZED_REF_AUDIO_KEY,
-    PRECOMPUTED_REF_CODE_KEY,
-    PRECOMPUTED_REF_IDS_KEY,
-    PRECOMPUTED_TEXT_IDS_KEY,
-    Qwen3TTSPromptEmbedsBuilder,
-)
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    _NORMALIZED_REF_AUDIO_KEY,
+    _PRECOMPUTED_REF_CODE_KEY,
+    _PRECOMPUTED_REF_IDS_KEY,
+    _PRECOMPUTED_TEXT_IDS_KEY,
     Qwen3TTSTalkerForConditionalGeneration,
 )
 
 
-def _make_minimal_talker(
-    tts_pad_embed: torch.Tensor | None = None,
-    *,
-    build_prompt_embeds=None,
-):
-    """Construct a bare ``Qwen3TTSTalkerForConditionalGeneration`` for preprocess tests.
-
-    Skips ``__init__`` and ``load_weights`` (so we don't need real config /
-    checkpoint files) and pre-populates the minimum set of attributes the
-    preprocess / decode paths read: the persistent ``tts_pad`` buffer and a
-    stub :class:`Qwen3TTSPromptEmbedsBuilder` with a configurable
-    ``build_prompt_embeds`` (defaults to raising, since most preprocess
-    tests don't exercise the builder).
-    """
+def _make_minimal_talker():
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
-    if tts_pad_embed is None:
-        tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
-    model._tts_pad_embed = tts_pad_embed
-
-    def _default_raise(**_kwargs):
-        raise AssertionError("build_prompt_embeds was not stubbed in this test")
-
-    model._prompt_builder = SimpleNamespace(
-        build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise,
-    )
+    model._ref_audio_artifact_cache_max_entries = 256
+    model._ref_audio_artifact_cache = OrderedDict()
     return model
 
 
-def _make_minimal_builder(
-    *,
-    config: SimpleNamespace | None = None,
-    talker_config: SimpleNamespace | None = None,
-    tts_pad_embed: torch.Tensor | None = None,
-):
-    """Construct a bare :class:`Qwen3TTSPromptEmbedsBuilder` instance.
-
-    Bypasses ``__init__`` so we don't need a real model checkpoint; tests
-    populate the attributes they actually exercise.
-    """
-    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
-    builder._config = config or SimpleNamespace(
-        tts_bos_token_id=100,
-        tts_eos_token_id=101,
-        tts_pad_token_id=102,
-    )
-    builder._talker_config = talker_config or SimpleNamespace(
-        codec_nothink_id=10,
-        codec_think_bos_id=11,
-        codec_think_eos_id=12,
-        codec_think_id=13,
-        codec_language_id={},
-        codec_pad_id=7,
-        codec_bos_id=8,
-        num_code_groups=2,
-        spk_is_dialect={},
-    )
-    builder._model_path = ""
-    builder._text_embedding = None
-    builder._text_projection = lambda x: x
-    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
-    builder._residual_code_embeddings = lambda: []
-    builder._speaker_encoder = None
-    builder._tts_pad_embed_buffer = (
-        tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16)
-    )
-    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw(
-        AssertionError("encode_ref_audio_batch_fn not stubbed")
-    )
-    builder._speaker_cache = None
-    builder._text_tokenizer = None
-    builder._ref_audio_artifact_cache_max_entries = 256
-    builder._ref_audio_artifact_cache = OrderedDict()
-    builder._resampler_cache = OrderedDict()
-    builder._resampler_cache_max = 16
-    return builder
-
-
 def test_single_token_prefill_uses_prefill_path():
+    model = _make_minimal_talker()
     full_prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(3, 4)
     trailing_text = torch.ones((2, 4), dtype=torch.float32)
+    tts_pad = torch.full((1, 4), 0.5, dtype=torch.float32)
     ref_code = torch.arange(32, dtype=torch.long).reshape(2, 16)
 
     def fake_build_prompt_embeds(*, task_type, info_dict):
-        return full_prompt_embeds, trailing_text, 2, ref_code
+        return full_prompt_embeds, trailing_text, tts_pad, 2, ref_code
 
-    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
+    model._build_prompt_embeds = fake_build_prompt_embeds
 
     input_ids = torch.tensor([123], dtype=torch.long)
     out_ids, out_embeds, update = model.preprocess(
@@ -121,22 +50,22 @@ def test_single_token_prefill_uses_prefill_path():
     assert update["meta"]["talker_text_offset"] == 0
     assert update["meta"]["ref_code_len"] == 2
     assert torch.equal(update["embed"]["prefill"], full_prompt_embeds)
-    # ``tts_pad`` is no longer round-tripped through ``info_dict`` — it is
-    # precomputed once into ``model._tts_pad_embed`` at load time.
-    assert "tts_pad" not in update["embed"]
+    assert torch.equal(update["embed"]["tts_pad"], tts_pad)
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text)
     assert torch.equal(update["codes"]["ref"], ref_code)
     assert update["codes"]["audio"].shape == (1, 16)
 
 
 def test_single_token_prefill_can_be_inferred_from_token_progress():
+    model = _make_minimal_talker()
     full_prompt_embeds = torch.arange(8, dtype=torch.float32).reshape(2, 4)
     trailing_text = torch.ones((1, 4), dtype=torch.float32)
+    tts_pad = torch.zeros((1, 4), dtype=torch.float32)
 
     def fake_build_prompt_embeds(*, task_type, info_dict):
-        return full_prompt_embeds, trailing_text, None, None
+        return full_prompt_embeds, trailing_text, tts_pad, None, None
 
-    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
+    model._build_prompt_embeds = fake_build_prompt_embeds
 
     out_ids, out_embeds, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -153,8 +82,7 @@ def test_single_token_prefill_can_be_inferred_from_token_progress():
 
 
 def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
-    model = _make_minimal_talker(tts_pad_embed=tts_pad)
+    model = _make_minimal_talker()
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -162,6 +90,7 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
     model.embed_input_ids = fake_embed_input_ids
     trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     out_ids, out_embeds, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -169,6 +98,7 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
         text=["hello"],
         task_type=["CustomVoice"],
         hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
         meta={"talker_text_offset": 1},
         _omni_is_prefill=False,
         _omni_num_computed_tokens=2,
@@ -185,8 +115,7 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
 
 
 def test_decode_advances_trailing_text_offset_across_multiple_steps():
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
-    model = _make_minimal_talker(tts_pad_embed=tts_pad)
+    model = _make_minimal_talker()
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -195,6 +124,7 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
     trailing_text = torch.arange(8, dtype=torch.float32).reshape(2, 4)
     state_tail = trailing_text
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
     meta = {"talker_text_offset": 0}
     seen_steps = []
 
@@ -205,6 +135,7 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
             text=["hello"],
             task_type=["CustomVoice"],
             hidden_states={"trailing_text": state_tail, "last": last_hidden},
+            embed={"tts_pad": tts_pad},
             meta=meta,
             _omni_is_prefill=False,
             _omni_num_computed_tokens=2,
@@ -223,8 +154,7 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
 
 
 def test_decode_compacts_long_trailing_text_after_large_offset():
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
-    model = _make_minimal_talker(tts_pad_embed=tts_pad)
+    model = _make_minimal_talker()
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -232,6 +162,7 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     model.embed_input_ids = fake_embed_input_ids
     trailing_text = torch.arange(130 * 4, dtype=torch.float32).reshape(130, 4)
     last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     _, _, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -239,6 +170,7 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
         text=["hello"],
         task_type=["CustomVoice"],
         hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
         meta={"talker_text_offset": 64},
         _omni_is_prefill=False,
         _omni_num_computed_tokens=2,
@@ -251,8 +183,7 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
 
 
 def test_decode_batch_preprocess_matches_decode_state_updates():
-    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
-    model = _make_minimal_talker(tts_pad_embed=tts_pad)
+    model = _make_minimal_talker()
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
@@ -262,6 +193,7 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     trailing_b = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100
     last_a = torch.full((4,), 2.0, dtype=torch.float32)
     last_b = torch.full((4,), 3.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
 
     out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
         input_ids=torch.tensor([101, 202], dtype=torch.long),
@@ -270,12 +202,14 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
                 "text": ["hello"],
                 "task_type": ["Base"],
                 "hidden_states": {"trailing_text": trailing_a, "last": last_a},
+                "embed": {"tts_pad": tts_pad},
                 "meta": {"talker_text_offset": 1},
             },
             {
                 "text": ["world"],
                 "task_type": ["CustomVoice"],
                 "hidden_states": {"trailing_text": trailing_b, "last": last_b},
+                "embed": {"tts_pad": tts_pad},
                 "meta": {"talker_text_offset": 2},
             },
         ],
@@ -294,49 +228,53 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
 
 
-def _stub_text_embedding(device_param: torch.nn.Parameter):
-    """Build a lambda that emulates ``nn.Embedding`` for the ``_device()`` helper.
-
-    ``Qwen3TTSPromptEmbedsBuilder._device()`` calls
-    ``next(self._text_embedding.parameters()).device``, so the stub needs a
-    ``parameters()`` attribute *and* must be callable as an embedding.
-    """
-    fn = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)  # noqa: E731
-    fn.parameters = lambda: iter([device_param])  # type: ignore[attr-defined]
-    return fn
-
-
 def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
-    builder = _make_minimal_builder()
+    model = _make_minimal_talker()
     device_param = torch.nn.Parameter(torch.empty(0))
-    builder._text_embedding = _stub_text_embedding(device_param)
-    builder._text_projection = lambda embeds: embeds
-    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    builder._text_tokenizer = FakeTokenizer()
-
-    builder._generate_icl_prompt = lambda **kwargs: (
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
         torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
         torch.ones((1, 4), device=kwargs["ref_code"].device),
     )
 
     normalize_calls = []
     ref_audio = np.arange(1024, dtype=np.float32)
-    builder.normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
+    model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
 
     ref_audio_ids = []
-    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (
+    model._encode_ref_audio_batch = lambda wavs, sr, *, device: (
         [ref_audio_ids.append(id(w)) or torch.ones((2, 2), dtype=torch.long) for w in wavs]
     )
-    builder.extract_speaker_embedding = lambda wav, _sr: (
+    model._extract_speaker_embedding = lambda wav, _sr: (
         ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, ref_code_len, ref_code = builder.build_prompt_embeds(
+    _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds(
         task_type="Base",
         info_dict={
             "text": ["hello"],
@@ -353,11 +291,11 @@ def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
 
 
 def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
-    builder = _make_minimal_builder()
+    model = _make_minimal_talker()
     wav1 = np.arange(2048, dtype=np.float32)
     wav2 = np.arange(3072, dtype=np.float32)
     normalize_calls = []
-    builder.normalize_ref_audio = lambda raw: (
+    model._normalize_ref_audio = lambda raw: (
         normalize_calls.append(raw)
         or (
             wav1 if raw == "a.wav" else wav2,
@@ -374,7 +312,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
             torch.full((3, 2), 22, dtype=torch.long),
         ]
 
-    builder._encode_ref_audio_batch_fn = _fake_encode_batch
+    model._encode_ref_audio_batch = _fake_encode_batch
 
     class FakeTextTokenizer:
         def __init__(self):
@@ -385,7 +323,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
             return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
 
     text_tok = FakeTextTokenizer()
-    builder._text_tokenizer = text_tok
+    model._get_tokenizer = lambda: text_tok
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -403,7 +341,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
         },
     }
 
-    builder.preprocess_batch(
+    model.preprocess_batch(
         req_ids=["r1", "r2"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
@@ -415,30 +353,30 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
     assert wavs[0] is wav1
     assert wavs[1] is wav2
     assert sr == 16000
-    assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
-    assert torch.equal(buf["r2"]["codes"][PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
-    assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav1
-    assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2
+    assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
+    assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1
+    assert buf["r2"][_NORMALIZED_REF_AUDIO_KEY][0] is wav2
     assert len(text_tok.calls) == 2
-    assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
-    assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
 
 
 def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
-    builder = _make_minimal_builder()
+    model = _make_minimal_talker()
     wav = np.arange(2048, dtype=np.float32)
-    builder.normalize_ref_audio = lambda raw: (wav, 16000)
-    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (_ for _ in ()).throw(
-        AssertionError("singleton should not call encode_ref_audio_batch")
+    model._normalize_ref_audio = lambda raw: (wav, 16000)
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+        AssertionError("singleton should not load speech tokenizer")
     )
 
     class FakeTextTokenizer:
         def __call__(self, texts, *, padding=False):
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
-    builder._text_tokenizer = FakeTextTokenizer()
+    model._get_tokenizer = lambda: FakeTextTokenizer()
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -449,27 +387,22 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit
         }
     }
 
-    builder.preprocess_batch(
+    model.preprocess_batch(
         req_ids=["r1"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
     )
 
-    assert buf["r1"][NORMALIZED_REF_AUDIO_KEY][0] is wav
-    assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
-    assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
 
 
 def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exists():
-    builder = _make_minimal_builder()
-    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
-
-    class _RaisingTokenizer:
-        def __call__(self, *_args, **_kwargs):
-            raise AssertionError("tokenizer not expected")
-
-    builder._text_tokenizer = _RaisingTokenizer()
-    builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
+    model = _make_minimal_talker()
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    model._get_tokenizer = lambda: (_ for _ in ()).throw(AssertionError("tokenizer not expected"))
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
         AssertionError("speech tokenizer not expected")
     )
     buf = {
@@ -483,53 +416,70 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
         }
     }
 
-    builder.preprocess_batch(
+    model.preprocess_batch(
         req_ids=["r1"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
     )
 
-    assert PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
-    assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+    assert _PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
+    assert _NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
 
 
 def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
-    builder = _make_minimal_builder()
+    model = _make_minimal_talker()
     device_param = torch.nn.Parameter(torch.empty(0))
-    builder._text_embedding = _stub_text_embedding(device_param)
-    builder._text_projection = lambda embeds: embeds
-    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    builder._text_tokenizer = FakeTokenizer()
-    builder._generate_icl_prompt = lambda **kwargs: (
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
         torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
         torch.ones((1, 4), device=kwargs["ref_code"].device),
     )
 
     ref_audio = np.arange(2048, dtype=np.float32)
     ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
-    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
-    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
+    model._encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
         AssertionError("serial encode not expected")
     )
     speaker_wav_ids = []
-    builder.extract_speaker_embedding = lambda wav, _sr: (
+    model._extract_speaker_embedding = lambda wav, _sr: (
         speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
+    _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds(
         task_type="Base",
         info_dict={
             "text": ["hello"],
             "ref_audio": ["ref.wav"],
             "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
             "non_streaming_mode": [False],
-            "codes": {PRECOMPUTED_REF_CODE_KEY: ref_code},
-            NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
+            "codes": {_PRECOMPUTED_REF_CODE_KEY: ref_code},
+            _NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
         },
     )
 

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py

@@ -1,129 +0,0 @@
-"""Regression tests for `estimate_prompt_len_from_additional_information`.
-
-Pins the 2D `voice_clone_prompt.ref_code` shape behaviour. Applying the
-singleton-batch unwrapper `_first(...)` to that value strips its outer
-dimension and reports `len(ref_code) == num_codebooks` instead of
-`num_frames`, which silently truncates `inputs_embeds` downstream in
-`_build_prompt_embeds`.
-"""
-
-import pytest
-
-from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-    Qwen3TTSPromptEmbedsBuilder,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _fake_tokenize(text: str, **_kwargs):
-    return [0] * (8 + max(1, len(text.split())))
-
-
-def test_estimate_prompt_len_uses_full_ref_code_length() -> None:
-    num_frames = 318
-    num_codebooks = 8
-    info = {
-        "task_type": ["Base"],
-        "text": ["hello"],
-        "ref_text": ["world"],
-        "voice_clone_prompt": [
-            {
-                "ref_spk_embedding": [0.0] * 512,
-                "ref_code": [[0] * num_codebooks for _ in range(num_frames)],
-                "icl_mode": True,
-            }
-        ],
-        "non_streaming_mode": [True],
-        "language": ["English"],
-    }
-
-    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
-        additional_information=info,
-        task_type="Base",
-        tokenize_prompt=_fake_tokenize,
-        codec_language_id={"english": 0},
-        spk_is_dialect=None,
-    )
-
-    # codec_lens = 1 + num_frames = 319; plus text-side and codec-prefix
-    # terms ~20. Would be ~30 if `_first` collapses ref_code to its first row.
-    assert est > 100, f"got {est}; expected ~339. Did `_first(ref_code)` collapse the 2D list again?"
-
-
-def test_estimate_prompt_len_handles_1d_ref_code() -> None:
-    num_frames = 50
-    info = {
-        "task_type": ["Base"],
-        "text": ["hello"],
-        "ref_text": ["world"],
-        "voice_clone_prompt": [
-            {
-                "ref_spk_embedding": [0.0] * 512,
-                "ref_code": list(range(num_frames)),
-                "icl_mode": True,
-            }
-        ],
-        "non_streaming_mode": [True],
-        "language": ["English"],
-    }
-
-    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
-        additional_information=info,
-        task_type="Base",
-        tokenize_prompt=_fake_tokenize,
-        codec_language_id={"english": 0},
-        spk_is_dialect=None,
-    )
-
-    assert est > 50, f"got {est}; 1D ref_code must contribute its own length"
-
-
-def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() -> None:
-    """`voice_clone_prompt.ref_ids` as a list-of-int (the natural pre-tokenized shape)
-    must be read as a sequence, not unwrapped to its first element. Without the fix
-    `_first` collapsed it to an int, the `isinstance(list)` / `hasattr("shape")`
-    branches fell through, and `ref_ids_len` defaulted to 0 — silently dropping the
-    reference-token contribution to the prefix length budget.
-    """
-
-    def make_info(num_ref_ids: int) -> dict:
-        return {
-            "task_type": ["Base"],
-            "text": ["hello"],
-            "ref_text": ["world"],
-            "voice_clone_prompt": [
-                {
-                    "ref_spk_embedding": [0.0] * 512,
-                    "ref_code": [[0] * 8 for _ in range(200)],
-                    "ref_ids": list(range(num_ref_ids)),
-                    "icl_mode": True,
-                }
-            ],
-            "non_streaming_mode": [True],
-            "language": ["English"],
-        }
-
-    est_30 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
-        additional_information=make_info(30),
-        task_type="Base",
-        tokenize_prompt=_fake_tokenize,
-        codec_language_id={"english": 0},
-        spk_is_dialect=None,
-    )
-    est_50 = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
-        additional_information=make_info(50),
-        task_type="Base",
-        tokenize_prompt=_fake_tokenize,
-        codec_language_id={"english": 0},
-        spk_is_dialect=None,
-    )
-
-    # Each additional ref_id contributes exactly one token to the prefix budget
-    # (the estimator strips a fixed 5-token tail from ref_ids regardless of length).
-    # Under the previous bug `_first` collapsed both lists to a single int, so both
-    # estimates would have been identical (delta = 0).
-    assert est_50 - est_30 == 20, (
-        f"got delta={est_50 - est_30}; expected 20 (= 50 - 30 extra ref_ids). "
-        f"Did `_first(ref_ids)` collapse the list to its first int?"
-    )

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -20,7 +20,7 @@ def _streaming_context() -> SimpleNamespace:
 
 
 def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNamespace) -> None:
-    inc_p, inc_o = q3._get_streaming_talker_tokens(
+    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
         "r1",
         [1, 2],
         [10, 11],
@@ -28,11 +28,13 @@ def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNam
     )
     assert inc_p == [1, 2]
     assert inc_o == [10, 11]
+    assert merged == [1, 2, 10, 11]
+    assert thinker_in == [1, 2]
 
 
 def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_context: SimpleNamespace) -> None:
     q3._get_streaming_talker_tokens("r2", [1, 2], [10, 11], streaming_context=_streaming_context)
-    inc_p, inc_o = q3._get_streaming_talker_tokens(
+    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
         "r2",
         [1, 2, 3, 4],
         [10, 11, 12, 13],
@@ -40,12 +42,14 @@ def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_conte
     )
     assert inc_p == [3, 4]
     assert inc_o == [12, 13]
+    assert merged == [1, 2, 10, 3, 4, 12, 13]
+    assert thinker_in == [1, 2, 10, 3, 4]
 
 
 def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     _streaming_context: SimpleNamespace,
 ) -> None:
-    inc_p, inc_o = q3._get_streaming_talker_tokens(
+    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
         "r3",
         [1, 2, 3, 4, 5, 6],
         [10],
@@ -54,6 +58,8 @@ def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     )
     assert inc_p == [1, 2, 3, 4]
     assert inc_o == [10]
+    assert merged == [1, 2, 3, 4, 10]
+    assert thinker_in == [1, 2, 3, 4]
 
 
 def test_get_streaming_talker_tokens_clear_state(_streaming_context: SimpleNamespace) -> None:
@@ -76,72 +82,6 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex
     assert state.talker2code2wav_last_seq_len == 0
 
 
-def test_streaming_input_prefill_chunk_is_cached() -> None:
-    transfer_manager = SimpleNamespace(_pending_streaming_prefills={})
-    request = SimpleNamespace(
-        external_req_id="rt-1",
-        output_token_ids=[100],
-        all_token_ids=[151644, 872, 100],
-        prompt_token_ids=[151644, 872],
-        resumable=True,
-        additional_information=None,
-    )
-    thinker_emb = torch.ones(2, 3)
-    thinker_hid = torch.full((2, 3), 2.0)
-
-    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
-        False,
-        request,
-        thinker_emb,
-        thinker_hid,
-        transfer_manager,
-    )
-
-    assert payload is None
-    cached = transfer_manager._pending_streaming_prefills["rt-1"]
-    assert torch.equal(cached["embed"]["prefill"], thinker_emb)
-    assert torch.equal(cached["hidden_states"]["output"], thinker_hid)
-    assert cached["ids"]["all"] == request.all_token_ids
-    assert cached["ids"]["prompt"] == request.prompt_token_ids
-
-
-def test_streaming_input_prefill_flushes_with_next_decode_chunk() -> None:
-    transfer_manager = SimpleNamespace(
-        _pending_streaming_prefills={
-            "rt-2": {
-                "embed": {"prefill": torch.ones(2, 3)},
-                "hidden_states": {"output": torch.full((2, 3), 2.0)},
-                "ids": {"all": [151644, 872, 100], "prompt": [151644, 872]},
-            }
-        }
-    )
-    request = SimpleNamespace(
-        external_req_id="rt-2",
-        output_token_ids=[101],
-        all_token_ids=[151644, 872, 100, 101],
-        prompt_token_ids=[151644, 872],
-        resumable=True,
-        additional_information=None,
-    )
-    thinker_emb = torch.full((1, 3), 3.0)
-    thinker_hid = torch.full((1, 3), 4.0)
-
-    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
-        False,
-        request,
-        thinker_emb,
-        thinker_hid,
-        transfer_manager,
-    )
-
-    assert payload is not None
-    assert payload.embed.prefill.shape == (3, 3)
-    assert payload.hidden_states.output.shape == (3, 3)
-    assert payload.ids.all == [151644, 872, 100]
-    assert payload.ids.prompt == [151644, 872]
-    assert "rt-2" not in transfer_manager._pending_streaming_prefills
-
-
 def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
     request = SimpleNamespace(
         request_id="codec",

tests/test_config_factory.py

@@ -1580,74 +1580,6 @@ class TestMingFlashOmniPipeline:
         assert len(stages) == 1
         assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
 
-    def test_image_pipeline_registered(self):
-        p = _PIPELINE_REGISTRY.get("ming_flash_omni_image")
-        assert p is not None
-        assert p.model_arch == "MingFlashOmniForConditionalGeneration"
-        assert len(p.stages) == 2
-        assert p.validate() == []
-
-    def test_image_thinker_stage(self):
-        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
-        assert s.model_stage == "thinker"
-        assert s.execution_type == StageExecutionType.LLM_AR
-        assert s.input_sources == ()
-        assert s.final_output is False
-        assert s.owns_tokenizer is True
-        assert s.requires_multimodal_data is True
-        # Image variant exports hidden states for the diffusion stage.
-        assert s.engine_output_type == "latent"
-        assert s.hf_config_name == "thinker_config"
-        assert s.sampling_constraints["detokenize"] is False
-        assert s.prompt_expand_func is not None
-
-    def test_image_dit_stage(self):
-        s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
-        assert s.model_stage == "dit"
-        assert s.execution_type == StageExecutionType.DIFFUSION
-        assert s.input_sources == (0,)
-        assert s.final_output is True
-        assert s.final_output_type == "image"
-        assert s.hf_config_name == "image_gen_config"
-        assert s.model_arch == "MingImagePipeline"
-        assert s.custom_process_input_func is not None
-
-    def test_image_processor_wiring_resolves(self):
-        """The prompt_expand_func and custom_process_input_func strings must point to real callables."""
-        thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0)
-        dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1)
-        for ref in (thinker.prompt_expand_func, dit.custom_process_input_func):
-            module_path, _, attr = ref.rpartition(".")
-            module = importlib.import_module(module_path)
-            assert callable(getattr(module, attr))
-
-    def test_image_yaml_loads_and_merges(self):
-        """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
-        if not deploy_path.exists():
-            pytest.skip("ming_flash_omni_image deploy yaml not found")
-
-        deploy = load_deploy_config(deploy_path)
-        assert len(deploy.stages) == 2
-        assert deploy.async_chunk is False
-        assert deploy.pipeline == "ming_flash_omni_image"
-        assert deploy.connectors is not None
-        assert "shared_memory_connector" in deploy.connectors
-
-        pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"]
-        stages = merge_pipeline_deploy(pipeline, deploy)
-        assert len(stages) == 2
-        # Stage 0 thinker: AR worker that emits latents.
-        assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration"
-        assert stages[0].yaml_engine_args["engine_output_type"] == "latent"
-        assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is False
-        assert stages[0].yaml_extras["prompt_expand_func"] is not None
-        # Stage 1 dit: diffusion stage with MingImagePipeline.
-        assert stages[1].yaml_engine_args["model_arch"] == "MingImagePipeline"
-        assert stages[1].custom_process_input_func is not None
-        assert stages[1].final_output is True
-        assert stages[1].final_output_type == "image"
-
 
 class TestBaseConfigInheritance:
     """Test deploy YAML base_config inheritance."""

tests/worker/test_omni_gpu_model_runner.py

@@ -161,11 +161,11 @@ def _make_runner_for_mimo(req_id="r_mimo"):
 
 
 def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
-    # `_talker_mtp_forward` calls `current_omni_platform.set_forward_context`,
-    # which would otherwise dispatch to the real device implementation.
+    # Patch the module-level `set_forward_context` symbol used inside
+    # OmniGPUModelRunner._talker_mtp_forward.
     import vllm_omni.worker.gpu_model_runner as mod  # Must be the same module that defines OmniGPUModelRunner
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
 
@@ -199,7 +199,7 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
 def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
 
@@ -215,7 +215,7 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
 def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
@@ -237,7 +237,7 @@ def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker
 def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
@@ -281,7 +281,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
 def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
@@ -322,7 +322,7 @@ def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
     import vllm_omni.worker.gpu_model_runner as mod
 
-    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
 

vllm_omni/config/pipeline_registry.py

@@ -65,11 +65,6 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.bagel.pipeline",
         "BAGEL_SINGLE_STAGE_PIPELINE",
     ),
-    # Lance (ByteDance) — BAGEL-lineage unified AR+diffusion, single-stage DiT.
-    "lance": (
-        "vllm_omni.model_executor.models.lance.pipeline",
-        "LANCE_PIPELINE",
-    ),
     "glm_image": (
         "vllm_omni.model_executor.models.glm_image.pipeline",
         "GLM_IMAGE_PIPELINE",
@@ -122,10 +117,6 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
         "MING_FLASH_OMNI_THINKER_ONLY_PIPELINE",
     ),
-    "ming_flash_omni_image": (
-        "vllm_omni.model_executor.models.ming_flash_omni.pipeline",
-        "MING_FLASH_OMNI_IMAGE_PIPELINE",
-    ),
     "moss_tts_nano": (
         "vllm_omni.model_executor.models.moss_tts_nano.pipeline",
         "MOSS_TTS_NANO_PIPELINE",

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -13,7 +13,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -426,19 +426,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
                 finish_reason = request.get_finished_reason()
-                is_segment_finished = True
+                is_segment_finished = request.is_finished() and request.resumable
                 finished = self._handle_stopped_request(request)
-                if not finished:
-                    # for streaming input request only
-                    if self.chunk_transfer_adapter:
-                        if self.vllm_config.model_config.stage_id != 0:
-                            # Downstream async-chunk stages receive real payloads from the
-                            # connector. This update only resumes polling for the next segment.
-                            self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
-                    request.discard_latest_async_tokens = True
-                    request.num_output_placeholders = 0
-                    request.spec_token_ids = []
-                    request._output_token_ids.clear()
                 if finished:
                     kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
@@ -482,7 +471,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
+                    self.chunk_transfer_adapter.save_async(pooler_output, request)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
@@ -633,36 +622,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
-        outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
-        if outstanding_async_tokens > 0:
-            # Async scheduling may already have sampled the previous
-            # segment's next token. Drop that late token instead of
-            # appending it to the new streaming segment.
-            session.discard_latest_async_tokens = True
-            session.num_computed_tokens -= session.num_output_placeholders
-            session.num_output_placeholders = 0
-            session.spec_token_ids = []
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
-            if self.vllm_config.model_config.stage_id != 0:
-                # Downstream async-chunk stages receive real payloads from the
-                # connector. This update only resumes polling for the next segment.
-                self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
-                # Do not replace prompt/additional_information here; the next
-                # upstream chunk will populate them in chunk transfer adapter.
-                session.arrival_time = update.arrival_time
-                session.sampling_params = update.sampling_params
-                if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
-                    self.num_waiting_for_streaming_input -= 1
-                session.status = RequestStatus.WAITING
-                if session in self.skipped_waiting:
-                    self.skipped_waiting.remove_requests((session,))
-                    self._enqueue_waiting_request(session)
+        if self.vllm_config.model_config.stage_id != 0:
+            self._replace_session_with_streaming_update(session, update)
 
-                if self.log_stats:
-                    session.record_event(EngineCoreEventType.QUEUED)
-                return
-        super()._update_request_as_session(session, update)
+        else:
+            super()._update_request_as_session(session, update)
 
     def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
         # TODO(wzliu)! for offline mode, we should not end process until all data is transferred

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -108,8 +108,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
             # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
-                if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks(
-                    request.request_id
+                if (
+                    self.chunk_transfer_adapter is not None
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
                 ):
                     self._pending_finish_reqs.append(request)
                 req_index += 1
@@ -159,7 +160,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
-                if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id):
+                if request.request_id in self.chunk_transfer_adapter.finished_requests:
                     self.waiting.pop_request()
                     self._pending_finish_reqs.append(request)
                     continue
@@ -488,7 +489,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
                 or (
                     self.chunk_transfer_adapter is not None
-                    and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id)
+                    and request.request_id in self.chunk_transfer_adapter.finished_requests
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
@@ -506,10 +507,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     routed_experts = model_runner_output.routed_experts_dict[req_id]
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
-                if not finished:
-                    # for streaming input request only
-                    if self.chunk_transfer_adapter:
-                        self.chunk_transfer_adapter.segment_finished_requests.discard(req_id)
                 if finished:
                     kv_transfer_params = self._free_request(request)
                     if self.chunk_transfer_adapter is not None:
@@ -672,26 +669,4 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         Do not expend prompt id using update.
         """
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
-        session._output_token_ids.clear()
-        session._all_token_ids.clear()
-        new_prompt = update.prompt_token_ids or ()
-        session._all_token_ids.extend(new_prompt)
-        session.num_computed_tokens = 0
-        session.prompt_token_ids = update.prompt_token_ids or ()
-        session.additional_information = update.additional_information or None
-        # Update block hashes for the new tokens.
-        session.update_block_hashes()
-        session.num_prompt_tokens = len(session.prompt_token_ids)
-        session.arrival_time = update.arrival_time
-        session.sampling_params = update.sampling_params
-        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
-            self.num_waiting_for_streaming_input -= 1
-        session.status = RequestStatus.WAITING
-        if session in self.skipped_waiting:
-            self.skipped_waiting.remove_requests((session,))
-            self._enqueue_waiting_request(session)
-
-        if self.log_stats:
-            session.record_event(EngineCoreEventType.QUEUED)
+        self._replace_session_with_streaming_update(session, update)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,7 +5,8 @@ from typing import Any
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.request import RequestStatus
+from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 
 from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
 
@@ -145,3 +146,28 @@ class OmniSchedulerMixin:
             finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {},
             pending_input_registrations=pending_input_registrations,
         )
+
+    def _replace_session_with_streaming_update(
+        self,
+        session: Request,
+        update: StreamingUpdate,
+    ) -> None:
+        """For streaming input: Replace an existing streaming session payload with the latest update."""
+        session._output_token_ids.clear()
+        session._all_token_ids.clear()
+        new_prompt = update.prompt_token_ids or ()
+        session._all_token_ids.extend(new_prompt)
+        session.num_computed_tokens = 0
+        session.prompt_token_ids = update.prompt_token_ids or ()
+        session.additional_information = update.additional_information or None
+        # Update block hashes for the new tokens.
+        session.update_block_hashes()
+        session.num_prompt_tokens = len(session.prompt_token_ids)
+        session.arrival_time = update.arrival_time
+        session.sampling_params = update.sampling_params
+        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+            self.num_waiting_for_streaming_input -= 1
+        session.status = RequestStatus.WAITING
+
+        if self.log_stats:
+            session.record_event(EngineCoreEventType.QUEUED)

vllm_omni/data_entry_keys.py

@@ -59,7 +59,6 @@ class Ids(TypedDict, total=False):
 
 class OmniPayloadMeta(TypedDict, total=False):
     finished: torch.Tensor
-    is_segment_finished: torch.Tensor
     stream_finished: torch.Tensor
     req_id: list[str]
     left_context_size: int
@@ -143,7 +142,6 @@ class IdsStruct(_StructBase):
 
 class MetaStruct(_StructBase):
     finished: torch.Tensor | None = None
-    is_segment_finished: torch.Tensor | None = None
     stream_finished: torch.Tensor | None = None
     req_id: list[str] | None = None
     left_context_size: int | None = None

vllm_omni/deploy/lance.yaml

@@ -1,20 +0,0 @@
-# Lance (ByteDance) — single-stage diffusion pipeline selector.
-#
-# Lance's HF config.json is descriptive metadata (no ``model_type``),
-# so the loader can't auto-detect from the model dir alone.  This YAML
-# only tells the loader "use the Lance pipeline" — one stage, no multi-
-# stage orchestration.
-
-pipeline: lance
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    max_num_batched_tokens: 32768
-    max_num_seqs: 1
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    devices: "0"
-    default_sampling_params:
-      seed: 42

vllm_omni/deploy/ming_flash_omni_image.yaml

@@ -1,63 +0,0 @@
-# Ming-flash-omni-2.0 deploy: thinker + image-generation (text-to-image / img2img).
-#
-# Stage 0: Thinker — multimodal understanding -> final-layer hidden states.
-#          TP=4 across GPUs 0-3.
-# Stage 1: DiT (MingImagePipeline) — condition encoder + FlowMatchEuler
-#          denoise + VAE decode. Diffusion worker, runs eager on GPU 4.
-#
-# Cross-stage transport: shared_memory_connector (intra-node).
-#
-# Fields omitted from a stage fall back to StageDeployConfig dataclass
-# defaults (see vllm_omni/config/stage_config.py).
-#   * Stage 0 max_num_seqs=2 leaves room for the CFG-text companion that
-#     `expand_cfg_prompts` (pipeline-level) spawns when a request carries
-#     a negative prompt — positive and negative hidden states both reach
-#     the diffusion stage.
-#   * Image-gen warmup adds latency on top of the thinker;
-#     pass `--stage-init-timeout <seconds>` if the default 300s is tight.
-pipeline: ming_flash_omni_image
-async_chunk: false
-
-trust_remote_code: true
-distributed_executor_backend: mp
-enable_prefix_caching: false
-
-stages:
-  - stage_id: 0
-    # max_num_seqs=2: parent + 1 CFG-text companion (negative prompt).
-    max_num_seqs: 2
-    gpu_memory_utilization: 0.85
-    enforce_eager: true
-    tensor_parallel_size: 4
-    devices: "0,1,2,3"
-    output_connectors:
-      to_stage_1: shared_memory_connector
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 1  # image-gen request: prefill only, no AR decode
-      seed: 42
-
-  - stage_id: 1
-    enforce_eager: true
-    devices: "4"
-    input_connectors:
-      from_stage_0: shared_memory_connector
-    default_sampling_params:
-      seed: 42
-
-connectors:
-  shared_memory_connector:
-    name: SharedMemoryConnector
-    extra:
-      shm_threshold_bytes: 65536  # 64KB threshold
-
-edges:
-  - from: 0
-    to: 1
-    window_size: -1
-    max_inflight: 1
-
-# Ming-flash-omni-2.0 has only been manually verified on CUDA devices
-platforms: {}

vllm_omni/deploy/qwen3_tts.yaml

@@ -64,7 +64,6 @@ stages:
       top_p: 1.0
 
   - stage_id: 1
-    # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
     max_num_seqs: 10
     gpu_memory_utilization: 0.3
     enforce_eager: false