[Rebase] Rebase to vllm releases/v0.22.0 (#3891)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: nainiu258 <cperfect02@163.com>
Signed-off-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Signed-off-by: JasonJ2021 <jasonj@zju.edu.cn>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: skf <54565339+skf-1999@users.noreply.github.com>
Co-authored-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Co-authored-by: Jason <72191212+JasonJ2021@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

.buildkite/test-merge.yml

@@ -107,7 +107,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 35
+        timeout_in_minutes: 60
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         agents:

.buildkite/test-ready.yml

@@ -183,7 +183,7 @@ steps:
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 15m bash -c "
+        timeout 20m bash -c "
                   pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
         "
     agents:

.gitignore

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

benchmarks/tts/model_configs.yaml

@@ -14,6 +14,7 @@ models:
     supported_tasks: [default_voice, voice_design]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       default_voice:
         voice: Vivian
@@ -27,6 +28,7 @@ models:
     supported_tasks: [voice_clone]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       voice_clone:
         task_type: Base
@@ -36,6 +38,7 @@ models:
     trust_remote_code: true
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       voice_clone: {}
       default_voice:

docker/Dockerfile.ci

@@ -1,5 +1,8 @@
+# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
+# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
+# be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.21.0
+ARG VLLM_BASE_TAG=v0.22.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docs/serving/speech_api.md

@@ -381,6 +381,44 @@ are cached in-process with a shared LRU so repeated requests with the same
 all TTS model types; deleting a voice invalidates every model-type slot at
 once.
 
+### Precomputed Custom Voices
+
+Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
+Generate a directory containing `custom_voice_manifest.json` plus one
+`.safetensors` file per voice, then set the pipeline-wide deploy config field:
+
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+
+Qwen3-TTS profiles are created with:
+
+```bash
+python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
+  --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+  --voice-name alice \
+  --ref-audio /path/to/reference.wav \
+  --ref-text "Original transcript of the reference audio" \
+  --mode icl \
+  --output-dir /path/to/custom_voices
+```
+
+VoxCPM2 profiles are created with:
+
+```bash
+python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
+  --model openbmb/VoxCPM2 \
+  --voice-name alice \
+  --ref-audio /path/to/reference.wav \
+  --mode ref_continuation \
+  --prompt-text "Original transcript of the reference audio" \
+  --output-dir /path/to/custom_voices
+```
+
+Only profiles whose safetensors payload can be loaded and validated are exposed
+by `GET /v1/audio/voices`. Valid precomputed voices can be used in
+`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
+
 **Configuration (environment variables):**
 
 | Variable | Default | Description |

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -380,6 +380,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
 ```
 Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
 
+### Precomputed custom voices
+For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
+```bash
+python qwen3_tts/precompute_custom_voice.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --voice-name alice \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio" \
+    --mode icl \
+    --output-dir /path/to/custom_voices
+```
+`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+Then start the server with that config and call the Speech API with only the voice name:
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
+
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
+    --output alice.wav
+```
+
 ### Streaming PCM
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
@@ -419,7 +444,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
 
 ### Notes
-- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
+- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
+- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
 
 ---
@@ -448,6 +474,23 @@ python voxcpm2/openai_speech_client.py \
 ```
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
 
+### Precomputed custom voices
+For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
+```bash
+python voxcpm2/precompute_custom_voice.py \
+    --model openbmb/VoxCPM2 \
+    --voice-name alice \
+    --ref-audio /path/to/reference.wav \
+    --mode ref_continuation \
+    --prompt-text "Original transcript of the reference audio" \
+    --output-dir /path/to/custom_voices
+```
+Add the output directory to the deploy config:
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
+
 ### Gradio demo (gapless streaming via AudioWorklet)
 ```bash
 python voxcpm2/gradio_demo.py

docs/user_guide/quantization/gguf.md

@@ -74,7 +74,7 @@ Online:
 vllm serve Qwen/Qwen-Image \
   --omni \
   --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
+  --diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
 ```
 
 ## Parameters

examples/offline_inference/lance/README.md

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

examples/offline_inference/lance/end2end.py

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

examples/online_serving/lance/README.md

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

examples/online_serving/lance/openai_chat_client.py

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

examples/online_serving/lance/run_server.sh

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

examples/online_serving/text_to_speech/README.md

@@ -412,6 +412,31 @@ curl -X POST http://localhost:8091/v1/audio/voices \
 ```
 Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
 
+### Precomputed custom voices
+For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
+```bash
+python qwen3_tts/precompute_custom_voice.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --voice-name alice \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio" \
+    --mode icl \
+    --output-dir /path/to/custom_voices
+```
+`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+Then start the server with that config and call the Speech API with only the voice name:
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
+
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
+    --output alice.wav
+```
+
 ### Streaming PCM
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
@@ -447,7 +472,8 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
 
 ### Notes
-- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
+- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
+- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
 
 ---
@@ -474,6 +500,23 @@ python voxcpm2/openai_speech_client.py \
 ```
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
 
+### Precomputed custom voices
+For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
+```bash
+python voxcpm2/precompute_custom_voice.py \
+    --model openbmb/VoxCPM2 \
+    --voice-name alice \
+    --ref-audio /path/to/reference.wav \
+    --mode ref_continuation \
+    --prompt-text "Original transcript of the reference audio" \
+    --output-dir /path/to/custom_voices
+```
+Add the output directory to the deploy config:
+```yaml
+custom_voice_dir: /path/to/custom_voices
+```
+After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
+
 ### Gradio demo (gapless streaming via AudioWorklet)
 ```bash
 python voxcpm2/gradio_demo.py

examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py

@@ -0,0 +1,243 @@
+"""Pre-compute Qwen3-TTS custom voice profiles.
+
+The generated directory can be passed to the server via
+``custom_voice_dir`` in ``vllm_omni/deploy/qwen3_tts.yaml``. Requests can then
+use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni.utils.custom_voice_io import safe_voice_stem  # noqa: E402
+
+MANIFEST_NAME = "custom_voice_manifest.json"
+_SPEAKER_ENCODER_SAMPLE_RATE = 24000
+_SPEAKER_ENCODER_NUM_MELS = 128
+_SPEAKER_ENCODER_N_FFT = 1024
+_SPEAKER_ENCODER_HOP_SIZE = 256
+_SPEAKER_ENCODER_WIN_SIZE = 1024
+_SPEAKER_ENCODER_FMIN = 0
+_SPEAKER_ENCODER_FMAX = 12000
+
+
+def _resolve_model_dir(model: str) -> str:
+    if os.path.isdir(model):
+        return model
+    from transformers.utils.hub import cached_file
+
+    config_path = cached_file(model, "config.json")
+    if config_path is None:
+        raise FileNotFoundError(f"Could not resolve config.json for {model}")
+    return os.path.dirname(config_path)
+
+
+def _read_audio_mono(path: str) -> tuple[np.ndarray, int]:
+    import soundfile as sf
+
+    wav, sr = sf.read(path, dtype="float32")
+    if wav.ndim > 1:
+        wav = wav.mean(axis=-1)
+    return np.asarray(wav, dtype=np.float32), int(sr)
+
+
+def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
+    if sr == target_sr:
+        return wav.astype(np.float32)
+    from vllm.multimodal.audio import AudioResampler
+
+    return AudioResampler(target_sr=target_sr).resample(wav.astype(np.float32), orig_sr=int(sr))
+
+
+def _load_speaker_encoder(model_dir: str, device: torch.device):
+    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 Qwen3TTSSpeakerEncoder
+
+    config = Qwen3TTSConfig.from_pretrained(model_dir)
+    encoder = Qwen3TTSSpeakerEncoder(config.speaker_encoder_config)
+    index_path = Path(model_dir) / "model.safetensors.index.json"
+    state: dict[str, torch.Tensor] = {}
+    if index_path.exists():
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+        weight_map = index.get("weight_map", {})
+        shard_names = sorted({v for k, v in weight_map.items() if k.startswith("speaker_encoder.")})
+        shard_paths = [Path(model_dir) / name for name in shard_names]
+    else:
+        shard_paths = sorted(Path(model_dir).glob("model*.safetensors"))
+    for shard in shard_paths:
+        with safe_open(str(shard), framework="pt", device="cpu") as f:
+            for key in f.keys():
+                if key.startswith("speaker_encoder."):
+                    state[key.removeprefix("speaker_encoder.")] = f.get_tensor(key)
+    if not state:
+        raise RuntimeError(
+            f"No speaker_encoder.* tensors found under {model_dir}. "
+            "Use a Qwen3-TTS Base checkpoint for pre-computing custom voice profiles."
+        )
+    encoder.load_state_dict(state)
+    encoder.to(device=device, dtype=torch.bfloat16).eval()
+    return config, encoder
+
+
+def _speaker_embedding(
+    encoder: torch.nn.Module,
+    speaker_encoder_config: object,
+    wav: np.ndarray,
+    sr: int,
+    device: torch.device,
+) -> torch.Tensor:
+    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import mel_spectrogram
+
+    target_sr = int(getattr(speaker_encoder_config, "sample_rate", _SPEAKER_ENCODER_SAMPLE_RATE))
+    mel_dim = int(getattr(speaker_encoder_config, "mel_dim", _SPEAKER_ENCODER_NUM_MELS))
+    if target_sr != _SPEAKER_ENCODER_SAMPLE_RATE or mel_dim != _SPEAKER_ENCODER_NUM_MELS:
+        raise ValueError(
+            "Unsupported Qwen3-TTS speaker encoder mel config for precompute: "
+            f"sample_rate={target_sr}, mel_dim={mel_dim}. "
+            f"Expected sample_rate={_SPEAKER_ENCODER_SAMPLE_RATE}, mel_dim={_SPEAKER_ENCODER_NUM_MELS} "
+            "to match the serving-time speaker embedding path."
+        )
+    wav = _resample(wav, sr, target_sr)
+    wav_t = torch.from_numpy(wav).to(device=device, dtype=torch.float32)
+    mels = mel_spectrogram(
+        wav_t.unsqueeze(0),
+        n_fft=_SPEAKER_ENCODER_N_FFT,
+        num_mels=_SPEAKER_ENCODER_NUM_MELS,
+        sampling_rate=_SPEAKER_ENCODER_SAMPLE_RATE,
+        hop_size=_SPEAKER_ENCODER_HOP_SIZE,
+        win_size=_SPEAKER_ENCODER_WIN_SIZE,
+        fmin=_SPEAKER_ENCODER_FMIN,
+        fmax=_SPEAKER_ENCODER_FMAX,
+    ).transpose(1, 2)
+    with torch.inference_mode():
+        return encoder(mels.to(device=device, dtype=torch.bfloat16))[0].float().cpu().contiguous()
+
+
+def _ref_code(model_dir: str, wav: np.ndarray, sr: int, device: torch.device) -> torch.Tensor:
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+    tok = Qwen3TTSTokenizer.from_pretrained(
+        str(Path(model_dir) / "speech_tokenizer"),
+        torch_dtype=torch.bfloat16,
+    )
+    try:
+        del tok.model.decoder
+        tok.model.decoder = None
+        tok.model.encoder.to(device)
+        tok.device = device
+    except Exception:
+        tok.device = device
+    with torch.inference_mode():
+        enc = tok.encode(wav, sr=int(sr), return_dict=True)
+    codes = getattr(enc, "audio_codes", None)
+    if isinstance(codes, list):
+        codes = codes[0] if codes else None
+    if not isinstance(codes, torch.Tensor):
+        raise RuntimeError("Speech tokenizer did not return audio_codes")
+    if codes.ndim == 3:
+        codes = codes[0]
+    return codes.to(dtype=torch.int32, device="cpu").contiguous()
+
+
+def _load_manifest(output_dir: Path, model: str, hidden_size: int) -> dict[str, Any]:
+    path = output_dir / MANIFEST_NAME
+    if path.exists():
+        return json.loads(path.read_text(encoding="utf-8"))
+    return {
+        "schema_version": 1,
+        "model_type": "qwen3_tts",
+        "model": model,
+        "hidden_size": hidden_size,
+        "voices": {},
+    }
+
+
+def _write_voice(
+    *,
+    model: str,
+    output_dir: Path,
+    voice_name: str,
+    ref_audio: str,
+    ref_text: str | None,
+    mode: str,
+    speaker_description: str | None,
+    device: torch.device,
+) -> None:
+    model_dir = _resolve_model_dir(model)
+    config, encoder = _load_speaker_encoder(model_dir, device)
+    wav, sr = _read_audio_mono(ref_audio)
+    if wav.size < 1024:
+        raise ValueError(f"Reference audio too short: {wav.size} samples")
+
+    tensors: dict[str, torch.Tensor] = {
+        "speaker_embedding": _speaker_embedding(encoder, config.speaker_encoder_config, wav, sr, device),
+    }
+    if mode == "icl":
+        if not ref_text or not ref_text.strip():
+            raise ValueError("--ref-text is required for --mode icl")
+        tensors["ref_code"] = _ref_code(model_dir, wav, sr, device)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    filename = f"{safe_voice_stem(voice_name)}.safetensors"
+    save_file(tensors, str(output_dir / filename))
+
+    hidden_size = int(getattr(config.talker_config, "hidden_size", tensors["speaker_embedding"].numel()))
+    manifest = _load_manifest(output_dir, model, hidden_size)
+    entry: dict[str, Any] = {
+        "name": voice_name,
+        "file": filename,
+        "mode": mode,
+        "embedding_dim": int(tensors["speaker_embedding"].numel()),
+    }
+    if "ref_code" in tensors:
+        entry["ref_code_length"] = int(tensors["ref_code"].shape[0])
+    if ref_text:
+        entry["ref_text"] = ref_text
+    if speaker_description:
+        entry["speaker_description"] = speaker_description
+    manifest.setdefault("voices", {})[voice_name] = entry
+    (output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
+    print(f"Wrote {output_dir / filename}")
+    print(f"Updated {output_dir / MANIFEST_NAME}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Pre-compute Qwen3-TTS custom voice profile")
+    parser.add_argument("--model", required=True, help="Qwen3-TTS Base model path or Hugging Face ID")
+    parser.add_argument("--voice-name", required=True)
+    parser.add_argument("--ref-audio", required=True)
+    parser.add_argument("--ref-text", default=None)
+    parser.add_argument("--mode", choices=["xvec", "icl"], default="xvec")
+    parser.add_argument("--speaker-description", default=None)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    _write_voice(
+        model=args.model,
+        output_dir=Path(args.output_dir),
+        voice_name=args.voice_name,
+        ref_audio=args.ref_audio,
+        ref_text=args.ref_text,
+        mode=args.mode,
+        speaker_description=args.speaker_description,
+        device=torch.device(args.device),
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py

@@ -0,0 +1,128 @@
+"""Pre-compute VoxCPM2 custom voice profiles.
+
+The generated directory can be passed to the server via
+``custom_voice_dir`` in ``vllm_omni/deploy/voxcpm2.yaml``. Requests can then
+use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni.utils.custom_voice_io import safe_voice_stem  # noqa: E402
+
+MANIFEST_NAME = "custom_voice_manifest.json"
+
+
+def _load_tts(model: str, device: torch.device):
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_import_utils import import_voxcpm2_core
+
+    VoxCPM = import_voxcpm2_core()
+    native = VoxCPM.from_pretrained(model, load_denoiser=False, optimize=False)
+    return native.tts_model.to(device).eval()
+
+
+def _load_manifest(output_dir: Path, model: str) -> dict[str, Any]:
+    path = output_dir / MANIFEST_NAME
+    if path.exists():
+        return json.loads(path.read_text(encoding="utf-8"))
+    return {
+        "schema_version": 1,
+        "model_type": "voxcpm2",
+        "model": model,
+        "voices": {},
+    }
+
+
+def _write_voice(
+    *,
+    model: str,
+    output_dir: Path,
+    voice_name: str,
+    ref_audio: str,
+    prompt_text: str | None,
+    mode: str,
+    speaker_description: str | None,
+    device: torch.device,
+) -> None:
+    if mode in ("continuation", "ref_continuation") and not prompt_text:
+        raise ValueError("--prompt-text is required for continuation/ref_continuation modes")
+
+    tts = _load_tts(model, device)
+    tensors: dict[str, torch.Tensor] = {}
+    with torch.inference_mode():
+        if mode in ("reference", "ref_continuation"):
+            tensors["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right").float().cpu().contiguous()
+        if mode in ("continuation", "ref_continuation"):
+            tensors["audio_feat"] = tts._encode_wav(ref_audio, padding_mode="left").float().cpu().contiguous()
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    filename = f"{safe_voice_stem(voice_name)}.safetensors"
+    save_file(tensors, str(output_dir / filename))
+
+    manifest = _load_manifest(output_dir, model)
+    entry: dict[str, Any] = {
+        "name": voice_name,
+        "file": filename,
+        "mode": mode,
+    }
+    if "ref_audio_feat" in tensors:
+        entry["ref_audio_feat_len"] = int(tensors["ref_audio_feat"].shape[0])
+    if "audio_feat" in tensors:
+        entry["audio_feat_len"] = int(tensors["audio_feat"].shape[0])
+    if prompt_text:
+        entry["prompt_text"] = prompt_text
+    if speaker_description:
+        entry["speaker_description"] = speaker_description
+
+    manifest.setdefault("voices", {})[voice_name] = entry
+    (output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
+    print(f"Wrote {output_dir / filename}")
+    print(f"Updated {output_dir / MANIFEST_NAME}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Pre-compute VoxCPM2 custom voice profile")
+    parser.add_argument("--model", default="openbmb/VoxCPM2", help="VoxCPM2 model path or Hugging Face ID")
+    parser.add_argument("--voice-name", required=True)
+    parser.add_argument("--ref-audio", required=True)
+    parser.add_argument(
+        "--prompt-text",
+        default=None,
+        help="Transcript of ref audio for continuation/ref_continuation modes",
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["reference", "continuation", "ref_continuation"],
+        default="reference",
+    )
+    parser.add_argument("--speaker-description", default=None)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = parser.parse_args()
+
+    _write_voice(
+        model=args.model,
+        output_dir=Path(args.output_dir),
+        voice_name=args.voice_name,
+        ref_audio=args.ref_audio,
+        prompt_text=args.prompt_text,
+        mode=args.mode,
+        speaker_description=args.speaker_description,
+        device=torch.device(args.device),
+    )
+
+
+if __name__ == "__main__":
+    main()

recipes/ByteDance/Lance.md

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

recipes/README.md

@@ -29,7 +29,6 @@ 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 |
@@ -48,6 +47,7 @@ recipes/
 | [`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

requirements/common.txt

@@ -1,4 +1,10 @@
 # Common dependencies for all platforms
+# transformers: keep aligned with upstream vLLM constraint.
+# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
+# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list
+# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed
+# _adjust_missing_and_unexpected_keys from list+list to set|set.
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -223,8 +223,12 @@ def _get_open_port() -> int:
         return s.getsockname()[1]
 
 
-def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None:
-    """Block until the given host:port accepts connections or timeout expires."""
+def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
+    """Block until the given host:port accepts connections or timeout expires.
+
+    If *proc* is provided, also monitors the process; raises RuntimeError
+    immediately if the server process exits before the port becomes available.
+    """
     start = time.time()
     while time.time() - start < timeout:
         try:
@@ -234,6 +238,10 @@ def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None:
                     return
         except Exception:
             pass
+        if proc is not None:
+            ret = proc.poll()
+            if ret is not None:
+                raise RuntimeError(f"Server process exited with code {ret} before port {host}:{port} became ready")
         time.sleep(2)
     raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s")
 
@@ -331,7 +339,7 @@ class DiffusionServer:
             env=env,
             cwd=str(Path(__file__).parent.parent.parent.parent),
         )
-        _wait_for_port(self.host, self.port)
+        _wait_for_port(self.host, self.port, proc=self.proc)
         print(f"DiffusionServer ready on {self.host}:{self.port}")
 
     def __enter__(self):

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -5,7 +5,8 @@
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
             "use_omni": false,
             "extra_cli_args": [
-                "--no-enable-prefix-caching"
+                "--no-enable-prefix-caching",
+                "--skip-mm-profiling"
             ]
         },
         "benchmark_params": [

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -17,7 +17,7 @@ AUDIO_MODEL = {
 }
 
 IMAGE_VIDEO_MODELS = {
-    "riverclouds/qwen_image_random": {"cuda": 2500, "rocm": 2100},
+    "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
 }
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}

tests/diffusion/test_diffusion_model_runner.py

@@ -301,8 +301,4 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
                 "vllm_config",
             ),
         ),
-        ("ir_op_priority", None),
-        ("enable_torch_wrap", vllm_config.compilation_config.ir_enable_torch_wrap),
-        ("enable_torch_wrap_exit", vllm_config.compilation_config.ir_enable_torch_wrap),
-        ("ir_op_priority_exit", None),
     ], ERROR_MESSAGE

tests/diffusion/test_stage_diffusion_proc.py

@@ -85,7 +85,7 @@ class MockOmniRequestOutput:
 BASE_HEIGHT = 512
 BASE_WIDTH = 512
 BASE_INFER_STEPS = 10
-DELAY_BASE = 2
+DELAY_BASE = 0.01
 
 
 class MockDiffusionEngine:
@@ -93,7 +93,7 @@ class MockDiffusionEngine:
         def simulate_step_delay(height, width, num_inference_steps) -> float:
             return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS)
 
-        DELAY_BASE = 2
+        DELAY_BASE = 0.01
         delay_scale = simulate_step_delay(
             request.sampling_params.height, request.sampling_params.width, request.sampling_params.num_inference_steps
         )
@@ -139,7 +139,7 @@ async def test_proc_process_request_with_batching_async_output():
     assert len(results) == len(test_requests)
     base_time = DELAY_BASE
     time_gap_std = DELAY_BASE * 2 * 2 * 1  # height/width/steps infer time scale
-    eps = 0.5
+    eps = 0.1
     for i, (res, elapsed_time) in enumerate(results):
         assert res.request_id == test_requests[i]["request_id"]
         assert isinstance(res, MockOmniRequestOutput)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -703,6 +703,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         kv_connector_output=None,
         cudagraph_stats=None,
         req_id_to_index={"req-s1": 0},
+        routed_experts=None,
         routed_experts_dict=None,
     )
 
@@ -789,6 +790,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         cudagraph_stats=None,
         req_id_to_index={"req-ar": 0},
         kv_extracted_req_ids=None,
+        routed_experts=None,
         routed_experts_dict=None,
     )
 

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -68,13 +68,13 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)  # PUB/SUB slow-joiner
 
     client = OmniCoordClientForStage(
-        router_addr,
+        coordinator.router_zmq_addr,
         "tcp://stage:coalesce",
         "tcp://stage:coalesce-out",
         0,
@@ -131,7 +131,7 @@ def test_omni_coordinator_registration_broadcast():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     # ZMQ PUB/SUB slow-joiner: allow SUB to connect before clients register.
@@ -139,9 +139,9 @@ def test_omni_coordinator_registration_broadcast():
 
     # Create 3 stage clients; each auto-registers on init.
     clients = [
-        OmniCoordClientForStage(router_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0),
-        OmniCoordClientForStage(router_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0),
-        OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
     ]
 
     msg = _wait_for_replica_list(sub, expected_count=3)
@@ -183,19 +183,19 @@ def test_omni_coordinator_heartbeat_timeout_handling():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)
 
     # A and B: real clients that send heartbeats every 5s.
-    client_a = OmniCoordClientForStage(router_addr, "tcp://stage:a", "tcp://stage:a-out", 0)
-    client_b = OmniCoordClientForStage(router_addr, "tcp://stage:b", "tcp://stage:b-out", 0)
+    client_a = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:a", "tcp://stage:a-out", 0)
+    client_b = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:b", "tcp://stage:b-out", 0)
 
     # C: raw DEALER that sends only registration, no heartbeat.
     dealer_ctx = zmq.Context.instance()
     dealer_c = dealer_ctx.socket(zmq.DEALER)
-    dealer_c.connect(router_addr)
+    dealer_c.connect(coordinator.router_zmq_addr)
     reg_event = {
         "input_addr": "tcp://stage:c",
         "output_addr": "tcp://stage:c-out",
@@ -254,12 +254,12 @@ def test_omni_coordinator_replica_shutdown_handling():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)  # PUB/SUB slow-joiner
 
-    client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
+    client = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
 
     msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None

tests/e2e/online_serving/test_lance.py

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

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -97,7 +97,7 @@ def _get_diffusion_feature_cases(model: str):
                     "--vae-patch-parallel-size",
                     "2",
                     "--vae-use-tiling",
-                    "--quantization-config",
+                    "--diffusion-quantization-config",
                     '{"method":"fp8"}',
                 ],
             ),

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,9 +1,13 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
+import base64
+import hashlib
+import io
 import json
 import logging
 import os
 import struct
+import wave
 from inspect import Signature, signature
 from pathlib import Path
 from types import SimpleNamespace
@@ -170,6 +174,27 @@ def create_mock_audio_output_for_test(
     )
 
 
+def _write_custom_voice_manifest(root: Path, *, model_type: str, voices: dict) -> None:
+    payload = {
+        "schema_version": 1,
+        "model_type": model_type,
+        "voices": voices,
+    }
+    (root / "custom_voice_manifest.json").write_text(json.dumps(payload), encoding="utf-8")
+
+
+def _wav_data_url(samples: np.ndarray, sample_rate: int) -> str:
+    pcm = (np.clip(samples, -1.0, 1.0) * 32767).astype("<i2")
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wav:
+        wav.setnchannels(1)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(pcm.tobytes())
+    encoded = base64.b64encode(buf.getvalue()).decode("ascii")
+    return f"data:audio/wav;base64,{encoded}"
+
+
 @pytest.fixture
 def test_app(mocker: MockerFixture):
     # Mock the engine client
@@ -972,6 +997,233 @@ class TestTTSMethods:
         params = speech_server._build_tts_params(req)
         assert "voice_clone_prompt" not in params
 
+    @pytest.mark.asyncio
+    async def test_resolve_ref_audio_reuses_decoded_audio_for_same_source(self, speech_server):
+        wav = np.linspace(-0.5, 0.5, 48000, dtype=np.float32)
+        ref_audio = _wav_data_url(wav, 24000)
+        speech_server.model_config.allowed_local_media_path = ""
+        speech_server.model_config.allowed_media_domains = None
+
+        first = await speech_server._resolve_ref_audio(ref_audio)
+        second = await speech_server._resolve_ref_audio(ref_audio)
+
+        assert first[1] == 24000
+        assert second[1] == 24000
+        assert first[0] is second[0]
+        assert first[0][0] == pytest.approx(float(wav[0]), abs=1e-4)
+        assert speech_server._get_resolved_ref_audio_artifact_key(
+            ref_audio
+        ) == speech_server._make_ref_audio_artifact_cache_key(np.asarray(first[0], dtype=np.float32), 24000)
+
+    def test_precomputed_qwen3_voice_infers_base_without_ref_audio(self, speech_server):
+        """Precomputed Qwen3 voices are reusable by name without per-request ref_audio."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "qwen3_tts",
+                "mode": "icl",
+                "ref_text": "reference transcript",
+                "ref_code_length": 3,
+            }
+        }
+        speech_server.supported_speakers = {"alice"}
+
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is None
+        assert req.task_type == "Base"
+
+        params = speech_server._build_tts_params(req)
+        assert params["task_type"] == ["Base"]
+        assert params["speaker"] == ["alice"]
+        assert params["x_vector_only_mode"] == [False]
+        assert params["ref_text"] == ["reference transcript"]
+        assert params["ref_code_length"] == [3]
+        assert "ref_audio" not in params
+
+    def test_uploaded_qwen3_voice_wins_over_same_named_precomputed_voice(self, speech_server, tmp_path):
+        """Qwen3 uploaded voices should take precedence over same-name precomputed voices."""
+        from safetensors.torch import save_file
+
+        uploaded_path = tmp_path / "alice.safetensors"
+        save_file({"speaker_embedding": torch.tensor([0.1] * 4)}, str(uploaded_path))
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.uploaded_speakers_dir = tmp_path
+        speech_server.uploaded_speakers = {
+            "alice": {
+                "name": "alice",
+                "created_at": 123,
+                "file_path": str(uploaded_path),
+                "embedding_source": "direct",
+            }
+        }
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "qwen3_tts",
+                "mode": "icl",
+                "ref_text": "precomputed transcript",
+                "ref_code_length": 3,
+            }
+        }
+
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is None
+
+        params = speech_server._build_tts_params(req)
+        assert params["speaker"] == ["alice"]
+        assert params["voice_created_at"] == [123]
+        assert params["task_type"] == ["Base"]
+        assert "voice_clone_prompt" in params
+        assert "ref_code_length" not in params
+        assert params.get("ref_text") != ["precomputed transcript"]
+
+    def test_precomputed_qwen3_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
+        """Manifest entries are not supported speakers unless their safetensors load."""
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="qwen3_tts",
+            voices={"Alice": {"file": "missing.safetensors", "mode": "xvec", "embedding_dim": 4}},
+        )
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                custom_voice_dir=str(tmp_path),
+                talker_config=SimpleNamespace(hidden_size=4),
+            )
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "alice" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is not None
+
+    def test_precomputed_qwen3_icl_without_ref_code_is_not_registered(self, speech_server, tmp_path):
+        """Qwen3 ICL profiles without ref_code cannot be exposed by the API layer."""
+        from safetensors.torch import save_file
+
+        save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="qwen3_tts",
+            voices={
+                "Alice": {
+                    "file": "alice.safetensors",
+                    "mode": "icl",
+                    "ref_text": "reference transcript",
+                    "embedding_dim": 4,
+                }
+            },
+        )
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                custom_voice_dir=str(tmp_path),
+                talker_config=SimpleNamespace(hidden_size=4),
+            )
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "alice" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Alice")
+        assert speech_server._validate_tts_request(req) is not None
+
+    def test_precomputed_voxcpm2_missing_safetensors_is_not_registered(self, speech_server, tmp_path):
+        """VoxCPM2 must not advertise a manifest-only voice that cannot hit prompt cache."""
+        _write_custom_voice_manifest(
+            tmp_path,
+            model_type="voxcpm2",
+            voices={"Bob": {"file": "missing.safetensors", "mode": "reference", "ref_audio_feat_len": 2}},
+        )
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(custom_voice_dir=str(tmp_path))
+        )
+
+        profiles = speech_server._load_precomputed_speakers()
+        assert profiles == {}
+
+        speech_server.precomputed_speakers = profiles
+        speech_server.supported_speakers = set(profiles)
+        assert "bob" not in speech_server.supported_speakers
+        req = OpenAICreateSpeechRequest(input="Hello", voice="Bob")
+        assert speech_server._validate_tts_request(req) is not None
+        assert speech_server._validate_tts_request(OpenAICreateSpeechRequest(input="Hello")) is None
+
+    def test_prepare_voxcpm2_rejects_supported_speaker_without_voice_profile(self, speech_server, mocker):
+        """VoxCPM2 named voices must be uploaded or precomputed on the real request path."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"bob", "default"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {}
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value="generator")
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={"prompt_token_ids": [1], "additional_information": {}}
+        )
+
+        with pytest.raises(ValueError, match="Invalid voice 'bob'"):
+            asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Bob")))
+
+        speech_server._build_voxcpm2_prompt.assert_not_awaited()
+        speech_server.engine_client.generate.assert_not_called()
+
+    def test_prepare_voxcpm2_accepts_default_voice(self, speech_server, mocker):
+        """VoxCPM2 default voice preserves the built-in zero-shot request path."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"default"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {}
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={"prompt_token_ids": [1], "additional_information": {}}
+        )
+
+        asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="default")))
+
+        speech_server._build_voxcpm2_prompt.assert_awaited_once()
+        speech_server.engine_client.generate.assert_called_once()
+
+    def test_prepare_voxcpm2_precomputed_voice_sets_model_cache_key(self, speech_server, mocker):
+        """VoxCPM2 precomputed voices must carry voice metadata to the model cache lookup."""
+        speech_server._tts_model_type = "voxcpm2"
+        speech_server.supported_speakers = {"alice"}
+        speech_server.uploaded_speakers = {}
+        speech_server.precomputed_speakers = {
+            "alice": {
+                "name": "Alice",
+                "model_type": "voxcpm2",
+                "mode": "reference",
+                "ref_audio_feat_len": 2,
+            }
+        }
+        speech_server.engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        speech_server.engine_client.generate = mocker.MagicMock(return_value=iter(()))
+        speech_server._build_voxcpm2_prompt = mocker.AsyncMock(
+            return_value={
+                "prompt_token_ids": [1],
+                "additional_information": {
+                    "voice_profile": speech_server.precomputed_speakers["alice"],
+                },
+            }
+        )
+
+        asyncio.run(speech_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="Hello", voice="Alice")))
+
+        prompt = speech_server.engine_client.generate.call_args.kwargs["prompt"]
+        additional = prompt["additional_information"]
+        assert additional["voice_name"] == "alice"
+        assert additional["voice_created_at"] == 0
+
     def test_build_tts_params(self, speech_server):
         """Test TTS parameter building."""
         req = OpenAICreateSpeechRequest(input="Hello", voice="Ryan", language="English")
@@ -2497,6 +2749,95 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
+        """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""
+        wav_list = [0.0] * 48000
+        artifact_key = "a" * 40
+        ref_audio = "data:audio/wav;base64,same"
+        qwen3_tts_server._put_resolved_ref_audio(
+            hashlib.sha1(ref_audio.encode("utf-8")).hexdigest(),
+            wav_list,
+            24000,
+            artifact_key,
+        )
+        qwen3_tts_server._ref_audio_model_artifact_ready.add(artifact_key)
+        qwen3_tts_server._codec_frame_rate = 25.0
+        qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))}
+        qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace(
+            codec_language_id={},
+            spk_is_dialect={},
+        )
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="Base",
+            ref_audio=ref_audio,
+            ref_text="reference",
+        )
+        request_id, _generator, tts_params = asyncio.run(
+            qwen3_tts_server._prepare_speech_generation(request, request_id="req-hot")
+        )
+
+        assert request_id == "req-hot"
+        assert "ref_audio" not in tts_params
+        assert tts_params["_qwen3_tts_ref_audio_cache_key"] == [artifact_key]
+        assert tts_params["ref_code_length"] == [50]
+        prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
+        assert prompt["additional_information"] is tts_params
+
+    def test_qwen3_ref_audio_artifact_ready_is_evicted_with_resolve_cache(self, qwen3_tts_server):
+        qwen3_tts_server._ref_audio_resolve_cache_max_entries = 1
+        qwen3_tts_server._ref_audio_resolve_cache_max_bytes = 1_000_000
+
+        qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a")
+        qwen3_tts_server._ref_audio_model_artifact_ready.add("artifact-a")
+        qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b")
+
+        assert "artifact-a" not in qwen3_tts_server._ref_audio_model_artifact_ready
+        assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()}
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_error(self, qwen3_tts_server):
+        async def failing_generator():
+            raise ValueError("boom")
+            yield  # pragma: no cover
+
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = "artifact-fail"
+
+        with pytest.raises(ValueError, match="boom"):
+            await anext(qwen3_tts_server._generate_audio_chunks(failing_generator(), "req-fail"))
+
+        assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-fail" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
+    @pytest.mark.asyncio
+    async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server):
+        async def pcm_generator():
+            yield SimpleNamespace(
+                multimodal_output={
+                    "audio": torch.zeros(16, dtype=torch.float32),
+                    "sr": 24000,
+                }
+            )
+            await asyncio.sleep(0)
+
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = "artifact-close"
+
+        stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close")
+        assert await anext(stream)
+        await stream.aclose()
+
+        assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-close" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
+    def test_qwen3_ref_audio_artifact_ready_requires_live_resolve_cache_entry(self, qwen3_tts_server):
+        qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = "artifact-evicted"
+
+        qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-evicted")
+
+        assert "req-evicted" not in qwen3_tts_server._request_ref_audio_artifact_keys
+        assert "artifact-evicted" not in qwen3_tts_server._ref_audio_model_artifact_ready
+
     def test_shutdown_is_idempotent(self, mocker: MockerFixture):
         """Calling shutdown() twice should not raise."""
         mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())

tests/entrypoints/test_stage_utils.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 from pytest_mock import MockerFixture
 
-from vllm_omni.entrypoints.stage_utils import set_stage_devices
+from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices
 
 
 def _make_dummy_torch(call_log):
@@ -130,3 +130,90 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
+
+
+# ---- _map_device_list unit tests ----
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_idempotency():
+    """Device IDs already in the visible set are returned as-is (idempotency)."""
+    result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"])
+    assert result == ["0", "1"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_idempotency_single():
+    """Single device ID in visible set passes through."""
+    result = _map_device_list(1, ["1"], ["0", "1"])
+    assert result == ["1"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_mapping():
+    """Device IDs < num_visible not in set are treated as indices."""
+    result = _map_device_list(0, ["0", "1"], ["6", "7", "8", "9"])
+    assert result == ["6", "7"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback():
+    """Device IDs >= num_visible are treated as physical IDs and passed through."""
+    result = _map_device_list(1, ["2"], ["0", "1"])
+    assert result == ["2"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback_multi():
+    """Multiple device IDs >= num_visible all pass through as physical IDs."""
+    result = _map_device_list(1, ["2", "3"], ["0", "1"])
+    assert result == ["2", "3"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback_mixed():
+    """When some devices are >= num_visible but some < num_visible, index mapping applies for the subset."""
+    # "1" < num_visible(2) so index mapping applies: visible[1] = "5"
+    # "2" >= num_visible(2) so it's dropped via the partial mapping path
+    result = _map_device_list(1, ["1", "2"], ["0", "5"])
+    assert result == ["5"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_with_gaps():
+    """Index-based mapping works with non-contiguous visible device lists."""
+    result = _map_device_list(0, ["0", "1"], ["4", "5", "7"])
+    assert result == ["4", "5"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_mapping_no_idempotency():
+    """Device IDs < num_visible not in visible set map via index even when no ID matches literally."""
+    # "0","1" with visible ["2","3"]: idempotency check fails (none in set),
+    # but index mapping succeeds: visible[0]="2", visible[1]="3"
+    result = _map_device_list(0, ["0", "1"], ["2", "3"])
+    assert result == ["2", "3"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_partial_mapping():
+    """When only a subset can map, returns the mapped subset (no error)."""
+    result = _map_device_list(0, ["0", "1", "2"], ["5", "6"])
+    assert result == ["5", "6"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_raises_on_non_numeric():
+    """Non-numeric device IDs raise ValueError."""
+    with pytest.raises(ValueError, match="must be non-negative integers"):
+        _map_device_list(0, ["a"], ["0", "1"])

tests/helpers/env.py

@@ -160,9 +160,41 @@ def _print_gpu_processes() -> None:
     print("=" * 80)
 
 
+def _cleanup_stale_device_locks() -> None:
+    """Remove stale device-initialization lock files whose recorded PID is dead.
+
+    Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a
+    crashed / killed test run and block subsequent orchestrator startups.
+    """
+    import glob as _glob
+
+    for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"):
+        try:
+            with open(lock_file) as fh:
+                content = fh.read().strip()
+            if not content:
+                continue
+            pid = int(content)
+        except (OSError, ValueError):
+            continue
+
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            print(f"  Removing stale device lock {lock_file} (PID {pid} is dead)")
+            try:
+                os.unlink(lock_file)
+            except OSError:
+                pass
+        except PermissionError:
+            pass
+
+
 def run_pre_test_cleanup() -> None:
     print("Pre-test GPU status:")
 
+    _cleanup_stale_device_locks()
+
     num_gpus = current_omni_platform.device_count()
     if num_gpus > 0:
         try:

tests/helpers/mark.py

@@ -21,7 +21,7 @@ def cuda_marks(*, res: str, num_cards: int):
         return marks
     test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards)
 
-    test_skipif = pytest.mark.skipif_cuda(
+    test_skipif = pytest.mark.skipif(
         not current_platform.is_cuda() or (current_platform.device_count() < num_cards),
         reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
     )
@@ -52,7 +52,7 @@ def xpu_marks(*, res: str, num_cards: int):
         return marks
     test_distributed = pytest.mark.distributed_xpu(num_cards=num_cards)
 
-    test_skipif = pytest.mark.skipif_xpu(
+    test_skipif = pytest.mark.skipif(
         not current_platform.is_xpu() or (current_platform.device_count() < num_cards),
         reason=f"Need at least {num_cards} XPUs to run the test.",
     )

tests/helpers/runtime.py

@@ -305,40 +305,132 @@ class OmniServer:
             time.sleep(2)
         raise RuntimeError(f"Server failed to start within {max_wait} seconds")
 
+    @staticmethod
+    def _reap_zombie(proc: "psutil.Process") -> bool:
+        """Reap a zombie child process via ``os.waitpid``.
+
+        ``psutil.Process.wait()`` uses ``pidfd_open`` + ``poll()``, which
+        never fires for a zombie (the zombie is already dead and will not
+        change state).  Since the test process is the parent, we can reap
+        the zombie directly with ``os.waitpid(pid, os.WNOHANG)`` and
+        retrieve its exit code.
+
+        Returns True if the process was a zombie and was reaped.
+        """
+        if proc.status() != psutil.STATUS_ZOMBIE:
+            return False
+        try:
+            os.waitpid(proc.pid, os.WNOHANG)
+        except ChildProcessError:
+            pass
+        return True
+
+    @staticmethod
+    def _wait_or_reap(proc: "psutil.Process", timeout: int) -> None:
+        """Wait for *proc* to exit, handling zombie state transparently.
+
+        When the process is already a zombie (e.g. orchestrator thread
+        hung during shutdown and the main process exited), ``pidfd_open``
+        + ``poll()`` inside ``psutil`` will never see a state change.
+        Fall back to ``os.waitpid`` to reap the zombie.
+        """
+        if OmniServer._reap_zombie(proc):
+            return
+        try:
+            proc.wait(timeout=timeout)
+        except psutil.TimeoutExpired:
+            OmniServer._reap_zombie(proc)
+
     def _kill_process_tree(self, pid):
+        """Kill the process tree rooted at *pid*.
+
+        Terminate the parent **first** so the OmniServer can gracefully shut
+        down its stage-engine children through the orchestrator.  This avoids
+        the ``subprocess died unexpectedly`` ERROR that the APIServer monitor
+        thread logs when children are killed before the parent, which in turn
+        can cause CI watchdogs to false-trigger on the upstream ``Shutdown
+        initiated`` message.
+
+        When the parent does not exit within the grace period (e.g. CPU-
+        offloaded workers stuck in CUDA D-state), the method falls back to
+        killing children first so the parent can be reaped cleanly.
+        """
         try:
             parent = psutil.Process(pid)
             children = parent.children(recursive=True)
-            all_pids = [pid] + [child.pid for child in children]
-
-            for child in children:
-                try:
-                    child.terminate()
-                except psutil.NoSuchProcess:
-                    pass
-
-            _, still_alive = psutil.wait_procs(children, timeout=10)
-
-            for child in still_alive:
-                try:
-                    child.kill()
-                except psutil.NoSuchProcess:
-                    pass
 
+            # 1. Terminate the parent first — let it run its graceful
+            #    shutdown cascade (orchestrator → stage pools → engine cores).
             try:
                 parent.terminate()
-                parent.wait(timeout=10)
-            except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+            except psutil.NoSuchProcess:
+                pass
+
+            # 2. Give the parent time to shut down its children cleanly.
+            parent_exited = False
+            try:
+                parent.wait(timeout=15)
+                parent_exited = True
+            except psutil.NoSuchProcess:
+                parent_exited = True
+            except psutil.TimeoutExpired:
+                parent_exited = OmniServer._reap_zombie(parent)
+
+            if not parent_exited:
+                # Parent is stuck — children (e.g. CPU-offloaded CFG workers)
+                # are likely in uninterruptible sleep.  Kill children first
+                # so the parent can be reaped without lingering as a zombie.
+                for child in children:
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+                psutil.wait_procs(children, timeout=5)
                 try:
                     parent.kill()
                 except psutil.NoSuchProcess:
                     pass
+                OmniServer._wait_or_reap(parent, timeout=5)
+            else:
+                # Parent exited cleanly — clean up any remaining children.
+                for child in children:
+                    try:
+                        if child.is_running():
+                            child.terminate()
+                    except psutil.NoSuchProcess:
+                        pass
 
+                gone, still_alive = psutil.wait_procs(children, timeout=10)
+
+                for child in still_alive:
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                try:
+                    if parent.is_running() and not OmniServer._reap_zombie(parent):
+                        parent.kill()
+                        parent.wait(timeout=10)
+                except psutil.NoSuchProcess:
+                    pass
+
+            # 3. Final sweep — ``kill -9`` anything that escaped.
             time.sleep(1)
-            alive_processes = []
-            for check_pid in all_pids:
-                if psutil.pid_exists(check_pid):
-                    alive_processes.append(check_pid)
+            alive_processes: list[int] = []
+            for child in children:
+                try:
+                    if child.is_running():
+                        alive_processes.append(child.pid)
+                except psutil.NoSuchProcess:
+                    pass
+            # Only count the parent as alive if it is NOT a zombie
+            # (zombies are already dead — just waiting to be reaped).
+            try:
+                if parent.is_running() and parent.status() != psutil.STATUS_ZOMBIE:
+                    alive_processes.append(parent.pid)
+            except psutil.NoSuchProcess:
+                pass
 
             if alive_processes:
                 print(f"Warning: Processes still alive: {alive_processes}")

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -25,6 +25,7 @@ class _FakeDecoder(nn.Module):
         self.total_upsample = total_upsample
         self.decode_calls: list[dict[str, int]] = []
         self.batched_decode_calls: list[dict[str, int]] = []
+        self.decode_codes: list[torch.Tensor] = []
         self.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
     def to(self, *args, **kwargs):
@@ -37,6 +38,7 @@ class _FakeDecoder(nn.Module):
         chunk_size: int = 300,
         left_context_size: int = 25,
     ) -> torch.Tensor:
+        self.decode_codes.append(codes.detach().cpu().clone())
         self.decode_calls.append(
             {
                 "chunk_size": chunk_size,
@@ -181,6 +183,136 @@ def test_forward_trims_trailing_padding_without_context():
     torch.testing.assert_close(audio, expected)
 
 
+def test_forward_reuses_cached_ref_context_for_followup_chunk():
+    model = _make_model()
+
+    first_codes = torch.tensor(
+        [
+            9,
+            8,
+            1,
+            2,
+            9,
+            8,
+            3,
+            4,
+        ],
+        dtype=torch.long,
+    )
+    model.forward(
+        input_ids=first_codes,
+        runtime_additional_information=[
+            {
+                "meta": {
+                    "left_context_size": 2,
+                    "ref_context_size": 2,
+                    "ref_context_request_id": "rid",
+                    "ref_context_included": True,
+                }
+            }
+        ],
+    )
+
+    followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
+    model.forward(
+        input_ids=followup_codes,
+        runtime_additional_information=[
+            {
+                "meta": {
+                    "left_context_size": 3,
+                    "ref_context_size": 2,
+                    "ref_context_request_id": "rid",
+                    "ref_context_included": False,
+                    "finished": torch.tensor(True),
+                }
+            }
+        ],
+    )
+
+    torch.testing.assert_close(
+        model.decoder.decode_codes[-1],
+        torch.tensor(
+            [
+                [
+                    [9, 8, 5, 6, 7],
+                    [9, 8, 15, 16, 17],
+                ]
+            ],
+            dtype=torch.long,
+        ),
+    )
+    assert "rid" not in model._ref_context_cache
+
+
+def test_forward_fails_fast_when_ref_context_cache_is_missing():
+    model = _make_model()
+
+    followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
+    with pytest.raises(ValueError, match="Missing Qwen3-TTS ref context cache"):
+        model.forward(
+            input_ids=followup_codes,
+            runtime_additional_information=[
+                {
+                    "meta": {
+                        "left_context_size": 3,
+                        "ref_context_size": 2,
+                        "ref_context_request_id": "rid",
+                        "ref_context_included": False,
+                    }
+                }
+            ],
+        )
+
+
+def test_ref_context_cache_evicts_lru_entries_when_requests_abort():
+    model = _make_model()
+    model._ref_context_cache_max_entries = 2
+
+    first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
+    for rid in ("a", "b", "c"):
+        model.forward(
+            input_ids=first_codes,
+            runtime_additional_information=[
+                {
+                    "meta": {
+                        "left_context_size": 2,
+                        "ref_context_size": 2,
+                        "ref_context_request_id": rid,
+                        "ref_context_included": True,
+                    }
+                }
+            ],
+        )
+
+    assert list(model._ref_context_cache) == ["b", "c"]
+    assert model._ref_context_cache_bytes == sum(model._tensor_nbytes(t) for t in model._ref_context_cache.values())
+
+
+def test_ref_context_cache_evicts_to_byte_cap():
+    model = _make_model()
+    model._ref_context_cache_max_entries = 100
+    model._ref_context_cache_max_bytes = 33
+
+    first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long)
+    for rid in ("a", "b"):
+        model.forward(
+            input_ids=first_codes,
+            runtime_additional_information=[
+                {
+                    "meta": {
+                        "left_context_size": 2,
+                        "ref_context_size": 2,
+                        "ref_context_request_id": rid,
+                        "ref_context_included": True,
+                    }
+                }
+            ],
+        )
+
+    assert list(model._ref_context_cache) == ["b"]
+    assert model._ref_context_cache_bytes == 32
+
+
 def test_connector_codec_chunking_does_not_override_decode_chunking():
     model = _make_model(
         async_chunk=True,

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -2,6 +2,7 @@ from collections import OrderedDict
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 import torch
 
 from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
@@ -9,6 +10,7 @@ from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
     PRECOMPUTED_REF_CODE_KEY,
     PRECOMPUTED_REF_IDS_KEY,
     PRECOMPUTED_TEXT_IDS_KEY,
+    REF_AUDIO_CACHE_KEY,
     Qwen3TTSPromptEmbedsBuilder,
 )
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
@@ -493,6 +495,41 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
     assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
 
 
+def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_without_normalize():
+    builder = _make_minimal_builder()
+    ref_code = torch.full((2, 2), 5, dtype=torch.long)
+    builder.put_ref_audio_artifacts("same-ref", ref_code=ref_code)
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
+        AssertionError("speech tokenizer not expected")
+    )
+
+    class FakeTextTokenizer:
+        def __call__(self, texts, *, padding=False):
+            return {"input_ids": [[7, 8, 9] for _ in texts]}
+
+    builder._text_tokenizer = FakeTextTokenizer()
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+            REF_AUDIO_CACHE_KEY: ["same-ref"],
+        }
+    }
+
+    builder.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], ref_code)
+    assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+
+
 def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
     builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
@@ -536,3 +573,112 @@ def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
     assert speaker_wav_ids == [id(ref_audio)]
     assert ref_code_len == 2
     assert torch.equal(out_ref_code, ref_code)
+
+
+def test_voice_clone_prompt_artifacts_do_not_resolve_ref_audio_for_cache_fill():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+    builder._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
+    builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("speaker embedding not expected")
+    )
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+
+    _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],
+            "voice_clone_prompt": {
+                "ref_code": ref_code,
+                "ref_spk_embedding": torch.ones(4, dtype=torch.bfloat16),
+                "icl_mode": True,
+            },
+        },
+    )
+
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)
+
+
+def test_ref_audio_artifact_only_uses_cache_without_ref_audio_payload():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+    builder._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
+    builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("speaker embedding not expected")
+    )
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+    builder.put_ref_audio_artifacts(
+        "same-ref",
+        ref_code=ref_code,
+        ref_spk_embedding=torch.ones(4, dtype=torch.bfloat16),
+    )
+
+    _prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+            REF_AUDIO_CACHE_KEY: ["same-ref"],
+        },
+    )
+
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)
+
+
+def test_ref_audio_artifact_only_cache_miss_fails_fast():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+
+    with pytest.raises(RuntimeError, match="artifact cache miss"):
+        builder.build_prompt_embeds(
+            task_type="Base",
+            info_dict={
+                "text": ["hello"],
+                "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+                "non_streaming_mode": [False],
+                REF_AUDIO_CACHE_KEY: ["missing-ref"],
+            },
+        )

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py

@@ -79,6 +79,27 @@ def test_estimate_prompt_len_handles_1d_ref_code() -> None:
     assert est > 50, f"got {est}; 1D ref_code must contribute its own length"
 
 
+def test_estimate_prompt_len_uses_ref_code_length_without_ref_audio() -> None:
+    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information={
+            "task_type": ["Base"],
+            "text": ["hello"],
+            "ref_text": ["reference transcript"],
+            "x_vector_only_mode": [False],
+            "ref_code_length": [6],
+            "_qwen3_tts_ref_audio_cache_key": ["same-ref"],
+            "language": ["English"],
+        },
+        task_type="Base",
+        tokenize_prompt=lambda _text: list(range(10)),
+        codec_language_id=None,
+        spk_is_dialect=None,
+        estimate_ref_code_len=lambda _ref_audio: None,
+    )
+
+    assert est == 15
+
+
 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

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -269,11 +269,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
 
     assert payload is not None
     assert payload.meta.left_context_size == 2
+    assert payload.meta.ref_context_size == 2
+    assert payload.meta.ref_context_request_id == rid
+    assert payload.meta.ref_context_included is True
     assert len(payload.codes.audio) == _Q * 12
 
 
-def test_ref_code_context_applies_to_all_streaming_chunks():
-    """ref_code is prepended as decoder context on every chunk, not just the first."""
+def test_followup_ref_code_context_is_sent_as_metadata_handle():
+    """Follow-up chunks keep full ref context semantically without resending it."""
     tm = _tm()
     rid = "r-ref2"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(35)]
@@ -289,9 +292,13 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
     )
 
     assert payload is not None
-    # ref_code (2 frames) prepended as left context on second chunk too
+    # ref_code (2 frames) is represented in metadata so Code2Wav can restore it
+    # from its request-local cache. It must not be resent in codes.audio.
     assert payload.meta.left_context_size == 10 + 2
-    assert len(payload.codes.audio) == _Q * (35 + 2)
+    assert payload.meta.ref_context_size == 2
+    assert payload.meta.ref_context_request_id == rid
+    assert payload.meta.ref_context_included is False
+    assert len(payload.codes.audio) == _Q * 35
 
 
 def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():

tests/test_arg_utils.py

@@ -164,7 +164,8 @@ def test_split_mixed_real_world():
     # orchestrator-only keys never reach engine
     assert "stage_init_timeout" not in engine
     assert "deploy_config" not in engine
-    assert "async_chunk" not in engine
+    # async_chunk is in SHARED_FIELDS so it flows to both
+    assert "async_chunk" in engine
 
 
 # ============================================================================
@@ -172,24 +173,47 @@ def test_split_mixed_real_world():
 # ============================================================================
 
 
-def test_user_typed_unclassified_warns(caplog):
+def test_user_typed_unclassified_warns(monkeypatch):
     """If the user types a flag we can't route, warn — don't silently drop."""
-    raw = {"bogus_flag": "value", "max_num_seqs": 64}
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    buf = io.StringIO()
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.WARNING)
+    vllm_logger.addHandler(handler)
+    try:
+        raw = {"bogus_flag": "value", "max_num_seqs": 64}
         split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"})
-    assert any("bogus_flag" in rec.message for rec in caplog.records), (
-        f"Expected warning mentioning 'bogus_flag', got: {[rec.message for rec in caplog.records]}"
-    )
+        assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}"
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
-def test_unclassified_without_user_typed_silent(caplog):
+def test_unclassified_without_user_typed_silent():
     """Without user_typed, unclassified keys drop silently (argparse defaults
     for server flags shouldn't spam logs on every launch)."""
-    raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    buf = io.StringIO()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.WARNING)
+    vllm_logger.addHandler(handler)
+    try:
+        raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
         split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None)
-    # No warnings because we don't know these were user-typed.
-    assert not any("host" in rec.message or "port" in rec.message for rec in caplog.records)
+        assert "host" not in buf.getvalue()
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
 # ============================================================================
@@ -290,7 +314,7 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     # Spot-check expected entries
     assert "stage_init_timeout" in blacklist
     assert "deploy_config" in blacklist
-    assert "async_chunk" in blacklist
+    assert "async_chunk" not in blacklist  # now in SHARED_FIELDS
     # Shared fields must NOT appear — they flow to both orchestrator and engine
     assert "model" not in blacklist
     assert "log_stats" not in blacklist
@@ -339,18 +363,31 @@ def test_ambiguous_field_strict_raises():
         split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True)
 
 
-def test_ambiguous_field_non_strict_routes_to_orchestrator(caplog):
+def test_ambiguous_field_non_strict_routes_to_orchestrator():
     """strict=False logs ERROR but routes the ambiguous field to orchestrator."""
 
     @dataclass
     class _AmbiguousEngine:
         deploy_config: str | None = None
 
-    with caplog.at_level(logging.ERROR, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    buf = io.StringIO()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.ERROR)
+    vllm_logger.addHandler(handler)
+    try:
         orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False)
-    assert orch.deploy_config == "x"
-    assert "deploy_config" not in engine
-    assert any("both OrchestratorArgs" in r.message for r in caplog.records)
+        assert orch.deploy_config == "x"
+        assert "deploy_config" not in engine
+        assert "both OrchestratorArgs" in buf.getvalue()
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
 # Sentinel-default precedence invariants (#3035)
@@ -426,24 +463,20 @@ def test_help_text_preserves_default_after_nullify():
 
 _OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset(
     {
-        "model_stage",
         "model_arch",
         "engine_output_type",
         "hf_config_name",
         "custom_process_next_stage_input_func",
         "subtalker_sampling_params",
-        "async_chunk",
         "omni_kv_config",
         "quantization_config",
         "worker_type",
         "task_type",
         "worker_cls",
-        "enable_sleep_mode",
         "omni_master_address",
         "omni_master_port",
         "stage_configs_path",
         "output_modalities",
-        "log_stats",
         "custom_pipeline_args",
     }
 )

tests/test_config_factory.py

@@ -1027,6 +1027,9 @@ class TestDeployConfigLoading:
             "async_scheduling",
             "compilation_config",
             "config_format",
+            "custom_voice_dir",
+            "data_parallel_size",
+            "devices",
             "disable_hybrid_kv_cache_manager",
             "enable_flashinfer_autotune",
             "enforce_eager",
@@ -1068,6 +1071,32 @@ class TestDeployConfigLoading:
             f"added={actual_fields - expected_fields}, removed={expected_fields - actual_fields}"
         )
 
+    def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path):
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+        deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml"
+        custom_voice_dir = tmp_path / "voices"
+        deploy_path.write_text(
+            f"""
+async_chunk: true
+custom_voice_dir: {custom_voice_dir}
+stages:
+  - stage_id: 0
+    devices: "0"
+  - stage_id: 1
+    devices: "0"
+""",
+            encoding="utf-8",
+        )
+
+        deploy = load_deploy_config(deploy_path)
+        pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert deploy.custom_voice_dir == str(custom_voice_dir)
+        assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
+
     def test_load_qwen3_omni_moe_deploy_config(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         deploy = load_deploy_config(deploy_path)

tests/test_tts_custom_voice_profiles.py

@@ -0,0 +1,283 @@
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _write_manifest(root, *, model_type: str, voices: dict) -> None:
+    payload = {
+        "schema_version": 1,
+        "model_type": model_type,
+        "voices": voices,
+    }
+    (root / "custom_voice_manifest.json").write_text(json.dumps(payload), encoding="utf-8")
+
+
+def test_qwen3_custom_voice_profiles_warm_speaker_cache(tmp_path):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+        Qwen3TTSTalkerForConditionalGeneration,
+    )
+
+    save_file(
+        {
+            "speaker_embedding": torch.arange(4, dtype=torch.float32),
+            "ref_code": torch.arange(6, dtype=torch.int32).reshape(3, 2),
+        },
+        str(tmp_path / "alice.safetensors"),
+    )
+    _write_manifest(
+        tmp_path,
+        model_type="qwen3_tts",
+        voices={
+            "Alice": {
+                "file": "alice.safetensors",
+                "mode": "icl",
+                "ref_text": "reference transcript",
+                "embedding_dim": 4,
+            }
+        },
+    )
+
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    model.config = SimpleNamespace(
+        custom_voice_dir=str(tmp_path),
+        speaker_encoder_config=SimpleNamespace(enc_dim=4),
+    )
+    model._speaker_cache = SpeakerEmbeddingCache()
+
+    model._load_custom_voice_profiles()
+
+    key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_icl", created_at=0)
+    cached = model._speaker_cache.get(key)
+    assert cached is not None
+    assert cached["icl_mode"] is True
+    assert cached["ref_text"] == "reference transcript"
+    torch.testing.assert_close(cached["ref_spk_embedding"], torch.arange(4, dtype=torch.float32))
+    torch.testing.assert_close(cached["ref_code"], torch.arange(6, dtype=torch.int32).reshape(3, 2))
+
+
+def test_qwen3_icl_profile_without_ref_code_is_not_downgraded(tmp_path):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+        Qwen3TTSTalkerForConditionalGeneration,
+    )
+
+    save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
+    _write_manifest(
+        tmp_path,
+        model_type="qwen3_tts",
+        voices={
+            "Alice": {
+                "file": "alice.safetensors",
+                "mode": "icl",
+                "ref_text": "reference transcript",
+                "embedding_dim": 4,
+            }
+        },
+    )
+
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    model.config = SimpleNamespace(
+        custom_voice_dir=str(tmp_path),
+        speaker_encoder_config=SimpleNamespace(enc_dim=4),
+    )
+    model._speaker_cache = SpeakerEmbeddingCache()
+
+    model._load_custom_voice_profiles()
+
+    icl_key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_icl", created_at=0)
+    xvec_key = model._speaker_cache.make_cache_key("alice", model_type="qwen3_tts_xvec", created_at=0)
+    assert model._speaker_cache.get(icl_key) is None
+    assert model._speaker_cache.get(xvec_key) is None
+
+
+def test_qwen3_prompt_len_uses_precomputed_ref_code_length():
+    from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+        Qwen3TTSPromptEmbedsBuilder,
+    )
+
+    prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information={
+            "text": ["hello"],
+            "task_type": ["Base"],
+            "speaker": ["alice"],
+            "x_vector_only_mode": [False],
+            "ref_code_length": [6],
+            "ref_text": ["reference transcript"],
+        },
+        task_type="Base",
+        tokenize_prompt=lambda _text: list(range(10)),
+        codec_language_id=None,
+        spk_is_dialect=None,
+        estimate_ref_code_len=lambda _ref_audio: None,
+    )
+
+    assert prompt_len == 15
+
+
+def test_custom_voice_profile_file_must_stay_under_manifest_dir(tmp_path):
+    from vllm_omni.utils.speaker_cache import iter_custom_voice_profiles
+
+    _write_manifest(
+        tmp_path,
+        model_type="qwen3_tts",
+        voices={"Alice": {"file": "../alice.safetensors", "mode": "xvec"}},
+    )
+
+    assert iter_custom_voice_profiles(tmp_path, expected_model_type="qwen3_tts") == []
+
+
+def test_custom_voice_profiles_drop_stale_provenance_metadata(tmp_path):
+    from vllm_omni.utils.speaker_cache import iter_custom_voice_profiles
+
+    save_file({"speaker_embedding": torch.arange(4, dtype=torch.float32)}, str(tmp_path / "alice.safetensors"))
+    _write_manifest(
+        tmp_path,
+        model_type="qwen3_tts",
+        voices={
+            "Alice": {
+                "file": "alice.safetensors",
+                "mode": "xvec",
+                "audio_codec_config_hash": "old-hash",
+                "audio_codec_version": "old-codec",
+                "model_revision": "old-revision",
+            }
+        },
+    )
+
+    profiles = iter_custom_voice_profiles(tmp_path, expected_model_type="qwen3_tts")
+
+    assert len(profiles) == 1
+    assert "audio_codec_config_hash" not in profiles[0]
+    assert "audio_codec_version" not in profiles[0]
+    assert "model_revision" not in profiles[0]
+
+
+def test_voxcpm2_custom_voice_profiles_warm_full_prompt_cache(tmp_path):
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+        VoxCPM2TalkerForConditionalGeneration,
+    )
+
+    save_file(
+        {
+            "ref_audio_feat": torch.ones(2, 3, 4, dtype=torch.float32),
+            "audio_feat": torch.full((5, 3, 4), 2.0, dtype=torch.float32),
+        },
+        str(tmp_path / "bob.safetensors"),
+    )
+    _write_manifest(
+        tmp_path,
+        model_type="voxcpm2",
+        voices={
+            "Bob": {
+                "file": "bob.safetensors",
+                "mode": "ref_continuation",
+                "prompt_text": "prompt transcript",
+                "ref_audio_feat_len": 2,
+                "audio_feat_len": 5,
+            }
+        },
+    )
+
+    model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
+    model.config = SimpleNamespace(custom_voice_dir=str(tmp_path))
+    model._speaker_cache = SpeakerEmbeddingCache()
+
+    model._load_custom_voice_profiles()
+
+    key = model._speaker_cache.make_cache_key("bob", model_type="voxcpm2", created_at=0)
+    cached = model._speaker_cache.get(key)
+    assert cached is not None
+    assert cached["mode"] == "ref_continuation"
+    assert cached["prompt_text"] == "prompt transcript"
+    torch.testing.assert_close(cached["ref_audio_feat"], torch.ones(2, 3, 4))
+    torch.testing.assert_close(cached["audio_feat"], torch.full((5, 3, 4), 2.0))
+
+
+def test_voxcpm2_custom_voice_profiles_skip_unloadable_manifest_entry(tmp_path):
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+        VoxCPM2TalkerForConditionalGeneration,
+    )
+
+    _write_manifest(
+        tmp_path,
+        model_type="voxcpm2",
+        voices={
+            "Bob": {
+                "file": "missing.safetensors",
+                "mode": "reference",
+                "ref_audio_feat_len": 2,
+            }
+        },
+    )
+
+    model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
+    model.config = SimpleNamespace(custom_voice_dir=str(tmp_path))
+    model._speaker_cache = SpeakerEmbeddingCache()
+
+    model._load_custom_voice_profiles()
+
+    key = model._speaker_cache.make_cache_key("bob", model_type="voxcpm2", created_at=0)
+    assert model._speaker_cache.get(key) is None
+
+
+def test_voxcpm2_precomputed_voice_cache_miss_raises_before_zero_shot():
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+        VoxCPM2TalkerForConditionalGeneration,
+    )
+
+    model = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
+    model.config = SimpleNamespace(bos_token_id=1)
+    model._speaker_cache = SpeakerEmbeddingCache()
+    model._get_multichar_zh_split = lambda: {}
+    model._get_or_create_state = lambda _req_id: SimpleNamespace()
+
+    with pytest.raises(ValueError, match="not loaded in the model cache"):
+        model.preprocess(
+            torch.tensor([1, 2], dtype=torch.long),
+            None,
+            request_id="req-1",
+            text_token_ids=[[2]],
+            voice_name=["bob"],
+            voice_created_at=0,
+            voice_profile={"mode": "reference"},
+        )
+
+
+def test_voxcpm2_prompt_can_account_for_precomputed_profile_lengths():
+    from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import build_voxcpm2_prompt
+
+    class _Tokenizer:
+        bos_token_id = 1
+        unk_token_id = 0
+
+        def get_vocab(self):
+            return {}
+
+        def encode(self, text, add_special_tokens=True):
+            ids = [10 + i for i, _ in enumerate(text.split())]
+            return [self.bos_token_id, *ids] if add_special_tokens else ids
+
+    prompt = build_voxcpm2_prompt(
+        hf_config=SimpleNamespace(audio_vae_config={"sample_rate": 48000, "encoder_rates": [2]}, patch_size=3),
+        tokenizer=_Tokenizer(),
+        split_map={},
+        text="hello world",
+        voice_profile={
+            "mode": "ref_continuation",
+            "ref_audio_feat_len": 2,
+            "audio_feat_len": 5,
+            "prompt_text": "prompt words",
+        },
+    )
+
+    # text(2) + audio_start(1) + ref prefix(2 + ref_len) + prompt_text(2) + prompt_audio_len(5)
+    assert len(prompt["prompt_token_ids"]) == 14
+    additional = prompt["additional_information"]
+    assert additional["voice_profile"]["mode"] == "ref_continuation"
+    assert additional["prompt_text"] == ["prompt words"]

tools/nightly/generate_nightly_perf_html.py

@@ -250,7 +250,7 @@ def _collect_diffusion_records(input_dir: str) -> list[dict[str, Any]]:
 
 
 def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
+    by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True)
     return sorted(
         by_date_desc,
         key=lambda r: (
@@ -264,8 +264,8 @@ def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
 
 
 def _sort_diffusion_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
-    return sorted(by_date_desc, key=lambda r: (r.get("test_name") or ""))
+    by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True)
+    return sorted(by_date_desc, key=lambda r: r.get("test_name") or "")
 
 
 def _ensure_parent_dir(path: str) -> None:

vllm_omni/config/pipeline_registry.py

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

vllm_omni/config/stage_config.py

@@ -516,6 +516,7 @@ class DeployConfig:
     enable_chunked_prefill: bool | None = None
     data_parallel_size: int | None = None
     pipeline_parallel_size: int | None = None
+    custom_voice_dir: str | None = None
 
 
 _STAGE_RESERVED_KEYS = frozenset(
@@ -693,6 +694,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
         "enable_chunked_prefill",
         "data_parallel_size",
         "pipeline_parallel_size",
+        "custom_voice_dir",
     ):
         if name in raw_dict:
             kwargs[name] = raw_dict[name]
@@ -814,6 +816,7 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = (
     "enable_chunked_prefill",
     "data_parallel_size",
     "pipeline_parallel_size",
+    "custom_voice_dir",
 )
 
 

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -24,9 +24,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
     uses_full_payload_input_coordinator,
 )
+from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.outputs import OmniConnectorOutput
 
@@ -321,7 +323,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     req = self.requests.get(req_id)
                     if req is not None and not req.is_finished():
                         outputs[req.client_index].append(
-                            EngineCoreOutput(
+                            OmniEngineCoreOutput(
                                 request_id=req_id,
                                 new_token_ids=[],
                                 kv_transfer_params={"kv_ready": True},
@@ -417,11 +419,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped = True
 
             if stopped:
-                if (
-                    model_runner_output.routed_experts_dict is not None
-                    and req_id in model_runner_output.routed_experts_dict
-                ):
-                    routed_experts = model_runner_output.routed_experts_dict[req_id]
+                if model_runner_output.routed_experts is not None:
+                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
 
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
@@ -452,7 +451,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+            if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs:
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -463,7 +462,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
                         finish_reason=finish_reason,
@@ -501,7 +500,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=request.request_id,
                         new_token_ids=[],
                         finish_reason=request.get_finished_reason(),
@@ -633,6 +632,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
+        session.discard_latest_async_tokens = False
         outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -29,9 +29,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
     uses_full_payload_input_coordinator,
 )
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
+from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
 
 logger = init_logger(__name__)
@@ -499,11 +501,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 stopped = True
 
             if stopped:
-                if (
-                    model_runner_output.routed_experts_dict is not None
-                    and req_id in model_runner_output.routed_experts_dict
-                ):
-                    routed_experts = model_runner_output.routed_experts_dict[req_id]
+                if model_runner_output.routed_experts is not None:
+                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 if not finished:
@@ -524,7 +523,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped_running_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+            if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs:
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if new_token_ids and self.structured_output_manager.should_advance(request):
@@ -543,7 +542,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
                         finish_reason=finish_reason,
@@ -580,7 +579,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         getattr(request, "external_req_id", None),
                     )
             outputs[request.client_index].append(
-                EngineCoreOutput(
+                OmniEngineCoreOutput(
                     request_id=request.request_id,
                     new_token_ids=[],
                     finish_reason=finish_reason,
@@ -607,7 +606,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=request.request_id,
                         new_token_ids=[],
                         finish_reason=request.get_finished_reason(),

vllm_omni/core/sched/utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared utilities for omni schedulers."""
+
+import numpy as np
+from vllm.v1.outputs import RoutedExpertsLists
+
+
+def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None:
+    """Extract per-request routed experts from RoutedExpertsLists using slot_mapping.
+
+    Matches upstream RoutedExpertsManager.get() pattern — filters routing_data
+    rows whose slot_mapping entries belong to this request's block_table.
+    """
+    if routed_experts is None:
+        return None
+    slots = getattr(request, "block_table", None)
+    if slots is None:
+        return None
+    slot_set = set(slots)
+    mask = np.isin(routed_experts.slot_mapping, list(slot_set))
+    data = routed_experts.routing_data[mask]
+    return data if data.size > 0 else None

vllm_omni/data_entry_keys.py

@@ -77,6 +77,9 @@ class OmniPayloadMeta(TypedDict, total=False):
     decode_flag: bool
     codec_streaming: bool
     ref_code_len: int
+    ref_context_size: int
+    ref_context_request_id: str
+    ref_context_included: bool
     talker_prefill_offset: int
 
 
@@ -161,6 +164,9 @@ class MetaStruct(_StructBase):
     decode_flag: bool | None = None
     codec_streaming: bool | None = None
     ref_code_len: int | None = None
+    ref_context_size: int | None = None
+    ref_context_request_id: str | None = None
+    ref_context_included: bool | None = None
     talker_prefill_offset: int | None = None
     codec_chunk_frames: int | None = None
     codec_left_context_frames: int | None = None

vllm_omni/deploy/lance.yaml

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

vllm_omni/diffusion/data.py

@@ -847,21 +847,6 @@ class OmniDiffusionConfig:
         self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
         self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
 
-    @staticmethod
-    def _looks_like_lance_subfolder(model: str | None) -> bool:
-        """Return True when ``--model`` points at a Lance per-component subfolder.
-
-        Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and
-        ``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may
-        reasonably hand the AR-style sub-checkpoint path directly.  The
-        ``LancePipeline`` constructor knows to walk up to the repo root from
-        either subfolder name.
-        """
-        if not model:
-            return False
-        base = os.path.basename(str(model).rstrip("/"))
-        return base in {"Lance_3B", "Lance_3B_Video"}
-
     def enrich_config(self) -> None:
         """Load model metadata from HuggingFace and populate config fields.
 
@@ -915,15 +900,6 @@ class OmniDiffusionConfig:
             else:
                 cfg = get_hf_file_to_dict("config.json", self.model)
                 if cfg is None:
-                    # Lance ships its top-level config.json one directory above
-                    # the per-checkpoint subfolders (``Lance_3B/`` or
-                    # ``Lance_3B_Video/``).  Try to recover that case before
-                    # raising.
-                    if self._looks_like_lance_subfolder(self.model):
-                        self.model_class_name = "LancePipeline"
-                        self.set_tf_model_config(TransformerConfig())
-                        self.update_multimodal_support()
-                        return
                     raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
 
                 self.set_tf_model_config(TransformerConfig.from_dict(cfg))
@@ -934,19 +910,6 @@ class OmniDiffusionConfig:
                     self.model_class_name = "BagelPipeline"
                     self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
-                elif (
-                    model_type == "lance"
-                    or "LancePipeline" in architectures
-                    or cfg.get("model_name") == "Lance"
-                    or self._looks_like_lance_subfolder(self.model)
-                ):
-                    # Lance ships a non-HF top-level config.json (model_name only)
-                    # plus per-component subfolders; resolve to the Lance pipeline.
-                    # Also accept --model pointing directly at the ``Lance_3B`` or
-                    # ``Lance_3B_Video`` subfolder by walking up to the repo root.
-                    self.model_class_name = "LancePipeline"
-                    self.set_tf_model_config(TransformerConfig())
-                    self.update_multimodal_support()
                 elif model_type == "neo_chat":
                     self.model_class_name = "SenseNovaU1Pipeline"
                     self.tf_model_config = TransformerConfig()

vllm_omni/diffusion/diffusion_engine.py

@@ -867,7 +867,7 @@ class DiffusionEngine:
                 worker_thread.join(timeout=10)
             if worker_thread.is_alive():
                 logger.warning(
-                    "Worker thread did not terminate within 10s; scheduler and executor shutdown will be skipped."
+                    "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred."
                 )
                 return
             else:

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -64,6 +64,7 @@ import contextlib
 import errno
 import logging
 import os
+import time
 from collections.abc import Iterable, Iterator
 
 logger = logging.getLogger(__name__)
@@ -112,9 +113,43 @@ def _safe_repo_filename(model: str) -> str:
     return model.replace("/", "__").replace(os.sep, "__") + ".lock"
 
 
+def _dotfile_lock_acquire(lock_dir: str, model: str, timeout: float = 300.0, poll_interval: float = 0.5) -> bool:
+    """Acquire an exclusive lock via atomic directory creation.
+
+    ``os.makedirs(path, exist_ok=False)`` is atomic on POSIX filesystems
+    (including Lustre and NFS), making it a reliable fallback when
+    ``fcntl.flock`` is unsupported (``ENOLCK`` on FSx/Lustre mounts
+    without the ``flock`` mount option).
+
+    Returns ``True`` if the lock was acquired, ``False`` on timeout.
+    The caller must call ``os.rmdir(lock_path)`` to release.
+    """
+    lock_path = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
+
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            os.makedirs(lock_path, exist_ok=False)
+            logger.info("Acquired dotfile prefetch lock for %s at %s", model, lock_path)
+            return True
+        except FileExistsError:
+            if time.monotonic() >= deadline:
+                logger.warning(
+                    "Timed out waiting for dotfile prefetch lock %s after %.0fs; proceeding unlocked",
+                    lock_path,
+                    timeout,
+                )
+                return False
+            time.sleep(poll_interval)
+
+
 @contextlib.contextmanager
 def _repo_prefetch_lock(model: str) -> Iterator[None]:
-    """Hold an exclusive ``fcntl.flock`` keyed on the HF repo id.
+    """Hold an exclusive lock keyed on the HF repo id.
+
+    Tries ``fcntl.flock`` first; falls back to an atomic directory-creation
+    dotfile lock when flock is unsupported (e.g. FSx for Lustre without the
+    ``flock`` mount option, which returns ``ENOLCK``).
 
     Why we still need a node-wide lock on top of ``snapshot_download``'s
     own per-blob ``.lock`` files:
@@ -131,71 +166,93 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]:
     * Multiple unrelated processes on the same node (concurrent CI jobs,
       multiple ``OmniServer`` instances, every ``DiffusionWorker``
       subprocess from the multiproc executor) can enter this code path
-      in parallel. ``flock`` on a per-repo file makes the whole
+      in parallel. An exclusive lock makes the whole
       ``snapshot_download(...)`` call a true node-wide critical
       section, so the *first* entrant fully populates the cache and
       every subsequent ``from_pretrained`` sees a warm, complete tree.
 
-    Best-effort: if we cannot create / acquire the lock (no ``fcntl``,
-    read-only FS, NFS without flock support) we log and run the download
-    anyway. Worst case behaviour reverts to "snapshot_download +
-    transformers v5 race" which is exactly the pre-fix state.
+    Best-effort: if we cannot acquire any lock (no ``fcntl``, read-only
+    FS, dotfile lock timeout) we log and run the download anyway. Worst
+    case behaviour reverts to "snapshot_download + transformers v5 race"
+    which is exactly the pre-fix state.
     """
+    lock_dir = None
+    lock_path = None
+    dotfile_held = None
+    fd = None
+    flock_held = False
+
+    # --- fcntl.flock path ---
     try:
         import fcntl  # type: ignore[import-not-found]
     except ImportError:  # pragma: no cover - non-POSIX (Windows)
-        logger.debug("fcntl unavailable on this platform; running prefetch unlocked")
-        yield
-        return
+        fcntl = None
 
-    try:
-        lock_dir = _node_lock_dir()
-    except OSError as exc:
-        logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
-        yield
-        return
-
-    lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
-
-    try:
-        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
-    except OSError as exc:
-        logger.warning("Could not open prefetch lock %s (%s); proceeding without flock", lock_path, exc)
-        yield
-        return
-
-    locked = False
-    try:
+    if fcntl is not None:
         try:
-            fcntl.flock(fd, fcntl.LOCK_EX)
-            locked = True
-            logger.debug("Acquired prefetch flock for %s at %s", model, lock_path)
+            lock_dir = _node_lock_dir()
         except OSError as exc:
-            # FS without flock support (ENOLCK / EOPNOTSUPP / EINVAL).
-            # EACCES on some Lustre configs without ``flock`` mount opt.
-            if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
-                logger.warning(
-                    "fcntl.flock not supported on '%s' (%s: %s); running prefetch unlocked",
-                    lock_path,
-                    type(exc).__name__,
-                    exc,
-                )
-            else:
-                raise
+            logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
+            fcntl = None  # force dotfile fallback
+
+    if fcntl is not None and lock_dir is not None:
+        lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
+        try:
+            fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
+        except OSError as exc:
+            logger.warning("Could not open prefetch lock %s (%s); falling back to dotfile lock", lock_path, exc)
+            fd = None
+
+        if fd is not None:
+            try:
+                fcntl.flock(fd, fcntl.LOCK_EX)
+                flock_held = True
+                logger.info("Acquired flock prefetch lock for %s at %s", model, lock_path)
+            except OSError as exc:
+                if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
+                    logger.warning(
+                        "fcntl.flock not supported on '%s' (%s: %s); falling back to dotfile lock",
+                        lock_path,
+                        type(exc).__name__,
+                        exc,
+                    )
+                    os.close(fd)
+                    fd = None
+                else:
+                    raise
+
+    # --- dotfile fallback ---
+    if not flock_held:
+        if lock_dir is None:
+            try:
+                lock_dir = _node_lock_dir()
+            except OSError as exc:
+                logger.warning("Could not allocate lock dir (%s); running prefetch unlocked", exc)
+                yield
+                return
+
+        if _dotfile_lock_acquire(lock_dir, model):
+            dotfile_held = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
+
+    try:
         yield
     finally:
-        if locked:
+        if flock_held and fd is not None:
             with contextlib.suppress(OSError):
                 fcntl.flock(fd, fcntl.LOCK_UN)
-        with contextlib.suppress(OSError):
-            os.close(fd)
+        if fd is not None:
+            with contextlib.suppress(OSError):
+                os.close(fd)
+        if dotfile_held is not None:
+            with contextlib.suppress(OSError):
+                os.rmdir(dotfile_held)
 
 
 def prefetch_subfolders(
     model: str,
     subfolders: Iterable[str],
     *,
-    local_files_only: bool = False,
+    local_files_only: bool | None = None,
     include_root_metadata: bool = True,
 ) -> None:
     """Materialise ``model``'s ``subfolders`` in the HF cache before loading.
@@ -206,17 +263,21 @@ def prefetch_subfolders(
         subfolders: Iterable of subfolder names (e.g. ``["text_encoder",
             "vae"]``) whose contents need to be fully present before any
             worker calls ``from_pretrained(subfolder=...)``.
-        local_files_only: When True, skip the prefetch entirely. The caller
-            has explicitly promised the cache is already populated (as happens
-            for local model checkouts), so hitting the network would defeat
-            the intent and may fail in air-gapped environments.
+        local_files_only: When ``True``, skip the prefetch entirely.
+            When ``None`` (default), auto-detect: skip if *model* is a
+            local directory, run otherwise.
         include_root_metadata: When True, also pull ``*.json`` at the repo
             root so ``model_index.json`` / ``config.json`` resolution during
             ``from_pretrained`` also hits a warm cache.
     """
+    if local_files_only is None:
+        local_files_only = os.path.isdir(model)
+
     if local_files_only or not model or os.path.isdir(model):
         return
 
+    logger.info("Prefetching %s subfolders: %s", model, list(subfolders))
+
     try:
         from huggingface_hub import snapshot_download
     except ImportError:  # pragma: no cover - huggingface_hub is a hard dep
@@ -251,6 +312,7 @@ def prefetch_subfolders(
                 repo_id=model,
                 allow_patterns=allow_patterns,
             )
+        logger.info("Prefetch complete for %s", model)
     except Exception as exc:
         # Best-effort: propagate only via logging. The subsequent
         # ``from_pretrained`` call will raise a clearer, call-site-specific
@@ -312,3 +374,31 @@ def _looks_like_auth_error(exc: BaseException) -> bool:
     # only load-bearing signal.
     msg = str(exc).lower()
     return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
+
+
+def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0):
+    """Call *load_fn* with retry on the transformers v5 shard-resolution race.
+
+    When the prefetch lock cannot be acquired (e.g. flock unsupported on
+    the filesystem and dotfile lock times out), ``from_pretrained`` may
+    still hit the ``cached_files`` race. This wrapper retries with
+    exponential backoff when the OSError message matches the specific
+    "does not appear to have a file named" pattern.
+    """
+    for attempt in range(max_retries):
+        try:
+            return load_fn()
+        except OSError as exc:
+            if "does not appear to have a file" not in str(exc):
+                raise
+            if attempt == max_retries - 1:
+                raise
+            delay = base_delay * (attempt + 1)
+            logger.warning(
+                "from_pretrained failed with shard-resolution race (%s); retrying in %.1fs (attempt %d/%d)",
+                exc,
+                delay,
+                attempt + 1,
+                max_retries,
+            )
+            time.sleep(delay)

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -106,16 +106,7 @@ class BagelRotaryEmbedding(nn.Module):
     linear, dynamic-NTK, YaRN, …), we delegate the ``inv_freq`` /
     ``attention_scaling`` computation to HF's ``ROPE_INIT_FUNCTIONS`` so
     that the frequency basis and scaling factor are identical to the
-    original checkpoint.
-
-    For Qwen2.5-VL-style multimodal RoPE (``rope_scaling.rope_type == "mrope"``)
-    the ``inv_freq`` basis is the standard default-rope one; the difference
-    is that position ids are 3-D ``(t, h, w)`` per token and the
-    ``mrope_section`` describes how the head dimension is split across axes.
-    This module accepts either 2-D scalar position ids ``(B, S)`` or 3-D
-    multimodal position ids ``(B, 3, S)`` and dispatches accordingly so the
-    same module works for both BAGEL (1-D rope) and Lance (Qwen2.5-VL mrope).
-    This module has no learnable parameters.
+    original checkpoint.  This module has no learnable parameters.
     """
 
     def __init__(self, config):
@@ -129,17 +120,10 @@ class BagelRotaryEmbedding(nn.Module):
         rope_type = (
             rope_scaling.get("rope_type") or rope_scaling.get("type") or rope_parameters.get("rope_type") or "default"
         )
-        # Cache mrope_section for the forward-time section-split.
-        self._mrope_section: list[int] | None = None
-        if rope_type == "mrope":
-            section = rope_scaling.get("mrope_section") or rope_parameters.get("mrope_section")
-            if section is None:
-                raise ValueError("rope_scaling.rope_type == 'mrope' requires 'mrope_section'.")
-            self._mrope_section = list(section)
 
-        if rope_type in ("default", "mrope"):
-            # mrope shares the default sinusoidal frequency basis; the
-            # multimodal split happens at forward time, not at init.
+        if rope_type == "default":
+            # transformers>=5.0 removed the 'default' entry from
+            # ROPE_INIT_FUNCTIONS; use the plain sinusoidal formula.
             rope_theta = (
                 getattr(config, "rope_theta", None)
                 or rope_parameters.get("rope_theta")
@@ -163,36 +147,11 @@ class BagelRotaryEmbedding(nn.Module):
 
         Args:
             x: Input tensor (only used for dtype inference).
-            position_ids: Either 2-D scalar ``(batch_size, seq_len)`` for plain
-                1-D RoPE, or 3-D multimodal ``(batch_size, 3, seq_len)`` for
-                Qwen2.5-VL-style mRoPE.  The latter is auto-detected from
-                ``position_ids.ndim``.
+            position_ids: Position indices, shape (batch_size, seq_len).
 
         Returns:
             cos, sin: Rotary embeddings, each of shape (batch_size, seq_len, dim).
         """
-        if position_ids.ndim == 3 and self._mrope_section is not None:
-            # multimodal path: position_ids is (B, 3, S) with rows = (t, h, w).
-            # Compute per-axis frequencies, then assemble the per-section
-            # rotary basis matching Qwen2-VL's ``apply_multimodal_rotary_pos_emb``.
-            B = position_ids.shape[0]
-            inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(B, 3, -1, 1)
-            position_ids_expanded = position_ids[:, :, None, :].float()
-            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)
-            # ``freqs`` is (B, 3, S, head_dim/2); double along the last axis
-            # to get the full ``head_dim`` rotary basis.
-            emb = torch.cat((freqs, freqs), dim=-1)  # (B, 3, S, head_dim)
-            cos_per_axis = emb.cos() * self.attention_scaling
-            sin_per_axis = emb.sin() * self.attention_scaling
-            # ``mrope_section`` (e.g. [16, 24, 24] for Qwen2.5-VL) sums to
-            # head_dim/2; doubled it sums to head_dim and cycles axis = i % 3.
-            sec_full = self._mrope_section * 2
-            cos_split = cos_per_axis.split(sec_full, dim=-1)
-            sin_split = sin_per_axis.split(sec_full, dim=-1)
-            cos = torch.cat([c[:, i % 3] for i, c in enumerate(cos_split)], dim=-1)
-            sin = torch.cat([s[:, i % 3] for i, s in enumerate(sin_split)], dim=-1)
-            return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
-
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
         freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
@@ -582,12 +541,8 @@ class PackedAttentionMoT(nn.Module):
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)
-        # Pre-merge code cast to float32 before q_norm/k_norm — bf16
-        # RMSNorm accumulation drift compounds across segmented prefill
-        # (x2t builds the cache via 3+ sequential forward_cache_update_*
-        # calls), producing degenerate token repetition in generate_text.
-        q = self.q_norm(q.to(torch.float32))
-        k = self.k_norm(k.to(torch.float32))
+        q = self.q_norm(q)
+        k = self.k_norm(k)
 
         cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
         q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
@@ -602,52 +557,16 @@ class PackedAttentionMoT(nn.Module):
             cache_v = past_key_values.value_cache[self.layer_idx]
             full_k = torch.cat([cache_k, k], dim=0)
             full_v = torch.cat([cache_v, v], dim=0)
-            cache_len = cache_k.shape[0]
         else:
             full_k = k
             full_v = v
-            cache_len = 0
 
-        if is_causal and cache_len > 0:
-            # PyTorch SDPA's ``is_causal=True`` with ``Q != K`` uses
-            # ``tril(diagonal=0)`` — top-left aligned, so for Q=1 against a
-            # long cache only ``q_0 -> k_0`` is unmasked (degenerate
-            # generate_text output).  Build the correct bottom-right
-            # aligned mask manually: ``mask[i, j] = -inf if j > cache_len
-            # + i`` so each new query attends to all of ``cache + self
-            # up to i``.  Bypass ``DiffusionAttention`` entirely — its
-            # ``_maybe_reshape_attn_mask`` helper only handles 2-D
-            # ``(B, K)`` shape, not ``(Q, K)``, so the reshape path
-            # silently drops the per-query mask for Q>1.
-            Q_len = q.shape[0]
-            K_len = full_k.shape[0]
-            arange_q = torch.arange(Q_len, device=q.device).unsqueeze(1)
-            arange_k = torch.arange(K_len, device=q.device).unsqueeze(0)
-            mask = arange_k > (cache_len + arange_q)  # (Q, K) bool
-            attn_bias = torch.zeros(Q_len, K_len, dtype=q.dtype, device=q.device)
-            attn_bias.masked_fill_(mask, float("-inf"))
-            # Permute to (B=1, H, S, D) for SDPA.
-            q_4d = q.unsqueeze(0).permute(0, 2, 1, 3)
-            k_4d = full_k.unsqueeze(0).permute(0, 2, 1, 3)
-            v_4d = full_v.unsqueeze(0).permute(0, 2, 1, 3)
-            attn_out_4d = torch.nn.functional.scaled_dot_product_attention(
-                q_4d,
-                k_4d,
-                v_4d,
-                attn_mask=attn_bias.unsqueeze(0).unsqueeze(0),  # (1, 1, Q, K)
-                dropout_p=0.0,
-                is_causal=False,
-                scale=1.0 / (self.head_dim**0.5),
-                enable_gqa=(self.num_heads != self.num_kv_heads),
-            )
-            attn_out = attn_out_4d.permute(0, 2, 1, 3)
-        else:
-            attn = self.attn_causal if is_causal else self.attn_noncausal
-            attn_out = attn(
-                q.unsqueeze(0),
-                full_k.unsqueeze(0),
-                full_v.unsqueeze(0),
-            )
+        attn = self.attn_causal if is_causal else self.attn_noncausal
+        attn_out = attn(
+            q.unsqueeze(0),
+            full_k.unsqueeze(0),
+            full_v.unsqueeze(0),
+        )
 
         attn_out = attn_out.squeeze(0).reshape(-1, self.q_size)
         attn_out, _ = self.o_proj(attn_out)
@@ -1695,26 +1614,10 @@ class Bagel(nn.Module):
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
-        # Lance i2v: tokens to freeze at their initial (encoded-image)
-        # value throughout the denoise loop.  Matches upstream's
-        # ``mse_loss_indexes``-exclusion behaviour from PR #33: cond
-        # positions are never updated and get ``timestep=0``.
-        frame_condition_token_indexes: torch.LongTensor | None = None,
     ):
         x_t = packed_init_noises
-        # Snapshot the pinned subtensor BEFORE the denoise loop touches
-        # x_t; the cond positions in packed_init_noises hold the
-        # VAE-encoded conditioning latent that must be preserved verbatim.
-        pinned_x_t = None
-        if frame_condition_token_indexes is not None:
-            frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long()
-            pinned_x_t = x_t[frame_condition_token_indexes].clone()
 
-        # Use num_timesteps + 1 sample points so we get `num_timesteps` denoise
-        # steps after dropping the terminal t=0 (which has no dt).  Upstream
-        # Lance / BAGEL both use this convention; without the +1 we silently
-        # run one fewer denoise iteration than the user asked for.
-        timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device)
+        timesteps = torch.linspace(1, 0, num_timesteps, device=x_t.device)
         timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps)
         dts = timesteps[:-1] - timesteps[1:]
         timesteps = timesteps[:-1]
@@ -1766,12 +1669,6 @@ class Bagel(nn.Module):
                 trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-            if frame_condition_token_indexes is not None:
-                # Cond positions stay at t=0 (clean signal).  Matches upstream
-                # PR #33 lance.py line 1605:
-                #     timestep[current_vae_mse_indexes_local_in_vae] = t
-                # (cond positions remain at the ``torch.zeros`` init value).
-                timestep[frame_condition_token_indexes] = 0.0
                 in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1]
                 cfg_text_scale_ = cfg_text_scale if in_cfg_window else 1.0
                 cfg_img_scale_ = cfg_img_scale if in_cfg_window else 1.0
@@ -1835,12 +1732,6 @@ class Bagel(nn.Module):
                 trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-            if frame_condition_token_indexes is not None:
-                # Cond positions stay at t=0 (clean signal).  Matches upstream
-                # PR #33 lance.py line 1605:
-                #     timestep[current_vae_mse_indexes_local_in_vae] = t
-                # (cond positions remain at the ``torch.zeros`` init value).
-                timestep[frame_condition_token_indexes] = 0.0
                 v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
@@ -1892,11 +1783,7 @@ class Bagel(nn.Module):
                 "num_branches": num_branches,
                 "seq_len": seq_len,
                 "batched_query_lens": packed_seqlens.repeat(num_branches),
-                # Concat along the *sequence* (last) dim — works for both
-                # scalar 1-D position ids ``(S,)`` and Lance's 3-D mRoPE
-                # ``(3, S)``.  Default ``dim=0`` collapses ``(3, S) + (3, S)``
-                # to ``(6, S)`` and breaks ``BagelRotaryEmbedding``.
-                "batched_position_ids": torch.cat(branches_pid, dim=-1),
+                "batched_position_ids": torch.cat(branches_pid),
                 "batched_text_indexes": torch.cat(
                     [packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
                 ),
@@ -1911,12 +1798,6 @@ class Bagel(nn.Module):
 
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-            if frame_condition_token_indexes is not None:
-                # Cond positions stay at t=0 (clean signal).  Matches upstream
-                # PR #33 lance.py line 1605:
-                #     timestep[current_vae_mse_indexes_local_in_vae] = t
-                # (cond positions remain at the ``torch.zeros`` init value).
-                timestep[frame_condition_token_indexes] = 0.0
             if t > cfg_interval[0] and t <= cfg_interval[1]:
                 cfg_text_scale_ = cfg_text_scale
                 cfg_img_scale_ = cfg_img_scale
@@ -1947,12 +1828,6 @@ class Bagel(nn.Module):
                     trajectory_log_probs.append(out.log_prob)
             else:
                 x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
-                if pinned_x_t is not None:
-                    # i2v: restore cond positions to their encoded-image
-                    # latent.  Matches upstream PR #33 lance.py line 1712:
-                    #     x_t[mse_indexes] = x_t[mse_indexes] - v_t[mse_indexes] * dts[i]
-                    # (cond positions are excluded from the update).
-                    x_t[frame_condition_token_indexes] = pinned_x_t
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
                 trajectory_timesteps.append(timesteps[i])
@@ -1984,7 +1859,6 @@ class Bagel(nn.Module):
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
-        frame_condition_token_indexes: torch.LongTensor | None = None,
     ):
         """CFG parallel denoising loop: each rank computes one CFG branch.
 
@@ -2044,12 +1918,6 @@ class Bagel(nn.Module):
 
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
-            if frame_condition_token_indexes is not None:
-                # Cond positions stay at t=0 (clean signal).  Matches upstream
-                # PR #33 lance.py line 1605:
-                #     timestep[current_vae_mse_indexes_local_in_vae] = t
-                # (cond positions remain at the ``torch.zeros`` init value).
-                timestep[frame_condition_token_indexes] = 0.0
             use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
 
             if use_cfg_this_step:
@@ -2202,7 +2070,7 @@ class Bagel(nn.Module):
             packed_sequence = packed_text_embedding.new_zeros((int(local_seqlens.sum()), self.hidden_size))
             packed_sequence[local_text_indexes] = packed_text_embedding
 
-            # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
+            assert timestep.unique().shape[0] == 1
             packed_pos_embed = self.latent_pos_embed(local_vae_pos_ids)
             local_timestep = timestep[: local_x_t.shape[0]]
             packed_timestep_embeds = self.time_embedder(local_timestep)
@@ -2239,7 +2107,7 @@ class Bagel(nn.Module):
         packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
         packed_sequence[packed_text_indexes] = packed_text_embedding
 
-        # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
+        assert timestep.unique().shape[0] == 1
         packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
         packed_timestep_embeds = self.time_embedder(timestep)
         x_t_emb = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
@@ -2291,7 +2159,7 @@ class Bagel(nn.Module):
         packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
         packed_sequence[packed_text_indexes] = packed_text_embedding
 
-        # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
+        assert timestep.unique().shape[0] == 1
         packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
         packed_timestep_embeds = self.time_embedder(timestep)
         x_t = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -355,7 +355,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
 
         gen_params = BagelGenParams(
             num_timesteps=int(req.sampling_params.num_inference_steps or 50),
-            timestep_shift=float(extra_args.get("timestep_shift", 3.0)),
+            timestep_shift=3.0,
             cfg_text_scale=cfg_text_scale,
             cfg_img_scale=cfg_img_scale,
             cfg_interval=cfg_interval,
@@ -820,11 +820,6 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         custom = {}
         if think_text is not None:
             custom["think_text"] = think_text
-        # Mirror the PIL image into ``custom_output`` so callers reading via
-        # the orchestrator IPC boundary (which strips the bare ``output``
-        # field) can still recover the result.  ``video_frames`` already
-        # uses this pattern.
-        custom["image"] = img
 
         return DiffusionOutput(
             output=img,

vllm_omni/diffusion/models/lance/__init__.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance (ByteDance) diffusion model components.
-
-Lance is a unified autoregressive + diffusion multimodal model on a
-Qwen2.5-VL-3B backbone.  Architecturally it is the BAGEL family (ByteDance
-Mixture-of-Transformers): the released ``Lance_3B`` checkpoint uses the exact
-same ``*_moe_gen`` MoT weight layout as BAGEL, plus ``vae2llm`` / ``llm2vae`` /
-``time_embedder`` / ``latent_pos_embed`` connectors.  The deltas vs BAGEL are:
-
-  * backbone is Qwen2.5-VL (mRoPE) instead of Qwen2,
-  * understanding ViT is Qwen2.5-VL vision (not SigLIP), loaded from the base
-    ``Qwen/Qwen2.5-VL-3B-Instruct`` rather than from the Lance checkpoint,
-  * VAE is Wan2.2 (reused from the vLLM-Omni WAN path) instead of the BAGEL AE,
-  * video path adds 3D latent position embeddings (follow-up; this module
-    implements the image path first).
-
-Because Lance is BAGEL-lineage, the transformer core is reused verbatim from
-``vllm_omni.diffusion.models.bagel.bagel_transformer`` and only the pipeline
-wiring (ViT / VAE / checkpoint layout) is specialized here.
-"""
-
-from vllm_omni.diffusion.models.lance.pipeline_lance import (
-    LancePipeline,
-    get_lance_post_process_func,
-)
-
-__all__ = [
-    "LancePipeline",
-    "get_lance_post_process_func",
-]

vllm_omni/diffusion/models/lance/prompts.py

@@ -1,121 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance chat / system prompts.
-
-Matches the upstream Lance training distribution (see
-``bytedance/Lance/data/system_prompt_render.py`` and
-``data/common.py::generate_system_prompt``):
-
-  * each task has a specific system prompt describing what to attend to;
-  * the prompt is wrapped in the Qwen chat template
-    ``<|im_start|>system\\n…<|im_end|>\\n<|im_start|>user\\n…<|im_end|>\\n``
-    ``<|im_start|>assistant\\n``;
-  * vision tokens are framed by ``<|vision_start|><|video_pad|><|vision_end|>``
-    (upstream uses ``<|video_pad|>`` even for image type by default).
-
-These are pure string formatting helpers — the runtime pipeline still does
-the actual VAE / ViT prefill separately.
-"""
-
-from __future__ import annotations
-
-# Task-specific system prompts (verbatim from upstream's
-# ``generate_system_prompt`` — first variant of each pool).
-SYSTEM_PROMPTS: dict[tuple[str, str], str] = {
-    # (task, vision_type) -> system prompt
-    ("t2i", "image"): (
-        "Describe the image by detailing the color, quantity, text, shape, "
-        "size, texture, spatial relationships of the objects and background:"
-    ),
-    ("t2v", "video"): (
-        "Describe the video by detailing the color, quantity, visible text, "
-        "shape, size, texture, spatial relationships and motion/camera "
-        "movements of the objects and background:"
-    ),
-    # i2v reuses the t2v system prompt (matches upstream Lance PR #33's
-    # ``data/common.py`` where ``system_prompt_type in ("t2v", "i2v")``
-    # selects the same prompt pool).
-    ("i2v", "video"): (
-        "Describe the video by detailing the color, quantity, visible text, "
-        "shape, size, texture, spatial relationships and motion/camera "
-        "movements of the objects and background:"
-    ),
-    ("image_edit", "image"): (
-        "Describe the key features of the input image (color, shape, size, "
-        "texture, objects, background), then explain how the user's text "
-        "instruction should alter or modify the image. Generate a new image "
-        "that meets the user's requirements while maintaining consistency "
-        "with the original input where appropriate."
-    ),
-    ("video_edit", "video"): (
-        "Describe the key features of the input video (color, shape, size, "
-        "texture, objects, background), then explain how the user's text "
-        "instruction should alter or modify the video. Generate a new video "
-        "that meets the user's requirements while maintaining consistency "
-        "with the original input where appropriate."
-    ),
-    ("x2t_image", "image"): (
-        "Generate a detailed and accurate description of the image, including all the key moments and visual details."
-    ),
-    ("x2t_video", "video"): (
-        "Generate a detailed and accurate description of the video, including all the key moments and visual details."
-    ),
-}
-
-# Vision-content placeholders (Qwen2.5-VL convention).  Upstream uses
-# ``<|video_pad|>`` for image inputs too unless ``force_video_pad=False``
-# (counter-intuitively named; see ``render_qwenvl_prompt`` upstream).
-VISION_START = "<|vision_start|>"
-VISION_END = "<|vision_end|>"
-VIDEO_PAD = "<|video_pad|>"
-IMAGE_PAD = "<|image_pad|>"
-
-
-def render_lance_prompt(
-    task: str,
-    user_text: str,
-    *,
-    vision_token: str | None = None,
-    system_prompt: str | None = None,
-) -> str:
-    """Render a Lance-compatible user-side prompt.
-
-    ``task`` is one of ``t2i``, ``t2v``, ``image_edit``, ``video_edit``,
-    ``x2t_image``, ``x2t_video``.  ``vision_token`` should be the visual
-    placeholder to embed inside the user message (e.g. ``VIDEO_PAD`` or a
-    ``VISION_START + VIDEO_PAD + VISION_END`` block); pass ``None`` for
-    text-only inputs.  ``system_prompt`` overrides the default task system
-    prompt — required for x2t QA examples whose upstream JSON carries a
-    per-example instruction (e.g. ``"Look at the image carefully and answer
-    the question."``); without it x2t falls back to the caption-style
-    default and the model describes instead of answering.
-
-    Returns a single string ready to be tokenized; no further wrapping is
-    needed by the caller.
-    """
-    # ``t2v`` / ``x2t_video`` / ``video_edit`` use video system prompts;
-    # the rest are image.
-    vision_type = "video" if task in {"t2v", "i2v", "x2t_video", "video_edit"} else "image"
-    sys_prompt = (
-        system_prompt if system_prompt else SYSTEM_PROMPTS.get((task, vision_type), "You are a helpful assistant.")
-    )
-
-    if vision_token is None:
-        user_msg = user_text
-    else:
-        # Upstream concatenates the vision block immediately before the text.
-        user_msg = f"{vision_token}{user_text}"
-
-    return (
-        f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n"
-    )
-
-
-__all__ = [
-    "IMAGE_PAD",
-    "SYSTEM_PROMPTS",
-    "VIDEO_PAD",
-    "VISION_END",
-    "VISION_START",
-    "render_lance_prompt",
-]

vllm_omni/diffusion/models/lance/wan_vae.py

@@ -1,870 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# Ported from ByteDance Lance upstream (https://github.com/bytedance/Lance,
-# modeling/vae/wan/{vae2_2.py,model.py}). Upstream copyright:
-#
-#   Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
-#   Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
-#   Licensed under the Apache License, Version 2.0.
-#
-# Lance ships its VAE as a raw ``Wan2.2_VAE.pth`` whose state dict is keyed for
-# the upstream ``WanVAE_`` nn.Module here (not for diffusers' ``AutoencoderKLWan``).
-# We port that module verbatim so we can load the .pth directly, and wrap it with
-# :class:`LanceWanVAE` which exposes BAGEL's 4-D ``encode(BCHW)/decode(BCHW)``
-# surface (and a 5-D path for the video checkpoint).
-"""Wan2.2 VAE used by Lance, ported from upstream so ``Wan2.2_VAE.pth`` loads
-natively without state-dict surgery."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-CACHE_T = 2
-
-
-class CausalConv3d(nn.Conv3d):
-    """Causal 3D conv with feature-map caching across temporal chunks."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._padding = (
-            self.padding[2],
-            self.padding[2],
-            self.padding[1],
-            self.padding[1],
-            2 * self.padding[0],
-            0,
-        )
-        self.padding = (0, 0, 0)
-
-    def forward(self, x, cache_x=None):
-        padding = list(self._padding)
-        if cache_x is not None and self._padding[4] > 0:
-            cache_x = cache_x.to(x.device)
-            x = torch.cat([cache_x, x], dim=2)
-            padding[4] -= cache_x.shape[2]
-        x = F.pad(x, padding)
-        return super().forward(x)
-
-
-class RMS_norm(nn.Module):
-    def __init__(self, dim, channel_first=True, images=True, bias=False):
-        super().__init__()
-        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
-        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
-        self.channel_first = channel_first
-        self.scale = dim**0.5
-        self.gamma = nn.Parameter(torch.ones(shape))
-        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
-
-    def forward(self, x):
-        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
-
-
-class Upsample(nn.Upsample):
-    def forward(self, x):
-        return super().forward(x.float()).type_as(x)
-
-
-class Resample(nn.Module):
-    def __init__(self, dim, mode):
-        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
-        super().__init__()
-        self.dim = dim
-        self.mode = mode
-
-        if mode == "upsample2d":
-            self.resample = nn.Sequential(
-                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, dim, 3, padding=1),
-            )
-        elif mode == "upsample3d":
-            self.resample = nn.Sequential(
-                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
-                nn.Conv2d(dim, dim, 3, padding=1),
-            )
-            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
-        elif mode == "downsample2d":
-            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
-        elif mode == "downsample3d":
-            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
-            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
-        else:
-            self.resample = nn.Identity()
-
-    def forward(self, x, feat_cache=None, feat_idx=[0]):
-        b, c, t, h, w = x.size()
-        if self.mode == "upsample3d":
-            if feat_cache is not None:
-                idx = feat_idx[0]
-                if feat_cache[idx] is None:
-                    feat_cache[idx] = "Rep"
-                    feat_idx[0] += 1
-                else:
-                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
-                        cache_x = torch.cat(
-                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                            dim=2,
-                        )
-                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
-                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
-                    if feat_cache[idx] == "Rep":
-                        x = self.time_conv(x)
-                    else:
-                        x = self.time_conv(x, feat_cache[idx])
-                    feat_cache[idx] = cache_x
-                    feat_idx[0] += 1
-                    x = x.reshape(b, 2, c, t, h, w)
-                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
-                    x = x.reshape(b, c, t * 2, h, w)
-        t = x.shape[2]
-        x = rearrange(x, "b c t h w -> (b t) c h w")
-        x = self.resample(x)
-        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
-
-        if self.mode == "downsample3d":
-            if feat_cache is not None:
-                idx = feat_idx[0]
-                if feat_cache[idx] is None:
-                    feat_cache[idx] = x.clone()
-                    feat_idx[0] += 1
-                else:
-                    cache_x = x[:, :, -1:, :, :].clone()
-                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
-                    feat_cache[idx] = cache_x
-                    feat_idx[0] += 1
-        return x
-
-
-class ResidualBlock(nn.Module):
-    def __init__(self, in_dim, out_dim, dropout=0.0):
-        super().__init__()
-        self.in_dim = in_dim
-        self.out_dim = out_dim
-        self.residual = nn.Sequential(
-            RMS_norm(in_dim, images=False),
-            nn.SiLU(),
-            CausalConv3d(in_dim, out_dim, 3, padding=1),
-            RMS_norm(out_dim, images=False),
-            nn.SiLU(),
-            nn.Dropout(dropout),
-            CausalConv3d(out_dim, out_dim, 3, padding=1),
-        )
-        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
-
-    def forward(self, x, feat_cache=None, feat_idx=[0]):
-        h = self.shortcut(x)
-        for layer in self.residual:
-            if isinstance(layer, CausalConv3d) and feat_cache is not None:
-                idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                    cache_x = torch.cat(
-                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                        dim=2,
-                    )
-                x = layer(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
-                feat_idx[0] += 1
-            else:
-                x = layer(x)
-        return x + h
-
-
-class AttentionBlock(nn.Module):
-    """Single-head causal self-attention over spatial tokens, per frame."""
-
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-        self.norm = RMS_norm(dim)
-        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
-        self.proj = nn.Conv2d(dim, dim, 1)
-        nn.init.zeros_(self.proj.weight)
-
-    def forward(self, x):
-        identity = x
-        b, c, t, h, w = x.size()
-        x = rearrange(x, "b c t h w -> (b t) c h w")
-        x = self.norm(x)
-        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
-        x = F.scaled_dot_product_attention(q, k, v)
-        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
-        x = self.proj(x)
-        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
-        return x + identity
-
-
-def _patchify(x, patch_size):
-    if patch_size == 1:
-        return x
-    if x.dim() == 4:
-        return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
-    if x.dim() == 5:
-        return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size)
-    raise ValueError(f"Invalid input shape: {x.shape}")
-
-
-def _unpatchify(x, patch_size):
-    if patch_size == 1:
-        return x
-    if x.dim() == 4:
-        return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
-    if x.dim() == 5:
-        return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size)
-    return x
-
-
-class AvgDown3D(nn.Module):
-    def __init__(self, in_channels, out_channels, factor_t, factor_s=1):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-        assert in_channels * self.factor % out_channels == 0
-        self.group_size = in_channels * self.factor // out_channels
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
-        x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
-        B, C, T, H, W = x.shape
-        x = x.view(
-            B,
-            C,
-            T // self.factor_t,
-            self.factor_t,
-            H // self.factor_s,
-            self.factor_s,
-            W // self.factor_s,
-            self.factor_s,
-        )
-        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
-        x = x.view(B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s)
-        x = x.view(B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s)
-        return x.mean(dim=2)
-
-
-class DupUp3D(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.factor_t = factor_t
-        self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-        assert out_channels * self.factor % in_channels == 0
-        self.repeats = out_channels * self.factor // in_channels
-
-    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
-        x = x.repeat_interleave(self.repeats, dim=1)
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            self.factor_t,
-            self.factor_s,
-            self.factor_s,
-            x.size(2),
-            x.size(3),
-            x.size(4),
-        )
-        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
-        x = x.view(
-            x.size(0),
-            self.out_channels,
-            x.size(2) * self.factor_t,
-            x.size(4) * self.factor_s,
-            x.size(6) * self.factor_s,
-        )
-        if first_chunk:
-            x = x[:, :, self.factor_t - 1 :, :, :]
-        return x
-
-
-class Down_ResidualBlock(nn.Module):
-    def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False):
-        super().__init__()
-        self.avg_shortcut = AvgDown3D(
-            in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1
-        )
-        downsamples = []
-        for _ in range(mult):
-            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
-            in_dim = out_dim
-        if down_flag:
-            mode = "downsample3d" if temperal_downsample else "downsample2d"
-            downsamples.append(Resample(out_dim, mode=mode))
-        self.downsamples = nn.Sequential(*downsamples)
-
-    def forward(self, x, feat_cache=None, feat_idx=[0]):
-        x_copy = x.clone()
-        for module in self.downsamples:
-            x = module(x, feat_cache, feat_idx)
-        return x + self.avg_shortcut(x_copy)
-
-
-class Up_ResidualBlock(nn.Module):
-    def __init__(self, in_dim, out_dim, dropout, mult, temporal_upsample=False, up_flag=False):
-        super().__init__()
-        if up_flag:
-            self.avg_shortcut = DupUp3D(
-                in_dim, out_dim, factor_t=2 if temporal_upsample else 1, factor_s=2 if up_flag else 1
-            )
-        else:
-            self.avg_shortcut = None
-        upsamples = []
-        for _ in range(mult):
-            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
-            in_dim = out_dim
-        if up_flag:
-            mode = "upsample3d" if temporal_upsample else "upsample2d"
-            upsamples.append(Resample(out_dim, mode=mode))
-        self.upsamples = nn.Sequential(*upsamples)
-
-    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
-        x_main = x.clone()
-        for module in self.upsamples:
-            x_main = module(x_main, feat_cache, feat_idx)
-        if self.avg_shortcut is not None:
-            return x_main + self.avg_shortcut(x, first_chunk)
-        return x_main
-
-
-class Encoder3d(nn.Module):
-    def __init__(
-        self,
-        dim=128,
-        z_dim=4,
-        dim_mult=(1, 2, 4, 4),
-        num_res_blocks=2,
-        attn_scales=(),
-        temperal_downsample=(True, True, False),
-        dropout=0.0,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.z_dim = z_dim
-        self.dim_mult = list(dim_mult)
-        self.num_res_blocks = num_res_blocks
-        self.attn_scales = list(attn_scales)
-        self.temperal_downsample = list(temperal_downsample)
-
-        dims = [dim * u for u in [1] + self.dim_mult]
-        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
-
-        downsamples = []
-        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
-            t_down_flag = self.temperal_downsample[i] if i < len(self.temperal_downsample) else False
-            downsamples.append(
-                Down_ResidualBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    dropout=dropout,
-                    mult=num_res_blocks,
-                    temperal_downsample=t_down_flag,
-                    down_flag=i != len(self.dim_mult) - 1,
-                )
-            )
-        self.downsamples = nn.Sequential(*downsamples)
-
-        self.middle = nn.Sequential(
-            ResidualBlock(out_dim, out_dim, dropout),
-            AttentionBlock(out_dim),
-            ResidualBlock(out_dim, out_dim, dropout),
-        )
-        self.head = nn.Sequential(
-            RMS_norm(out_dim, images=False),
-            nn.SiLU(),
-            CausalConv3d(out_dim, z_dim, 3, padding=1),
-        )
-
-    def forward(self, x, feat_cache=None, feat_idx=[0]):
-        if feat_cache is not None:
-            idx = feat_idx[0]
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                cache_x = torch.cat(
-                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                    dim=2,
-                )
-            x = self.conv1(x, feat_cache[idx])
-            feat_cache[idx] = cache_x
-            feat_idx[0] += 1
-        else:
-            x = self.conv1(x)
-
-        for layer in self.downsamples:
-            if feat_cache is not None:
-                x = layer(x, feat_cache, feat_idx)
-            else:
-                x = layer(x)
-
-        for layer in self.middle:
-            if isinstance(layer, ResidualBlock) and feat_cache is not None:
-                x = layer(x, feat_cache, feat_idx)
-            else:
-                x = layer(x)
-
-        for layer in self.head:
-            if isinstance(layer, CausalConv3d) and feat_cache is not None:
-                idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                    cache_x = torch.cat(
-                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                        dim=2,
-                    )
-                x = layer(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
-                feat_idx[0] += 1
-            else:
-                x = layer(x)
-        return x
-
-
-class Decoder3d(nn.Module):
-    def __init__(
-        self,
-        dim=128,
-        z_dim=4,
-        dim_mult=(1, 2, 4, 4),
-        num_res_blocks=2,
-        attn_scales=(),
-        temporal_upsample=(False, True, True),
-        dropout=0.0,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.z_dim = z_dim
-        self.dim_mult = list(dim_mult)
-        self.num_res_blocks = num_res_blocks
-        self.attn_scales = list(attn_scales)
-        self.temporal_upsample = list(temporal_upsample)
-
-        dims = [dim * u for u in [self.dim_mult[-1]] + self.dim_mult[::-1]]
-        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
-        self.middle = nn.Sequential(
-            ResidualBlock(dims[0], dims[0], dropout),
-            AttentionBlock(dims[0]),
-            ResidualBlock(dims[0], dims[0], dropout),
-        )
-
-        upsamples = []
-        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
-            t_up_flag = self.temporal_upsample[i] if i < len(self.temporal_upsample) else False
-            upsamples.append(
-                Up_ResidualBlock(
-                    in_dim=in_dim,
-                    out_dim=out_dim,
-                    dropout=dropout,
-                    mult=num_res_blocks + 1,
-                    temporal_upsample=t_up_flag,
-                    up_flag=i != len(self.dim_mult) - 1,
-                )
-            )
-        self.upsamples = nn.Sequential(*upsamples)
-        self.head = nn.Sequential(
-            RMS_norm(out_dim, images=False),
-            nn.SiLU(),
-            CausalConv3d(out_dim, 12, 3, padding=1),
-        )
-
-    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
-        if feat_cache is not None:
-            idx = feat_idx[0]
-            cache_x = x[:, :, -CACHE_T:, :, :].clone()
-            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                cache_x = torch.cat(
-                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                    dim=2,
-                )
-            x = self.conv1(x, feat_cache[idx])
-            feat_cache[idx] = cache_x
-            feat_idx[0] += 1
-        else:
-            x = self.conv1(x)
-
-        for layer in self.middle:
-            if isinstance(layer, ResidualBlock) and feat_cache is not None:
-                x = layer(x, feat_cache, feat_idx)
-            else:
-                x = layer(x)
-
-        for layer in self.upsamples:
-            if feat_cache is not None:
-                x = layer(x, feat_cache, feat_idx, first_chunk)
-            else:
-                x = layer(x)
-
-        for layer in self.head:
-            if isinstance(layer, CausalConv3d) and feat_cache is not None:
-                idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
-                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
-                    cache_x = torch.cat(
-                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
-                        dim=2,
-                    )
-                x = layer(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
-                feat_idx[0] += 1
-            else:
-                x = layer(x)
-        return x
-
-
-def _count_conv3d(model):
-    return sum(1 for m in model.modules() if isinstance(m, CausalConv3d))
-
-
-class WanVAE_(nn.Module):
-    """Upstream Wan2.2 VAE module — encoder3d/decoder3d sandwich with 2x patchify
-    on input. State-dict-compatible with ``Wan2.2_VAE.pth``."""
-
-    def __init__(
-        self,
-        dim=160,
-        dec_dim=256,
-        z_dim=16,
-        dim_mult=(1, 2, 4, 4),
-        num_res_blocks=2,
-        attn_scales=(),
-        temperal_downsample=(True, True, False),
-        dropout=0.0,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.z_dim = z_dim
-        self.dim_mult = list(dim_mult)
-        self.num_res_blocks = num_res_blocks
-        self.attn_scales = list(attn_scales)
-        self.temperal_downsample = list(temperal_downsample)
-        self.temporal_upsample = list(temperal_downsample)[::-1]
-
-        self.encoder = Encoder3d(
-            dim, z_dim * 2, self.dim_mult, num_res_blocks, self.attn_scales, self.temperal_downsample, dropout
-        )
-        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
-        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
-        self.decoder = Decoder3d(
-            dec_dim, z_dim, self.dim_mult, num_res_blocks, self.attn_scales, self.temporal_upsample, dropout
-        )
-
-    def encode(self, x, scale):
-        self.clear_cache()
-        x = _patchify(x, patch_size=2)
-        t = x.shape[2]
-        iter_ = 1 + (t - 1) // 4
-        out = None
-        for i in range(iter_):
-            self._enc_conv_idx = [0]
-            if i == 0:
-                out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
-            else:
-                out_ = self.encoder(
-                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
-                    feat_cache=self._enc_feat_map,
-                    feat_idx=self._enc_conv_idx,
-                )
-                out = torch.cat([out, out_], 2)
-        mu, log_var = self.conv1(out).chunk(2, dim=1)
-        if isinstance(scale[0], torch.Tensor):
-            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
-        else:
-            mu = (mu - scale[0]) * scale[1]
-        self.clear_cache()
-        return mu, log_var
-
-    def decode(self, z, scale):
-        self.clear_cache()
-        if isinstance(scale[0], torch.Tensor):
-            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
-        else:
-            z = z / scale[1] + scale[0]
-        iter_ = z.shape[2]
-        x = self.conv2(z)
-        out = None
-        for i in range(iter_):
-            self._conv_idx = [0]
-            if i == 0:
-                out = self.decoder(
-                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True
-                )
-            else:
-                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
-                out = torch.cat([out, out_], 2)
-        out = _unpatchify(out, patch_size=2)
-        self.clear_cache()
-        return out
-
-    def clear_cache(self):
-        self._conv_num = _count_conv3d(self.decoder)
-        self._conv_idx = [0]
-        self._feat_map = [None] * self._conv_num
-        self._enc_conv_num = _count_conv3d(self.encoder)
-        self._enc_conv_idx = [0]
-        self._enc_feat_map = [None] * self._enc_conv_num
-
-
-# ----------------------------------------------------------------------------
-# Wan2.2 VAE wrapper + BAGEL-surface adapter
-# ----------------------------------------------------------------------------
-
-
-# Per-channel mean/std for the 48-channel Wan2.2 latent space (upstream constants).
-_WAN22_LATENT_MEAN: tuple[float, ...] = (
-    -0.2289,
-    -0.0052,
-    -0.1323,
-    -0.2339,
-    -0.2799,
-    0.0174,
-    0.1838,
-    0.1557,
-    -0.1382,
-    0.0542,
-    0.2813,
-    0.0891,
-    0.1570,
-    -0.0098,
-    0.0375,
-    -0.1825,
-    -0.2246,
-    -0.1207,
-    -0.0698,
-    0.5109,
-    0.2665,
-    -0.2108,
-    -0.2158,
-    0.2502,
-    -0.2055,
-    -0.0322,
-    0.1109,
-    0.1567,
-    -0.0729,
-    0.0899,
-    -0.2799,
-    -0.1230,
-    -0.0313,
-    -0.1649,
-    0.0117,
-    0.0723,
-    -0.2839,
-    -0.2083,
-    -0.0520,
-    0.3748,
-    0.0152,
-    0.1957,
-    0.1433,
-    -0.2944,
-    0.3573,
-    -0.0548,
-    -0.1681,
-    -0.0667,
-)
-_WAN22_LATENT_STD: tuple[float, ...] = (
-    0.4765,
-    1.0364,
-    0.4514,
-    1.1677,
-    0.5313,
-    0.4990,
-    0.4818,
-    0.5013,
-    0.8158,
-    1.0344,
-    0.5894,
-    1.0901,
-    0.6885,
-    0.6165,
-    0.8454,
-    0.4978,
-    0.5759,
-    0.3523,
-    0.7135,
-    0.6804,
-    0.5833,
-    1.4146,
-    0.8986,
-    0.5659,
-    0.7069,
-    0.5338,
-    0.4889,
-    0.4917,
-    0.4069,
-    0.4999,
-    0.6866,
-    0.4093,
-    0.5709,
-    0.6065,
-    0.6415,
-    0.4944,
-    0.5726,
-    1.2042,
-    0.5458,
-    1.6887,
-    0.3971,
-    1.0600,
-    0.3943,
-    0.5537,
-    0.5444,
-    0.4089,
-    0.7468,
-    0.7744,
-)
-
-
-class LanceWanVAE(nn.Module):
-    """Wan2.2 VAE wrapped for BAGEL's pipeline.
-
-    Exposes BAGEL's image-VAE surface — ``encode(BCHW) -> BC_zHW`` and
-    ``decode(BC_zHW) -> BCHW`` — by treating each image as a 1-frame video clip.
-    A 5-D ``encode_video``/``decode_video`` path is also provided for the
-    ``Lance_3B_Video`` checkpoint.
-
-    Construction is lazy: the heavy ``WanVAE_`` and ``Wan2.2_VAE.pth`` are not
-    materialized until first use. Once built, the inner module is registered as
-    a submodule so ``self.parameters()``, ``self.to(device)`` and ``vae_dtype =
-    next(vae.parameters()).dtype`` (used by BAGEL's decode path) all behave.
-    """
-
-    z_channels: int = 48
-    downsample_spatial: int = 16
-    downsample_temporal: int = 4
-
-    def __init__(
-        self,
-        vae_path: str,
-        dtype: torch.dtype = torch.bfloat16,
-        device: torch.device | str | None = None,
-        *,
-        lazy: bool = False,
-    ):
-        super().__init__()
-        self._vae_path = vae_path
-        self._dtype = dtype
-        self._device = torch.device(device) if device is not None else None
-        self._built = False
-        # Registered lazily once the .pth is loaded:
-        # self.model: WanVAE_   (set in _ensure_built)
-        if not lazy:
-            # BAGEL's decode path inspects ``next(vae.parameters()).dtype`` before
-            # the first encode/decode call, so we have to register the submodule
-            # eagerly when a real device is provided.  ``lazy=True`` is reserved
-            # for config-only / no-GPU paths.
-            self._ensure_built()
-
-    def _ensure_built(self) -> None:
-        if self._built:
-            return
-        device = self._device or (
-            next(self.parameters(), torch.empty(0)).device
-            if any(True for _ in self.parameters())
-            else torch.device("cpu")
-        )
-        # Wan2.2 VAE config from upstream Wan2_2_VAE defaults:
-        #   z_dim=48 (channels), c_dim=160 (encoder width), dec_dim=256 (decoder width),
-        #   dim_mult=[1,2,4,4], temperal_downsample=[False,True,True].
-        model = WanVAE_(
-            dim=160,
-            dec_dim=256,
-            z_dim=self.z_channels,
-            dim_mult=(1, 2, 4, 4),
-            num_res_blocks=2,
-            attn_scales=(),
-            temperal_downsample=(False, True, True),
-            dropout=0.0,
-        )
-        logger.info("Loading Wan2.2 VAE state dict from %s", self._vae_path)
-        state = torch.load(self._vae_path, map_location="cpu", weights_only=True)
-        missing, unexpected = model.load_state_dict(state, strict=False)
-        if missing:
-            logger.warning("Wan2.2 VAE missing keys (%d): %s", len(missing), missing[:5])
-        if unexpected:
-            logger.warning("Wan2.2 VAE unexpected keys (%d): %s", len(unexpected), unexpected[:5])
-        model = model.to(device=device, dtype=self._dtype).eval()
-        model.requires_grad_(False)
-        # Register as submodule so .to() / .parameters() pick it up.
-        self.model = model
-        # Per-channel normalization tensors (matches upstream Wan2_2_VAE.scale = [mean, 1/std]).
-        mean = torch.tensor(_WAN22_LATENT_MEAN, dtype=self._dtype, device=device)
-        inv_std = torch.tensor([1.0 / s for s in _WAN22_LATENT_STD], dtype=self._dtype, device=device)
-        # Buffers, so .to(device) follows along.
-        self.register_buffer("_latent_mean", mean, persistent=False)
-        self.register_buffer("_latent_inv_std", inv_std, persistent=False)
-        self._built = True
-
-    @torch.inference_mode()
-    def encode_video(self, video: torch.Tensor, *, use_sample: bool = True) -> torch.Tensor:
-        """Encode a 5-D clip ``[B, 3, T, H, W]`` -> latent ``[B, 48, t, h, w]``."""
-        self._ensure_built()
-        video = video.to(self.model.encoder.conv1.weight.dtype)
-        mu, log_var = self.model.encode(video, [self._latent_mean, self._latent_inv_std])
-        if use_sample:
-            std = torch.exp(0.5 * log_var)
-            mu = mu + std * torch.randn_like(std)
-        return mu
-
-    @torch.inference_mode()
-    def decode_video(self, latent: torch.Tensor) -> torch.Tensor:
-        """Decode a 5-D latent ``[B, 48, t, h, w]`` -> video ``[B, 3, T, H, W]``."""
-        self._ensure_built()
-        latent = latent.to(self.model.decoder.conv1.weight.dtype)
-        out = self.model.decode(latent, [self._latent_mean, self._latent_inv_std])
-        return out.clamp_(-1.0, 1.0)
-
-    # ----- BAGEL image-VAE surface (4-D) -------------------------------- #
-
-    @torch.inference_mode()
-    def encode(self, padded_images: torch.Tensor) -> torch.Tensor:
-        """``[B, 3, H, W]`` -> ``[B, 48, H/16, W/16]`` (single-frame image path).
-
-        Each image is wrapped as a 1-frame clip, encoded, and the temporal axis
-        is squeezed back out so the result matches BAGEL's iteration pattern.
-        """
-        if padded_images.dim() == 4:
-            video = padded_images.unsqueeze(2)  # B,C,1,H,W
-            squeeze_time = True  # image path: caller iterates 3-D latents
-        elif padded_images.dim() == 5:
-            video = padded_images
-            # video path: caller expects 4-D latents (C,T,H,W), even when T==1
-            # (e.g. video2video with a 1-frame reference clip).  Squeezing here
-            # would drop the temporal axis and break the downstream indexing
-            # ``latent[:, :t_lat, ...]`` in ``forward_cache_update_vae``.
-            squeeze_time = False
-        else:
-            raise ValueError(f"LanceWanVAE.encode expects 4-D BCHW or 5-D BCTHW, got {tuple(padded_images.shape)}")
-        latent = self.encode_video(video, use_sample=True)
-        # B,48,t,h,w -> B,48,h,w when t == 1 (image path only)
-        if squeeze_time and latent.shape[2] == 1:
-            return latent.squeeze(2)
-        return latent
-
-    @torch.inference_mode()
-    def decode(self, latent: torch.Tensor) -> torch.Tensor:
-        """``[B, 48, h, w]`` -> ``[B, 3, H, W]`` (single-frame image path)."""
-        if latent.dim() == 4:
-            latent = latent.unsqueeze(2)  # B,48,1,h,w
-        elif latent.dim() != 5:
-            raise ValueError(f"LanceWanVAE.decode expects 4-D BCHW or 5-D BCTHW, got {tuple(latent.shape)}")
-        video = self.decode_video(latent)  # B,3,T,H,W
-        if video.shape[2] == 1:
-            return video.squeeze(2)
-        return video
-
-
-def build_wan22_vae(vae_path: str, dtype: torch.dtype = torch.bfloat16, device=None) -> LanceWanVAE:
-    """Convenience factory: lazy-construct a :class:`LanceWanVAE` adapter."""
-    return LanceWanVAE(vae_path=vae_path, dtype=dtype, device=device)
-
-
-__all__ = ["LanceWanVAE", "WanVAE_", "build_wan22_vae"]

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -275,21 +275,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         self.device = get_local_device()
         model = od_config.model
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            )
         )
         # Qwen2.5-VL ships a vision tower that text-to-image does not use.
         # Drop it while the model is still on CPU, before moving to GPU, so

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -27,7 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -243,22 +243,23 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         model = od_config.model
 
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
 
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
@@ -266,8 +267,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -25,7 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
@@ -208,7 +208,7 @@ class QwenImageEditPlusPipeline(
         model = od_config.model
 
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # Defend against a transformers v5 multi-worker race where a peer
         # worker's partially-written shard set makes our from_pretrained
@@ -219,15 +219,16 @@ class QwenImageEditPlusPipeline(
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
 
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
@@ -236,8 +237,8 @@ class QwenImageEditPlusPipeline(
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -24,7 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
     AutoencoderKLQwenImage,
@@ -216,29 +216,30 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
         self.device = get_local_device()
         model = od_config.model
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         # modules keep same as transformers & diffusers
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
         )
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         # modules re-implemented for vLLM-Omni

vllm_omni/diffusion/registry.py

@@ -141,11 +141,6 @@ _DIFFUSION_MODELS = {
         "pipeline_bagel",
         "BagelPipeline",
     ),
-    "LancePipeline": (
-        "lance",
-        "pipeline_lance",
-        "LancePipeline",
-    ),
     "MingImagePipeline": (
         "ming_flash_omni",
         "pipeline_ming_imagegen",
@@ -476,7 +471,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
-    "LancePipeline": "get_lance_post_process_func",
     "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",

vllm_omni/distributed/omni_coordinator/omni_coordinator.py

@@ -43,17 +43,20 @@ class OmniCoordinator:
             heartbeat_timeout: Seconds before a replica is considered
                 unhealthy if no heartbeat / update is received.
         """
-        self._router_zmq_addr = router_zmq_addr
-        self._pub_zmq_addr = pub_zmq_addr
         self._heartbeat_timeout = heartbeat_timeout
 
         # Dedicated ZMQ context for this coordinator instance.
         self._ctx = zmq.Context()
         self._router = self._ctx.socket(zmq.ROUTER)
-        self._router.bind(self._router_zmq_addr)
+        self._router.bind(router_zmq_addr)
+        # Recover the actual bound address (port=0 → real port).
+        self.router_zmq_addr = self._router.getsockopt_string(zmq.LAST_ENDPOINT)
+        self._pub_zmq_addr = pub_zmq_addr  # keep original for internal use
 
         self._pub = self._ctx.socket(zmq.PUB)
-        self._pub.bind(self._pub_zmq_addr)
+        self._pub.bind(pub_zmq_addr)
+        # Recover the actual bound address (port=0 → real port).
+        self.pub_zmq_addr = self._pub.getsockopt_string(zmq.LAST_ENDPOINT)
 
         self._replicas: dict[str, ReplicaInfo] = {}
         self._lock = threading.Lock()

vllm_omni/engine/arg_utils.py

@@ -144,7 +144,8 @@ class OmniEngineArgs(EngineArgs):
     force_cutlass_fp8: bool | None = None
     worker_type: str | None = None
     task_type: str | None = None
-    worker_cls: str = None
+    worker_cls: str = None  # type: ignore[assignment]  # Upstream default is "auto"; omni resolves
+    # in __post_init__ based on worker_type (ar/generation), so None is safe here.
     enable_sleep_mode: bool = False
     omni: bool = False
 
@@ -444,6 +445,48 @@ class OrchestratorArgs:
     # === Pre-built Objects ===
     parallel_config: Any = None
 
+    # === Diffusion model config ===
+    num_gpus: int | None = None
+    model_class_name: str | None = None
+    diffusion_load_format: str | None = None
+    diffusers_load_kwargs: str = "{}"
+    diffusers_call_kwargs: str = "{}"
+    ulysses_degree: int | None = None
+    ulysses_mode: str = "strict"
+    ring_degree: int | None = None
+    diffusion_quantization_config: str | None = None
+    use_hsdp: bool = False
+    hsdp_shard_size: int = -1
+    hsdp_replicate_size: int = 1
+    diffusion_attention_backend: str | None = None
+    diffusion_attention_config: str | None = None
+    cache_backend: str = "none"
+    cache_config: str | None = None
+    enable_cache_dit_summary: bool = False
+    step_execution: bool = False
+    vae_use_slicing: bool = False
+    vae_use_tiling: bool = False
+    enable_multithread_weight_load: bool = True
+    num_weight_load_threads: int = 4
+    enable_cpu_offload: bool = False
+    enable_layerwise_offload: bool = False
+    boundary_ratio: float | None = None
+    flow_shift: float | None = None
+    diffusion_kv_cache_dtype: str | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
+    cfg_parallel_size: int = 1
+    vae_patch_parallel_size: int = 1
+    default_sampling_params: str | None = None
+    max_generated_image_size: int | None = None
+    tts_max_instructions_length: int | None = None
+    enable_diffusion_pipeline_profiler: bool = False
+    enable_ar_profiler: bool = False
+    auxiliary_text_encoder: str | None = None
+    log_file: str | None = None
+    replica_id: int | None = None
+    omni_replica_address: str | None = None
+
     # === Multi-stage guards ===
     # --tokenizer is captured by the orchestrator and forwarded to stages
     # only when the stage does not define tokenizer/tokenizer_subdir itself.
@@ -459,6 +502,8 @@ SHARED_FIELDS: frozenset[str] = frozenset(
         "stage_id",  # orch: route (headless); engine: identity
         "log_stats",  # both want the flag
         "stage_configs_path",  # orch: load legacy YAML; engine: may reference for validation
+        "async_chunk",  # orch: read from CLI, redistribute; engine: per-stage flag
+        "tokenizer",  # orch: detect model type; engine: tokenization
     }
 )
 

vllm_omni/engine/async_omni_engine.py

@@ -2124,7 +2124,7 @@ class AsyncOmniEngine:
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
                         )
-                quantization_config = kwargs.get("quantization_config")
+                quantization_config = kwargs.get("diffusion_quantization_config")
                 if quantization_config is not None:
                     if (
                         not hasattr(cfg.engine_args, "quantization_config")

vllm_omni/engine/output_processor.py

@@ -1,11 +1,7 @@
 from typing import Any
 
-import numpy as np
 import torch
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    split_routed_experts,
-)
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
@@ -172,7 +168,6 @@ class OmniRequestState(RequestState):
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
         kv_transfer_params: dict[str, Any] | None = None,
-        routed_experts: np.ndarray | None = None,
     ) -> OmniRequestOutput | PoolingRequestOutput | None:
         """Create a request output from generation results.
 
@@ -199,7 +194,6 @@ class OmniRequestState(RequestState):
                 finish_reason,
                 stop_reason,
                 kv_transfer_params,
-                routed_experts,
             )
 
         finished = finish_reason is not None
@@ -234,16 +228,7 @@ class OmniRequestState(RequestState):
 
         external_req_id = self.external_req_id
 
-        # Split routing data into prompt and generation portions, matching
-        # upstream make_request_output behaviour.
-        prompt_routed_experts = None
-        gen_routed_experts = None
-        if routed_experts is not None:
-            prompt_len = len(self.prompt_token_ids) if self.prompt_token_ids else 0
-            num_gen = self.detokenizer.num_output_tokens() if self.detokenizer is not None else None
-            prompt_routed_experts, gen_routed_experts = split_routed_experts(routed_experts, prompt_len, num_gen)
-
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, gen_routed_experts)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
 
         if self.parent_req is None:
             outputs = [output]
@@ -258,7 +243,6 @@ class OmniRequestState(RequestState):
             outputs,
             finished,
             kv_transfer_params,
-            prompt_routed_experts,
         )
 
     def _new_completion_output(
@@ -266,10 +250,11 @@ class OmniRequestState(RequestState):
         token_ids: list[int],
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
-        routed_experts: np.ndarray | None = None,
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
+        # Note: upstream _new_completion_output no longer accepts routed_experts
+        # as a parameter; it reads from ``self.routed_experts_chunks`` internally.
+        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
 
         # Inter-stage processors need the full cumulative token sequence.
         # In DELTA mode, base_output.token_ids only has the latest step's

vllm_omni/engine/stage_engine_core_proc.py

@@ -101,6 +101,11 @@ class StageEngineCoreProc(EngineCoreProc):
             set_death_signal(signal.SIGTERM)
             set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}")
             decorate_logs()
+            # Workaround for flashinfer/jit-cache version mismatch in CI.
+            # The parent process handles this gracefully via ring_globals.py,
+            # but the subprocess hits an unprotected import in TopKTopPSampler.
+            # Setting this env var allows the same graceful fallback to work.
+            os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
             os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))
 
             engine_core = StageEngineCoreProc(

vllm_omni/engine/stage_init_utils.py

@@ -824,6 +824,10 @@ def build_vllm_config(
     if upgraded is not vllm_config.quant_config:
         vllm_config = replace(vllm_config, quant_config=upgraded)
 
+    custom_voice_dir = engine_args_dict.get("custom_voice_dir")
+    if custom_voice_dir:
+        setattr(vllm_config.model_config.hf_config, "custom_voice_dir", custom_voice_dir)
+
     return vllm_config, executor_class
 
 
@@ -856,6 +860,45 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
     return input_processor
 
 
+def _cleanup_stale_lock_if_dead(lock_file: str) -> bool:
+    """If *lock_file* exists and its recorded PID is dead, unlink the file.
+
+    Returns ``True`` if the stale lock was cleaned up (caller should retry),
+    ``False`` otherwise (lock holder appears alive, or file could not be read).
+    """
+    try:
+        with open(lock_file) as fh:
+            content = fh.read().strip()
+        if not content:
+            return False
+        pid = int(content)
+    except (OSError, ValueError):
+        return False
+
+    # Check whether the PID is still alive.
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        # PID does not exist — stale lock.
+        logger.info(
+            "Removing stale device lock %s (PID %s is dead)",
+            lock_file,
+            pid,
+        )
+        try:
+            os.unlink(lock_file)
+            return True
+        except OSError:
+            logger.debug("Failed to unlink stale lock %s", lock_file)
+            return False
+    except PermissionError:
+        # PID exists but we cannot signal it (different user) — treat as alive.
+        return False
+
+    # PID is alive — legitimate lock holder.
+    return False
+
+
 def acquire_device_locks(
     stage_id: int,
     engine_args_dict: dict[str, Any],
@@ -934,6 +977,7 @@ def acquire_device_locks(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
+            already_cleaned_stale = False  # only try stale cleanup once per device
 
             while not lock_acquired:
                 try:
@@ -948,6 +992,11 @@ def acquire_device_locks(
                         logger.debug("Acquired exclusive lock for device %s", device_id)
                     except BlockingIOError:
                         os.close(lock_fd)
+                        # Detect and clean stale locks from dead processes.
+                        if not already_cleaned_stale:
+                            already_cleaned_stale = True
+                            if _cleanup_stale_lock_if_dead(lock_file):
+                                continue  # retry flock immediately
                         if time.time() - wait_start > stage_init_timeout:
                             logger.warning(
                                 "Timeout waiting for device %s initialization lock, proceeding anyway",

vllm_omni/engine/stage_pool.py

@@ -22,6 +22,7 @@ from vllm_omni.engine.stage_client import (
     StagePoolDiffusionClient,
     StagePoolLLMClient,
 )
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.metrics.stats import StageStats
 from vllm_omni.metrics.utils import count_tokens_from_outputs
@@ -533,6 +534,10 @@ class StagePool:
     ) -> int:
         """Submit a stage-entry request into this pool."""
         params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id]
+        # Convert plain vllm SamplingParams for single-stage diffusion models
+        # that receive sampling params from the user/caller directly.
+        if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams):
+            params = OmniDiffusionSamplingParams()
         submit_kwargs = dict(submit_kwargs or {})
         if self.stage_type == "diffusion":
             replica_id = await self._pick_or_select(
@@ -593,6 +598,8 @@ class StagePool:
     ) -> int:
         """Submit a streaming update to an already admitted request."""
         params = req_state.sampling_params_list[self.stage_id]
+        if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams):
+            params = OmniDiffusionSamplingParams()
         replica_id = self.get_bound_replica_id(request_id)
         if replica_id is None or self.clients[replica_id] is None:
             replica_id = await self._pick_or_select(request_id)

vllm_omni/entrypoints/async_omni.py

@@ -295,6 +295,15 @@ class AsyncOmni(EngineClient, OmniBase):
             # Start final output dispatcher on the first call to generate()
             self._final_output_handler()
 
+            # Forward bare sampling_params (e.g. from /v1/completions) as the stage-0 entry.
+            if sampling_params_list is None and sampling_params is not None:
+                if self.num_stages == 1:
+                    sampling_params_list = [sampling_params]
+                else:
+                    default = list(self.default_sampling_params_list)
+                    default[0] = sampling_params
+                    sampling_params_list = default
+
             # Expand sampling params for PD disaggregation (user may provide N-1 params)
             if (
                 sampling_params_list is not None
@@ -919,8 +928,14 @@ class AsyncOmni(EngineClient, OmniBase):
         for result in results:
             if isinstance(result, dict) and result.get("todo"):
                 continue
-            if isinstance(result, list):
-                merged.update(result)
+            if isinstance(result, (list, set)):
+                for item in result:
+                    if isinstance(item, (list, set)):
+                        merged.update(item)
+                    elif isinstance(item, int):
+                        merged.add(item)
+            elif isinstance(result, int):
+                merged.add(result)
         return sorted(merged)
 
     async def pin_lora(self, adapter_id: int) -> bool:

vllm_omni/entrypoints/cli/serve.py

@@ -436,7 +436,7 @@ class OmniServeCommand(CLISubcommand):
             "Equivalent to setting DiffusionParallelConfig.ring_degree.",
         )
         omni_config_group.add_argument(
-            "--quantization-config",
+            "--diffusion-quantization-config",
             type=json.loads,
             default=None,
             help=(

vllm_omni/entrypoints/omni_base.py

@@ -463,7 +463,7 @@ class OmniBase(PDDisaggregationMixin):
             peak_memory_mb=peak_memory_mb,
         )
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
         logger.info("[%s] Shutting down", self.__class__.__name__)
         self._shutdown_base()
 

vllm_omni/entrypoints/openai/api_server.py

@@ -55,13 +55,8 @@ from vllm.entrypoints.openai.engine.protocol import (
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.orca_metrics import metrics_header
-from vllm.entrypoints.openai.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
 from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config
-from vllm.entrypoints.openai.speech_to_text.serving import (
-    OpenAIServingTranscription,
-    OpenAIServingTranslation,
-)
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
@@ -74,6 +69,13 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
+from vllm.entrypoints.speech_to_text.transcription.serving import (
+    OpenAIServingTranscription,
+)
+from vllm.entrypoints.speech_to_text.translation.serving import (
+    OpenAIServingTranslation,
+)
 from vllm.entrypoints.utils import (
     create_error_response,
     load_aware_call,
@@ -355,7 +357,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     warnings_module.filterwarnings("ignore", message=".*PydanticSerializationUnexpectedValue.*", category=UserWarning)
 
     # Add process-specific prefix to stdout and stderr.
-    decorate_logs("APIServer")
+    decorate_logs("APIServer", skip_if_decorated=True)
 
     listen_address, sock = setup_openai_server(args)
 
@@ -646,6 +648,19 @@ async def omni_init_app_state(
 
     state.vllm_config = vllm_config
 
+    # Propagate enable_in_reasoning to the API-server process. The engine core
+    # runs in a separate process, so the contextvar that backs
+    # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers
+    # call `get_enable_structured_outputs_in_reasoning()` during request
+    # handling and need to see the real flag, otherwise they silently fall
+    # back to False and mismatch the engine-side bitmask gating.
+    if vllm_config is not None:
+        from vllm.tool_parsers.structural_tag_registry import (
+            set_enable_structured_outputs_in_reasoning,
+        )
+
+        set_enable_structured_outputs_in_reasoning(vllm_config.structured_outputs_config.enable_in_reasoning)
+
     # Get supported tasks
     supported_tasks: set[str] = {"generate"}
     if hasattr(engine_client, "get_supported_tasks"):
@@ -1531,6 +1546,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                 prompt=request.prompt,
                 extra_body=extra_body,
                 request_id=f"img_gen-{random_uuid()}",
+                raw_request=raw_request,
             )
             if isinstance(generation_result, ErrorResponse):
                 return JSONResponse(
@@ -1937,6 +1953,7 @@ async def edit_images(
                 output_format=output_format,
                 output_compression=output_compression,
                 size=size_str,
+                raw_request=raw_request,
             )
             if stream and not isinstance(generation_result, ErrorResponse):
                 return StreamingResponse(

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -9,8 +9,8 @@ from uuid import uuid4
 
 import numpy as np
 from vllm.entrypoints.openai.engine.protocol import UsageInfo
-from vllm.entrypoints.openai.realtime.connection import RealtimeConnection as VllmRealtimeConnection
-from vllm.entrypoints.openai.realtime.protocol import TranscriptionDelta, TranscriptionDone
+from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection
+from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2574,6 +2574,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_format: str = "png",
         output_compression: int = 100,
         size: str = "auto",
+        raw_request: Request | None = None,
     ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]:
         """Generate diffusion images and return raw images plus generation stats."""
         if request_id is None:
@@ -2635,6 +2636,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     output_format=output_format,
                     output_compression=output_compression,
                     size=size,
+                    raw_request=raw_request,
                 )
 
             result = None
@@ -2659,6 +2661,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         peak_memory_mb = result.peak_memory_mb
         cot_output = None
 
+        req_out = getattr(result, "request_output", None)
+        if req_out:
+            prompt_obj = getattr(req_out, "prompt", None)
+            if isinstance(prompt_obj, dict):
+                extra = prompt_obj.get("extra", {})
+                if isinstance(extra, dict):
+                    ar_text = extra.get("ar_generated_text")
+                    if isinstance(ar_text, str) and ar_text.strip():
+                        cot_output = ar_text
+
         req_out = getattr(result, "request_output", None)
         if req_out:
             prompt_obj = getattr(req_out, "prompt", None)
@@ -2679,6 +2691,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_format: str,
         output_compression: int,
         size: str,
+        raw_request: Request | None = None,
     ) -> AsyncIterator[str]:
         """Yield image edit SSE chunks from multi-stage diffusion outputs."""
         created = int(time.time())
@@ -2730,7 +2743,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             logger.error("EngineDeadError during streaming image edit: %s", exc)
             data = self.create_streaming_error_response(exc)
             yield f"data: {data}\n\n"
-            terminate_if_errored(engine=self.engine_client)
+            if raw_request is not None:
+                terminate_if_errored(
+                    server=raw_request.app.state.server,
+                    engine=self.engine_client,
+                )
+            else:
+                logger.warning(
+                    "[OmniOpenAIServingChat] Engine dead during streaming image edit, "
+                    "but cannot terminate server (no raw_request context)."
+                )
         except Exception as exc:
             logger.exception("Streaming image edit failed: %s", exc)
             chunk = ImageEditStreamError(

vllm_omni/entrypoints/openai/serving_speech.py

@@ -8,6 +8,7 @@ import os
 import re
 import struct
 import time
+from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 from http import HTTPStatus
 from pathlib import Path
@@ -55,7 +56,11 @@ from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
     create_instruction as ming_create_instruction,
 )
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.speaker_cache import get_speaker_cache
+from vllm_omni.utils.speaker_cache import (
+    get_speaker_cache,
+    iter_custom_voice_profiles,
+    load_validated_profile_tensors,
+)
 
 logger = init_logger(__name__)
 
@@ -107,6 +112,9 @@ _TTS_LANGUAGES: set[str] = {
 }
 _REF_AUDIO_MIN_DURATION = 1.0  # seconds
 _REF_AUDIO_MAX_DURATION = 30.0  # seconds
+_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256
+_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024
+_QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key"
 _TTS_MAX_INSTRUCTIONS_LENGTH = 500
 _TTS_MAX_NEW_TOKENS_MIN = 1
 _TTS_MAX_NEW_TOKENS_MAX = 4096
@@ -261,8 +269,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             logger.warning("Invalid SPEAKER_MAX_UPLOADED=%r; using default 1000", _raw_cap)
             self._max_uploaded_speakers = 1000
         self.uploaded_speakers: dict[str, dict] = {}
+        self.precomputed_speakers: dict[str, dict[str, Any]] = {}
         self.supported_speakers: set[str] = set()
         self._ref_audio_data_url_cache: dict[str, str] = {}
+        self._ref_audio_resolve_cache: OrderedDict[str, tuple[list[float], int, int, str]] = OrderedDict()
+        self._ref_audio_resolve_cache_bytes = 0
+        self._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES
+        self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES
+        self._ref_audio_model_artifact_ready: set[str] = set()
+        self._request_ref_audio_artifact_keys: dict[str, str] = {}
         self._speaker_cache = get_speaker_cache()
         self._last_upload_ts = 0
         self._upload_lock = asyncio.Lock()
@@ -395,6 +410,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         # Determine TTS model type or None
         self._tts_model_type = self._detect_tts_model_type()
+        self.precomputed_speakers = self._load_precomputed_speakers()
 
         # GLM-TTS lazy-cached resources (populated on first GLM-TTS request)
         self._glm_tts_text_tokenizer: object | None = None
@@ -405,6 +421,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         # Merge built-in speakers into the set initialized by _init_speaker_storage.
         self.supported_speakers |= self._load_supported_speakers()
+        self.supported_speakers |= set(self.precomputed_speakers)
         self._tts_tokenizer = None
         self._voxcpm2_tokenizer = None
         self._voxcpm2_split_map: dict[int, list[int]] = {}
@@ -570,6 +587,45 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "glm_tts"
         return None
 
+    def _get_custom_voice_dir(self) -> str | None:
+        try:
+            value = getattr(self.engine_client.model_config.hf_config, "custom_voice_dir", None)
+        except AttributeError:
+            return None
+        if isinstance(value, os.PathLike):
+            return os.fspath(value)
+        if isinstance(value, str) and value:
+            return value
+        return None
+
+    def _load_precomputed_speakers(self) -> dict[str, dict[str, Any]]:
+        """Load precomputed voice names from ``custom_voice_dir`` for API validation."""
+        if self._tts_model_type not in ("qwen3_tts", "voxcpm2"):
+            return {}
+        custom_voice_dir = self._get_custom_voice_dir()
+        if not custom_voice_dir:
+            return {}
+
+        profiles: dict[str, dict[str, Any]] = {}
+        qwen3_embedding_dim = self._get_qwen_tts_expected_speaker_embedding_dim()
+        for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type=self._tts_model_type):
+            tensors = load_validated_profile_tensors(
+                profile,
+                expected_model_type=self._tts_model_type,
+                qwen3_embedding_dim=qwen3_embedding_dim,
+            )
+            if tensors is None:
+                continue
+            profiles[profile["voice_name_lower"]] = profile
+        if profiles:
+            logger.info(
+                "Loaded %d precomputed %s voice profile(s) from %s",
+                len(profiles),
+                self._tts_model_type,
+                custom_voice_dir,
+            )
+        return profiles
+
     def _compute_max_instructions_length(self) -> int:
         """Compute max instructions length with precedence: CLI > stage config > default.
 
@@ -748,11 +804,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self._voxcpm2_encode("")  # lazy-init tokenizer + split_map
         ref_audio = None
         ref_sr = None
+        voice_profile = None
         if request.ref_audio is not None:
             ref_audio, ref_sr = await self._resolve_ref_audio(request.ref_audio)
         elif uploaded_ref is not None:
             wav_np, ref_sr = uploaded_ref
             ref_audio = wav_np.tolist()
+        elif request.voice is not None:
+            voice_profile = self.precomputed_speakers.get(request.voice.lower())
         return build_voxcpm2_prompt(
             hf_config=self.engine_client.model_config.hf_config,
             tokenizer=self._voxcpm2_tokenizer,
@@ -761,6 +820,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             ref_audio=ref_audio,
             ref_sr=ref_sr,
             ref_text=request.ref_text,
+            voice_profile=voice_profile,
         )
 
     def _load_uploaded_audio(self, voice_name: str) -> tuple[np.ndarray, int] | None:
@@ -1205,12 +1265,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self._tts_model_type == "cosyvoice3":
             return self._validate_cosyvoice3_request(request)
         if self._tts_model_type == "voxcpm2":
-            if request.max_new_tokens is not None:
-                if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
-                    return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
-                if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
-                    return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
-            return None  # VoxCPM2 accepts any text input
+            return self._validate_voxcpm2_request(request)
         if self._tts_model_type == "ming_flash_omni_tts":
             return self._validate_ming_tts_request(request)
         if self._tts_model_type == "moss_tts_nano":
@@ -1219,6 +1274,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return self._validate_glm_tts_request(request)
         return self._validate_qwen_tts_request(request)
 
+    def _validate_voxcpm2_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        """Validate VoxCPM2 request parameters. Returns error message or None."""
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+
+        if request.voice is not None:
+            request.voice = request.voice.lower()
+            available_voices = set(self.uploaded_speakers) | set(self.precomputed_speakers) | {"default"}
+            if request.voice not in available_voices:
+                supported = ", ".join(sorted(available_voices)) or "none"
+                return f"Invalid voice '{request.voice}'. Supported: {supported}"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+            if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
+                return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
+
+        return None
+
     def _voxcpm2_encode(self, text: str) -> list[int]:
         """Tokenize text for VoxCPM2, splitting multichar Chinese tokens."""
         from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
@@ -1315,11 +1390,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # Infer Base task when ref_audio or ref_text is provided without explicit task_type.
         if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None):
             request.task_type = "Base"
-        task_type = request.task_type or "CustomVoice"
 
         # Normalize voice to lowercase for case-insensitive matching
         if request.voice is not None:
             request.voice = request.voice.lower()
+            if request.task_type is None and request.voice in self.precomputed_speakers:
+                request.task_type = "Base"
+        task_type = request.task_type or "CustomVoice"
 
         # Validate input is not empty
         if not request.input or not request.input.strip():
@@ -1379,6 +1456,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     file_path = Path(speaker_info["file_path"])
                     if not file_path.exists():
                         return f"Data file for uploaded speaker '{request.voice}' not found on disk"
+                elif voice_lower in self.precomputed_speakers:
+                    profile = self.precomputed_speakers[voice_lower]
+                    mode = str(profile.get("mode") or "xvec").lower()
+                    ref_text = request.ref_text or profile.get("ref_text")
+                    if mode == "icl" and (not isinstance(ref_text, str) or not ref_text.strip()):
+                        return (
+                            f"Precomputed voice '{request.voice}' uses ICL mode but has no ref_text in "
+                            "the request or custom voice manifest"
+                        )
                 else:
                     # need ref_audio for built-in speaker
                     if request.ref_audio is None:
@@ -1677,6 +1763,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         URLs, ``data:`` base64 URIs, and ``file:`` local paths (the latter
         gated by ``--allowed-local-media-path``).
         """
+        cache_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest()
+        cached = self._ref_audio_resolve_cache.get(cache_key)
+        if cached is not None:
+            self._ref_audio_resolve_cache.move_to_end(cache_key)
+            wav_list, sr, _, _ = cached
+            logger.debug(
+                "Resolved ref_audio from cache: samples=%d sr=%d duration_s=%.3f",
+                len(wav_list),
+                sr,
+                len(wav_list) / sr if sr > 0 else 0.0,
+            )
+            return wav_list, sr
+
         # In diffusion mode, model_config may not be available
         if self._diffusion_mode:
             connector = MediaConnector()
@@ -1686,11 +1785,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 allowed_local_media_path=model_config.allowed_local_media_path,
                 allowed_media_domains=model_config.allowed_media_domains,
             )
+        fetch_start_s = time.perf_counter()
         wav_np, sr = await connector.fetch_audio_async(ref_audio_str)
+        fetch_decode_ms = (time.perf_counter() - fetch_start_s) * 1000.0
         wav_np = np.asarray(wav_np, dtype=np.float32)
         if wav_np.ndim > 1:
             wav_np = np.mean(wav_np, axis=-1)
         sr = int(sr)
+        artifact_key = self._make_ref_audio_artifact_cache_key(wav_np, sr)
         duration = len(wav_np) / sr if sr > 0 else 0.0
         if duration < _REF_AUDIO_MIN_DURATION:
             raise ValueError(
@@ -1702,7 +1804,83 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 f"Reference audio too long ({duration:.1f}s). "
                 f"Maximum {_REF_AUDIO_MAX_DURATION:.0f}s supported — use a shorter clip."
             )
-        return wav_np.tolist(), sr
+        tolist_start_s = time.perf_counter()
+        wav_list = wav_np.tolist()
+        tolist_ms = (time.perf_counter() - tolist_start_s) * 1000.0
+        logger.debug(
+            "Resolved ref_audio: fetch_decode_ms=%.3f tolist_ms=%.3f samples=%d sr=%d duration_s=%.3f",
+            fetch_decode_ms,
+            tolist_ms,
+            len(wav_np),
+            sr,
+            duration,
+        )
+        self._put_resolved_ref_audio(cache_key, wav_list, sr, artifact_key)
+        return wav_list, sr
+
+    @staticmethod
+    def _make_ref_audio_artifact_cache_key(wav: np.ndarray, sr: int) -> str:
+        wav_f32 = wav.astype(np.float32, copy=False).reshape(-1)
+        h = hashlib.sha1()
+        h.update(int(sr).to_bytes(4, byteorder="little", signed=False))
+        h.update(int(wav_f32.size).to_bytes(8, byteorder="little", signed=False))
+        h.update(wav_f32.tobytes(order="C"))
+        return h.hexdigest()
+
+    def _get_resolved_ref_audio_artifact_key(self, ref_audio_str: str) -> str | None:
+        source_key = hashlib.sha1(ref_audio_str.encode("utf-8")).hexdigest()
+        cached = self._ref_audio_resolve_cache.get(source_key)
+        if cached is None:
+            return None
+        self._ref_audio_resolve_cache.move_to_end(source_key)
+        return cached[3]
+
+    def _put_resolved_ref_audio(self, cache_key: str, wav_list: list[float], sr: int, artifact_key: str) -> None:
+        if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0:
+            return
+        # Approximate list[float] storage. CPython float objects add per-element
+        # overhead, so max_entries remains the hard cache cap.
+        size = len(wav_list) * 40
+        if size > self._ref_audio_resolve_cache_max_bytes:
+            return
+        previous = self._ref_audio_resolve_cache.pop(cache_key, None)
+        if previous is not None:
+            self._ref_audio_resolve_cache_bytes -= previous[2]
+            if previous[3] != artifact_key:
+                self._discard_ref_audio_artifact_ready_if_unreferenced(previous[3])
+        self._ref_audio_resolve_cache[cache_key] = (wav_list, int(sr), size, artifact_key)
+        self._ref_audio_resolve_cache_bytes += size
+        while len(self._ref_audio_resolve_cache) > self._ref_audio_resolve_cache_max_entries:
+            _, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
+            self._ref_audio_resolve_cache_bytes -= old_size
+            self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
+        while self._ref_audio_resolve_cache_bytes > self._ref_audio_resolve_cache_max_bytes:
+            _, (_, _, old_size, old_artifact_key) = self._ref_audio_resolve_cache.popitem(last=False)
+            self._ref_audio_resolve_cache_bytes -= old_size
+            self._discard_ref_audio_artifact_ready_if_unreferenced(old_artifact_key)
+
+    def _discard_ref_audio_artifact_ready_if_unreferenced(self, artifact_key: str) -> None:
+        if artifact_key and all(entry[3] != artifact_key for entry in self._ref_audio_resolve_cache.values()):
+            self._ref_audio_model_artifact_ready.discard(artifact_key)
+
+    def _qwen3_tts_can_use_ref_audio_artifact_only(self, tts_params: dict[str, Any], artifact_key: str | None) -> bool:
+        if self._tts_model_type != "qwen3_tts":
+            return False
+        if not artifact_key or artifact_key not in self._ref_audio_model_artifact_ready:
+            return False
+        return (tts_params.get("task_type") or ["CustomVoice"])[0] == "Base"
+
+    def _track_ref_audio_artifact_warmup(self, request_id: str, artifact_key: str | None) -> None:
+        if artifact_key:
+            self._request_ref_audio_artifact_keys[request_id] = artifact_key
+
+    def _mark_ref_audio_artifact_ready_for_request(self, request_id: str) -> None:
+        artifact_key = self._request_ref_audio_artifact_keys.pop(request_id, None)
+        if artifact_key and any(entry[3] == artifact_key for entry in self._ref_audio_resolve_cache.values()):
+            self._ref_audio_model_artifact_ready.add(artifact_key)
+
+    def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None:
+        self._request_ref_audio_artifact_keys.pop(request_id, None)
 
     async def _generate_audio_chunks(
         self,
@@ -1733,6 +1911,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         first_chunk = True
         first_audio_chunk_s: float | None = None
         stream_start_s = request_start_s if request_start_s is not None else time.perf_counter()
+        artifact_ready = False
 
         try:
             async for res in generator:
@@ -1787,6 +1966,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     if first_audio_chunk_s is None:
                         first_audio_chunk_s = time.perf_counter()
                     yield self.create_audio(audio_obj).audio_data
+            self._mark_ref_audio_artifact_ready_for_request(request_id)
+            artifact_ready = True
             total_ms = (time.perf_counter() - stream_start_s) * 1000.0
             if first_audio_chunk_s is not None:
                 first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0
@@ -1840,6 +2021,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             )
             logger.exception("Streaming speech generation failed for %s: %s", request_id, e)
             raise
+        finally:
+            if not artifact_ready:
+                self._discard_ref_audio_artifact_warmup(request_id)
 
     @staticmethod
     def _extract_audio_output(res) -> tuple[dict | None, str | None]:
@@ -1893,13 +2077,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         # Speaker (voice)
         if request.voice is not None:
+            voice_lower = request.voice.lower()
             params["speaker"] = [request.voice]
-            params["voice_created_at"] = [self._voice_created_at(request.voice.lower())]
+            params["voice_created_at"] = [self._voice_created_at(voice_lower)]
 
             # Uploaded voices use task_type="Base" (CustomVoice requires built-in spk_id).
             # If ref_text was provided at upload time, use in-context cloning; otherwise x_vector only.
-            if request.voice.lower() in self.uploaded_speakers and request.ref_audio is None:
-                speaker_info = self.uploaded_speakers[request.voice.lower()]
+            if voice_lower in self.uploaded_speakers and request.ref_audio is None:
+                speaker_info = self.uploaded_speakers[voice_lower]
 
                 # Check if this voice was uploaded with a pre-computed embedding.
                 # Populate request.speaker_embedding so the existing code path
@@ -1907,6 +2092,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 embedding = self._get_uploaded_speaker_embedding(request.voice)
                 if embedding is not None:
                     request.speaker_embedding = embedding
+                    params["speaker"] = [voice_lower]
                     params["task_type"] = ["Base"]
                     logger.info("Auto-set speaker_embedding for uploaded voice: %s", request.voice)
                 else:
@@ -1914,6 +2100,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     if not audio_data:
                         raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing or corrupted")
                     stored_ref_text = speaker_info.get("ref_text")
+                    params["speaker"] = [voice_lower]
                     params["ref_audio"] = [audio_data]
                     params["task_type"] = ["Base"]
                     if stored_ref_text:
@@ -1924,6 +2111,19 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     logger.info(
                         "Auto-set ref_audio for uploaded voice: %s (icl=%s)", request.voice, bool(stored_ref_text)
                     )
+            elif voice_lower in self.precomputed_speakers and request.ref_audio is None:
+                profile = self.precomputed_speakers[voice_lower]
+                mode = str(profile.get("mode") or "xvec").lower()
+                params["speaker"] = [voice_lower]
+                params["task_type"] = ["Base"]
+                params["x_vector_only_mode"] = [mode != "icl"]
+                ref_text = request.ref_text or profile.get("ref_text")
+                if isinstance(ref_text, str) and ref_text.strip():
+                    params["ref_text"] = [ref_text]
+                ref_code_length = profile.get("ref_code_length")
+                if mode == "icl" and ref_code_length:
+                    params["ref_code_length"] = [int(ref_code_length)]
+                logger.info("Using precomputed Qwen3-TTS custom voice profile: %s (mode=%s)", voice_lower, mode)
 
         elif params["task_type"][0] == "CustomVoice":
             params["speaker"] = ["Vivian"]  # Default for CustomVoice
@@ -2334,6 +2534,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self.engine_client.errored:
             raise self.engine_client.dead_error
 
+        request_id = request_id or f"speech-{random_uuid()}"
+        qwen3_ref_audio_warmup_artifact_key: str | None = None
+
         # If this is a streaming request, we need to coerce
         # cumulative outputs to delta outputs; this ensures
         # we don't emit redundant MM data & drain after emitting.
@@ -2412,12 +2615,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         elif self._tts_model_type == "voxcpm2":
             # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
             # raw waveform tuple for prefill-length accounting, not a base64 data URL.
+            validation_error = self._validate_voxcpm2_request(request)
+            if validation_error:
+                raise ValueError(validation_error)
+
             uploaded_ref: tuple[np.ndarray, int] | None = None
             if request.voice:
                 voice_lower = request.voice.lower()
-                if voice_lower not in self.uploaded_speakers and voice_lower not in self.supported_speakers:
-                    all_voices = sorted(self.uploaded_speakers.keys() | self.supported_speakers)
-                    raise ValueError(f"Invalid voice '{request.voice}'. Supported: {', '.join(all_voices) or 'none'}")
                 if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
                     if self.uploaded_speakers[voice_lower].get("embedding_source") == "direct":
                         raise ValueError(
@@ -2430,7 +2634,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             tts_params = {}
             if request.voice:
                 voice_lower = request.voice.lower()
-                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
+                if voice_lower in self.uploaded_speakers or voice_lower in self.precomputed_speakers:
                     additional = prompt.setdefault("additional_information", {})
                     additional["voice_name"] = voice_lower
                     additional["voice_created_at"] = self._voice_created_at(voice_lower)
@@ -2478,7 +2682,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     ref_audio_source = tts_params["ref_audio"][0]
                 if ref_audio_source is not None and isinstance(ref_audio_source, str):
                     wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
-                    tts_params["ref_audio"] = [[wav_list, sr]]
+                    artifact_key = self._get_resolved_ref_audio_artifact_key(ref_audio_source)
+                    if self._tts_model_type == "qwen3_tts" and artifact_key:
+                        tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key]
+                    ref_code_length = self._estimate_ref_code_len([wav_list, sr])
+                    if self._tts_model_type == "qwen3_tts" and ref_code_length is not None:
+                        tts_params["ref_code_length"] = [int(ref_code_length)]
+                    if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key):
+                        logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key)
+                    else:
+                        tts_params["ref_audio"] = [[wav_list, sr]]
+                        if self._tts_model_type == "qwen3_tts":
+                            qwen3_ref_audio_warmup_artifact_key = artifact_key
 
                 ph_len = await self._estimate_prompt_len_async(tts_params)
                 prompt = tokens_input(prompt_token_ids=[1] * ph_len)
@@ -2506,7 +2721,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             tts_params = {}
             prompt = {"prompt": request.input}
 
-        request_id = request_id or f"speech-{random_uuid()}"
         if self._is_fish_speech:
             model_type = "fish_speech"
         elif self._tts_model_type == "covo_audio":
@@ -2643,6 +2857,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             sampling_params_list=sampling_params_list,
             output_modalities=["audio"],
         )
+        self._track_ref_audio_artifact_warmup(request_id, qwen3_ref_audio_warmup_artifact_key)
         return request_id, generator, tts_params
 
     async def _generate_pcm_chunks(self, generator, request_id: str):
@@ -2657,8 +2872,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest):
         """Yield raw PCM bytes for a speech request as soon as chunks are decoded."""
         request_id, generator, _ = await self._prepare_speech_generation(request)
-        async for chunk in self._generate_pcm_chunks(generator, request_id):
-            yield chunk
+        try:
+            async for chunk in self._generate_pcm_chunks(generator, request_id):
+                yield chunk
+        finally:
+            self._discard_ref_audio_artifact_warmup(request_id)
 
     async def _generate_audio_bytes(
         self,
@@ -2667,101 +2885,108 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id: str | None = None,
     ) -> tuple[bytes | str, str]:
         request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id)
+        artifact_ready = False
 
-        # MOSS-TTS-Nano emits delta chunks per yield (single-stage,
-        # async_chunk=false). The engine surfaces each yield as its own
-        # RequestOutput, so we need to accumulate across the async-for loop —
-        # final_output alone only carries the last (often empty) sentinel.
-        is_moss = self._tts_model_type == "moss_tts_nano"
-        moss_chunks: list[Any] = []
-        moss_sample_rate: int | None = None
+        try:
+            # MOSS-TTS-Nano emits delta chunks per yield (single-stage,
+            # async_chunk=false). The engine surfaces each yield as its own
+            # RequestOutput, so we need to accumulate across the async-for loop —
+            # final_output alone only carries the last (often empty) sentinel.
+            is_moss = self._tts_model_type == "moss_tts_nano"
+            moss_chunks: list[Any] = []
+            moss_sample_rate: int | None = None
 
-        final_output: OmniRequestOutput | None = None
-        async for res in generator:
-            final_output = res
-            if not is_moss:
-                continue
-            try:
-                step_audio, step_key = self._extract_audio_output(res)
-            except Exception:
-                continue
-            if step_key is None:
-                continue
-            chunk = step_audio[step_key]
-            candidates = chunk if isinstance(chunk, list) else [chunk]
-            for cand in candidates:
-                if hasattr(cand, "numel") and cand.numel() > 0:
-                    moss_chunks.append(cand)
-            sr_step = step_audio.get("sr")
-            if sr_step is not None:
-                sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step
-                moss_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step)
+            final_output: OmniRequestOutput | None = None
+            async for res in generator:
+                final_output = res
+                if not is_moss:
+                    continue
+                try:
+                    step_audio, step_key = self._extract_audio_output(res)
+                except Exception:
+                    continue
+                if step_key is None:
+                    continue
+                chunk = step_audio[step_key]
+                candidates = chunk if isinstance(chunk, list) else [chunk]
+                for cand in candidates:
+                    if hasattr(cand, "numel") and cand.numel() > 0:
+                        moss_chunks.append(cand)
+                sr_step = step_audio.get("sr")
+                if sr_step is not None:
+                    sr_val_step = sr_step[-1] if isinstance(sr_step, list) and sr_step else sr_step
+                    moss_sample_rate = int(sr_val_step.item()) if hasattr(sr_val_step, "item") else int(sr_val_step)
 
-        if final_output is None:
-            raise ValueError("No output generated from the model.")
+            if final_output is None:
+                raise ValueError("No output generated from the model.")
 
-        audio_output, audio_key = self._extract_audio_output(final_output)
-        if audio_key is None:
-            raise ValueError("TTS model did not produce audio output.")
+            audio_output, audio_key = self._extract_audio_output(final_output)
+            if audio_key is None:
+                raise ValueError("TTS model did not produce audio output.")
 
-        audio_tensor = audio_output[audio_key]
-        sr_raw = audio_output.get("sr", 24000)
-        sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
-        sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+            audio_tensor = audio_output[audio_key]
+            sr_raw = audio_output.get("sr", 24000)
+            sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+            sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
 
-        if is_moss:
-            # Prefer the engine's own consolidated audio when present. After the
-            # vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so
-            # final_output already carries the full concatenated waveform; the
-            # delta-accumulator below is kept as a fallback for DELTA-style
-            # engines that surface chunks one yield at a time.
-            if isinstance(audio_tensor, list):
-                non_empty_final = [c for c in audio_tensor if hasattr(c, "numel") and c.numel() > 0]
-                final_audio = torch.cat(non_empty_final, dim=-1) if non_empty_final else None
-            elif hasattr(audio_tensor, "numel") and audio_tensor.numel() > 0:
-                final_audio = audio_tensor
-            else:
-                final_audio = None
+            if is_moss:
+                # Prefer the engine's own consolidated audio when present. After the
+                # vllm 0.20 rebase non-stream requests resolve to FINAL_ONLY, so
+                # final_output already carries the full concatenated waveform; the
+                # delta-accumulator below is kept as a fallback for DELTA-style
+                # engines that surface chunks one yield at a time.
+                if isinstance(audio_tensor, list):
+                    non_empty_final = [c for c in audio_tensor if hasattr(c, "numel") and c.numel() > 0]
+                    final_audio = torch.cat(non_empty_final, dim=-1) if non_empty_final else None
+                elif hasattr(audio_tensor, "numel") and audio_tensor.numel() > 0:
+                    final_audio = audio_tensor
+                else:
+                    final_audio = None
 
-            if final_audio is not None:
-                audio_tensor = final_audio
-            elif moss_chunks:
-                audio_tensor = torch.cat(moss_chunks, dim=-1)
-            else:
-                audio_tensor = np.zeros((0,), dtype=np.float32)
-            if moss_sample_rate is not None:
-                sample_rate = moss_sample_rate
-        elif isinstance(audio_tensor, list):
-            async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
-            if async_chunk:
-                non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
-                audio_tensor = (
-                    torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
-                )
-            else:
-                audio_history = audio_tensor
-                audio_tensor = np.zeros((0,), dtype=np.float32)
-                # Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
-                for candidate in reversed(audio_history):
-                    if candidate.numel() > 0:
-                        audio_tensor = candidate
-                        break
-        if hasattr(audio_tensor, "float"):
-            audio_tensor = audio_tensor.float().detach().cpu().numpy()
+                if final_audio is not None:
+                    audio_tensor = final_audio
+                elif moss_chunks:
+                    audio_tensor = torch.cat(moss_chunks, dim=-1)
+                else:
+                    audio_tensor = np.zeros((0,), dtype=np.float32)
+                if moss_sample_rate is not None:
+                    sample_rate = moss_sample_rate
+            elif isinstance(audio_tensor, list):
+                async_chunk = bool(getattr(self.engine_client.model_config, "async_chunk", False))
+                if async_chunk:
+                    non_empty_chunks = [candidate for candidate in audio_tensor if candidate.numel() > 0]
+                    audio_tensor = (
+                        torch.cat(non_empty_chunks, dim=-1) if non_empty_chunks else np.zeros((0,), dtype=np.float32)
+                    )
+                else:
+                    audio_history = audio_tensor
+                    audio_tensor = np.zeros((0,), dtype=np.float32)
+                    # Non-async Qwen3-TTS returns cumulative history snapshots, so keep the latest non-empty tensor.
+                    for candidate in reversed(audio_history):
+                        if candidate.numel() > 0:
+                            audio_tensor = candidate
+                            break
+            if hasattr(audio_tensor, "float"):
+                audio_tensor = audio_tensor.float().detach().cpu().numpy()
 
-        if audio_tensor.ndim > 1:
-            audio_tensor = audio_tensor.squeeze()
+            if audio_tensor.ndim > 1:
+                audio_tensor = audio_tensor.squeeze()
 
-        audio_obj = CreateAudio(
-            audio_tensor=audio_tensor,
-            sample_rate=sample_rate,
-            response_format=request.response_format or "wav",
-            speed=request.speed or 1.0,
-            stream_format=request.stream_format,
-            base64_encode=base64_encode,
-        )
-        audio_response: AudioResponse = self.create_audio(audio_obj)
-        return audio_response.audio_data, audio_response.media_type
+            audio_obj = CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=sample_rate,
+                response_format=request.response_format or "wav",
+                speed=request.speed or 1.0,
+                stream_format=request.stream_format,
+                base64_encode=base64_encode,
+            )
+            audio_response: AudioResponse = self.create_audio(audio_obj)
+            self._mark_ref_audio_artifact_ready_for_request(request_id)
+            artifact_ready = True
+            return audio_response.audio_data, audio_response.media_type
+        finally:
+            if not artifact_ready:
+                self._discard_ref_audio_artifact_warmup(request_id)
 
     async def _create_diffusion_speech(
         self,

vllm_omni/entrypoints/stage_utils.py

@@ -114,6 +114,24 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list:
         return list(device_list)
 
     logical_ids = [int(device) for device in device_list]
+
+    # If ALL logical device IDs are >= num_visible, none can be 0-based indices
+    # into the visible list — they must be physical device IDs.  This happens
+    # during multi-replica stage initialization where split_devices_for_replicas
+    # assigns physical device IDs (e.g. "2") that are not part of the current
+    # process's CUDA_VISIBLE_DEVICES but are valid physical GPUs.  Pass them
+    # through directly so the replica subprocess inherits the correct device.
+    if logical_ids and all(d >= num_visible for d in logical_ids):
+        logger.warning(
+            "Stage %s has device IDs %s, none of which are < the visible device count %d "
+            "(visible=%s). Treating them as physical device IDs and passing through.",
+            stage_id,
+            device_list,
+            num_visible,
+            visible_device_list,
+        )
+        return list(device_list)
+
     mapped_devices = [visible_device_list[idx] for idx in logical_ids if idx < num_visible]
     mapping_pairs = [
         f"{logical_id}->{visible_device_list[logical_id]}" for logical_id in logical_ids if logical_id < num_visible

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import replace
 from functools import partial
 from threading import Lock
 
@@ -365,6 +364,12 @@ class CosyVoice3Model(
         else:
             raise ValueError(f"Model stage not supported {self.model_stage}")
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig:
         """Create VllmConfig for the inner Qwen2 LLM.
 
@@ -542,17 +547,21 @@ class CosyVoice3Model(
         if self.model_stage != "cosyvoice3_talker":
             return None
 
-        sampler = getattr(self, "_talker_sampler", None)
-        if sampler is None:
-            sampler = Sampler()
-            self._talker_sampler = sampler
-
         if not self._cosyvoice3_ras_enabled(sampling_metadata):
+            sampler = getattr(self, "_talker_sampler", None)
+            if sampler is None:
+                sampler = Sampler()
+                self._talker_sampler = sampler
             return sampler(logits=logits, sampling_metadata=sampling_metadata)
 
         logits = logits.to(torch.float32)
-        sampling_for_processors = replace(sampling_metadata, no_penalties=True)
-        logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False)
+        # Apply logits processors directly — RAS handles its own repetition
+        # logic.  We avoid instantiating Sampler() here because its import
+        # chain pulls in flashinfer / GPU deps that fail in CPU-only tests.
+        if sampling_metadata.allowed_token_ids_mask is not None:
+            logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf"))
+        for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
+            logits = processor.apply(logits)
 
         sampling_cfg = dict(self.config.llm.get("sampling", {}))
         default_top_p = float(sampling_cfg.get("top_p", 0.8))

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -243,6 +243,12 @@ class CovoAudioForConditionalGeneration(
             else lambda *args, **kwargs: None
         )
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def fused_thinker_talker_preprocess(
         self,
         input_ids: torch.Tensor | None,

vllm_omni/model_executor/models/lance/__init__.py

@@ -1,9 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance (ByteDance) — vLLM-Omni model_executor package.
-
-Single-stage topology: a self-contained diffusion stage (``LancePipeline``
-defined in ``vllm_omni.diffusion.models.lance``) that owns the Qwen2-MoT LLM,
-the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous to
-``bagel_single_stage``. The ``PipelineConfig`` lives in :mod:`.pipeline`.
-"""

vllm_omni/model_executor/models/lance/pipeline.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance pipeline topology.
-
-Single-stage DiT — self-contained diffusion stage that handles all modalities
-(text2img, image_edit, x2t_image, video) internally via its own Qwen2-MoT
-LLM, Qwen2.5-VL ViT, Wan2.2 VAE, and tokenizer. Mirrors ``bagel_single_stage``.
-"""
-
-from vllm_omni.config.stage_config import (
-    PipelineConfig,
-    StageExecutionType,
-    StagePipelineConfig,
-)
-
-LANCE_PIPELINE = PipelineConfig(
-    model_type="lance",
-    model_arch="LancePipeline",
-    hf_architectures=(),
-    stages=(
-        StagePipelineConfig(
-            stage_id=0,
-            model_stage="dit",
-            execution_type=StageExecutionType.DIFFUSION,
-            input_sources=(),
-            final_output=True,
-            final_output_type="image",
-        ),
-    ),
-)

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py

@@ -143,6 +143,12 @@ class MingFlashOmniForConditionalGeneration(
                 f"For image generation use stage_type: diffusion with model_arch: MingImagePipeline."
             )
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def forward(
         self,
         input_ids: torch.Tensor,

vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py

@@ -400,6 +400,13 @@ class MossTTSNanoForGeneration(nn.Module):
         runtime_additional_information: list[dict[str, Any]] | None = None,
         **kwargs: Any,
     ) -> OmniOutput:
+        if self._lm is None or self._audio_tokenizer is None:
+            # load_weights() loads the real LM from HuggingFace.  When
+            # load_format=dummy the vllm DummyModelLoader skips calling
+            # model.load_weights(), so we do it lazily on first forward.
+            # Must happen before _make_dummy_hidden so self._device is set.
+            self.load_weights([])
+
         sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000)
         sr_tensor = torch.tensor(sr, dtype=torch.int32)
         empty = torch.zeros((0,), dtype=torch.float32)
@@ -418,9 +425,6 @@ class MossTTSNanoForGeneration(nn.Module):
                 },
             )
 
-        if self._lm is None or self._audio_tokenizer is None:
-            raise RuntimeError("MOSS-TTS-Nano model not loaded.  Was load_weights() called?")
-
         outputs: list[torch.Tensor] = []
         srs: list[torch.Tensor] = []
         last_chunk_flags: list[bool] = []

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -91,18 +91,17 @@ class Qwen2_5OmniForConditionalGeneration(
 
         elif self.model_stage == "talker":
             multimodal_config.skip_mm_profiling = True
-            # register the process function for the talker stage
             self.has_preprocess = True
             self.set_custom_preprocess(self.talker_preprocess)
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
-            self.talker = init_vllm_registered_model(
-                vllm_config=vllm_config,
-                prefix=maybe_prefix(prefix, "talker"),
-                hf_config=talker_config,
-                # Use registry architecture key
-                architectures=["Qwen2_5OmniTalkerModel"],
-            )
+            with self._mark_language_model(vllm_config=vllm_config):
+                self.talker = init_vllm_registered_model(
+                    vllm_config=vllm_config,
+                    prefix=maybe_prefix(prefix, "talker"),
+                    hf_config=talker_config,
+                    architectures=["Qwen2_5OmniTalkerModel"],
+                )
             self.model = self.talker
             self.token2wav = None
             # set suppress start id according to token2wav
@@ -130,12 +129,13 @@ class Qwen2_5OmniForConditionalGeneration(
             self.token2wav_config = getattr(config, "token2wav_config", None)
             self.token2wav = None
             if self.token2wav_config is not None:
-                self.token2wav = init_vllm_registered_model(
-                    vllm_config=vllm_config,
-                    prefix=maybe_prefix(prefix, "token2wav"),
-                    hf_config=self.token2wav_config,
-                    architectures=["Qwen2_5OmniToken2WavModel"],
-                )
+                with self._mark_language_model(vllm_config=vllm_config):
+                    self.token2wav = init_vllm_registered_model(
+                        vllm_config=vllm_config,
+                        prefix=maybe_prefix(prefix, "token2wav"),
+                        hf_config=self.token2wav_config,
+                        architectures=["Qwen2_5OmniToken2WavModel"],
+                    )
             # voice resources (loaded on demand)
             self._token2wav_conds: dict[str, torch.Tensor] = {}
             self._token2wav_ref_mels: dict[str, torch.Tensor] = {}
@@ -149,6 +149,18 @@ class Qwen2_5OmniForConditionalGeneration(
             (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None
         )
 
+    def get_language_model(self) -> nn.Module:
+        """Return the language model for MoE detection in upstream load_model.
+
+        Upstream GPUModelRunner.load_model() calls get_language_model() to find
+        the MoE model inside VLM wrappers. For thinker (which IS MoE), delegate
+        to the child. For talker/code2wav (not MoE, not SupportsMultiModal),
+        return self.model directly to avoid NotImplementedError.
+        """
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     # -------------------- Device utilities --------------------
     @staticmethod
     def _module_device(module: nn.Module) -> torch.device:

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -278,6 +278,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             return self.model.sampler
         return Sampler()
 
+    def get_language_model(self) -> torch.nn.Module:
+        """Delegate to the active stage's language model for upstream MoE resolution."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def embed_input_ids(
         self,
         input_ids: torch.Tensor,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -64,6 +64,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
+    SupportsLoRA,
     SupportsMRoPE,
     SupportsMultiModal,
     SupportsPP,
@@ -1071,6 +1072,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     nn.Module,
     SupportsMultiModal,
     SupportsPP,
+    SupportsLoRA,
     SupportsMRoPE,
     Qwen3OmniMoeConditionalGenerationMixin,
     SupportsTranscription,
@@ -1089,6 +1091,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             "k_proj",
             "v_proj",
         ],
+        "attn.qkv": [
+            "attn.q",
+            "attn.k",
+            "attn.v",
+        ],
         "gate_up_proj": [
             "gate_proj",
             "up_proj",
@@ -1111,7 +1118,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.vllm_config = vllm_config  # needed for torch compile forward context
-        thinker_config: Qwen3OmniMoeThinkerConfig = vllm_config.model_config.hf_config
+        hf_config = vllm_config.model_config.hf_config
+        thinker_config: Qwen3OmniMoeThinkerConfig = getattr(hf_config, "thinker_config", None) or hf_config
         quant_config = vllm_config.quant_config
         multimodal_config = vllm_config.model_config.multimodal_config
         self.config = thinker_config
@@ -1660,6 +1668,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         self,
         input_tokens: list[int],
         mm_features: list[MultiModalFeatureSpec],
+        **kwargs,
     ) -> tuple[torch.Tensor, int]:
         """Compute M-RoPE input positions using mm_features directly."""
         seq_len = len(input_tokens)

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -757,6 +757,14 @@ class Qwen3TTSPromptEmbedsBuilder:
             ref_audio_list = info_dict.get("ref_audio")
             if not isinstance(ref_audio_list, list) or not ref_audio_list:
                 continue
+            cache_key_from_serving = first_value(info_dict.get(REF_AUDIO_CACHE_KEY))
+            if isinstance(cache_key_from_serving, str) and cache_key_from_serving:
+                cached = self.get_ref_audio_artifacts(cache_key_from_serving)
+                if cached is not None:
+                    cached_ref_code = coerce_ref_code_tensor(cached.get("ref_code"), device=device)
+                    if isinstance(cached_ref_code, torch.Tensor):
+                        info_dict.setdefault("codes", {})[PRECOMPUTED_REF_CODE_KEY] = cached_ref_code
+                        continue
             try:
                 wav, sr = self.normalize_ref_audio(ref_audio_list[0])
             except Exception:
@@ -1029,6 +1037,19 @@ class Qwen3TTSPromptEmbedsBuilder:
             xvec_only = bool((info_dict.get("x_vector_only_mode") or [False])[0])
             in_context_mode = not xvec_only
             voice_clone_prompt = _normalize_voice_clone_prompt(info_dict.get("voice_clone_prompt"))
+            if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
+                icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
+                if isinstance(icl_flag, bool):
+                    in_context_mode = icl_flag
+                    xvec_only = not in_context_mode
+            prompt_has_ref_code = voice_clone_prompt is not None and self._has_ref_code_like(
+                voice_clone_prompt.get("ref_code")
+            )
+            prompt_spk = voice_clone_prompt.get("ref_spk_embedding") if voice_clone_prompt is not None else None
+            prompt_has_speaker = isinstance(prompt_spk, torch.Tensor) or (
+                isinstance(prompt_spk, (list, np.ndarray)) and len(prompt_spk) > 0
+            )
+            needs_ref_audio_cache_key = not (prompt_has_speaker and (not in_context_mode or prompt_has_ref_code))
             ref_audio_wav: np.ndarray | None = None
             ref_audio_sr: int | None = None
 
@@ -1056,7 +1077,8 @@ class Qwen3TTSPromptEmbedsBuilder:
             # from ref_audio: x_vector_only_mode, precomputed ref_code
             # with no ref_spk_embedding, voice_clone_prompt.icl_mode=False.
             ref_audio_cache_key = info_dict.pop(REF_AUDIO_CACHE_KEY, None)
-            if not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key):
+            ref_audio_cache_key = first_value(ref_audio_cache_key)
+            if needs_ref_audio_cache_key and not (isinstance(ref_audio_cache_key, str) and ref_audio_cache_key):
                 normalized = info_dict.get(NORMALIZED_REF_AUDIO_KEY)
                 wav_peek: np.ndarray | None = None
                 sr_peek: int | None = None
@@ -1075,13 +1097,25 @@ class Qwen3TTSPromptEmbedsBuilder:
             cached_artifacts = None
             if isinstance(ref_audio_cache_key, str) and ref_audio_cache_key:
                 cached_artifacts = self.get_ref_audio_artifacts(ref_audio_cache_key)
+            ref_audio_list = info_dict.get("ref_audio")
+            has_ref_audio_payload = isinstance(ref_audio_list, list) and bool(ref_audio_list)
+            artifact_only = (
+                isinstance(ref_audio_cache_key, str) and bool(ref_audio_cache_key) and not has_ref_audio_payload
+            )
+            if artifact_only and cached_artifacts is None:
+                raise RuntimeError(
+                    "Qwen3-TTS ref_audio artifact cache miss for artifact-only request; "
+                    "retry with ref_audio or precompute a custom voice profile."
+                )
 
             # Speaker cache: only for uploaded (named) speakers
             _speaker_cache_key = None
+            _cache_lookup_voice = None
             if voice_clone_prompt is None and self._speaker_cache is not None:
                 _speaker_list = info_dict.get("speaker")
                 if isinstance(_speaker_list, list) and _speaker_list:
                     _voice_name = str(_speaker_list[0]).lower()
+                    _cache_lookup_voice = _voice_name
                     # Per-mode namespace — xvec and icl produce different artifacts
                     # for the same voice, so they must not share a cache slot.
                     _mode = "xvec" if xvec_only else "icl"
@@ -1103,15 +1137,18 @@ class Qwen3TTSPromptEmbedsBuilder:
                             "ref_code": ref_code_cached,
                             "ref_spk_embedding": ref_spk_embed_cached,
                             "icl_mode": _cached.get("icl_mode"),
+                            "ref_text": _cached.get("ref_text"),
                         }
                         _speaker_cache_key = None  # hit → don't store again
 
-            # Official implementation may pass `voice_clone_prompt.icl_mode`.
-            if voice_clone_prompt is not None and "icl_mode" in voice_clone_prompt:
-                icl_flag = _as_singleton(voice_clone_prompt.get("icl_mode"))
-                if isinstance(icl_flag, bool):
-                    in_context_mode = icl_flag
-                    xvec_only = not in_context_mode
+            if voice_clone_prompt is None and _speaker_cache_key is not None:
+                ref_audio_list = info_dict.get("ref_audio")
+                if not isinstance(ref_audio_list, list) or not ref_audio_list:
+                    raise ValueError(
+                        f"Qwen3-TTS speaker '{_cache_lookup_voice}' was requested without ref_audio, "
+                        "but no precomputed cache entry was loaded"
+                    )
+
             ref_code = None
             if voice_clone_prompt is not None:
                 ref_code = _as_singleton(voice_clone_prompt.get("ref_code"))
@@ -1136,6 +1173,8 @@ class Qwen3TTSPromptEmbedsBuilder:
                 if isinstance(cached_ref_code, torch.Tensor):
                     ref_code_t = cached_ref_code
                     ref_code_len = int(ref_code_t.shape[0])
+            if ref_code_t is None and in_context_mode and artifact_only:
+                raise RuntimeError("Qwen3-TTS ref_audio artifact cache entry is missing ref_code.")
             if ref_code_t is None and in_context_mode:
                 # Compute ref_code from ref_audio if not provided.
                 wav_np, sr = _get_ref_audio()
@@ -1161,6 +1200,8 @@ class Qwen3TTSPromptEmbedsBuilder:
                     .to(device=input_ids.device, dtype=torch.bfloat16)
                     .view(1, 1, -1)
                 )
+            elif artifact_only:
+                raise RuntimeError("Qwen3-TTS ref_audio artifact cache entry is missing ref_spk_embedding.")
             else:
                 wav_np, sr = _get_ref_audio()
                 speaker_embed = self.extract_speaker_embedding(wav_np, sr).view(1, 1, -1)
@@ -1207,6 +1248,8 @@ class Qwen3TTSPromptEmbedsBuilder:
                     )
                 if ref_ids is None:
                     ref_text = _as_singleton(info_dict.get("ref_text"))
+                    if (not isinstance(ref_text, str) or not ref_text.strip()) and voice_clone_prompt is not None:
+                        ref_text = _as_singleton(voice_clone_prompt.get("ref_text"))
                     if isinstance(ref_text, str) and ref_text.strip():
                         ref_ids = tok(
                             build_ref_text(ref_text),
@@ -1489,6 +1532,13 @@ class Qwen3TTSPromptEmbedsBuilder:
                     except Exception:
                         ref_code_len = None
 
+                if ref_code_len is None:
+                    ref_code_length = _first(info.get("ref_code_length"), None)
+                    try:
+                        if ref_code_length is not None:
+                            ref_code_len = int(ref_code_length)  # type: ignore[arg-type]
+                    except (TypeError, ValueError):
+                        ref_code_len = None
                 if ref_code_len is None and estimate_ref_code_len is not None:
                     ref_code_len = estimate_ref_code_len(info.get("ref_audio"))
                 if ref_code_len is None:
@@ -1511,6 +1561,10 @@ class Qwen3TTSPromptEmbedsBuilder:
 
                     if ref_ids is None:
                         ref_text = _first(info.get("ref_text"), "")
+                        if (not isinstance(ref_text, str) or not ref_text.strip()) and isinstance(
+                            voice_clone_prompt, dict
+                        ):
+                            ref_text = _first(voice_clone_prompt.get("ref_text"), "")
                         if not isinstance(ref_text, str) or not ref_text.strip():
                             raise ValueError(
                                 "Base in-context non-streaming requires `ref_text` or tokenized `ref_ids`."

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from collections import Counter
+from collections import Counter, OrderedDict
 from collections.abc import Iterable
 from typing import Any
 
@@ -24,6 +24,9 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
 
 logger = init_logger(__name__)
 
+_REF_CONTEXT_CACHE_MAX_ENTRIES = 4096
+_REF_CONTEXT_CACHE_MAX_BYTES = 64 * 1024 * 1024
+
 
 def _codec_ids_from_payload_or_input(
     input_ids: torch.Tensor,
@@ -100,6 +103,10 @@ class Qwen3TTSCode2Wav(nn.Module):
         self._output_sample_rate = int(tok_config.output_sample_rate)
         self._total_upsample = int(self.decoder.total_upsample)
         self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0)
+        self._ref_context_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
+        self._ref_context_cache_bytes = 0
+        self._ref_context_cache_max_entries = _REF_CONTEXT_CACHE_MAX_ENTRIES
+        self._ref_context_cache_max_bytes = _REF_CONTEXT_CACHE_MAX_BYTES
 
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
         # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner.
@@ -205,6 +212,49 @@ class Qwen3TTSCode2Wav(nn.Module):
         self._batch_stats_actual_frames.update(actual_frames)
         self._batch_stats_bucket_groups[(group_size, bucket_frames)] += 1
 
+    @staticmethod
+    def _tensor_nbytes(tensor: torch.Tensor) -> int:
+        return int(tensor.numel() * tensor.element_size())
+
+    def _evict_ref_context_cache_if_needed(self) -> None:
+        evicted = 0
+        while len(self._ref_context_cache) > self._ref_context_cache_max_entries:
+            _, cached = self._ref_context_cache.popitem(last=False)
+            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
+            evicted += 1
+        while self._ref_context_cache_bytes > self._ref_context_cache_max_bytes and len(self._ref_context_cache) > 1:
+            _, cached = self._ref_context_cache.popitem(last=False)
+            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
+            evicted += 1
+        if evicted:
+            logger.debug(
+                "Evicted %d Qwen3-TTS ref context cache entries; entries=%d bytes=%d",
+                evicted,
+                len(self._ref_context_cache),
+                self._ref_context_cache_bytes,
+            )
+
+    def _cache_ref_context(self, request_id: str, tensor: torch.Tensor) -> None:
+        previous = self._ref_context_cache.pop(request_id, None)
+        if previous is not None:
+            self._ref_context_cache_bytes -= self._tensor_nbytes(previous)
+        cached = tensor.detach().contiguous()
+        self._ref_context_cache[request_id] = cached
+        self._ref_context_cache.move_to_end(request_id)
+        self._ref_context_cache_bytes += self._tensor_nbytes(cached)
+        self._evict_ref_context_cache_if_needed()
+
+    def _get_ref_context(self, request_id: str) -> torch.Tensor | None:
+        cached = self._ref_context_cache.get(request_id)
+        if cached is not None:
+            self._ref_context_cache.move_to_end(request_id)
+        return cached
+
+    def _pop_ref_context(self, request_id: str) -> None:
+        cached = self._ref_context_cache.pop(request_id, None)
+        if cached is not None:
+            self._ref_context_cache_bytes -= self._tensor_nbytes(cached)
+
     def log_decode_batch_stats(self) -> None:
         if not self._batch_stats_enabled or self._batch_stats_requests == 0:
             return
@@ -268,6 +318,32 @@ class Qwen3TTSCode2Wav(nn.Module):
         valid_codes_qf: list[torch.Tensor] = []
         valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
+        ref_context_size = [0] * len(request_ids_list)
+        ref_context_request_ids: list[str | None] = [None] * len(request_ids_list)
+        ref_context_included = [False] * len(request_ids_list)
+        finished_flags = [False] * len(request_ids_list)
+
+        def _meta_int(value: Any) -> int:
+            if isinstance(value, list):
+                value = value[0] if value else 0
+            if isinstance(value, torch.Tensor):
+                value = value.reshape(-1)[0].item() if value.numel() > 0 else 0
+            return int(value or 0)
+
+        def _meta_str(value: Any) -> str | None:
+            if isinstance(value, list):
+                value = value[0] if value else None
+            if value is None:
+                return None
+            return str(value)
+
+        def _meta_bool(value: Any) -> bool:
+            if isinstance(value, list):
+                value = value[0] if value else False
+            if isinstance(value, torch.Tensor):
+                return bool(value.reshape(-1)[0].item()) if value.numel() > 0 else False
+            return bool(value)
+
         if runtime_infos:
             for i, info in enumerate(runtime_infos):
                 if i >= len(left_context_size):
@@ -276,7 +352,15 @@ class Qwen3TTSCode2Wav(nn.Module):
                     continue
                 meta = info.get("meta", {})
                 if "left_context_size" in meta:
-                    left_context_size[i] = meta["left_context_size"]
+                    left_context_size[i] = _meta_int(meta["left_context_size"])
+                if "ref_context_size" in meta:
+                    ref_context_size[i] = _meta_int(meta["ref_context_size"])
+                if "ref_context_request_id" in meta:
+                    ref_context_request_ids[i] = _meta_str(meta["ref_context_request_id"])
+                if "ref_context_included" in meta:
+                    ref_context_included[i] = _meta_bool(meta["ref_context_included"])
+                if "finished" in meta:
+                    finished_flags[i] = _meta_bool(meta["finished"])
         for i, req_ids in enumerate(request_ids_list):
             runtime_info = runtime_infos[i] if i < len(runtime_infos) else None
             req_ids = _codec_ids_from_payload_or_input(req_ids, runtime_info)
@@ -284,6 +368,7 @@ class Qwen3TTSCode2Wav(nn.Module):
                 parsed.append((0, 0))
                 continue
             ctx_frames = left_context_size[i]
+            ref_ctx_frames = ref_context_size[i]
             flat = req_ids
             n = flat.numel()
             if n == 0 or n % q != 0:
@@ -302,6 +387,25 @@ class Qwen3TTSCode2Wav(nn.Module):
             frames = n // q
             # [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F])
             codes_qf = flat.reshape(q, frames)
+            ref_req_id = ref_context_request_ids[i]
+            if ref_req_id is not None and ref_ctx_frames > 0:
+                if ref_context_included[i]:
+                    if frames < ref_ctx_frames:
+                        raise ValueError(
+                            "Qwen3-TTS ref context metadata says ref prefix is included, "
+                            f"but frames={frames} < ref_context_size={ref_ctx_frames}"
+                        )
+                    self._cache_ref_context(ref_req_id, codes_qf[:, :ref_ctx_frames])
+                else:
+                    cached_ref = self._get_ref_context(ref_req_id)
+                    if cached_ref is None:
+                        raise ValueError(
+                            "Missing Qwen3-TTS ref context cache for "
+                            f"request {ref_req_id!r}; first chunk must include ref_code"
+                        )
+                    cached_ref = cached_ref.to(device=codes_qf.device, dtype=codes_qf.dtype)
+                    codes_qf = torch.cat((cached_ref, codes_qf), dim=1)
+                    frames = int(codes_qf.shape[1])
             parsed.append((ctx_frames, frames))
             valid_codes_qf.append(codes_qf)
             valid_indices.append(i)
@@ -443,6 +547,10 @@ class Qwen3TTSCode2Wav(nn.Module):
             if wav.shape[0] > 0:
                 audios[idx] = wav.to(dtype=torch.float32).reshape(-1)
 
+        for req_id, finished in zip(ref_context_request_ids, finished_flags, strict=False):
+            if req_id is not None and finished:
+                self._pop_ref_context(req_id)
+
         return OmniOutput(
             text_hidden_states=None,
             multimodal_outputs={"model_outputs": audios, "sr": srs},

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -23,7 +23,11 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.utils.speaker_cache import get_speaker_cache
+from vllm_omni.utils.speaker_cache import (
+    get_speaker_cache,
+    iter_custom_voice_profiles,
+    load_validated_profile_tensors,
+)
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
 from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder
@@ -441,6 +445,46 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             encode_ref_audio_batch=self._encode_ref_audio_batch,
             speaker_cache=self._speaker_cache,
         )
+        self._load_custom_voice_profiles()
+
+    # -------------------- custom voice profiles --------------------
+
+    def _load_custom_voice_profiles(self) -> None:
+        """Preload offline Qwen3-TTS custom voice profiles into speaker cache."""
+        custom_voice_dir = getattr(self.config, "custom_voice_dir", None)
+        if not custom_voice_dir:
+            return
+
+        expected_dim = int(getattr(self.config.speaker_encoder_config, "enc_dim", 0) or 0)
+        loaded = 0
+        for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type="qwen3_tts"):
+            tensors = load_validated_profile_tensors(
+                profile,
+                expected_model_type="qwen3_tts",
+                qwen3_embedding_dim=expected_dim,
+            )
+            if tensors is None:
+                continue
+
+            speaker_embedding = tensors["speaker_embedding"].reshape(-1).contiguous().cpu()
+            mode = str(profile.get("mode") or "xvec").lower()
+            ref_code = tensors.get("ref_code")
+            artifacts: dict[str, Any] = {
+                "ref_spk_embedding": speaker_embedding,
+                "ref_code": ref_code.contiguous().cpu() if isinstance(ref_code, torch.Tensor) else None,
+                "icl_mode": mode == "icl",
+                "ref_text": profile.get("ref_text"),
+            }
+            key = self._speaker_cache.make_cache_key(
+                profile["voice_name_lower"],
+                model_type=f"qwen3_tts_{mode}",
+                created_at=0,
+            )
+            self._speaker_cache.put(key, artifacts)
+            loaded += 1
+
+        if loaded:
+            logger.info("Loaded %d precomputed Qwen3-TTS custom voice profile(s) from %s", loaded, custom_voice_dir)
 
     # -------------------- vLLM required hooks --------------------
 

vllm_omni/model_executor/models/registry.py

@@ -2,6 +2,7 @@ from vllm.model_executor.models.registry import (
     _VLLM_MODELS,
     _LazyRegisteredModel,
     _ModelRegistry,
+    _resolve_module_name,
 )
 
 _OMNI_MODELS = {
@@ -261,7 +262,7 @@ OmniModelRegistry = _ModelRegistry(
     {
         **{
             model_arch: _LazyRegisteredModel(
-                module_name=f"vllm.model_executor.models.{mod_relname}",
+                module_name=_resolve_module_name(mod_relname),
                 class_name=cls_name,
             )
             for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()

vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py

@@ -35,7 +35,11 @@ from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.utils.speaker_cache import get_speaker_cache
+from vllm_omni.utils.speaker_cache import (
+    get_speaker_cache,
+    iter_custom_voice_profiles,
+    load_validated_profile_tensors,
+)
 
 from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
 from .voxcpm2_import_utils import import_voxcpm2_core
@@ -97,6 +101,7 @@ def build_voxcpm2_prompt(
     ref_audio: Any | None = None,
     ref_sr: int | None = None,
     ref_text: str | None = None,
+    voice_profile: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     """Build a VoxCPM2 prefill prompt whose ``prompt_token_ids`` length matches
     the talker-side prefill length.
@@ -110,7 +115,24 @@ def build_voxcpm2_prompt(
         ids = ids[1:]
     prefill_len = len(ids) + 1  # + audio_start
     additional: dict[str, Any] = {"text_token_ids": [ids]}
-    if ref_audio is not None:
+    if voice_profile is not None and ref_audio is None:
+        mode = str(voice_profile.get("mode") or "reference").lower()
+        ref_audio_feat_len = int(voice_profile.get("ref_audio_feat_len") or 0)
+        audio_feat_len = int(voice_profile.get("audio_feat_len") or 0)
+        prompt_text = voice_profile.get("prompt_text") or voice_profile.get("ref_text")
+        additional["voice_profile"] = dict(voice_profile)
+
+        if mode in ("reference", "ref_continuation"):
+            prefill_len += ref_audio_feat_len + 2  # ref_start / ref_end
+        if mode in ("continuation", "ref_continuation"):
+            if isinstance(prompt_text, str) and prompt_text:
+                additional["prompt_text"] = [prompt_text]
+                prompt_ids = split_multichar_chinese(tokenizer.encode(prompt_text, add_special_tokens=True), split_map)
+                if prompt_ids and prompt_ids[0] == bos:
+                    prompt_ids = prompt_ids[1:]
+                prefill_len += len(prompt_ids)
+            prefill_len += audio_feat_len
+    elif ref_audio is not None:
         vae = hf_config.audio_vae_config
         patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"])
         ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples)
@@ -450,6 +472,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
         # Speaker cache for ref_audio_feat across requests
         self._speaker_cache = get_speaker_cache()
+        self._load_custom_voice_profiles()
 
         self._perf = _PerfTimer(enabled=_ENABLE_PROFILING)
         self._cfm_buffers: _CFMBufferManager | None = None
@@ -472,6 +495,51 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         # one-shot by design: fires at most once per process to avoid log spam.
         self._active_state_warned = False
 
+    # -------------------- custom voice profiles --------------------
+
+    @staticmethod
+    def _clone_prompt_cache(cache: dict[str, Any]) -> dict[str, Any]:
+        cloned: dict[str, Any] = {}
+        for key, value in cache.items():
+            cloned[key] = value.clone() if isinstance(value, torch.Tensor) else value
+        return cloned
+
+    def _load_custom_voice_profiles(self) -> None:
+        """Preload offline VoxCPM2 prompt caches into the shared speaker cache."""
+        custom_voice_dir = getattr(self.config, "custom_voice_dir", None)
+        if not custom_voice_dir:
+            return
+
+        loaded = 0
+        for profile in iter_custom_voice_profiles(custom_voice_dir, expected_model_type="voxcpm2"):
+            tensors = load_validated_profile_tensors(profile, expected_model_type="voxcpm2")
+            if tensors is None:
+                continue
+
+            ref_audio_feat = tensors.get("ref_audio_feat")
+            audio_feat = tensors.get("audio_feat")
+            mode = str(profile.get("mode") or "").lower()
+
+            prompt_cache: dict[str, Any] = {"mode": mode}
+            if ref_audio_feat is not None:
+                prompt_cache["ref_audio_feat"] = ref_audio_feat.contiguous().cpu()
+            if audio_feat is not None:
+                prompt_cache["audio_feat"] = audio_feat.contiguous().cpu()
+            prompt_text = profile.get("prompt_text") or profile.get("ref_text")
+            if isinstance(prompt_text, str) and prompt_text:
+                prompt_cache["prompt_text"] = prompt_text
+
+            key = self._speaker_cache.make_cache_key(
+                profile["voice_name_lower"],
+                model_type="voxcpm2",
+                created_at=0,
+            )
+            self._speaker_cache.put(key, prompt_cache)
+            loaded += 1
+
+        if loaded:
+            logger.info("Loaded %d precomputed VoxCPM2 custom voice profile(s) from %s", loaded, custom_voice_dir)
+
     @property
     def tts(self) -> nn.Module:
         return self._tts
@@ -1161,6 +1229,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                 prompt_audio = prompt_audio[0] if prompt_audio else None
             if isinstance(prompt_text, list):
                 prompt_text = prompt_text[0] if prompt_text else None
+            voice_profile = info_dict.get("voice_profile")
+            if isinstance(voice_profile, list):
+                voice_profile = voice_profile[0] if voice_profile else None
+            requires_precomputed_cache = isinstance(voice_profile, dict)
 
             state.prompt_cache = None
             voice_name = info_dict.get("voice_name")
@@ -1168,20 +1240,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                 voice_name = voice_name[0] if voice_name else None
             _created_at = int(info_dict.get("voice_created_at") or 0)
 
-            if ref_audio or (prompt_audio and prompt_text):
-                # Check speaker cache for reference-only mode
-                if voice_name and ref_audio and not prompt_audio:
-                    _cache_key = self._speaker_cache.make_cache_key(
-                        voice_name, model_type="voxcpm2", created_at=_created_at
-                    )
-                    cached = self._speaker_cache.get(_cache_key)
-                    if cached is not None:
+            if voice_name:
+                _cache_key = self._speaker_cache.make_cache_key(
+                    voice_name, model_type="voxcpm2", created_at=_created_at
+                )
+                cached = self._speaker_cache.get(_cache_key)
+                if cached is not None:
+                    if "mode" in cached:
+                        state.prompt_cache = self._clone_prompt_cache(cached)
+                    elif "ref_audio_feat" in cached:
                         state.prompt_cache = {
                             "mode": "reference",
                             "ref_audio_feat": cached["ref_audio_feat"].clone(),
                         }
-                        logger.debug("Speaker cache HIT for VoxCPM2 speaker '%s'", voice_name)
+                    logger.debug("Speaker cache HIT for VoxCPM2 speaker '%s'", voice_name)
 
+            if state.prompt_cache is None and requires_precomputed_cache:
+                requested_mode = voice_profile.get("mode") if isinstance(voice_profile, dict) else None
+                raise ValueError(
+                    f"Precomputed VoxCPM2 voice '{voice_name}' was accepted by serving but is not loaded "
+                    f"in the model cache (voice_created_at={_created_at}, mode={requested_mode})"
+                )
+
+            if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
                 if state.prompt_cache is None:
                     try:
                         state.prompt_cache = self._build_prompt_cache(

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -159,6 +159,13 @@ class VoxtralTTSForConditionalGeneration(
             raise ValueError("Invalid model stage")
 
     # -------------------- CUDA Graph for acoustic transformer --------------------
+
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def _enable_acoustic_transformer_cudagraph(self):
         """Initialize and capture CUDA graphs for compute_mm_logits."""
         if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"):

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -248,12 +248,29 @@ def talker2code2wav_async_chunk(
     left_context_size = max(0, end_index - context_length)
     window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
 
-    # Prepend a bounded ref_code tail as decoder context for every chunk so the
-    # vocoder keeps voice-clone speaker identity without making Stage1 shapes
-    # depend on full reference-audio length. The decoder is causal with sliding
-    # attention, so frames older than this context window cannot affect the
-    # emitted chunk. Use `.get()` (not `.pop()`) to keep ref_code for later chunks.
+    # Prepend the bounded ref_code tail to the first emitted chunk so Code2Wav
+    # can cache the same reference context that mainline would otherwise send on
+    # every chunk. Follow-up chunks send only a metadata handle; Code2Wav restores
+    # this context locally before decode.
     ref_code = request_payload.get(request_id)
+    ref_context_size = 0
+    ref_context_request_id: str | None = None
+    ref_context_included = False
+    if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+        if ref_code.ndim == 1:
+            num_quantizers = len(window_frames[0])
+            if ref_code.numel() % num_quantizers == 0:
+                ref_code = ref_code.reshape(-1, num_quantizers)
+            else:
+                logger.warning(
+                    "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
+                    ref_code.numel(),
+                    num_quantizers,
+                )
+                ref_code = None
+        elif ref_code.ndim != 2:
+            logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
+            ref_code = None
     if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
         ref_context = ref_code
         if ref_code_context_frames > 0 and int(ref_context.shape[0]) > ref_code_context_frames:
@@ -263,9 +280,15 @@ def talker2code2wav_async_chunk(
                 int(ref_context.shape[0]),
             )
             ref_context = ref_context[-ref_code_context_frames:]
-        ref_frames = ref_context.tolist()
-        window_frames = ref_frames + window_frames
-        left_context_size += len(ref_frames)
+        ref_context_size = int(ref_context.shape[0]) if ref_context.ndim > 1 else 0
+        if ref_context_size > 0:
+            ref_context_request_id = request_id
+            emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0))
+            if emitted_chunks <= 0:
+                ref_frames = ref_context.tolist()
+                window_frames = ref_frames + window_frames
+                ref_context_included = True
+            left_context_size += ref_context_size
 
     num_quantizers = len(window_frames[0])
     num_frames = len(window_frames)
@@ -274,12 +297,18 @@ def talker2code2wav_async_chunk(
         dtype=torch.long,
     )
 
+    meta = MetaStruct(
+        left_context_size=left_context_size,
+        finished=torch.tensor(finished, dtype=torch.bool),
+    )
+    if ref_context_size > 0 and ref_context_request_id is not None:
+        meta.ref_context_size = ref_context_size
+        meta.ref_context_request_id = ref_context_request_id
+        meta.ref_context_included = ref_context_included
+
     return OmniPayloadStruct(
         codes=CodesStruct(audio=code_predictor_codes),
-        meta=MetaStruct(
-            left_context_size=left_context_size,
-            finished=torch.tensor(finished, dtype=torch.bool),
-        ),
+        meta=meta,
         speaker=extract_speaker_from_request(request),
         language=extract_language_from_request(request),
     )

vllm_omni/patch.py

@@ -128,3 +128,25 @@ for module_name, module in list(sys.modules.items()):
         module.StreamingUpdate = OmniStreamingUpdate
     if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
         module.EngineCoreRequest = OmniEngineCoreRequest
+
+
+# Patch: add qwen3_omni_moe to vllm's chat template fallback registry.
+# Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone
+# chat_template.json (not in tokenizer_config.json).  transformers < 5.9.0
+# does not load this file, so the tokenizer has no chat_template attribute.
+# vllm's resolve_chat_template falls back to MODEL_TYPE_TO_CHAT_TEMPLATE
+# which has "qwen" but not "qwen3_omni_moe".  Register the same fallback.
+def _patch_chat_template_registry():
+    try:
+        from vllm.transformers_utils.chat_templates.registry import (
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK,
+            _get_qwen_chat_template_fallback,
+        )
+
+        if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = _get_qwen_chat_template_fallback
+    except ImportError:
+        pass
+
+
+_patch_chat_template_registry()

vllm_omni/platforms/cuda/platform.py

@@ -229,4 +229,4 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         if envs.VLLM_USE_OINK_OPS:
             rms_norm = ["oink"] + default
 
-        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
+        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)

vllm_omni/platforms/interface.py

@@ -11,6 +11,7 @@ from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import init_logger
 from vllm.platforms import Platform
+from vllm.platforms.interface import PlatformEnum
 
 logger = init_logger(__name__)
 
@@ -239,6 +240,7 @@ class OmniPlatform(Platform):
 
 class UnspecifiedOmniPlatform(OmniPlatform):
     _omni_enum = OmniPlatformEnum.UNSPECIFIED
+    _enum = PlatformEnum.UNSPECIFIED
     device_type = ""
 
     @classmethod

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -16,11 +16,6 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    extract_routed_experts_for_current_batch,
-    get_global_experts_capturer,
-    issue_routing_d2h_copy,
-)
 from vllm.sequence import IntermediateTensors
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import (
@@ -257,7 +252,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | IntermediateTensors | None:
         if self.vllm_config.model_config.enable_return_routed_experts:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                 capturer.finalize_pending_copy()
         if self.ascend_config.profiling_chunk_config.enabled:
@@ -728,12 +723,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             deferred_state_corrections_fn()
 
         if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"):
-            issue_routing_d2h_copy(
-                input_batch_req_ids=self.input_batch.req_ids,
-                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                positions=self.positions,
-                positions_cpu=self._positions_cpu,
-            )
+            self._omni_routed_experts_d2h(scheduler_output)
 
         return None
 
@@ -911,17 +901,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
             if self.speculative_config is not None:
                 self.finalize_kv_connector()
 
-        routed_experts_dict = None
+        routed_experts_lists = None
         if self.model_config.enable_return_routed_experts:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"):
-                routed_experts_dict = extract_routed_experts_for_current_batch(
-                    req_ids=req_ids_output_copy,
-                    requests=self.requests,
-                    req_id_to_index=self.input_batch.req_id_to_index,
-                    num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
-                    max_model_len=self.max_model_len,
-                )
+                routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
 
         #  -------------------------------------- Omni-new -------------------------------------------------
         engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
@@ -1042,9 +1026,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
             kv_connector_output=kv_connector_output,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
             cudagraph_stats=cudagraph_stats,
-            routed_experts_dict=routed_experts_dict,
         )
         model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
+        model_runner_output.routed_experts = routed_experts_lists
         #  -------------------------------------- Omni-new -------------------------------------------------
 
         if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"):