[Lance] Fix x2t cache-offset SDPA mask for Q != K

PyTorch SDPA's is_causal=True with Q != K uses tril(diagonal=0) which is
top-left aligned. For a single decode step (Q=1) against a long ViT+text
cache (K=574+), only q_0 -> k_0 is unmasked — producing degenerate output
("TheTheThe..." token loops).

Bypass DiffusionAttention (whose _maybe_reshape_attn_mask only handles
2-D (B, K) masks) and call torch.nn.functional.scaled_dot_product_attention
directly with an explicit 4-D (1, 1, Q, K) bottom-right-aligned causal mask
when cache_len > 0. Also cast q/k to fp32 before RMSNorm for precision.

Verified all 7 tasks on Lance:
- x2t_image:  "No" (valid VQA answer)
- x2t_video:  "The purple sphere ... right direc..." (valid VQA answer)
- t2i / i2v / t2v / image_edit / video_edit all produce expected media.

Signed-off-by: lishunyang12 <lishunyang12@163.com>

.buildkite/test-nightly.yml

@@ -459,6 +459,52 @@ 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
@@ -624,7 +670,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: 3
+                        nvidia.com/gpu: 2
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -987,60 +1033,6 @@ 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

.gitignore

@@ -270,3 +270,6 @@ 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

@@ -0,0 +1,108 @@
+#!/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/tts/bench_tts.py

@@ -130,6 +130,9 @@ 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,8 +32,11 @@ models:
         task_type: Base
 
   openbmb/VoxCPM2:
-    supported_tasks: [voice_clone]
+    supported_tasks: [voice_clone, default_voice]
+    trust_remote_code: true
     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,6 +19,8 @@ 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
 
@@ -185,6 +187,29 @@ 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)
@@ -235,6 +260,20 @@ 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_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** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 **Frame Interpolation Support**

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -329,6 +329,16 @@ 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=1024, help="Output image height.")
-    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
+    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("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
     parser.add_argument(
         "--bot-task",
@@ -229,7 +229,21 @@ def main():
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
+    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}"
+    )
     for sp in params_list:
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps

examples/offline_inference/lance/README.md

@@ -0,0 +1,52 @@
+# 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

@@ -0,0 +1,333 @@
+# 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/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.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,6 +37,10 @@ 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.
@@ -135,6 +139,8 @@ 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,6 +155,18 @@ 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()
@@ -182,6 +194,8 @@ 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 = {
@@ -225,6 +239,7 @@ 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,6 +271,14 @@ 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,6 +80,12 @@ 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()
 
@@ -122,7 +128,7 @@ def run_e2e():
     if mm_processor_kwargs:
         prompts["mm_processor_kwargs"] = mm_processor_kwargs
 
-    sampling_params_list = [OmniDiffusionSamplingParams()]
+    sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
 
     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.qwen3_tts_talker import (
-            Qwen3TTSTalkerForConditionalGeneration,
+        from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+            Qwen3TTSPromptEmbedsBuilder,
         )
 
         if model_name not in _cache:
@@ -124,7 +124,7 @@ def _estimate_prompt_len(
             except Exception:
                 return None
 
-        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+        return Qwen3TTSPromptEmbedsBuilder.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

@@ -0,0 +1,40 @@
+# 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

@@ -0,0 +1,184 @@
+#!/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

@@ -0,0 +1,15 @@
+#!/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 --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
+| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.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,34 +63,27 @@ 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 dual-stage YAML, which adds an image-gen stage behind the
-thinker:
+Launch with the image deploy YAML, which adds an image-gen stage behind
+the thinker:
 
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
     --stage-init-timeout 1800 \
-    --port 8188
+    --init-timeout 1800 \
+    --port 8091
 ```
 
 Then request image output by passing `"modalities": ["image"]`:
 
 ```bash
-curl http://127.0.0.1:8188/v1/chat/completions \
+curl -s http://127.0.0.1:8091/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
+      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
       "modalities": ["image"]
-    }' -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
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
 ```
 
 ### Optional knobs
@@ -109,7 +102,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:8188/v1/chat/completions \
+curl http://127.0.0.1:8091/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "Jonathan1909/Ming-flash-omni-2.0",
@@ -122,7 +115,7 @@ curl http://127.0.0.1:8188/v1/chat/completions \
             "byte5_text":["理解与生成统一"]}}}
       ],
       "messages": [{"role": "user", "content": "Draw a poster."}]
-    }' -o /tmp/ming_response.json
+    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
 ```
 
 ### img2img (reference image)
@@ -142,9 +135,9 @@ into the diffusion stage as `extra[reference_image]`:
 
 ### GPU layout
 
-The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
+The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
 (TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
-`runtime.devices` per stage to relocate; with fewer GPUs available, drop
+`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.
@@ -156,7 +149,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 (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
+| `["image"]` | Thinker + Imagegen (image deploy 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_dual.yaml runs the thinker on cards 0-3
+    # The shipped ming_flash_omni_image.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,10 +18,6 @@ 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 \
@@ -289,7 +285,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 **without** `async_chunk` if you use `/v1/realtime`, for example:
+Ensure the server is running, for example:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091

examples/online_serving/sensenova_u1/README.md

@@ -6,11 +6,21 @@
 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,10 +3,18 @@
 
 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"
+    --port "$PORT" \
+    "${EXTRA_ARGS[@]}"

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` | (offline only) | — | — | — |
+| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
 | 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,16 +321,29 @@ 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`, `--output`.
+The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--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,6 +12,8 @@ Examples:
 """
 
 import argparse
+import base64
+import os
 
 import httpx
 
@@ -19,6 +21,24 @@ 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 = {
@@ -26,16 +46,36 @@ 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...")
@@ -45,7 +85,6 @@ 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)
 
@@ -86,6 +125,30 @@ 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

@@ -0,0 +1,177 @@
+# 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

@@ -0,0 +1,112 @@
+# 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

@@ -0,0 +1,109 @@
+# 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,10 +29,13 @@ 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 |
@@ -44,6 +47,7 @@ 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

@@ -0,0 +1,107 @@
+# 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

@@ -0,0 +1,92 @@
+"""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_generation_scheduler_update_session.py

@@ -1,6 +1,6 @@
-"""Unit tests for OmniSchedulerMixin streaming session replacement.
+"""Unit tests for generation streaming session replacement.
 
-These tests pin the behavior of `_replace_session_with_streaming_update` against
+These tests pin the behavior of `_update_request_as_session` 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,19 +19,34 @@ 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_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
 # isort: on
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class _SchedulerStub(OmniSchedulerMixin):
-    """Minimal scheduler surface required by OmniSchedulerMixin."""
+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."""
 
     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:
@@ -71,7 +86,7 @@ class TestReplaceSessionWithStreamingUpdate:
 
         update = _make_update(prompt_token_ids=[40, 41, 42])
         sched.num_waiting_for_streaming_input = 3
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session._output_token_ids == []
         assert list(session._all_token_ids) == [40, 41, 42]
@@ -88,7 +103,7 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         update = _make_update(prompt_token_ids=None)
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session.prompt_token_ids == ()
         assert list(session._all_token_ids) == []
@@ -108,7 +123,7 @@ class TestReplaceSessionWithStreamingUpdate:
             pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)")
         update = replace(base, additional_information=None)
 
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
         assert session.additional_information is None
 
     def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None:
@@ -116,14 +131,14 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         sched.num_waiting_for_streaming_input = 5
-        sched._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(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._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(session, _make_update())
 
         assert session.events
         assert session.events[-1].type == EngineCoreEventType.QUEUED

tests/dfx/perf/scripts/run_benchmark.py

@@ -191,7 +191,16 @@ 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"}
+    exclude_keys = {
+        "request_rate",
+        "baseline",
+        "num_prompts",
+        "max_concurrency",
+        "task",
+        "enabled",
+        "eval_phase",
+        "trust_remote_code",
+    }
 
     for key, value in params.items():
         if key in exclude_keys or value is None:
@@ -207,6 +216,14 @@ 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_cfgp2.json

@@ -1,35 +1,33 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_fp8_cfgp2",
-        "description": "TP=2 Quantization=fp8 CfgP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_cfgp2",
+        "description": "TP=2 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_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.1035,
-                    "latency_p99": 9.9057,
-                    "peak_memory_mb_max": 66470,
-                    "peak_memory_mb_mean": 66470
+                    "throughput_qps": 0.21,
+                    "latency_p99": 4.7469,
+                    "peak_memory_mb_max": 101100,
+                    "peak_memory_mb_mean": 101100
                 }
             }
         ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -1,35 +1,33 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp2_fp8_sp2",
-        "description": "TP=2 Quantization=fp8 SP=2 baseline",
+        "test_name": "test_hunyuan_image_tp2_sp2",
+        "description": "TP=2 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,
-                "usp": 2,
-                "quantization": "fp8",
-                "distributed-executor-backend": "mp",
-                "enforce-eager": true,
+                "ulysses-degree": 2,
                 "enable-diffusion-pipeline-profiler": true
             }
         },
         "benchmark_params": [
             {
-                "name": "1024x1024_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.08,
-                    "latency_p99": 12.0731,
-                    "peak_memory_mb_max": 66314,
-                    "peak_memory_mb_mean": 66314
+                    "throughput_qps": 0.20,
+                    "latency_p99": 5.1025,
+                    "peak_memory_mb_max": 97402,
+                    "peak_memory_mb_mean": 97402
                 }
             }
         ]

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -1,34 +1,32 @@
 [
     {
-        "test_name": "test_hunyuan_image_tp4_fp8",
-        "description": "TP=4 Quantization=fp8 baseline",
+        "test_name": "test_hunyuan_image_tp4",
+        "description": "TP=4 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_steps50",
+                "name": "1024x1024_steps8",
                 "dataset": "random",
                 "task": "t2i",
                 "width": 1024,
                 "height": 1024,
-                "num-inference-steps": 50,
+                "num-inference-steps": 8,
                 "num-prompts": 10,
                 "max-concurrency": 1,
                 "skip-performance-assertion": true,
                 "baseline": {
-                    "throughput_qps": 0.08,
-                    "latency_p99": 13.1227,
-                    "peak_memory_mb_max": 46838,
-                    "peak_memory_mb_mean": 46838
+                    "throughput_qps": 0.21,
+                    "latency_p99": 4.9235,
+                    "peak_memory_mb_max": 57576,
+                    "peak_memory_mb_mean": 57576
                 }
             }
         ]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -17,7 +17,16 @@ 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"}
+    exclude_keys = {
+        "request_rate",
+        "baseline",
+        "num_prompts",
+        "max_concurrency",
+        "task",
+        "enabled",
+        "eval_phase",
+        "trust_remote_code",
+    }
     args = []
     for key, value in params.items():
         if key in exclude_keys or value is None:

tests/dfx/perf/tests/test_voxcpm2.json

@@ -0,0 +1,198 @@
+[
+  {
+    "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,10 +4,16 @@ 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
@@ -27,6 +33,16 @@ 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:
@@ -452,6 +468,61 @@ 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)
@@ -468,6 +539,29 @@ 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)
@@ -486,6 +580,315 @@ 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],
@@ -493,11 +896,18 @@ 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 ``pgrep -f``.
+    """Build a post-ready injector that kills processes matched by ordered ``grep_patterns``.
 
-    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).
+    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``.
 
     Args:
         grep_patterns: One pattern or an ordered list of patterns.
@@ -514,9 +924,45 @@ 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 pattern={pattern} signal={signal_name} limit={limit}",
+                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}",
                 flush=True,
             )
             pids = inject_process_kill(
@@ -535,26 +981,7 @@ def make_process_kill_fault_injector(
                 )
                 continue
             if filtered:
-                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)
+                _kill_and_wait(filtered, pattern, source="pgrep")
                 return
         logger.warning(
             "[reliability][process-kill] no process matched patterns=%s signal=%s limit=%s",

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -21,13 +21,19 @@ 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
@@ -97,7 +103,73 @@ FAULT_ERROR_KEYWORDS = (
     "500",
     "503",
 )
-RUNTIME_WORKER_PATTERN = "VLLM::Worker"
+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
 
 
 class _HasServeArgs(Protocol):
@@ -194,6 +266,55 @@ 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)
 
 
@@ -354,17 +475,7 @@ 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",
-    [
-        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",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
@@ -385,7 +496,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, FAULT_ERROR_KEYWORDS)
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
     else:
         pytest.fail("expected request failure after process-kill injection")
 
@@ -394,48 +505,22 @@ 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",
-    [
-        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",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     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 SIGKILL, /health→503, chat fails fast, concurrent chat does not hang."""
+    """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_*``.
+    """
     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."}],
@@ -444,20 +529,28 @@ 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=20)
+        ff_status, ff_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            payload,
+            timeout_sec=_PROCESS_KILL_WORKER_CHAT_FF_HTTP_TIMEOUT_SEC,
+        )
         elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
+            f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        )
         assert ff_status >= 500, (
             f"[process_kill fast_fail] expected server-side failure after fault, "
             f"got status={ff_status}, body={ff_body[:200]!r}"
         )
-    except Exception as exc:
-        ff_exc = exc
+    except Exception:
         elapsed = time.monotonic() - start
-        assert elapsed < 15, f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
+        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"
+        )
 
     payload_json = json.dumps(
         {
@@ -490,32 +583,195 @@ 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:
-            conc_debug.append(repr(exc))
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
             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#2327")
+@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")
 @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,28 +20,25 @@ 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_t2v_reliability_default",
+        "test_name": "wan22_i2v_reliability_default",
         "server_params": {
-            "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-            "server_args": [
-                "--num-gpus",
-                "1",
-                "--boundary-ratio",
-                "0.875",
-                "--flow-shift",
-                "5.0",
-                "--disable-log-stats",
-            ],
+            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         },
     }
 ]
@@ -73,9 +70,114 @@ 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
@@ -122,40 +224,15 @@ 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",
-    [
-        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",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_video_request_failure(
+def test_reliability_fault_process_kill_request_failure(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
-    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,
-    }
+    request_config = _video_request_config()
     try:
         openai_client_function.send_video_diffusion_request(request_config, request_num=1)
     except Exception as exc:
@@ -168,21 +245,11 @@ def test_reliability_fault_process_kill_video_request_failure(
 @pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
 @pytest.mark.parametrize(
     "fault_injector",
-    [
-        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",
-        ),
-    ],
+    WORKER_SIGNAL_FAULT_PARAMS,
     indirect=True,
 )
 @pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
-def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
+def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
     omni_server_after_fault_function,
     openai_client_function,
 ) -> None:
@@ -290,6 +357,37 @@ def test_reliability_fault_process_kill_video_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")
@@ -409,3 +507,127 @@ 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

@@ -0,0 +1,131 @@
+# 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,6 +19,7 @@ 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
@@ -313,3 +314,150 @@ 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

@@ -0,0 +1,529 @@
+# 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,6 +7,9 @@ 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]
 
@@ -36,6 +39,33 @@ 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()
@@ -77,3 +107,43 @@ 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/models/bagel/test_trajectory_recording.py

@@ -56,10 +56,7 @@ 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,
@@ -68,11 +65,8 @@ 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,6 +124,58 @@ 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/test_diffusion_engine_dummy_run.py

@@ -0,0 +1,57 @@
+# 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,6 +52,7 @@ 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,6 +269,7 @@ 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,6 +258,7 @@ 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."""
@@ -275,6 +276,7 @@ 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,9 +230,11 @@ 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})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
-    assert cleanup_calls == []  # no terminal cleanup; meta.finished is unobservable
+    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
 
 
 def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
@@ -248,12 +250,50 @@ 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})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
     assert connector.put.called
     sent_payload = connector.put.call_args.kwargs["data"]
     assert isinstance(sent_payload, OmniPayloadStruct)
-    assert sent_payload.meta is None  # confirms it's the empty struct on the wire
+    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
 
 
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
@@ -266,7 +306,9 @@ 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})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+    )
 
     assert len(cleanup_calls) == 1
     args, _ = cleanup_calls[0]

tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py

@@ -21,10 +21,12 @@ 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
@@ -150,3 +152,101 @@ 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,6 +7,9 @@ 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,
 )
@@ -14,9 +17,7 @@ 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 = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
+CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
 WORKER_EXTENSION_CLASS = (
     "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
 )

tests/e2e/offline_inference/test_sensenova_u1_img2img.py

@@ -167,3 +167,18 @@ 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,3 +156,17 @@ 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,6 +18,7 @@ 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]
 
@@ -25,7 +26,18 @@ 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_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+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"}}},
+)
 
 
 def _get_diffusion_feature_cases(model: str):
@@ -63,6 +75,7 @@ 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",
@@ -71,31 +84,33 @@ def _get_diffusion_feature_cases(model: str):
                 ],
             ),
             id="parallel_cfg_2",
-            marks=PARALLEL_FEATURE_MARKS,
+            marks=PARALLEL_2_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_FEATURE_MARKS,
+            marks=PARALLEL_2_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_FEATURE_MARKS,
+            marks=PARALLEL_2_FEATURE_MARKS,
         ),
         # Layerwise Offloading (single-card)
         pytest.param(
@@ -106,18 +121,19 @@ def _get_diffusion_feature_cases(model: str):
             id="single_card_layerwise_offload",
             marks=SINGLE_CARD_FEATURE_MARKS,
         ),
-        # Hybrid Sharded Data Parallel (4 GPUs)
+        # Hybrid Sharded Data Parallel (2 GPUs)
         pytest.param(
             OmniServerParams(
                 model=model,
+                stage_config_path=BAGEL_HSDP_2_DEPLOY,
                 server_args=[
                     "--use-hsdp",
                     "--hsdp-shard-size",
-                    "4",
+                    "2",
                 ],
             ),
-            id="parallel_hsdp_4",
-            marks=PARALLEL_FEATURE_MARKS,
+            id="parallel_hsdp_2",
+            marks=HSDP_2_FEATURE_MARKS,
         ),
     ]
 
@@ -140,7 +156,7 @@ def test_bagel(
     - Ulysses-SP (degree=2)
     - Ring-Attention (degree=2)
     - Layerwise Offloading
-    - Hybrid Sharded Data Parallel (size=4)
+    - Hybrid Sharded Data Parallel (size=2)
 
     Validation is delegated to assert_diffusion_response in tests/helpers/assertions.py,
     which checks output dimensions and basic correctness.
@@ -158,6 +174,7 @@ 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

@@ -0,0 +1,144 @@
+# 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,6 +74,41 @@ 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,7 +61,11 @@ test_params = [
             model=model,
             stage_config_path=default_path,
             use_stage_cli=True,
-            server_args=["--no-async-chunk"],
+            server_args=[
+                "--no-async-chunk",
+                "--stage-overrides",
+                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
+            ],
         ),
         id="default",
     ),

tests/engine/test_orchestrator.py

@@ -964,6 +964,7 @@ 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,7 +44,6 @@ 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 the talker's _build_prompt_embeds).
+# exercises the full code path through Qwen3TTSPromptEmbedsBuilder.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,3 +363,104 @@ 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,12 +1,19 @@
 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, pytest.mark.cpu]
+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")
 
 
 async def _noop(**kw):
@@ -59,6 +66,89 @@ 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 = []
@@ -91,22 +181,23 @@ def test_generate_accepts_request_after_repeated_cancellations():
                 await task
 
         assert len(await collect_outputs("after-cancel")) == 1
-        assert submitted_request_ids == [
-            "baseline",
-            "cancel-0",
-            "cancel-1",
-            "cancel-2",
-            "after-cancel",
-        ]
-        assert aborted_request_batches == [
-            ["cancel-0"],
-            ["cancel-1"],
-            ["cancel-2"],
-        ]
+
+        # 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)
 
     asyncio.run(run_test())
 
 
+@pytest.mark.cpu
 @pytest.mark.parametrize(
     "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
 )
@@ -132,3 +223,79 @@ 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,6 +16,7 @@ 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
@@ -165,6 +166,14 @@ 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(
@@ -623,16 +632,23 @@ 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"
     try:
-        app.request_states["req-1"] = object()
-        await app.abort("req-1")
+        app.request_states[req_id] = ClientRequestState(
+            request_id=req_id,
+            external_request_id=external_req_id,
+        )
+        await app.abort(external_req_id)
     finally:
         app.shutdown()
 
-    assert engine.aborted == [["req-1"]]
-    assert "req-1" not in app.request_states
+    assert engine.aborted == [[req_id]]
+    assert external_req_id 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
-    Qwen3TTSTalker.estimate_prompt_len_from_additional_information so that
-    the test directly observes the real _estimate_ref_code_len closure.
+    Qwen3TTSPromptEmbedsBuilder.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.qwen3_tts_talker import (
-        Qwen3TTSTalkerForConditionalGeneration,
+    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+        Qwen3TTSPromptEmbedsBuilder,
     )
 
     def _shortcut(**kwargs):
@@ -69,7 +69,7 @@ def estimator_cache(monkeypatch):
         return kwargs["estimate_ref_code_len"](ref_audio)
 
     monkeypatch.setattr(
-        Qwen3TTSTalkerForConditionalGeneration,
+        Qwen3TTSPromptEmbedsBuilder,
         "estimate_prompt_len_from_additional_information",
         staticmethod(_shortcut),
     )

tests/helpers/runtime.py

@@ -531,10 +531,6 @@ 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
@@ -1671,7 +1667,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"):
+        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"):
             if key in request_config:
                 extra_body[key] = request_config[key]
 
@@ -2147,8 +2143,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.qwen3_tts_talker import (
-                Qwen3TTSTalkerForConditionalGeneration,
+            from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+                Qwen3TTSPromptEmbedsBuilder,
             )
 
             if model_name not in _cache:
@@ -2160,7 +2156,7 @@ class OmniRunner:
 
             tok, tcfg = _cache[model_name]
             task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            return Qwen3TTSPromptEmbedsBuilder.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,19 +478,29 @@ 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)],
 )
-def test_snakebeta_triton_vs_eager(batch, channels, seq_len):
-    """Fused Triton SnakeBeta kernel must match eager PyTorch output."""
+@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."""
     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).eval()
-    x = torch.randn(batch, channels, seq_len, device=DEVICE)
+    snake = SnakeBeta(in_features=channels).to(DEVICE).to(dtype).eval()
+    x = torch.randn(batch, channels, seq_len, device=DEVICE, dtype=dtype)
 
     with torch.no_grad():
         eager_out = snake._eager_forward(x)
         triton_out = snake._triton_forward(x)
 
-    torch.testing.assert_close(triton_out, eager_out, atol=1e-5, rtol=1e-5)
+    assert triton_out.dtype == x.dtype
+    torch.testing.assert_close(triton_out, eager_out, atol=atol, rtol=rtol)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -4,34 +4,105 @@ 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():
+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).
+    """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
-    model._ref_audio_artifact_cache_max_entries = 256
-    model._ref_audio_artifact_cache = OrderedDict()
+    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,
+    )
     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, tts_pad, 2, ref_code
+        return full_prompt_embeds, trailing_text, 2, ref_code
 
-    model._build_prompt_embeds = fake_build_prompt_embeds
+    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
 
     input_ids = torch.tensor([123], dtype=torch.long)
     out_ids, out_embeds, update = model.preprocess(
@@ -50,22 +121,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)
-    assert torch.equal(update["embed"]["tts_pad"], tts_pad)
+    # ``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["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, tts_pad, None, None
+        return full_prompt_embeds, trailing_text, None, None
 
-    model._build_prompt_embeds = fake_build_prompt_embeds
+    model = _make_minimal_talker(build_prompt_embeds=fake_build_prompt_embeds)
 
     out_ids, out_embeds, update = model.preprocess(
         input_ids=torch.tensor([123], dtype=torch.long),
@@ -82,7 +153,8 @@ def test_single_token_prefill_can_be_inferred_from_token_progress():
 
 
 def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -90,7 +162,6 @@ 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),
@@ -98,7 +169,6 @@ 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,
@@ -115,7 +185,8 @@ def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
 
 
 def test_decode_advances_trailing_text_offset_across_multiple_steps():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -124,7 +195,6 @@ 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 = []
 
@@ -135,7 +205,6 @@ 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,
@@ -154,7 +223,8 @@ def test_decode_advances_trailing_text_offset_across_multiple_steps():
 
 
 def test_decode_compacts_long_trailing_text_after_large_offset():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
@@ -162,7 +232,6 @@ 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),
@@ -170,7 +239,6 @@ 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,
@@ -183,7 +251,8 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
 
 
 def test_decode_batch_preprocess_matches_decode_state_updates():
-    model = _make_minimal_talker()
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
+    model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
         return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
@@ -193,7 +262,6 @@ 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),
@@ -202,14 +270,12 @@ 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},
             },
         ],
@@ -228,53 +294,49 @@ 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():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
-    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={},
-    )
+    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)
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    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: (
+    builder._text_tokenizer = FakeTokenizer()
+
+    builder._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)
-    model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
+    builder.normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
 
     ref_audio_ids = []
-    model._encode_ref_audio_batch = lambda wavs, sr, *, device: (
+    builder._encode_ref_audio_batch_fn = lambda wavs, sr, *, device: (
         [ref_audio_ids.append(id(w)) or torch.ones((2, 2), dtype=torch.long) for w in wavs]
     )
-    model._extract_speaker_embedding = lambda wav, _sr: (
+    builder.extract_speaker_embedding = lambda wav, _sr: (
         ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds(
+    _prompt, _trailing, ref_code_len, ref_code = builder.build_prompt_embeds(
         task_type="Base",
         info_dict={
             "text": ["hello"],
@@ -291,11 +353,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():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     wav1 = np.arange(2048, dtype=np.float32)
     wav2 = np.arange(3072, dtype=np.float32)
     normalize_calls = []
-    model._normalize_ref_audio = lambda raw: (
+    builder.normalize_ref_audio = lambda raw: (
         normalize_calls.append(raw)
         or (
             wav1 if raw == "a.wav" else wav2,
@@ -312,7 +374,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
             torch.full((3, 2), 22, dtype=torch.long),
         ]
 
-    model._encode_ref_audio_batch = _fake_encode_batch
+    builder._encode_ref_audio_batch_fn = _fake_encode_batch
 
     class FakeTextTokenizer:
         def __init__(self):
@@ -323,7 +385,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()
-    model._get_tokenizer = lambda: text_tok
+    builder._text_tokenizer = text_tok
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -341,7 +403,7 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
         },
     }
 
-    model.preprocess_batch(
+    builder.preprocess_batch(
         req_ids=["r1", "r2"],
         model_intermediate_buffer=buf,
         device=torch.device("cpu"),
@@ -353,30 +415,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():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     wav = np.arange(2048, dtype=np.float32)
-    model._normalize_ref_audio = lambda raw: (wav, 16000)
-    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
-        AssertionError("singleton should not load speech tokenizer")
+    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")
     )
 
     class FakeTextTokenizer:
         def __call__(self, texts, *, padding=False):
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
-    model._get_tokenizer = lambda: FakeTextTokenizer()
+    builder._text_tokenizer = FakeTextTokenizer()
     buf = {
         "r1": {
             "task_type": ["Base"],
@@ -387,22 +449,27 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit
         }
     }
 
-    model.preprocess_batch(
+    builder.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():
-    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(
+    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(
         AssertionError("speech tokenizer not expected")
     )
     buf = {
@@ -416,70 +483,53 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
         }
     }
 
-    model.preprocess_batch(
+    builder.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():
-    model = _make_minimal_talker()
+    builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
-    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={},
-    )
+    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)
 
     class FakeTokenizer:
         def __call__(self, *_args, **_kwargs):
             return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
 
-    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: (
+    builder._text_tokenizer = FakeTokenizer()
+    builder._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)
-    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(
+    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(
         AssertionError("serial encode not expected")
     )
     speaker_wav_ids = []
-    model._extract_speaker_embedding = lambda wav, _sr: (
+    builder.extract_speaker_embedding = lambda wav, _sr: (
         speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
     )
 
-    _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds(
+    _prompt, _trailing, ref_code_len, out_ref_code = builder.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

@@ -0,0 +1,129 @@
+"""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, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r1",
         [1, 2],
         [10, 11],
@@ -28,13 +28,11 @@ 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, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r2",
         [1, 2, 3, 4],
         [10, 11, 12, 13],
@@ -42,14 +40,12 @@ 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, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r3",
         [1, 2, 3, 4, 5, 6],
         [10],
@@ -58,8 +54,6 @@ 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:
@@ -82,6 +76,72 @@ 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,6 +1580,74 @@ 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):
-    # Patch the module-level `set_forward_context` symbol used inside
-    # OmniGPUModelRunner._talker_mtp_forward.
+    # `_talker_mtp_forward` calls `current_omni_platform.set_forward_context`,
+    # which would otherwise dispatch to the real device implementation.
     import vllm_omni.worker.gpu_model_runner as mod  # Must be the same module that defines OmniGPUModelRunner
 
-    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "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, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "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, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "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, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "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, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "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, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
 
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
 

vllm_omni/config/pipeline_registry.py

@@ -65,6 +65,11 @@ _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",
@@ -117,6 +122,10 @@ _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 EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -426,8 +426,19 @@ 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 = request.is_finished() and request.resumable
+                is_segment_finished = True
                 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:
@@ -471,7 +482,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request)
+                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
@@ -622,11 +633,36 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
-        if self.vllm_config.model_config.stage_id != 0:
-            self._replace_session_with_streaming_update(session, update)
+        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)
 
-        else:
-            super()._update_request_as_session(session, update)
+                if self.log_stats:
+                    session.record_event(EngineCoreEventType.QUEUED)
+                return
+        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,9 +108,8 @@ 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 request.request_id in self.chunk_transfer_adapter.finished_requests
+                if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks(
+                    request.request_id
                 ):
                     self._pending_finish_reqs.append(request)
                 req_index += 1
@@ -160,7 +159,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 request.request_id in self.chunk_transfer_adapter.finished_requests:
+                if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id):
                     self.waiting.pop_request()
                     self._pending_finish_reqs.append(request)
                     continue
@@ -489,7 +488,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 request.request_id in self.chunk_transfer_adapter.finished_requests
+                    and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id)
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
@@ -507,6 +506,10 @@ 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:
@@ -669,4 +672,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         Do not expend prompt id using update.
         """
-        self._replace_session_with_streaming_update(session, 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)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,8 +5,7 @@ from typing import Any
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.engine import EngineCoreEventType
-from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
 
@@ -146,28 +145,3 @@ 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,6 +59,7 @@ 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
@@ -142,6 +143,7 @@ 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

@@ -0,0 +1,20 @@
+# 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

@@ -0,0 +1,63 @@
+# 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,6 +64,7 @@ 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

vllm_omni/diffusion/attention/backends/registry.py

@@ -59,6 +59,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
     TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
     SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
+    SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
 

vllm_omni/diffusion/attention/backends/sage_attn3.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+try:
+    from sageattn3 import sageattn3_blackwell  # noqa: F401
+except ImportError:
+    logger.warning(
+        "SageAttention3Backend is not available. Install `sageattn3` from "
+        "https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell"
+    )
+    raise ImportError
+
+
+# Wrapping sageattn3_blackwell as a torch.library custom op keeps it opaque to
+# torch.compile. Otherwise Dynamo graph-breaks on the raw pybind11 kernel and
+# Inductor fails scheduling with KeyError: 'op5'. The hasattr guard keeps this
+# idempotent across test re-imports that pop the module from sys.modules.
+if not hasattr(torch.ops.vllm_omni, "sageattn3_blackwell"):
+
+    @torch.library.custom_op("vllm_omni::sageattn3_blackwell", mutates_args=())
+    def _sageattn3_blackwell_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        is_causal: bool,
+    ) -> torch.Tensor:
+        from sageattn3 import sageattn3_blackwell as _kernel
+
+        return _kernel(query, key, value, is_causal=is_causal)
+
+    @_sageattn3_blackwell_op.register_fake
+    def _(query, key, value, is_causal):
+        return torch.empty_like(query)
+
+
+_sageattn3_blackwell_op = torch.ops.vllm_omni.sageattn3_blackwell
+
+
+class SageAttention3Backend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 128, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN_3"
+
+    @staticmethod
+    def get_impl_cls() -> type["SageAttention3Impl"]:
+        return SageAttention3Impl
+
+
+class SageAttention3Impl(AttentionImpl):
+    _warned_gqa_fallback_global: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None = None,
+    ) -> torch.Tensor:
+        query = query.transpose(1, 2).contiguous()
+        key = key.transpose(1, 2).contiguous()
+        value = value.transpose(1, 2).contiguous()
+
+        if key.shape[1] != query.shape[1]:
+            if query.shape[1] % key.shape[1] != 0:
+                raise ValueError(
+                    "GQA/MQA requires query heads to be a multiple of KV heads, "
+                    f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
+                )
+            if not type(self)._warned_gqa_fallback_global:
+                logger.warning("SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA.")
+                type(self)._warned_gqa_fallback_global = True
+            output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.causal,
+                dropout_p=self.dropout,
+                scale=self.softmax_scale,
+                enable_gqa=True,
+            )
+        else:
+            output = _sageattn3_blackwell_op(query, key, value, self.causal)
+
+        return output.transpose(1, 2).contiguous()

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -30,10 +30,6 @@ def _maybe_reshape_attn_mask(
       - broadcast_k: [batch_size, 1, 1, seq_len_k]
       - full_qk: [batch_size, 1, seq_len_q, seq_len_k]
     """
-    # Skip Attention Mask if all values are 1, `None` mask can speedup the computation
-    if attn_mask is not None and torch.all(attn_mask != 0):
-        attn_mask = None
-
     # Reshape Attention Mask
     # 2D [batch_size, seq_len_k] mask only.
     if (
@@ -89,6 +85,7 @@ class SDPAImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.requires_gqa = num_heads != num_kv_heads
         if backend_kwargs:
             logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
@@ -115,6 +112,7 @@ class SDPAImpl(AttentionImpl):
             dropout_p=0.0,
             is_causal=self.causal,
             scale=self.softmax_scale,
+            enable_gqa=self.requires_gqa,
         )
         out = output.permute(0, 2, 1, 3)
         return out

vllm_omni/diffusion/attention/parallel/ring.py

@@ -10,7 +10,7 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FLASH_ATTN
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -110,13 +110,12 @@ class RingParallelAttention:
         if backend_pref is not None:
             backend_pref = backend_pref.lower()
 
-        # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA
         # FP32 is not supported by Flash Attention, force SDPA
         if query.dtype == torch.float32:
             backend_pref = "sdpa"
-        elif not HAS_FA3 and not HAS_FLASH_ATTN:
+        elif not HAS_FA3 and not HAS_FLASH_ATTN and not HAS_AITER:
             if backend_pref != "sdpa":
-                logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.")
+                logger.warning_once("Flash Attention (FA2/FA3/AITER) is not available! Force enabling SDPA.")
             backend_pref = "sdpa"
 
         # Extract joint tensors
@@ -147,7 +146,13 @@ class RingParallelAttention:
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
         # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper)
-        attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA
+        # On ROCm, use AITER
+        if HAS_FA3:
+            attn_type = AttnType.FA3
+        elif HAS_AITER:
+            attn_type = AttnType.AITER
+        else:
+            attn_type = AttnType.FA
 
         return ring_flash_attn_func(
             query,

vllm_omni/diffusion/cache/__init__.py

@@ -5,12 +5,19 @@ Cache module for diffusion model inference acceleration.
 
 This module provides a unified cache backend system for different caching strategies:
 - TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
+- MagCache: Magnitude-based Cache for adaptive transformer caching
 - cache-dit: DBCache, SCM, and TaylorSeer caching strategies
 
 Cache backends are instantiated directly via their constructors and configured via OmniDiffusionConfig.
 """
 
 from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.prompt_embed_cache import (
+    PromptEmbedCache,
+    install_prompt_embed_cache,
+    resolve_prompt_embed_cache_config,
+    uninstall_prompt_embed_cache,
+)
 from vllm_omni.diffusion.cache.teacache import (
     CacheContext,
     TeaCacheConfig,
@@ -20,8 +27,12 @@ from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 
 __all__ = [
     "CacheBackend",
-    "TeaCacheConfig",
     "CacheContext",
     "TeaCacheBackend",
+    "TeaCacheConfig",
     "apply_teacache_hook",
+    "PromptEmbedCache",
+    "install_prompt_embed_cache",
+    "resolve_prompt_embed_cache_config",
+    "uninstall_prompt_embed_cache",
 ]

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -18,6 +18,7 @@ from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifie
 from cache_dit.caching.block_adapters import FakeDiffusionPipeline
 from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
+from cache_dit.caching.cache_blocks.pattern_3_4_5 import CachedBlocks_Pattern_3_4_5
 from cache_dit.caching.cache_contexts import BasicCacheConfig
 from cache_dit.caching.cache_contexts.cache_manager import CachedContextManager
 from vllm.logger import init_logger
@@ -1170,6 +1171,162 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
     return refresh_cache_context
 
 
+class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """
+    Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden
+    states during denoising.
+    """
+
+    @classmethod
+    def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool:
+        if kwargs.get("cache_dit_skip", False):
+            return False
+
+        # Prefix/text forwards either omit image_gen_indicators or update the
+        # DynamicCache. Denoising forwards are gen-only and use update_cache=False.
+        if kwargs.get("update_cache", True):
+            return False
+
+        exist_gen = kwargs.get("exist_gen")
+        exist_und = kwargs.get("exist_und")
+        if exist_gen is None or exist_und is None:
+            image_gen_indicators = kwargs.get("image_gen_indicators")
+            if image_gen_indicators is None:
+                return False
+            exist_gen = image_gen_indicators.any().item()
+            exist_und = (~image_gen_indicators).any().item()
+
+        return exist_gen and not exist_und
+
+    @staticmethod
+    def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+        kwargs = dict(kwargs)
+        kwargs.pop("cache_dit_skip", None)
+        return kwargs
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs):
+        block_kwargs = self._strip_cache_only_kwargs(kwargs)
+        if not self._is_denoising_call(kwargs):
+            hidden_states, new_encoder_hidden_states = self.call_blocks(
+                hidden_states,
+                *args,
+                **block_kwargs,
+            )
+            return self._process_forward_outputs(hidden_states, new_encoder_hidden_states)
+
+        return super().forward(hidden_states, *args, **block_kwargs)
+
+
+class SensenovaCachedAdapter(CachedAdapter):
+    """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks."""
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        SensenovaCachedBlocks(
+                            # 0. Transformer blocks configuration
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            # 1. Cache/Prune context configuration
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
+def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for SenseNova-U1 model.
+
+    Args:
+        pipeline: The SenseNova-U1 pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    transformer = getattr(getattr(pipeline, "language_model", None), "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.")
+
+    # Build DBCacheConfig
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on SenseNova-U1 decoder layers: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    SensenovaCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_3,
+            params_modifiers=[modifier],
+            check_forward_pattern=True,
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = pipeline.language_model.model
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for Flux.2-dev pipeline.
 
@@ -1587,6 +1744,7 @@ CUSTOM_DIT_ENABLERS.update(
         "LTX23Pipeline": enable_cache_for_ltx2,
         "LTX23ImageToVideoPipeline": enable_cache_for_ltx2,
         "BagelPipeline": enable_cache_for_bagel,
+        "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1,
         "GlmImagePipeline": enable_cache_for_glm_image,
         "Flux2Pipeline": enable_cache_for_flux2,
         "DreamIDOmniPipeline": enable_cache_for_dreamid_omni,

vllm_omni/diffusion/cache/magcache/__init__.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    MagCacheBlockHook,
+    MagCacheHeadHook,
+    MagCacheState,
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.cache.magcache.strategy import (
+    Flux2MagCacheStrategy,
+    FluxMagCacheStrategy,
+    MagCacheStrategy,
+    MagCacheStrategyRegistry,
+    get_strategy,
+    register_strategy,
+)
+
+__all__ = [
+    "Flux2MagCacheStrategy",
+    "FluxMagCacheStrategy",
+    "MagCacheBlockHook",
+    "MagCacheConfig",
+    "MagCacheHeadHook",
+    "MagCacheState",
+    "MagCacheStrategy",
+    "MagCacheStrategyRegistry",
+    "apply_mag_cache_hook",
+    "get_strategy",
+    "register_strategy",
+]

vllm_omni/diffusion/cache/magcache/backend.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache backend implementation.
+
+This module provides the MagCache backend that implements the CacheBackend
+interface using the hooks-based MagCache system.
+"""
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.cache.magcache.strategy import (
+    get_strategy,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+
+class MagCacheBackend(CacheBackend):
+    """
+    MagCache implementation using hooks.
+
+    MagCache (Magnitude-based Cache) is an adaptive caching technique that
+    speeds up diffusion inference by reusing transformer block computations
+    based on accumulated magnitude error between timesteps.
+
+    The backend applies MagCache hooks to the transformer which intercept the
+    forward pass and implement the caching logic transparently.
+
+    Example:
+        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
+        >>> from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+        >>> from vllm_omni.diffusion.cache.magcache.strategy import FluxMagCacheStrategy
+        >>> cache_config = DiffusionCacheConfig(
+        ...     mag_ratios=FluxMagCacheStrategy.FLUX_MAG_RATIOS,
+        ...     num_inference_steps=28,
+        ...     mag_threshold=0.24,
+        ...     mag_max_skip_steps=5,
+        ...     mag_retention_ratio=0.1,
+        ... )
+        >>> backend = MagCacheBackend(cache_config)
+        >>> backend.enable(pipeline)
+        >>> backend.refresh(pipeline, num_inference_steps=50)
+    """
+
+    def __init__(self, config: DiffusionCacheConfig):
+        super().__init__(config)
+        self._registered = False
+        self._magcache_config: MagCacheConfig | None = None
+        self._transformer_id: int | None = None
+
+    def enable(self, pipeline: Any) -> None:
+        """Enable MagCache on transformer using hooks.
+
+        This creates a MagCacheConfig from the backend's DiffusionCacheConfig
+        and applies the MagCache hook to the transformer.
+
+        Args:
+            pipeline: Diffusion pipeline instance. Extracts transformer and transformer_type:
+                     - transformer: pipeline.transformer
+                     - transformer_type: pipeline.transformer.__class__.__name__
+        """
+        transformer = pipeline.transformer
+        transformer_type = transformer.__class__.__name__
+
+        num_inference_steps = self.config.num_inference_steps or 28
+
+        mag_ratios = self.config.mag_ratios
+        strategy = None
+
+        if mag_ratios is None and not self.config.mag_calibrate:
+            strategy = get_strategy(transformer_type)
+            original_ratios = strategy.mag_ratios
+
+            if len(original_ratios) != num_inference_steps and hasattr(strategy, "nearest_interp"):
+                mag_ratios = strategy.nearest_interp(original_ratios, num_inference_steps)
+                logger.info(
+                    f"MagCache: Interpolated mag_ratios from {len(original_ratios)} to {num_inference_steps} steps"
+                )
+            else:
+                mag_ratios = original_ratios
+                if len(original_ratios) != num_inference_steps:
+                    logger.warning(
+                        f"MagCache: mag_ratios length ({len(original_ratios)}) != "
+                        f"num_inference_steps ({num_inference_steps}), "
+                        f"this may cause unexpected behavior"
+                    )
+
+            logger.info(f"MagCache: Using mag_ratios from {type(strategy).__name__}")
+
+        if mag_ratios is None and not self.config.mag_calibrate:
+            raise ValueError(
+                f"mag_ratios must be provided for MagCache. "
+                f"For {transformer_type}, you need to provide mag_ratios or run in calibrate mode."
+            )
+
+        self._magcache_config = MagCacheConfig(
+            transformer_type=transformer_type,
+            threshold=self.config.mag_threshold,
+            max_skip_steps=self.config.mag_max_skip_steps,
+            retention_ratio=self.config.mag_retention_ratio,
+            num_inference_steps=num_inference_steps,
+            mag_calibrate=self.config.mag_calibrate,
+            mag_ratios=mag_ratios if not self.config.mag_calibrate else None,
+        )
+        self._transformer_id = id(transformer)
+
+        apply_mag_cache_hook(transformer, self._magcache_config, strategy=strategy)
+
+        self._registered = True
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int) -> None:
+        """Refresh MagCache state for new generation.
+
+        Clears all cached residuals and resets counters/accumulators.
+        Should be called before each generation to ensure clean state.
+
+        Args:
+            pipeline: Diffusion pipeline instance. Extracts transformer via pipeline.transformer.
+            num_inference_steps: Number of inference steps for the current generation.
+                                May be used for cache context updates.
+        """
+        transformer = pipeline.transformer
+        current_transformer_id = id(transformer)
+
+        needs_re_register = False
+
+        if self._registered and hasattr(self, "_transformer_id"):
+            if current_transformer_id != self._transformer_id:
+                logger.warning(
+                    f"Transformer was replaced (id changed from {self._transformer_id} "
+                    f"to {current_transformer_id}), re-registering hooks"
+                )
+                needs_re_register = True
+
+        if not self._registered or needs_re_register:
+            self.enable(pipeline)
+            return
+
+        blocks_with_hooks = []
+
+        for name, submodule in transformer.named_children():
+            if not isinstance(submodule, torch.nn.ModuleList):
+                continue
+            for index, block in enumerate(submodule):
+                registry = getattr(block, "_hook_registry", None)
+                if registry is not None and len(registry._hooks) > 0:
+                    blocks_with_hooks.append((f"{name}.{index}", block, registry))
+
+        if not blocks_with_hooks:
+            logger.warning("No hooks found on transformer blocks, re-registering")
+            apply_mag_cache_hook(transformer, self._magcache_config)
+            self._transformer_id = current_transformer_id
+        else:
+            for name, block, registry in blocks_with_hooks:
+                for hook in registry._hooks.values():
+                    if hasattr(hook, "reset_state"):
+                        hook.reset_state(block)
+
+    def is_enabled(self) -> bool:
+        """Check if MagCache is enabled.
+
+        Returns:
+            True if enabled, False otherwise.
+        """
+        return self.enabled

vllm_omni/diffusion/cache/magcache/config.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class MagCacheConfig:
+    """
+    Configuration for MagCache applied to transformer models.
+
+    MagCache (Magnitude-based Cache) is an adaptive caching technique that speeds up
+    diffusion model inference by reusing transformer block computations based on
+    magnitude ratios between consecutive timesteps.
+
+    Reference: https://github.com/Zehong-Ma/MagCache
+
+    Args:
+        threshold: Accumulated error threshold. Higher = more aggressive skipping (faster, lower quality).
+            Default: 0.24
+        max_skip_steps: Max consecutive skip steps (K).
+            Default: 5
+        retention_ratio: Fraction of initial steps where skipping is disabled (stability).
+            Default: 0.1
+        num_inference_steps: Total inference steps. Required for retention step calculation.
+            Default: 28
+        mag_ratios: Pre-computed magnitude ratios per step. Calibrate or use strategy defaults.
+            Default: None
+        mag_calibrate: If True, runs without skipping and logs norm_ratios for calibration.
+            Default: False
+        transformer_type: Transformer class name for logging.
+            Default: "FluxTransformer2DModel"
+    """
+
+    threshold: float = 0.24
+    max_skip_steps: int = 5
+    retention_ratio: float = 0.1
+    num_inference_steps: int = 28
+    mag_ratios: torch.Tensor | list[float] | None = None
+    mag_calibrate: bool = False
+    transformer_type: str = "FluxTransformer2DModel"
+
+    def __post_init__(self) -> None:
+        """Validate and set default coefficients."""
+        if self.threshold <= 0:
+            raise ValueError(f"threshold must be positive, got {self.threshold}")
+
+        if self.max_skip_steps <= 0:
+            raise ValueError(f"max_skip_steps must be positive, got {self.max_skip_steps}")
+
+        if not 0 < self.retention_ratio < 1:
+            raise ValueError(f"retention_ratio must be in (0, 1), got {self.retention_ratio}")
+
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "num_inference_steps must be provided for MagCache. "
+                "This is required to determine retention steps and interpolate mag_ratios. "
+                "For Flux models, use num_inference_steps=28."
+            )
+
+        if self.num_inference_steps <= 0:
+            raise ValueError(f"num_inference_steps must be positive, got {self.num_inference_steps}")
+
+        if not self.mag_calibrate and self.mag_ratios is None:
+            raise ValueError(
+                "mag_ratios must be provided for MagCache inference because these ratios "
+                "are model-dependent. To get them for your model:\n"
+                "1. Initialize MagCacheConfig(mag_calibrate=True, ...)\n"
+                "2. Run inference on your model once.\n"
+                "3. Copy the printed ratios array and pass it to mag_ratios in the config.\n"
+                "For Flux models, you can import FLUX_MAG_RATIOS from vllm_omni.diffusion.cache.magcache.strategy."
+            )
+
+        if not self.mag_calibrate and self.mag_ratios is not None:
+            if not torch.is_tensor(self.mag_ratios):
+                self.mag_ratios = torch.tensor(self.mag_ratios)

vllm_omni/diffusion/cache/magcache/hook.py

@@ -0,0 +1,610 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Hook-based MagCache implementation for vLLM-Omni.
+
+This module implements a diffusers-style hook system for MagCache (Magnitude-based Cache),
+providing adaptive caching for diffusion model inference.
+
+MagCache speeds up inference by skipping transformer block computations when the accumulated
+magnitude error is below a threshold, reusing cached residuals instead.
+
+Based on: https://github.com/Zehong-Ma/MagCache
+Reference: diffusers/src/diffusers/hooks/mag_cache.py
+
+Architecture:
+- MagCacheStrategy: Model-specific strategy for preprocessing/postprocessing
+- MagCacheState: Per-step state tracking residuals and accumulated error
+- MagCacheHeadHook: Decides whether to skip based on accumulated error
+- MagCacheBlockHook: Computes and stores residuals at tail block
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
+from diffusers.utils.torch_utils import unwrap_module
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.state import MagCacheState
+from vllm_omni.diffusion.cache.magcache.strategy import MagCacheStrategy, MagCacheStrategyRegistry
+from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook, StateManager
+from vllm_omni.logger import init_logger
+
+logger = init_logger(__name__)
+
+_MAG_CACHE_LEADER_BLOCK_HOOK = "mag_cache_leader_block_hook"
+_MAG_CACHE_BLOCK_HOOK = "mag_cache_block_hook"
+
+
+class MagCacheHeadHook(ModelHook):
+    """Head block hook for MagCache - decides whether to skip computation."""
+
+    _HOOK_NAME = "mag_cache_head"
+
+    def __init__(
+        self,
+        state_manager: StateManager,
+        config: MagCacheConfig,
+        strategy: MagCacheStrategy | None = None,
+        is_tail: bool = False,
+    ):
+        super().__init__()
+        self.state_manager = state_manager
+        self.config = config
+        self._strategy = strategy
+        self._metadata = None
+        self._is_tail = is_tail
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        unwrapped_module = unwrap_module(module)
+        block_class = unwrapped_module.__class__
+
+        try:
+            self._metadata = TransformerBlockRegistry.get(block_class)
+        except ValueError:
+            if self._strategy is not None:
+                metadata = self._strategy.register_block_metadata(block_class)
+                if metadata is not None:
+                    TransformerBlockRegistry.register(model_class=block_class, metadata=metadata)
+                else:
+                    TransformerBlockRegistry.register(
+                        model_class=block_class,
+                        metadata=TransformerBlockMetadata(
+                            return_hidden_states_index=1,
+                            return_encoder_hidden_states_index=0,
+                        ),
+                    )
+            else:
+                TransformerBlockRegistry.register(
+                    model_class=block_class,
+                    metadata=TransformerBlockMetadata(
+                        return_hidden_states_index=1,
+                        return_encoder_hidden_states_index=0,
+                    ),
+                )
+            self._metadata = TransformerBlockRegistry.get(block_class)
+
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("magcache")
+
+        if hasattr(self._metadata, "hidden_states_argument_name"):
+            arg_name = self._metadata.hidden_states_argument_name
+        else:
+            arg_name = "hidden_states"
+        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+        state: MagCacheState = self.state_manager.get_state()
+        state.head_block_input = hidden_states
+
+        if state._is_first_step:
+            state.accumulated_ratio = 1.0
+            state.accumulated_err = 0.0
+            state.accumulated_steps = 0
+            state._is_first_step = False
+
+        should_compute = True
+
+        if self.config.mag_calibrate:
+            should_compute = True
+        else:
+            current_step = state.step_index
+            if current_step >= len(self.config.mag_ratios):
+                current_scale = 1.0
+            else:
+                current_scale = self.config.mag_ratios[current_step]
+
+            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)
+
+            if current_step >= retention_step:
+                state.accumulated_ratio *= current_scale
+                state.accumulated_steps += 1
+                state.accumulated_err += abs(1.0 - state.accumulated_ratio)
+
+                if (
+                    state.previous_residual is not None
+                    and state.accumulated_err <= self.config.threshold
+                    and state.accumulated_steps <= self.config.max_skip_steps
+                ):
+                    should_compute = False
+                else:
+                    state.accumulated_ratio = 1.0
+                    state.accumulated_steps = 0
+                    state.accumulated_err = 0.0
+
+        state.should_compute = should_compute
+
+        if not should_compute:
+            res = state.previous_residual
+
+            if isinstance(res, tuple):
+                res = tuple(r.to(hidden_states.device) for r in res)
+
+                if self._strategy is not None:
+                    original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                        "encoder_hidden_states", args, kwargs
+                    )
+                    if original_encoder_hidden_states.device != res[1].device:
+                        original_encoder_hidden_states = original_encoder_hidden_states.to(res[1].device)
+                    output, enc_output = self._strategy.apply_residual_tuple(
+                        hidden_states, original_encoder_hidden_states, res
+                    )
+                    ret_list = [None] * 2
+                    ret_list[self._metadata.return_hidden_states_index] = output
+                    ret_list[self._metadata.return_encoder_hidden_states_index] = enc_output
+                    return self.log_cache_hit(state, output, ret_list)
+                else:
+                    raise RuntimeError(
+                        f"MagCache residual is tuple but no strategy available for {self._metadata.transformer_type}. "
+                        f"Please register a MagCacheStrategy for this model."
+                    )
+            elif res.device != hidden_states.device:
+                res = res.to(hidden_states.device)
+
+            if self._strategy is not None:
+                output = self._strategy.apply_residual(hidden_states, res)
+            elif res.shape == hidden_states.shape:
+                output = hidden_states + res
+            elif (
+                hidden_states.ndim == 3
+                and res.ndim == 3
+                and hidden_states.shape[0] == res.shape[0]
+                and hidden_states.shape[2] == res.shape[2]
+            ):
+                diff = hidden_states.shape[1] - res.shape[1]
+                if diff > 0:
+                    output = hidden_states.clone()
+                    output[:, diff:, :] = output[:, diff:, :] + res
+                else:
+                    output = hidden_states + res
+            else:
+                output = hidden_states + res
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = output
+                ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
+                return self.log_cache_hit(state, output, ret_list)
+            else:
+                return self.log_cache_hit(state, output, None)
+        else:
+            output = self.fn_ref.original_forward(*args, **kwargs)
+            result = self.log_cache_miss(state, output)
+
+        if self._is_tail:
+            if isinstance(output, tuple):
+                out_hidden = output[self._metadata.return_hidden_states_index]
+            else:
+                out_hidden = output
+
+            in_hidden = state.head_block_input
+
+            if in_hidden is not None:
+                if self._strategy is not None:
+                    residual = self._strategy.compute_residual(out_hidden, in_hidden)
+                elif out_hidden.shape == in_hidden.shape:
+                    residual = out_hidden - in_hidden
+                else:
+                    residual = out_hidden
+
+                if self.config.mag_calibrate:
+                    self._perform_calibration_head(state, residual)
+
+                state.previous_residual = residual
+                self._advance_step_head(state)
+
+        return result
+
+    def _perform_calibration_head(
+        self,
+        state: MagCacheState,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        if self._strategy is not None:
+            ratio, std, cos_dis = self._strategy.compute_calibration_metrics(current_residual, state.previous_residual)
+        else:
+            if state.previous_residual is None:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+            else:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+
+        state.calibration_ratios.append(ratio)
+        state.norm_ratios.append(round(ratio, 5))
+        state.norm_stds.append(round(std, 5))
+        state.cos_dises.append(round(cos_dis, 5))
+
+    def _advance_step_head(self, state: MagCacheState) -> None:
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            if self.config.mag_calibrate:
+                logger.info("MagCache calibration complete.")
+                logger.info(f"norm_ratios: {state.norm_ratios}")
+                logger.info(f"norm_stds: {state.norm_stds}")
+                logger.info(f"cos_dises: {state.cos_dises}")
+                logger.info("Copy these values to DiffusionCacheConfig(mag_ratios=...) for production use")
+
+            state.step_index = 0
+            state.accumulated_ratio = 1.0
+            state.accumulated_steps = 0
+            state.accumulated_err = 0.0
+            state.previous_residual = None
+            state.calibration_ratios = []
+            state.norm_ratios = []
+            state.norm_stds = []
+            state.cos_dises = []
+            state._is_first_step = True
+
+    def log_cache_hit(self, state: MagCacheState, output, ret):
+        step = state.step_index
+        if state.previous_residual is not None:
+            if isinstance(state.previous_residual, tuple):
+                residual_shape = tuple(r.shape for r in state.previous_residual)
+            else:
+                residual_shape = state.previous_residual.shape
+        else:
+            residual_shape = "None"
+        logger.debug(
+            f"[MagCache][HEAD] STEP={step}: CACHE_HIT (err={state.accumulated_err:.6f}, "
+            f"steps_skipped={state.accumulated_steps}, residual_shape={residual_shape}"
+        )
+        return ret if ret is not None else output
+
+    def log_cache_miss(self, state: MagCacheState, output):
+        step = state.step_index
+        residual_norm = 0.0
+        if state.previous_residual is not None:
+            if isinstance(state.previous_residual, tuple):
+                residual_norm = sum(float(torch.norm(r).item()) for r in state.previous_residual)
+            else:
+                residual_norm = float(torch.norm(state.previous_residual).item())
+        logger.debug(
+            f"[MagCache][HEAD] STEP={step}: CACHE_MISS "
+            f"(err={state.accumulated_err:.6f}, acc_ratio={state.accumulated_ratio:.6f}, "
+            f"residual_norm={residual_norm:.6f}, threshold={self.config.threshold}, "
+            f"max_skip={self.config.max_skip_steps})"
+        )
+        return output
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.state_manager.reset()
+        return module
+
+
+class MagCacheBlockHook(ModelHook):
+    """Block hook for MagCache - computes residuals at tail block."""
+
+    _HOOK_NAME = "mag_cache_block"
+
+    def __init__(
+        self,
+        state_manager: StateManager,
+        is_tail: bool = False,
+        config: MagCacheConfig | None = None,
+        strategy: MagCacheStrategy | None = None,
+    ):
+        super().__init__()
+        self.state_manager = state_manager
+        self.is_tail = is_tail
+        self.config = config
+        self._strategy = strategy
+        self._metadata = None
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        unwrapped_module = unwrap_module(module)
+        block_class = unwrapped_module.__class__
+
+        try:
+            self._metadata = TransformerBlockRegistry.get(block_class)
+        except ValueError:
+            if self._strategy is not None:
+                metadata = self._strategy.register_block_metadata(block_class)
+                if metadata is not None:
+                    TransformerBlockRegistry.register(model_class=block_class, metadata=metadata)
+                else:
+                    TransformerBlockRegistry.register(
+                        model_class=block_class,
+                        metadata=TransformerBlockMetadata(
+                            return_hidden_states_index=1,
+                            return_encoder_hidden_states_index=0,
+                        ),
+                    )
+            else:
+                TransformerBlockRegistry.register(
+                    model_class=block_class,
+                    metadata=TransformerBlockMetadata(
+                        return_hidden_states_index=1,
+                        return_encoder_hidden_states_index=0,
+                    ),
+                )
+            self._metadata = TransformerBlockRegistry.get(block_class)
+
+        return module
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        self.state_manager.reset()
+        return module
+
+    @torch.compiler.disable
+    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
+        if self.state_manager._current_context is None:
+            self.state_manager.set_context("magcache")
+        state: MagCacheState = self.state_manager.get_state()
+
+        if not state.should_compute:
+            res = state.previous_residual
+            if res is None:
+                res = torch.zeros_like(args[0])
+
+            if hasattr(self._metadata, "hidden_states_argument_name"):
+                arg_name = self._metadata.hidden_states_argument_name
+            else:
+                arg_name = "hidden_states"
+            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
+
+            if self.is_tail:
+                self.advance_step(state)
+
+            if self._metadata.return_encoder_hidden_states_index is not None:
+                encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
+                    "encoder_hidden_states", args, kwargs
+                )
+
+                if self._strategy is not None:
+                    out_hidden, enc_out = self._strategy.apply_residual_tuple(hidden_states, encoder_hidden_states, res)
+                else:
+                    out_hidden = hidden_states + res
+                    enc_out = encoder_hidden_states
+
+                max_idx = max(
+                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
+                )
+                ret_list = [None] * (max_idx + 1)
+                ret_list[self._metadata.return_hidden_states_index] = out_hidden
+                ret_list[self._metadata.return_encoder_hidden_states_index] = enc_out
+                return tuple(ret_list)
+
+            if self._strategy is not None:
+                output = self._strategy.apply_residual(hidden_states, res)
+            else:
+                output = hidden_states + res
+
+            return output
+
+        output = self.fn_ref.original_forward(*args, **kwargs)
+
+        if self.is_tail:
+            if isinstance(output, tuple):
+                out_hidden = output[self._metadata.return_hidden_states_index]
+            else:
+                out_hidden = output
+
+            in_hidden = state.head_block_input
+
+            if in_hidden is None:
+                return output
+
+            if self._strategy is not None:
+                residual = self._strategy.compute_residual(output, in_hidden)
+            elif out_hidden.shape == in_hidden.shape:
+                residual = out_hidden - in_hidden
+            elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
+                residual = out_hidden - in_hidden
+            else:
+                residual = out_hidden
+
+            if self.config.mag_calibrate:
+                self.perform_calibration(state, residual)
+
+            state.previous_residual = residual
+            self.advance_step(state)
+
+            self.log_residual_computed(state, residual)
+
+        return output
+
+    def log_residual_computed(
+        self,
+        state: MagCacheState,
+        residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        step = state.step_index
+        if residual is None:
+            residual_norm = 0.0
+            residual_shape = "None"
+        elif isinstance(residual, tuple):
+            residual_norm = sum(float(torch.norm(r).item()) for r in residual)
+            residual_shape = tuple(r.shape for r in residual)
+        else:
+            residual_norm = float(torch.norm(residual).item())
+            residual_shape = residual.shape
+        logger.debug(
+            f"[MagCache][TAIL] STEP={step}: RESIDUAL_COMPUTED (norm={residual_norm:.6f}, shape={residual_shape})"
+        )
+
+    def perform_calibration(
+        self,
+        state: MagCacheState,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> None:
+        if self._strategy is not None:
+            ratio, std, cos_dis = self._strategy.compute_calibration_metrics(current_residual, state.previous_residual)
+        else:
+
+            def _get_norm(residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+                if isinstance(residual, tuple):
+                    return sum(torch.linalg.norm(r.float(), dim=-1) for r in residual)
+                return torch.linalg.norm(residual.float(), dim=-1)
+
+            if state.previous_residual is None:
+                ratio, std, cos_dis = 1.0, 0.0, 0.0
+            else:
+                curr_norm = _get_norm(current_residual)
+                prev_norm = _get_norm(state.previous_residual)
+                ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
+                std = (curr_norm / (prev_norm + 1e-8)).std().item()
+                cos_dis = (
+                    (
+                        1
+                        - F.cosine_similarity(
+                            current_residual.flatten(0, -2) if current_residual.ndim > 2 else current_residual,
+                            state.previous_residual.flatten(0, -2)
+                            if state.previous_residual.ndim > 2
+                            else state.previous_residual,
+                            dim=-1,
+                            eps=1e-8,
+                        )
+                    )
+                    .mean()
+                    .item()
+                )
+
+        state.calibration_ratios.append(ratio)
+        state.norm_ratios.append(round(ratio, 5))
+        state.norm_stds.append(round(std, 5))
+        state.cos_dises.append(round(cos_dis, 5))
+
+    def advance_step(self, state: MagCacheState) -> None:
+        state.step_index += 1
+        if state.step_index >= self.config.num_inference_steps:
+            if self.config.mag_calibrate:
+                logger.info("MagCache calibration complete.")
+                logger.info(f"norm_ratios: {state.norm_ratios}")
+                logger.info(f"norm_stds: {state.norm_stds}")
+                logger.info(f"cos_dises: {state.cos_dises}")
+                logger.info("Copy these values to DiffusionCacheConfig(mag_ratios=...) for production use")
+
+            state.step_index = 0
+            state.accumulated_ratio = 1.0
+            state.accumulated_steps = 0
+            state.accumulated_err = 0.0
+            state.previous_residual = None
+            state.calibration_ratios = []
+            state.norm_ratios = []
+            state.norm_stds = []
+            state.cos_dises = []
+            state._is_first_step = True
+
+
+def apply_mag_cache_hook(
+    module: torch.nn.Module,
+    config: MagCacheConfig,
+    strategy: MagCacheStrategy | None = None,
+) -> None:
+    """Apply MagCache optimization to a transformer module.
+
+    Args:
+        module: Transformer model to optimize (e.g., FluxTransformer2DModel)
+        config: MagCacheConfig specifying caching parameters
+        strategy: Optional strategy to use. If None, will be looked up from registry.
+    """
+    HookRegistry.check_if_exists_or_initialize(module)
+
+    transformer_type = config.transformer_type
+    if strategy is None:
+        strategy = MagCacheStrategyRegistry.get_if_exists(transformer_type)
+
+    if strategy is None:
+        logger.warning(
+            f"MagCache: No strategy found for '{transformer_type}'. "
+            f"Using default behavior. Available strategies: {list(MagCacheStrategyRegistry._registry.keys())}"
+        )
+    else:
+        strategy_name = type(strategy).__name__
+        logger.info(f"MagCache: Applying {strategy_name} for '{transformer_type}'")
+        if hasattr(type(strategy), "register_blocks"):
+            type(strategy).register_blocks()
+
+    state_manager = StateManager(MagCacheState, (), {})
+    remaining_blocks = []
+
+    for name, submodule in module.named_children():
+        if not isinstance(submodule, torch.nn.ModuleList):
+            continue
+        for index, block in enumerate(submodule):
+            remaining_blocks.append((f"{name}.{index}", block))
+
+    if not remaining_blocks:
+        logger.warning("MagCache: No transformer blocks found to apply hooks.")
+        return
+
+    if len(remaining_blocks) == 1:
+        name, block = remaining_blocks[0]
+        logger.info(f"MagCache: Applying Head+Tail Hook to single block '{name}'")
+        _apply_mag_cache_head_hook(block, state_manager, config, strategy, is_tail=True)
+        return
+
+    head_block_name, head_block = remaining_blocks.pop(0)
+    tail_block_name, tail_block = remaining_blocks.pop(-1)
+
+    logger.info(f"MagCache: Applying Head Hook to {head_block_name}")
+    _apply_mag_cache_head_hook(head_block, state_manager, config, strategy)
+
+    for name, block in remaining_blocks:
+        _apply_mag_cache_block_hook(block, state_manager, config, strategy=strategy)
+
+    logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
+    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True, strategy=strategy)
+
+
+def _apply_mag_cache_head_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    strategy: MagCacheStrategy | None = None,
+    is_tail: bool = False,
+) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
+
+    hook = MagCacheHeadHook(state_manager, config, strategy, is_tail=is_tail)
+    registry.register_hook(_MAG_CACHE_LEADER_BLOCK_HOOK, hook)
+
+
+def _apply_mag_cache_block_hook(
+    block: torch.nn.Module,
+    state_manager: StateManager,
+    config: MagCacheConfig,
+    is_tail: bool = False,
+    strategy: MagCacheStrategy | None = None,
+) -> None:
+    registry = HookRegistry.check_if_exists_or_initialize(block)
+
+    if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None:
+        registry.remove_hook(_MAG_CACHE_BLOCK_HOOK)
+
+    hook = MagCacheBlockHook(state_manager, is_tail, config, strategy)
+    registry.register_hook(_MAG_CACHE_BLOCK_HOOK, hook)

vllm_omni/diffusion/cache/magcache/state.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache state management.
+
+This module contains the MagCacheState class which manages the internal state
+for MagCache caching logic, including residuals, accumulated metrics, and step tracking.
+"""
+
+import torch
+
+
+class MagCacheState:
+    """State management for MagCache caching logic."""
+
+    def __init__(self) -> None:
+        """Initialize empty MagCache state."""
+        self.previous_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
+        self.head_block_input: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None
+        self.should_compute: bool = True
+        self.accumulated_ratio: float = 1.0
+        self.accumulated_err: float = 0.0
+        self.accumulated_steps: int = 0
+        self.step_index: int = 0
+        self.calibration_ratios: list[float] = []
+        self.norm_ratios: list[float] = []
+        self.norm_stds: list[float] = []
+        self.cos_dises: list[float] = []
+        self._is_first_step: bool = True
+
+    def reset(self) -> None:
+        """Reset all state variables for a new inference run."""
+        self.previous_residual = None
+        self.head_block_input = None
+        self.should_compute = True
+        self.accumulated_ratio = 1.0
+        self.accumulated_err = 0.0
+        self.accumulated_steps = 0
+        self.step_index = 0
+        self.calibration_ratios = []
+        self.norm_ratios = []
+        self.norm_stds = []
+        self.cos_dises = []
+        self._is_first_step = True

vllm_omni/diffusion/cache/magcache/strategy.py

@@ -0,0 +1,488 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+MagCache strategy definitions for different model architectures.
+
+This module provides model-specific strategies for MagCache, allowing easy
+extension to new models by implementing the MagCacheStrategy interface.
+
+Architecture:
+- MagCacheStrategy: Abstract base class defining the strategy interface
+- FluxMagCacheStrategy: Strategy for Flux (dual-stream) models
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+import torch
+from diffusers.hooks._helpers import TransformerBlockMetadata
+
+
+class MagCacheStrategy(ABC):
+    """
+    Abstract base class for MagCache strategies.
+
+    Each model architecture requires a specific strategy to handle:
+    - Residual computation (how to calculate the residual for caching)
+    - Residual application (how to apply cached residual)
+    - Model-specific magnitude ratios
+
+    Implement this class to add support for new model architectures.
+    """
+
+    @property
+    @abstractmethod
+    def mag_ratios(self) -> torch.Tensor:
+        """Return the default mag_ratios tensor for this model.
+
+        This tensor defines caching ratios for each transformer block.
+        Values should be calibrated for the specific model architecture.
+
+        Returns:
+            1D tensor of mag_ratios (one per transformer block).
+        """
+        pass
+
+    def register_block_metadata(self, block_class: type) -> TransformerBlockMetadata | None:
+        """Register model-specific transformer block metadata.
+
+        Override this method to provide custom metadata for transformer blocks
+        that have non-standard output formats (e.g., tuple returns).
+
+        Args:
+            block_class: The transformer block class to register.
+
+        Returns:
+            TransformerBlockMetadata if custom registration is needed, None otherwise.
+        """
+        return None
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute residual between block output and input.
+
+        Default implementation: output - head_input.
+        Override this method for models with non-standard output formats.
+
+        Args:
+            output: Output from transformer blocks.
+            head_input: Input to the first block.
+
+        Returns:
+            Residual tensor for caching.
+        """
+        return output - head_input
+
+    def apply_residual(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply cached residual to hidden states.
+
+        Default implementation: add residual to hidden_states.
+        This works for most model architectures.
+
+        Args:
+            hidden_states: Current hidden states.
+            residual: Cached residual to apply.
+
+        Returns:
+            Hidden states with residual added.
+        """
+        return hidden_states + residual
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply cached residual tuple to both hidden_states and encoder_hidden_states.
+
+        Default implementation: add residuals separately.
+        Override this method for models with specific residual application logic.
+
+        Args:
+            hidden_states: Current hidden states.
+            encoder_hidden_states: Current encoder hidden states.
+            residual: Tuple of (hidden_states_residual, encoder_hidden_states_residual).
+
+        Returns:
+            Tuple of (hidden_states, encoder_hidden_states) with residuals applied.
+        """
+        h_res, e_res = residual
+        return hidden_states + h_res, encoder_hidden_states + e_res
+
+    def compute_calibration_metrics(
+        self,
+        current_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        previous_residual: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> tuple[float, float, float]:
+        """Compute calibration metrics for mag_ratios generation.
+
+        Default implementation computes norm ratios and cosine dissimilarity.
+        Override this method for models with custom metric computation.
+
+        Args:
+            current_residual: Residual from the current step.
+            previous_residual: Residual from the previous step (None for first step).
+
+        Returns:
+            Tuple of (norm_ratio, norm_std, cos_dis):
+            - norm_ratio: Mean ratio of current to previous residual norms
+            - norm_std: Standard deviation of the norm ratios
+            - cos_dis: Mean cosine dissimilarity (1 - cosine_similarity)
+        """
+        import torch.nn.functional as F
+
+        if previous_residual is None:
+            return 1.0, 0.0, 0.0
+
+        curr_norm = torch.linalg.norm(current_residual.float(), dim=-1)
+        prev_norm = torch.linalg.norm(previous_residual.float(), dim=-1)
+
+        ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
+        std = (curr_norm / (prev_norm + 1e-8)).std().item()
+        cos_dis = (1 - F.cosine_similarity(current_residual, previous_residual, dim=-1, eps=1e-8)).mean().item()
+
+        return ratio, std, cos_dis
+
+    def get_calibration_metrics_names(self) -> tuple[str, str, str]:
+        """Return the names of calibration metrics for logging.
+
+        Returns:
+            Tuple of metric names in order: (norm_ratio_name, norm_std_name, cos_dis_name)
+        """
+        return ("norm_ratio", "norm_std", "cos_dis")
+
+
+class FluxMagCacheStrategy(MagCacheStrategy):
+    """
+    MagCache strategy for Flux (dual-stream) models.
+
+    Flux architecture:
+    - transformer blocks (dual-stream): image tokens and text tokens
+      processed independently with separate weights
+    - single transformer blocks (single-stream): concatenated sequence
+      (image + text tokens) shares the same group of weights
+    - Final norm_out and proj_out layers
+
+    This strategy provides:
+    - mag_ratios: Pre-computed magnitude ratios for Flux (28 steps)
+    - compute_residual: Handles tuple output format
+    - apply_residual_tuple: Handles decoder residual only
+    """
+
+    FLUX_MAG_RATIOS = torch.tensor(
+        [
+            1.0,
+            1.07313,
+            1.21035,
+            1.04432,
+            1.06818,
+            1.05547,
+            1.0183,
+            1.03405,
+            1.02574,
+            1.03042,
+            1.02739,
+            1.01955,
+            1.01585,
+            1.02439,
+            1.01154,
+            1.01377,
+            1.00994,
+            1.01444,
+            1.00839,
+            1.02269,
+            1.0007,
+            1.00714,
+            1.00484,
+            1.01381,
+            1.00426,
+            0.99764,
+            1.00778,
+            1.00233,
+        ]
+    )
+
+    @property
+    def mag_ratios(self) -> torch.Tensor:
+        """Return default mag_ratios for Flux model."""
+        return self.FLUX_MAG_RATIOS
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute residual for Flux output format (tuple or single tensor).
+
+        Flux single transformer blocks return a tuple, so we extract
+        the decoder output (index 1) before computing residual.
+        """
+        if isinstance(output, tuple):
+            decoder_output = output[1] if len(output) > 1 else output[0]
+            return decoder_output - head_input
+        else:
+            return output - head_input
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply residual tuple for Flux - only add decoder residual.
+
+        Flux has separate image and text processing, so the residual
+        is only applied to the decoder (image) branch.
+        """
+        if isinstance(residual, tuple):
+            decoder_residual = residual[1]
+        else:
+            decoder_residual = residual
+
+        output = hidden_states + decoder_residual
+        enc_output = encoder_hidden_states
+
+        return output, enc_output
+
+    @staticmethod
+    def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor:
+        """Interpolate mag_ratios to target length using nearest neighbor."""
+        src_length = len(src_array)
+        if target_length == 1:
+            return src_array[-1:]
+
+        scale = (src_length - 1) / (target_length - 1)
+        grid = torch.arange(target_length, device=src_array.device, dtype=torch.float32)
+        mapped_indices = torch.round(grid * scale).long()
+
+        return src_array[mapped_indices]
+
+
+class Flux2MagCacheStrategy(FluxMagCacheStrategy):
+    """MagCache strategy for Flux2 model.
+
+    Flux2 shares the same dual-stream architecture as Flux, but may have
+    different tensor shapes in some transformer blocks, requiring special
+    handling in residual computation.
+    """
+
+    FLUX2_MAG_RATIOS = torch.tensor(
+        [
+            1.0,
+            0.96528,
+            1.11559,
+            1.0565,
+            1.00425,
+            1.0805,
+            0.98616,
+            1.09289,
+            1.03196,
+            1.06679,
+            1.03941,
+            1.05375,
+            1.03128,
+            1.05349,
+            1.01983,
+            1.05535,
+            1.0662,
+            1.05748,
+            1.00318,
+            1.05222,
+            1.04556,
+            1.0506,
+            1.05058,
+            1.05219,
+            1.02025,
+            1.05052,
+            1.04143,
+            1.0498,
+        ]
+    )
+
+    @property
+    def mag_ratios(self) -> torch.Tensor:
+        """Return default mag_ratios for Flux2 model."""
+        return self.FLUX2_MAG_RATIOS
+
+    def register_block_metadata(self, block_class: type) -> TransformerBlockMetadata | None:
+        """Register Flux2-specific block metadata based on block type.
+
+        Flux2 has two block types with different output formats:
+        - Dual-stream (Flux2TransformerBlock): returns (encoder_hidden_states, hidden_states)
+        - Single-stream (Flux2SingleTransformerBlock): returns single tensor
+        """
+        class_name = block_class.__name__
+        is_single_stream = "Single" in class_name
+
+        if is_single_stream:
+            return TransformerBlockMetadata(
+                return_hidden_states_index=0,
+                return_encoder_hidden_states_index=None,
+            )
+        else:
+            return TransformerBlockMetadata(
+                return_hidden_states_index=1,
+                return_encoder_hidden_states_index=0,
+            )
+
+    def compute_residual(
+        self,
+        output: torch.Tensor,
+        head_input: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Compute residual for Flux2 with dual-stream support.
+
+        For dual-stream blocks, computes residual for both encoder and decoder branches.
+        For single-stream blocks: if shapes match, computes output - input; otherwise returns output.
+        """
+        if isinstance(output, tuple):
+            enc_output, dec_output = output[0], output[1]
+
+            if isinstance(head_input, tuple):
+                enc_head, dec_head = head_input[0], head_input[1]
+            else:
+                enc_head = head_input[:, : enc_output.shape[1], ...]
+                dec_head = head_input[:, enc_output.shape[1] :, ...]
+
+            enc_residual = enc_output - enc_head
+            dec_residual = dec_output - dec_head
+
+            return (enc_residual, dec_residual)
+        else:
+            if output.shape == head_input.shape:
+                return output - head_input
+            else:
+                return output
+
+    def apply_residual_tuple(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        residual: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply residual tuple for Flux2 with shape mismatch handling.
+
+        For Flux2 dual-stream blocks, handles the case where hidden_states
+        and decoder_residual may have different shapes.
+        """
+        if isinstance(residual, tuple):
+            decoder_residual = residual[1]
+        else:
+            decoder_residual = residual
+
+        if hidden_states.shape == decoder_residual.shape:
+            output = hidden_states + decoder_residual
+        else:
+            output = hidden_states
+        enc_output = encoder_hidden_states
+
+        return output, enc_output
+
+    def apply_residual(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply residual for Flux2.
+
+        For single-stream blocks: if shapes match, adds residual; otherwise returns input.
+        For dual-stream blocks: applies decoder residual only.
+        """
+        if isinstance(residual, tuple):
+            dec_residual = residual[1]
+            if hidden_states.shape == dec_residual.shape:
+                return hidden_states + dec_residual
+            else:
+                return hidden_states
+        else:
+            if residual.shape == hidden_states.shape:
+                return hidden_states + residual
+            else:
+                return hidden_states
+
+
+class MagCacheStrategyRegistry:
+    """Registry for MagCache strategies by transformer type."""
+
+    _registry: dict[str, MagCacheStrategy] = {}
+
+    @classmethod
+    def register(cls, name: str, strategy: MagCacheStrategy) -> None:
+        """Register a strategy with explicit name.
+
+        Args:
+            name: Transformer model type identifier (e.g., "FluxTransformer2DModel")
+            strategy: MagCacheStrategy instance
+        """
+        cls._registry[name] = strategy
+
+    @classmethod
+    def get(cls, transformer_type: str) -> MagCacheStrategy:
+        """Get strategy for given transformer type."""
+        if transformer_type not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(f"Unknown model type: '{transformer_type}'. Available types: {available}")
+        return cls._registry[transformer_type]
+
+    @classmethod
+    def get_if_exists(cls, transformer_type: str) -> MagCacheStrategy | None:
+        """Get strategy if exists, None otherwise."""
+        return cls._registry.get(transformer_type)
+
+
+MagCacheStrategyRegistry.register("FluxTransformer2DModel", FluxMagCacheStrategy())
+
+MagCacheStrategyRegistry.register("Flux2Transformer2DModel", Flux2MagCacheStrategy())
+
+
+def register_strategy(
+    transformer_cls_name: str,
+    strategy: MagCacheStrategy,
+) -> None:
+    """Register a MagCache strategy for a model type.
+
+    This allows extending MagCache support to new models without modifying
+    the core MagCache code.
+
+    Args:
+        transformer_cls_name: Transformer model type identifier (class name or type string)
+                               Must match pipeline.transformer.__class__.__name__
+        strategy: MagCacheStrategy instance for this model type
+
+    Example:
+        >>> class MyModelMagCacheStrategy(MagCacheStrategy):
+        ...     @property
+        ...     def mag_ratios(self):
+        ...         return torch.tensor([...])
+        >>> register_strategy("MyModelTransformer", MyModelMagCacheStrategy())
+    """
+    MagCacheStrategyRegistry.register(transformer_cls_name, strategy)
+
+
+def get_strategy(transformer_cls_name: str) -> MagCacheStrategy:
+    """Get strategy function for given transformer class.
+
+    This function looks up the strategy based on the exact transformer_cls_name string,
+    which should match the transformer type in the pipeline (i.e., pipeline.transformer.__class__.__name__).
+
+    Args:
+        transformer_cls_name: Transformer class name (e.g., "FluxTransformer2DModel")
+                              Must exactly match a registered strategy.
+
+    Returns:
+        MagCacheStrategy instance for the model
+
+    Raises:
+        ValueError: If model type not found in registry
+    """
+    return MagCacheStrategyRegistry.get(transformer_cls_name)

vllm_omni/diffusion/cache/prompt_embed_cache.py

@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Prompt-embedding cache for diffusion pipelines.
+
+Text encoders are the single most expensive per-request preprocessing step in
+most diffusion pipelines, and their inputs are frequently identical across
+requests (e.g. in GRPO-style rollouts the same prompt is submitted many times
+with different seeds to sample different images). Re-running the text encoder
+for each of those requests wastes compute and GPU memory.
+
+This module provides a small LRU cache plus a transparent wrapper for a
+pipeline's ``encode_prompt`` method. Because almost every diffusion pipeline in
+``vllm_omni/diffusion/models`` routes all text-encoder invocations through
+``self.encode_prompt(...)``, wrapping that single method is sufficient to
+cache results model-wide without per-pipeline edits.
+
+Design points:
+    * The wrapper is installed by :class:`DiffusionModelRunner` after the
+      pipeline has loaded, so each runner process owns its own cache.
+    * Cache keys are derived from the bound ``encode_prompt`` arguments. Only
+      inputs we can safely hash (``str`` / ``int`` / ``float`` / ``bool`` /
+      ``None`` / ``bytes`` / ``torch.device`` / ``torch.dtype`` / numpy
+      scalars / nested lists/tuples/dicts of those) participate in the key.
+      If any argument is a tensor, PIL image, or other non-trivial object, we
+      bypass the cache for that call to guarantee correctness.
+    * If a caller passes precomputed ``*_embeds`` into ``encode_prompt`` the
+      wrapper also bypasses the cache because the call is already short-circuit.
+    * Cache values are detached tensors. Downstream pipeline code typically
+      does non-inplace ops (``.repeat``, ``.view``, slicing); we do not clone
+      on hit since outputs are treated as read-only.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import os
+import threading
+from collections import OrderedDict
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Sentinel returned by :func:`_hashable` when an input cannot be hashed
+# safely (e.g. tensor / PIL image / custom object). Any such input disables
+# caching for that specific call.
+_NOT_HASHABLE = object()
+
+# Sentinel returned by :meth:`PromptEmbedCache.get` to distinguish a cache
+# miss from a legitimate cached value of ``None``.
+_CACHE_MISS = object()
+
+# Argument names whose presence of a non-None value means the caller is
+# providing precomputed embeddings. In every pipeline we inspected, passing
+# these makes ``encode_prompt`` a cheap passthrough, so caching it has no
+# benefit and just wastes memory on large tensor keys.
+_PRECOMPUTED_EMBED_ARGS = (
+    "prompt_embeds",
+    "negative_prompt_embeds",
+    "pooled_prompt_embeds",
+    "negative_pooled_prompt_embeds",
+    "prompt_embeds_mask",
+    "negative_prompt_embeds_mask",
+)
+
+
+def _hashable(obj: Any) -> Any:
+    """Convert *obj* into a hashable representation or return ``_NOT_HASHABLE``.
+
+    Only inputs whose textual / scalar identity fully determines the
+    ``encode_prompt`` output are considered safe. Tensors and PIL images (which
+    flow through e.g. image-edit pipelines) deliberately disable caching for
+    that call rather than risk stale results. Common value-like configuration
+    objects such as ``torch.device`` and ``torch.dtype`` are normalized to
+    stable string forms so they can participate in cache keys safely.
+    """
+    if obj is None or isinstance(obj, (str, int, float, bool, bytes)):
+        return obj
+    if isinstance(obj, torch.device):
+        return ("__torch_device__", str(obj))
+    if isinstance(obj, torch.dtype):
+        return ("__torch_dtype__", str(obj))
+    if type(obj).__module__.split(".", 1)[0] == "numpy" and hasattr(obj, "item"):
+        try:
+            return _hashable(obj.item())
+        except (ValueError, TypeError):
+            return _NOT_HASHABLE
+    if isinstance(obj, (list, tuple)):
+        out = []
+        for item in obj:
+            h = _hashable(item)
+            if h is _NOT_HASHABLE:
+                return _NOT_HASHABLE
+            out.append(h)
+        return ("__seq__", tuple(out))
+    if isinstance(obj, dict):
+        items = []
+        for k in sorted(obj.keys(), key=repr):
+            v = _hashable(obj[k])
+            if v is _NOT_HASHABLE:
+                return _NOT_HASHABLE
+            items.append((repr(k), v))
+        return ("__dict__", tuple(items))
+    # torch.Tensor, PIL.Image, np.ndarray, and any other object → not safe.
+    return _NOT_HASHABLE
+
+
+def _detach_output(value: Any) -> Any:
+    """Detach tensors in the return value so the cache never holds grad state."""
+    if isinstance(value, torch.Tensor):
+        return value.detach()
+    if isinstance(value, tuple):
+        detached = tuple(_detach_output(v) for v in value)
+        # Preserve namedtuple types so attribute access keeps working.
+        if hasattr(value, "_fields"):
+            return type(value)(*detached)
+        return detached
+    if isinstance(value, list):
+        return [_detach_output(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _detach_output(v) for k, v in value.items()}
+    return value
+
+
+class PromptEmbedCache:
+    """Thread-safe LRU cache for ``encode_prompt`` outputs.
+
+    The cache stores whatever ``encode_prompt`` returns (tensor, tuple of
+    tensors, ``None``, etc.). Lookup / insertion is O(1) amortized. Eviction
+    is least-recently-used.
+    """
+
+    def __init__(self, max_size: int = 32, enabled: bool = True) -> None:
+        if max_size <= 0:
+            raise ValueError(f"max_size must be positive, got {max_size}")
+        self.max_size = max_size
+        self.enabled = enabled
+        self._store: OrderedDict[Any, Any] = OrderedDict()
+        self._lock = threading.Lock()
+        self.hits = 0
+        self.misses = 0
+        self.bypassed = 0
+
+    def get(self, key: Any) -> Any:
+        """Return the cached value or ``_CACHE_MISS`` if absent.
+
+        Returning a sentinel (rather than ``None``) lets callers cache
+        legitimate ``None`` results from the wrapped function.
+        """
+        with self._lock:
+            if key not in self._store:
+                self.misses += 1
+                return _CACHE_MISS
+            self._store.move_to_end(key)
+            self.hits += 1
+            return self._store[key]
+
+    def put(self, key: Any, value: Any) -> None:
+        with self._lock:
+            self._store[key] = value
+            self._store.move_to_end(key)
+            while len(self._store) > self.max_size:
+                self._store.popitem(last=False)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._store.clear()
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            return {
+                "size": len(self._store),
+                "max_size": self.max_size,
+                "hits": self.hits,
+                "misses": self.misses,
+                "bypassed": self.bypassed,
+            }
+
+
+def _build_key(
+    signature: inspect.Signature,
+    model_tag: str,
+    args: tuple,
+    kwargs: dict,
+) -> Any:
+    """Return a hashable cache key or ``None`` if this call must bypass cache.
+
+    We bind the raw positional / keyword arguments to the method signature so
+    that keyword and positional invocations with the same underlying values
+    collide in the cache.
+    """
+    try:
+        bound = signature.bind_partial(*args, **kwargs)
+    except TypeError:
+        return None
+    bound.apply_defaults()
+
+    # Bypass when the caller supplies precomputed embeddings.
+    for name in _PRECOMPUTED_EMBED_ARGS:
+        if name in bound.arguments and bound.arguments[name] is not None:
+            return None
+
+    items: list[tuple[str, Any]] = []
+    for name in sorted(bound.arguments.keys()):
+        h = _hashable(bound.arguments[name])
+        if h is _NOT_HASHABLE:
+            return None
+        items.append((name, h))
+    return (model_tag, tuple(items))
+
+
+def install_prompt_embed_cache(
+    pipeline: Any,
+    *,
+    max_size: int = 32,
+    enabled: bool = True,
+    model_tag: str | None = None,
+) -> PromptEmbedCache | None:
+    """Wrap ``pipeline.encode_prompt`` so results are cached by argument identity.
+
+    Idempotent: calling twice on the same pipeline is a no-op and returns the
+    existing cache. Returns ``None`` if the pipeline has no ``encode_prompt``
+    method.
+    """
+    existing = getattr(pipeline, "_prompt_embed_cache", None)
+    if existing is not None:
+        return existing
+
+    encode_fn = getattr(pipeline, "encode_prompt", None)
+    if encode_fn is None or not callable(encode_fn):
+        logger.debug(
+            "Prompt-embed cache: pipeline %s has no encode_prompt(); skipping.",
+            type(pipeline).__name__,
+        )
+        return None
+
+    try:
+        signature = inspect.signature(encode_fn)
+    except (TypeError, ValueError) as e:
+        logger.warning(
+            "Prompt-embed cache: cannot introspect encode_prompt on %s (%s); skipping.",
+            type(pipeline).__name__,
+            e,
+        )
+        return None
+
+    cache = PromptEmbedCache(max_size=max_size, enabled=enabled)
+    tag = model_tag or type(pipeline).__name__
+
+    @functools.wraps(encode_fn)
+    def cached_encode_prompt(*args: Any, **kwargs: Any) -> Any:
+        if not cache.enabled:
+            return encode_fn(*args, **kwargs)
+        key = _build_key(signature, tag, args, kwargs)
+        if key is None:
+            with cache._lock:
+                cache.bypassed += 1
+            return encode_fn(*args, **kwargs)
+        hit = cache.get(key)
+        if hit is not _CACHE_MISS:
+            return hit
+        out = encode_fn(*args, **kwargs)
+        cache.put(key, _detach_output(out))
+        return out
+
+    cached_encode_prompt.__wrapped__ = encode_fn  # type: ignore[attr-defined]
+    # Use object.__setattr__ to bypass diffusers ModelMixin attribute checks.
+    try:
+        pipeline.encode_prompt = cached_encode_prompt  # type: ignore[assignment]
+    except Exception:
+        object.__setattr__(pipeline, "encode_prompt", cached_encode_prompt)
+    object.__setattr__(pipeline, "_prompt_embed_cache", cache)
+    object.__setattr__(pipeline, "_prompt_embed_cache_original_fn", encode_fn)
+
+    logger.info(
+        "Prompt-embed cache installed on %s (max_size=%d, enabled=%s).",
+        tag,
+        max_size,
+        enabled,
+    )
+    return cache
+
+
+def uninstall_prompt_embed_cache(pipeline: Any) -> None:
+    """Restore the original ``encode_prompt`` on *pipeline* if wrapped."""
+    original = getattr(pipeline, "_prompt_embed_cache_original_fn", None)
+    if original is None:
+        return
+    try:
+        pipeline.encode_prompt = original  # type: ignore[assignment]
+    except Exception:
+        object.__setattr__(pipeline, "encode_prompt", original)
+    for attr in ("_prompt_embed_cache", "_prompt_embed_cache_original_fn"):
+        if hasattr(pipeline, attr):
+            try:
+                delattr(pipeline, attr)
+            except Exception:
+                object.__setattr__(pipeline, attr, None)
+
+
+def resolve_prompt_embed_cache_config(
+    enable: bool | None = None,
+    max_size: int | None = None,
+) -> tuple[bool, int]:
+    """Combine explicit args with env-var overrides.
+
+    Environment variables (useful for quick enablement in GRPO jobs without
+    touching config files):
+
+        ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE`` (``1``/``0``/``true``/``false``)
+        ``OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE`` (positive int)
+    """
+    env_enable = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE")
+    if env_enable is not None:
+        parsed = env_enable.strip().lower()
+        if parsed in ("1", "true", "yes", "on"):
+            enable = True
+        elif parsed in ("0", "false", "no", "off"):
+            enable = False
+
+    env_size = os.environ.get("OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE")
+    if env_size is not None:
+        try:
+            env_size_int = int(env_size)
+            if env_size_int > 0:
+                max_size = env_size_int
+        except ValueError:
+            logger.warning(
+                "Ignoring non-integer OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE=%r.",
+                env_size,
+            )
+
+    return bool(enable), int(max_size or 32)