Fix LingBot diffusion alignment CI

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

.buildkite/test-merge.yml

@@ -108,10 +108,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        # The suite relaunches an OmniServer per test class (~30 model loads,
-        # incl. ten ~27.5 GiB Bagel loads); 60 min passed only ~1 in 3 runs
-        # even on main, and vLLM 0.25 weight loading needs ~10 min more.
-        timeout_in_minutes: 90
+        timeout_in_minutes: 60
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         agents:
@@ -149,9 +146,7 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Entrypoint Test with L4"
-        # 30 min left no headroom for a cold image pull (~11 min for new base
-        # layers) plus per-class engine relaunches; main passed at 26/30.
-        timeout_in_minutes: 45
+        timeout_in_minutes: 30
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
         agents:

.gitignore

@@ -275,8 +275,3 @@ tests/.ci_generated/
 
 # Local Lance precision-matching dump scripts (not part of the repo)
 lance_compare/
-
-# FP8 quantization test sample images: tests/diffusion/quantization/
-# test_quantization_fp8.py saves these into the repo root for manual
-# inspection; they are byproducts and must never be committed.
-test_*_fp8.png

benchmarks/diffusion/lingbot_video_parity.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import statistics
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import imageio.v3 as iio
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="LingBot dense parity and timing harness.")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
+    parser.add_argument("--official-repo", default="/tmp/lingbot-video")
+    parser.add_argument("--output-dir", default="/tmp/lingbot_dense_parity")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--runs", type=int, default=0, help="Optional in-process steady-state native runs.")
+    return parser.parse_args()
+
+
+def _run(cmd: list[str], cwd: Path | None = None) -> tuple[float, str]:
+    start = time.perf_counter()
+    proc = subprocess.run(cmd, cwd=cwd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    elapsed = time.perf_counter() - start
+    if proc.returncode != 0:
+        print(proc.stdout)
+        raise subprocess.CalledProcessError(proc.returncode, proc.args, output=proc.stdout)
+    return elapsed, proc.stdout
+
+
+def _read_video(path: Path) -> np.ndarray:
+    frames = iio.imread(path, plugin="pyav")
+    if frames.dtype != np.float32:
+        frames = frames.astype(np.float32) / 255.0
+    return frames
+
+
+def _compare(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
+    if a.shape != b.shape:
+        return {"shape_a": list(a.shape), "shape_b": list(b.shape), "shape_match": False}
+    diff = a - b
+    mse = float(np.mean(diff * diff))
+    mae = float(np.mean(np.abs(diff)))
+    psnr = float("inf") if mse == 0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
+    return {
+        "shape": list(a.shape),
+        "shape_match": True,
+        "mae": mae,
+        "mse": mse,
+        "psnr": psnr,
+    }
+
+
+def _extract_total_generation_seconds(log: str) -> float | None:
+    match = re.search(r"Total generation time:\s*([0-9.]+)\s*seconds", log)
+    return float(match.group(1)) if match else None
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _benchmark_native_requests(args: argparse.Namespace) -> dict[str, Any]:
+    load_start = time.perf_counter()
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+    load_seconds = time.perf_counter() - load_start
+    prompt = {"prompt": args.prompt}
+    latencies: list[float] = []
+    peak_mb = 0.0
+
+    for _ in range(args.runs):
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+        sampling_params = OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            num_frames=args.num_frames,
+            num_inference_steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            generator=generator,
+            output_type="pt",
+            extra_args={"shift": args.shift},
+        )
+        start = time.perf_counter()
+        output = omni.generate(prompt, sampling_params)
+        latencies.append(time.perf_counter() - start)
+        peak_mb = max(peak_mb, _peak_memory_mb(output))
+        del output
+    omni.close()
+
+    return {
+        "load_seconds": load_seconds,
+        "run_seconds": latencies,
+        "mean_seconds": statistics.mean(latencies) if latencies else None,
+        "median_seconds": statistics.median(latencies) if latencies else None,
+        "min_seconds": min(latencies) if latencies else None,
+        "max_seconds": max(latencies) if latencies else None,
+        "peak_mb": peak_mb,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    official_out = output_dir / "official_diffusers.mp4"
+    native_out = output_dir / "vllm_omni_native.mp4"
+
+    official_cmd = [
+        sys.executable,
+        str(Path(args.official_repo) / "scripts" / "inference.py"),
+        "--backend",
+        "diffusers",
+        "--model_dir",
+        args.model,
+        "--mode",
+        "t2v",
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(official_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num_frames",
+        str(args.num_frames),
+        "--steps",
+        str(args.steps),
+        "--guidance_scale",
+        str(args.guidance_scale),
+        "--shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+        "--quiet_progress",
+    ]
+    native_cmd = [
+        sys.executable,
+        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
+        "--model",
+        args.model,
+        "--prompt",
+        args.prompt,
+        "--output",
+        str(native_out),
+        "--height",
+        str(args.height),
+        "--width",
+        str(args.width),
+        "--num-frames",
+        str(args.num_frames),
+        "--num-inference-steps",
+        str(args.steps),
+        "--guidance-scale",
+        str(args.guidance_scale),
+        "--flow-shift",
+        str(args.shift),
+        "--seed",
+        str(args.seed),
+        "--fps",
+        str(args.fps),
+    ]
+
+    official_time, official_log = _run(official_cmd, cwd=Path(args.official_repo))
+    native_time, native_log = _run(native_cmd, cwd=Path.cwd())
+
+    result = {
+        "official_seconds": official_time,
+        "native_process_seconds": native_time,
+        "native_request_seconds": _extract_total_generation_seconds(native_log),
+        "official_output": str(official_out),
+        "native_output": str(native_out),
+        "accuracy": _compare(_read_video(official_out), _read_video(native_out)),
+        "official_log_tail": official_log.splitlines()[-20:],
+        "native_log_tail": native_log.splitlines()[-20:],
+    }
+    if args.runs > 0:
+        result["native_steady_state"] = _benchmark_native_requests(args)
+
+    result_path = output_dir / "comparison.json"
+    result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
+    print(json.dumps(result, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

docker/Dockerfile.ci

@@ -14,71 +14,8 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-
-# Install vLLM (CUDA 13.0 build) from precompiled wheel at the selected commit.
-# Must use direct URL because the wheel has a PEP 440 local version identifier
-# (e.g. +gaeee7ef93.cu130) which pip/uv refuse to install from a PEP 503 package index.
-# The cu130 wheel index lives at /<commit>/cu130/vllm/ but the actual wheel files
-# are stored at the top level /<commit>/<wheel>.
-ENV VLLM_PRECOMPILED_WHEEL_COMMIT=dd10e03f95f94edbea1975c67ace3a35ec9a8a40
-# Keep the base image's CUDA-12 toolkit libs on the loader path. NOTE: this does
-# NOT resolve libnvJitLink.so.13 (that CUDA-13 lib is not in /usr/local/cuda here) —
-# the .so.13 JIT linker is handled separately via ldconfig after the cu130 wheels
-# are installed (see the nvidia-cuda-nvjitlink-cu13 step below).
-ENV LD_LIBRARY_PATH=/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
-RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
-    from urllib.parse import urljoin, quote; \
-    index_url='https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/cu130/vllm/'; \
-    html=urllib.request.urlopen(index_url).read().decode(); \
-    m=re.search(r'href=\"([^\"]*x86_64\.whl)\"',html); \
-    assert m, f'No x86_64 vLLM wheel found at {index_url}'; \
-    print(quote(urljoin(index_url,m.group(1)), safe=':/+%'))") && \
-    echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
-    uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
-
-
-
 RUN uv pip install --system ".[dev]"
 
-RUN uv pip install --system --upgrade \
-        "flashinfer-python==0.6.13" \
-        "flashinfer-cubin==0.6.13"
-
-RUN uv pip install --system --upgrade \
-    "flashinfer-jit-cache==0.6.13" \
-    --index-url https://flashinfer.ai/whl/cu130
-
-# Force-reinstall torch AFTER flashinfer, because flashinfer --upgrade pulls in
-# a newer torch (2.12.x) whose libtorch_cuda.so lacks
-# at::cuda::getCurrentCUDABlasHandle, causing "undefined symbol" at import time.
-RUN uv pip install --system --force-reinstall \
-    "torch==2.11.0" "torchaudio==2.11.0" "torchvision==0.26.0" \
-    --torch-backend cu130
-
-# CUDA 13 JIT linker: libnvJitLink.so.13 is provided by the nvidia-nvjitlink
-# package that the cu130 torch/vllm wheels above pull into site-packages (the
-# v0.24.0 base image's /usr/local/cuda toolkit is CUDA 12, so the LD_LIBRARY_PATH
-# above never resolves the .so.13). Its lib dir is not on the dynamic-loader path,
-# so bitsandbytes / torch.compile worker subprocesses fail to dlopen it. Locate the
-# already-installed lib and register its dir with ldconfig (system-wide, survives
-# subprocess env resets). Only fall back to the pip wheel (correct name:
-# nvidia-nvjitlink-cu13, NOT nvidia-cuda-nvjitlink-cu13, which does not exist) if
-# the torch deps did not bring it in.
-RUN NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)" && \
-    if [ -z "$NVJITLINK_LIB" ]; then \
-        uv pip install --system "nvidia-nvjitlink-cu13" && \
-        NVJITLINK_LIB="$(find / -name 'libnvJitLink.so.13*' 2>/dev/null | head -1)"; \
-    fi && \
-    test -n "$NVJITLINK_LIB" && \
-    dirname "$NVJITLINK_LIB" > /etc/ld.so.conf.d/nvjitlink-cu13.conf && \
-    ldconfig && \
-    echo "Registered ${NVJITLINK_LIB} with ldconfig"
-
-# Re-pin numpy LAST. numba requires numpy <= 2.4, but the flashinfer and torch cu130
-# reinstalls above resolve numpy back up to 2.5, which breaks `import numba` at test
-# time ("Numba needs NumPy 2.4 or less"). This must be the final dependency op so
-# nothing can clobber it again.
-RUN uv pip install --system "numpy==2.2.6"
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docs/configuration/composable_parallel.md

@@ -1,239 +0,0 @@
-# Composable parallel strategies
-
-In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the existing stage configs. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy.
-
-Strategies are intentionally narrow: each one only writes the axes it declares, leaving every other engine arg untouched. This makes a strategy file a small, reusable overlay you can swap independently of the deploy YAML (see `configuration/stage_configs.md`).
-
-!!! note
-    Composable parallel is **opt-in**. Passing `--strategy-config <path/to/strategy.yaml>` overlays the derived sizing onto the registry-merged stages *before* any CLI overrides are applied. When `--strategy-config` is omitted, stages remain exactly as their deploy YAML produces them.
-
-## Top-level schema reference
-
-A strategy file has a single top-level mapping, `strategies`, keyed by the `model_stage` name of each stage you want to overlay. Each value is a list of axis declarations (see [StrategySpec fields](#strategyspec-fields) below).
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `strategies` | `dict[str, list]` | required | — | Maps `model_stage` name (e.g. `thinker`, `talker`, `code2wav`, `dit`) to a list of axis declarations. The key MUST be the stage name, not the integer `stage_id`. See `vllm_omni/config/composable_parallel/strategy_loader.py`. |
-
-### StrategySpec fields
-
-Each entry in a stage's list becomes one `StrategySpec` (see `vllm_omni/config/composable_parallel/spec.py`). The loader fills in `routing` and `aggregation` from per-kind defaults, so the file stays compact.
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `axis` | str | required | — | Mesh-axis kind. One of `tp`, `dp`, `pp`, `ep`, `stage_replica` (wired today) or `sp_ulysses`, `sp_ring`, `cfg`, `vae_pp`, `hsdp`, `stage_pp`, `cp` (reserved; declaring one raises `AxisTranslationError` at translate time). See `MeshAxisKind` in `vllm_omni/config/composable_parallel/spec.py`. |
-| `size` | int | required | — | Axis degree. Must be `> 0`. |
-| `routing` | str | optional | per-kind default | Stateless routing policy. Accepted only for `dp` / `stage_replica`; other kinds reject it. Values: `random`, `round_robin`, `least_queue`. (`hash` is recognised but currently raises `NotImplementedError`.) See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. |
-| `l1_owner` | str | optional | per-kind default | Overrides which layer owns request routing for this axis. `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. Recorded into `shard_extension`. |
-| `name` | str | optional | axis kind | Human-readable label for this spec, used in error messages. |
-
-Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs` — let advanced authors bind layer-walker hook IDs (`LayerHookSpec`) or kernel IDs (`KernelSpec`) onto an axis. These are not exposed by the YAML loader today and are intended for programmatic construction. See `vllm_omni/config/composable_parallel/spec.py` for the dataclasses.
-
-### CLI flags
-
-| Flag | Description |
-|------|-------------|
-| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. |
-| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). |
-| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. |
-
-!!! warning
-    `--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy.
-
-## Sub-schemas
-
-### `axis`
-
-The set of recognised axis kinds is the `MeshAxisKind` enum in `vllm_omni/config/composable_parallel/spec.py`:
-
-```yaml
-# Translatable today (see `_SUPPORTED_KINDS` in translator.py).
-axis: tp              # tensor parallel — sets tensor_parallel_size
-axis: dp              # engine data parallel — sets data_parallel_size (vLLM's intra-engine DP LB)
-axis: pp              # pipeline parallel — sets pipeline_parallel_size
-axis: ep              # dense expert parallel — sets enable_expert_parallel=True; size must equal tp*dp
-axis: stage_replica   # omni-coordinator-level replica fan-out — sets num_replicas + omni_lb_policy
-
-# Reserved (raise AxisTranslationError today; declared in the type system for forward compat).
-axis: sp_ulysses
-axis: sp_ring
-axis: cfg
-axis: vae_pp
-axis: hsdp
-axis: stage_pp
-axis: cp
-```
-
-### `routing`
-
-The full closed hierarchy lives in `vllm_omni/config/composable_parallel/routing.py` (`Broadcast`, `PartitionByHash`, `ShardSequence`, `ShardSpatial`, `RouteByExpert`, `PipelineMicrobatch`, `DuplicateWithCondTag`, `RouteByStage`). For YAML authors only `RouteByStage` is exposed as a `routing:` string, and only for `dp` / `stage_replica` axes (other axes have fixed routing — `tp`/`ep` broadcast, `pp` microbatch). Recognised policies:
-
-```yaml
-routing: random         # uniform-random across replicas/DP ranks
-routing: round_robin    # cycle through replicas/DP ranks (the default for dp/stage_replica)
-routing: least_queue    # send to the replica with the smallest backlog
-# routing: hash         # recognised but raises NotImplementedError; no key-stable balancer yet
-```
-
-### `aggregation`
-
-Aggregation is filled in automatically by the loader (`_default_aggregation` in `vllm_omni/config/composable_parallel/strategy_loader.py`) from the axis kind, so users do not author it:
-
-| Axis kind | Aggregation pattern |
-|-----------|---------------------|
-| `tp`, `ep` | `TakeRank` (all ranks agree; pick one) |
-| `dp`, ep (in some flows) | `Union` (disjoint replica results) |
-| `pp` | `StitchPipeline` |
-| `stage_replica` | `FanInByStage` |
-
-Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSpatial`, `StitchPipeline`, `Combine`, `FanInByStage` (see `vllm_omni/config/composable_parallel/aggregation.py`).
-
-## Precedence
-
-From highest to lowest, when the same axis is touched by multiple layers:
-
-1. **CLI overrides** (`--stage_N_tensor_parallel_size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
-2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`.
-3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/<model_type>.yaml`).
-4. **Parser defaults.**
-
-After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `--devices` or `--tensor-parallel-size` that conflicts with the strategy's world size cannot slip past the pre-spawn check.
-
-## Worked example
-
-### Example 1 — `strategy_tp2.yaml`
-
-Shard the `thinker` stage with TP=2 on top of the bundled Qwen2.5-Omni deploy YAML. Routing is fixed (broadcast) for TP, so no `routing:` field is needed.
-
-```yaml
-# examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml
-strategies:
-  thinker:
-    - axis: tp
-      size: 2
-```
-
-Launch:
-
-```bash
-python3 end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_tp2.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```
-
-Effective `OmniParallelConfig` per stage after the overlay (see `OmniParallelConfig` in `vllm_omni/config/composable_parallel/translator.py`):
-
-| Stage (`model_stage`) | `tensor_parallel_size` | `data_parallel_size` | `pipeline_parallel_size` | `stage_replica_size` | `omni_lb_policy` |
-|-----------------------|------------------------|----------------------|--------------------------|----------------------|------------------|
-| `thinker`             | **2** (from strategy)  | 1                    | 1                        | 1                    | —                |
-| `talker`              | 1                      | 1                    | 1                        | 1                    | —                |
-| `code2wav`            | 1                      | 1                    | 1                        | 1                    | —                |
-
-### Example 2 — `strategy_stage_replica.yaml`
-
-Replicate the `thinker` stage across two independent engines and let omni's StagePool round-robin requests across them.
-
-```yaml
-# examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml
-strategies:
-  thinker:
-    - axis: stage_replica
-      size: 2
-      routing: round_robin
-```
-
-Launch:
-
-```bash
-python3 end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_stage_replica.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```
-
-Effective per-stage config and the derived orchestrator knob:
-
-| Stage (`model_stage`) | `tensor_parallel_size` | `stage_replica_size` | `omni_lb_policy` (derived) |
-|-----------------------|------------------------|----------------------|----------------------------|
-| `thinker`             | 1                      | **2**                | **`round-robin`**          |
-| `talker`              | 1                      | 1                    | —                          |
-| `code2wav`            | 1                      | 1                    | —                          |
-
-The derived `omni_lb_policy = round-robin` is **not** written into any stage's engine args. Instead it is surfaced on `StrategyApplyResult.omni_lb_policy` and applied once at orchestrator construction by `AsyncOmniEngine._apply_strategy_lb_policy` (see `vllm_omni/engine/async_omni_engine.py`, lines 1107–1134). If you also pass `--omni-lb-policy round-robin` the values must match; passing a different value raises `Conflicting load-balancer policy`.
-
-## Per-field deep-dive
-
-### `name`
-
-Default: the axis kind (e.g. `"tp"`).
-
-Human-readable label that appears in `AxisTranslationError` and `StrategyApplyError` messages. Set it to disambiguate two stages declaring the same axis kind when scanning logs.
-
-### `mesh_axis` (`axis` + `size`)
-
-Default: no default — both fields are required.
-
-`axis` chooses the parallelism dimension; `size` is its degree. Five kinds are wired end-to-end today (`tp`, `dp`, `pp`, `ep`, `stage_replica`); the remaining kinds in `MeshAxisKind` are reserved and raise `AxisTranslationError` at translate time so a typo cannot silently become a no-op (see `_SUPPORTED_KINDS` in `vllm_omni/config/composable_parallel/translator.py`). For `ep`, the size MUST equal `tensor_parallel_size * data_parallel_size`.
-
-### `routing`
-
-Default: per-kind — `RouteByStage(round_robin)` for `dp` and `stage_replica`; `PipelineMicrobatch()` for `pp`; `Broadcast()` for `tp`/`ep` and any other kind.
-
-User-facing string only on `dp` and `stage_replica` axes. The four recognised values are `random`, `round_robin`, `least_queue`, `hash`; `hash` is rejected today with `NotImplementedError`. See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`.
-
-### `l1_owner`
-
-Default: per-kind — `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`.
-
-Overrides which layer owns request routing along this axis. `engine` means vLLM realizes routing intra-engine; `delegated` means the omni coordinator owns routing across independent replicas. Set this only when overriding the per-kind default in `_DEFAULT_L1_OWNER` (`vllm_omni/config/composable_parallel/translator.py`); invalid values fail fast in `_resolve_l1_owner`.
-
-### `stage_replica` and `omni_lb_policy`
-
-`stage_replica` is *not* a vLLM world dimension; it spins up `size` independent engine replicas of one stage, coordinated by the omni `StagePool` load balancer. The axis writes both:
-
-- The stage's `num_replicas` (a per-stage deploy knob).
-- The pipeline-wide `omni_lb_policy` string (an orchestrator-level knob).
-
-Because `omni_lb_policy` is read once at `AsyncOmniEngine` construction — not from per-stage engine args — only one value is allowed across the whole pipeline. If two stages declare conflicting policies, application raises (see `apply_strategy_specs` in `vllm_omni/config/composable_parallel/apply.py`, lines 253–260). To set this knob manually instead, use `--omni-lb-policy` on the CLI; see [the model stage docs](stage_configs.md) for how stage replicas plug into the orchestrator at runtime.
-
-### `model_stage` keys
-
-Default: no default — required for every entry under `strategies:`.
-
-Strategy keys MUST be `model_stage` *names* (strings), the same labels stage configs use for `engine_args.model_stage` (see [`model_stage` in `stage_configs.md`](stage_configs.md#engine_argsmodel_stage)). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`).
-
-## Common pitfalls
-
-!!! warning "Using `stage_id` integers as keys"
-    Strategy keys are `model_stage` names (str), not integer stage IDs. A file that uses an integer index will raise at apply time:
-
-    ```text
-    StrategyApplyError: strategy role model_stage=0 did not match any stage;
-    available stages: 'thinker'(id=0), 'talker'(id=1), 'code2wav'(id=2)
-    ```
-
-    Use the `model_stage` name (e.g. `thinker`) instead.
-
-!!! warning "Mixing CLI overrides with a strategy on the same axis"
-    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--stage_0_tensor_parallel_size`, `--stage-overrides '{"0": {...}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
-
-!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
-    Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
-
-    ```text
-    --strategy-config (<path>) was provided but model '<model>' resolves via the
-    legacy stage_configs YAML path, which does not support composable-parallel
-    strategies; the strategy is ignored. Use a registry-based model to apply it.
-    ```
-
-    Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy.
-
-!!! warning "Conflicting `omni_lb_policy` across stages"
-    `omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises:
-
-    ```text
-    StrategyApplyError: role 'talker' derives omni_lb_policy='least-queue-length'
-    but role 'thinker' already derived 'round-robin'. The load-balancer policy is
-    pipeline-wide; only one value is allowed.
-    ```
-
-    Make the policies match, drop one, or supply the value explicitly via `--omni-lb-policy`.

docs/configuration/pd_disaggregation.md

@@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
     runtime:
       devices: "2"
     engine_input_source: [1]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
 
   - stage_id: 3
     runtime:
@@ -131,7 +131,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
     engine_args:
       max_num_seqs: 1
     engine_input_source: [2]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
 ```
 
 Compared with the default Qwen3-Omni config:

docs/contributing/model/adding_omni_model.md

@@ -475,64 +475,92 @@ thinker_hidden_states = output.multimodal_output["24"]
 
 ### Implementation Example
 
-Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function:
-
-| Suffix | Role | Runs when |
-|--------|------|-----------|
-| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector |
-| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads |
-| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only |
+Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`:
 
 ```python
-# qwen3_omni.py (Thinker → Talker, non-async path)
+# qwen3_omni.py
 
-def thinker2talker_token_only(
-    source_outputs: list[Any],
+def thinker2talker(
+    stage_list: list[Any],
+    engine_input_source: list[int],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """Allocate talker prefill slots; bulk tensors arrive via the connector."""
-    ...
+    """
+    Process thinker outputs to create talker inputs.
+
+    Args:
+        stage_list: List of stage objects
+        engine_input_source: Source stage IDs (typically [0] for thinker)
+        prompt: Original prompt data
+
+    Returns:
+        List of OmniTokensPrompt for talker stage
+    """
+    source_stage_id = engine_input_source[0]
+    thinker_outputs = stage_list[source_stage_id].engine_outputs
+    talker_inputs = []
+
+    for thinker_output in thinker_outputs:
+        output = thinker_output.outputs[0]
+        # Extract thinker embeddings and hidden states
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
+
+        info = {
+            "thinker_prefill_embeddings": thinker_prefill_embeddings,
+            "thinker_hidden_states": thinker_hidden_states,
+            "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
+            "thinker_input_ids": thinker_output.prompt_token_ids,
+        }
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * computed_length,
+                additional_information=info,
+                multi_modal_data=None,
+            )
+        )
+
+    return talker_inputs
 
 
-def thinker2talker_full_payload(
-    transfer_manager: Any,
-    pooling_output: dict[str, Any],
-    request: OmniEngineCoreRequest,
-) -> dict[str, Any] | None:
-    """Pack accumulated thinker hidden states into OmniPayload for stage-1."""
-    ...
+def talker2code2wav(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """
+    Process talker outputs to create code2wav inputs.
+    """
+    source_stage_id = engine_input_source[0]
+    talker_outputs = stage_list[source_stage_id].engine_outputs
+    code2wav_inputs = []
 
+    for talker_output in talker_outputs:
+        output = talker_output.outputs[0]
+        # Extract codec codes
+        codec_codes = (
+            output.multimodal_output["code_predictor_codes"]
+            .to(torch.long)
+            .transpose(0, 1)
+            .cpu()
+            .to(torch.long)
+            .reshape(-1)
+            .tolist()
+        )
 
-def thinker2talker_async_chunk(
-    transfer_manager: Any,
-    multimodal_output: OmniPayload | dict[str, Any],
-    request: OmniEngineCoreRequest,
-    is_finished: bool = False,
-) -> OmniPayloadStruct | None:
-    """Stream thinker rows to talker while async_chunk is enabled."""
-    ...
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+            )
+        )
+
+    return code2wav_inputs
 ```
 
-Wire these in `pipeline.py`:
-
-```python
-StagePipelineConfig(
-    stage_id=0,
-    custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
-    async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
-    ...
-),
-StagePipelineConfig(
-    stage_id=1,
-    sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
-    ...
-),
-```
-
-Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract.
-
 ## Testing
 
 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).

docs/models/supported_models.md

@@ -46,6 +46,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
+| `LingBotVideoPipeline` | LingBot-Video dense T2V | `robbyant/lingbot-video-dense-1.3b` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

examples/offline_inference/qwen2_5_omni/README.md

@@ -60,24 +60,3 @@ If media file paths are not provided, the script will use default assets. Suppor
 - `mixed_modalities`: Audio + image + video
 - `use_audio_in_video`: Extract audio from video
 - `text`: Text-only query
-
-### Composable parallelism (strategy configs)
-
-You can shard or replicate a stage with a small composable-parallel
-`strategy.yaml` overlaid onto the **bundled default deploy config** -- no bespoke
-deploy YAML required. Supply the strategy with `--strategy-config` and any matching
-device layout via `--stage-overrides` (explicit kwargs forwarded to `Omni`).
-
-Both examples below require >= 3 GPUs (the thinker uses 2, talker + code2wav share a third):
-
-```bash
-# Tensor-parallel the thinker (TP=2)
-python end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_tp2.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-
-# Replicate the thinker across 2 engines (round-robin load balancing)
-python end2end.py --query-type text --num-prompts 6 \
-    --strategy-config strategy_stage_replica.yaml \
-    --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-```

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -553,39 +553,6 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help="Optional explicit deploy YAML (otherwise the bundled default is used).",
-    )
-    parser.add_argument(
-        "--strategy-config",
-        type=str,
-        default=None,
-        help=(
-            "Optional composable-parallel strategy.yaml mapping role -> parallel axes "
-            "(e.g. tp / stage_replica). Overlaid onto the merged stage configs."
-        ),
-    )
-    parser.add_argument(
-        "--stage-overrides",
-        type=str,
-        default=None,
-        help=(
-            "Optional JSON of per-stage overrides applied on top of the default "
-            'deploy config, e.g. \'{"0": {"devices": "0,1"}}\' to give the '
-            "thinker a 2-GPU pool. Lets you run a strategy on the bundled default "
-            "deploy config without writing a bespoke deploy YAML."
-        ),
-    )
-    parser.add_argument(
-        "--omni-lb-policy",
-        type=str,
-        default=None,
-        choices=["random", "round-robin", "least-queue-length"],
-        help="StagePool load-balancer policy for replicated stages (orchestrator-level knob).",
-    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml

@@ -1,22 +0,0 @@
-# Composable-parallel strategy for Qwen2.5-Omni.
-#
-# Replicate the thinker stage across 2 independent engines and let the omni
-# StagePool balance requests round-robin across them. The thinker's `devices`
-# entry acts as a per-replica template, so omni expands it to one GPU per replica.
-# The `routing: round_robin` below auto-derives the StagePool load-balancer policy,
-# so --omni-lb-policy is not needed.
-#
-# Rather than hand-write a bespoke deploy YAML, run on the bundled default deploy
-# config and pass the matching 2-GPU device template via --stage-overrides (explicit
-# kwargs): the thinker gets a two-GPU pool while talker + code2wav share a third
-# card. Requires >= 3 GPUs.
-#
-# Usage:
-#   python3 end2end.py --query-type text --num-prompts 6 \
-#       --strategy-config strategy_stage_replica.yaml \
-#       --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-strategies:
-  thinker:
-    - axis: stage_replica
-      size: 2
-      routing: round_robin

examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml

@@ -1,19 +0,0 @@
-# Composable-parallel strategy for Qwen2.5-Omni: tensor-parallel the thinker.
-#
-# Shard the thinker stage across 2 GPUs with tensor parallelism (TP=2). TP routing
-# is fixed (broadcast), so no `routing` / load-balancer policy is involved -- this
-# is the simplest axis and a good starting point for the mental model.
-#
-# The strategy sets tensor_parallel_size=2 on the thinker. Rather than hand-write a
-# bespoke deploy YAML, run on the bundled default deploy config and pass a matching
-# 2-GPU device layout via --stage-overrides (explicit kwargs), so the thinker gets a
-# two-GPU pool while talker + code2wav share a third card. Requires >= 3 GPUs.
-#
-# Usage:
-#   python3 end2end.py --query-type text --num-prompts 6 \
-#       --strategy-config strategy_tp2.yaml \
-#       --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}'
-strategies:
-  thinker:
-    - axis: tp
-      size: 2

examples/offline_inference/text_to_video/text_to_video_lingbot.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text-to-video with the LingBot dense preset.")
+    parser.add_argument("--model", default="/home/models/lingbot-video-dense-1.3b")
+    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--output", default="lingbot_video_output.mp4")
+    parser.add_argument("--height", type=int, default=192)
+    parser.add_argument("--width", type=int, default=320)
+    parser.add_argument("--num-frames", type=int, default=9)
+    parser.add_argument("--num-inference-steps", type=int, default=2)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--flow-shift", type=float, default=3.0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
+    parser.add_argument("--batch-cfg", action="store_true")
+    return parser.parse_args()
+
+
+def _peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _extract_video(output: Any) -> Any:
+    if isinstance(output, list):
+        output = output[0] if output else None
+    if isinstance(output, OmniRequestOutput):
+        if output.final_output_type != "image":
+            raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
+        if not output.images:
+            raise ValueError("No video frames found in OmniRequestOutput.")
+        return output.images[0] if len(output.images) == 1 else output.images
+    return output
+
+
+def _to_frame_array(video: Any) -> np.ndarray:
+    if isinstance(video, torch.Tensor):
+        tensor = video.detach().cpu()
+        if tensor.dim() == 5:
+            tensor = tensor[0] if tensor.shape[1] not in (3, 4) else tensor[0].permute(1, 2, 3, 0)
+        elif tensor.dim() == 4 and tensor.shape[0] in (3, 4):
+            tensor = tensor.permute(1, 2, 3, 0)
+        if tensor.is_floating_point():
+            if float(tensor.min()) < 0.0:
+                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
+            else:
+                tensor = tensor.clamp(0, 1)
+        return tensor.float().numpy()
+    if isinstance(video, np.ndarray):
+        array = video[0] if video.ndim == 5 else video
+        if np.issubdtype(array.dtype, np.integer):
+            return array.astype(np.float32) / 255.0
+        return array.astype(np.float32, copy=False)
+    if isinstance(video, list):
+        return np.stack([_to_frame_array(frame) for frame in video], axis=0)
+    return np.asarray(video, dtype=np.float32)
+
+
+def main() -> None:
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    omni = Omni(
+        model=args.model,
+        model_class_name="LingBotVideoPipeline",
+        flow_shift=args.flow_shift,
+        parallel_config=DiffusionParallelConfig(),
+    )
+
+    prompt: dict[str, str] = {"prompt": args.prompt}
+    if args.negative_prompt:
+        prompt["negative_prompt"] = args.negative_prompt
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=generator,
+        output_type=args.output_type,
+        extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
+    )
+
+    start = time.perf_counter()
+    output = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+    print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
+    peak_mb = _peak_memory_mb(output)
+    if peak_mb:
+        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
+
+    video = _to_frame_array(_extract_video(output))
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    from diffusers.utils import export_to_video
+
+    export_to_video(list(video), str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+    omni.close()
+
+
+if __name__ == "__main__":
+    main()

recipes/README.md

@@ -48,6 +48,7 @@ recipes/
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
+| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense text-to-video serving | 1x L20X / A100 / H100 |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
 | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |

recipes/Robbyant/LingBot-Video.md

@@ -0,0 +1,143 @@
+# LingBot-Video
+
+> Native dense text-to-video serving for `robbyant/lingbot-video-dense-1.3b`
+
+## Summary
+
+- Vendor: Robbyant
+- Model: `robbyant/lingbot-video-dense-1.3b`
+- Task: Text-to-video generation
+- Mode: Offline generation and online serving with the OpenAI-compatible `/v1/videos` API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to run the dense LingBot-Video checkpoint with
+vLLM-Omni's native pipeline. The runtime path does not import the upstream
+`lingbot_video` Python package; it loads the checkpoint components directly
+with the in-tree `LingBotVideoPipeline`, dense DiT transformer, shared FlowUniPC
+scheduler, Qwen3-VL text encoder, and Wan VAE.
+
+This recipe is intentionally text-to-video only. Text-to-image, image-conditioned
+video, TI2V, MoE checkpoints, and fused-expert kernels are not covered here.
+
+## References
+
+- Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
+- Upstream project: <https://github.com/Robbyant/lingbot-video>
+- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
+- Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)
+
+## Hardware Support
+
+This recipe documents the CUDA single-GPU dense checkpoint path. Multi-GPU
+parallelism, Cache-DiT, CPU offload, and the LingBot MoE path are not validated
+for this model in this PR.
+
+## GPU
+
+### 1 x NVIDIA L20X / A100 / H100
+
+The dense 1.3B checkpoint has been smoke-tested on a single NVIDIA L20X at a
+small validation shape (`192x320`, 9 frames, 2 steps). Use a larger GPU or lower
+resolution/frame count if your target workload runs out of memory.
+
+#### Offline T2V
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
+  --model robbyant/lingbot-video-dense-1.3b \
+  --prompt "a robotic arm picks up a red block" \
+  --output lingbot_t2v.mp4 \
+  --height 192 \
+  --width 320 \
+  --num-frames 9 \
+  --num-inference-steps 2 \
+  --guidance-scale 3.0 \
+  --flow-shift 3.0 \
+  --seed 42 \
+  --fps 24
+```
+
+#### Online Serving
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+vllm serve robbyant/lingbot-video-dense-1.3b \
+  --omni \
+  --model-class-name LingBotVideoPipeline \
+  --port 8091
+```
+
+After the server is ready, submit a text-to-video job:
+
+```bash
+create_response=$(curl -s http://localhost:8091/v1/videos \
+  -F "model=robbyant/lingbot-video-dense-1.3b" \
+  -F "prompt=a robotic arm picks up a red block" \
+  -F "width=320" \
+  -F "height=192" \
+  -F "num_frames=9" \
+  -F "fps=24" \
+  -F "num_inference_steps=2" \
+  -F "guidance_scale=3.0" \
+  -F "flow_shift=3.0" \
+  -F "seed=42")
+
+video_id=$(echo "${create_response}" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "${status}" = "completed" ]; then
+    break
+  fi
+  if [ "${status}" = "failed" ]; then
+    curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o lingbot_t2v.mp4
+```
+
+## Key Parameters
+
+| Parameter | Suggested smoke value | Notes |
+|-----------|-----------------------|-------|
+| `height` | `192` | Must be a multiple of 16 |
+| `width` | `320` | Must be a multiple of 16 |
+| `num_frames` | `9` | Must be `1` or `4n + 1`; this PR validates T2V with video outputs |
+| `num_inference_steps` | `2` | Use more steps for quality sweeps |
+| `guidance_scale` | `3.0` | CFG is active when this is greater than `1.0` |
+| `flow_shift` | `3.0` | Scheduler flow-shift; aliases the pipeline's internal `shift` |
+| `negative_prompt` | model default | Optional text describing artifacts to avoid |
+| `fps` | `24` | Output MP4 frame rate |
+
+## Validation
+
+Local dense smoke run:
+
+- Shape: 9 frames at `192x320`
+- Steps: 2
+- Request generation time: `0.2923s`
+- Peak reserved GPU memory: `14548 MiB`
+
+Local parity harness against the upstream repository:
+
+- Shape: `[9, 192, 320, 3]`
+- MAE: `0.0065238`
+- MSE: `0.00006650`
+- PSNR: `41.77 dB`
+- Native request time: `0.2875s`
+
+Do not treat these as production benchmarks; they are draft-PR validation
+numbers for the small dense smoke configuration.
+
+## Known Limitations
+
+- T2V only. T2I, I2V, and TI2V are not claimed by this PR.
+- Dense checkpoint only. MoE/fused-expert support belongs in a separate PR.
+- No Cache-DiT, TeaCache, tensor parallelism, sequence parallelism, CFG
+  parallelism, HSDP, CPU offload, or VAE patch parallelism validation yet.
+- Only one request per LingBot pipeline batch is currently supported.

tests/ar_diffusion/test_kv_cache.py

@@ -9,9 +9,14 @@ KV stack on CPU (block bookkeeping only, no GPU tensors).
 import pytest
 import torch
 from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
-from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.request import RequestStatus
 
+try:
+    from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+except ModuleNotFoundError:
+    KVCacheSpecRegistry = None
+    from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
 from vllm_omni.experimental.ar_diffusion.kv_cache import (
     ARDiffusionKVCache,
     ARDiffusionKVConfig,
@@ -48,7 +53,10 @@ def test_spec_registration_resolves_to_chunk_window_manager():
     # Without explicit registration the MRO walk would fall back to the parent
     # SlidingWindowManager; assert the subclass manager wins.
     spec = make_spec()
-    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    if KVCacheSpecRegistry is not None:
+        assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+    else:
+        assert spec_manager_map[type(spec)] is ChunkWindowManager
 
 
 def test_spec_kind_is_sliding_window():

tests/config/composable_parallel/test_apply.py

@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for applying strategy specs onto merged stage configs."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    RouteByStage,
-    StrategyApplyError,
-    StrategySpec,
-    TakeRank,
-    apply_strategy_specs,
-)
-from vllm_omni.config.stage_config import StageConfig
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def _stage(stage_id: int, model_stage: str, engine_args=None, runtime=None) -> StageConfig:
-    return StageConfig(
-        stage_id=stage_id,
-        model_stage=model_stage,
-        yaml_engine_args=dict(engine_args or {}),
-        yaml_runtime=dict(runtime or {"num_replicas": 1}),
-    )
-
-
-def _qwen_stages() -> list[StageConfig]:
-    return [
-        _stage(0, "thinker"),
-        _stage(1, "talker"),
-        _stage(2, "code2wav"),
-    ]
-
-
-def test_apply_tp_by_role():
-    stages = _qwen_stages()
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-    # untouched roles keep their config
-    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
-
-
-def test_apply_by_model_stage():
-    stages = _qwen_stages()
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_apply_stage_replica_sets_num_replicas_and_surfaces_lb():
-    stages = _qwen_stages()
-    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
-    assert stages[1].yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def test_only_declared_axes_are_written():
-    stages = _qwen_stages()
-    # strategy declares only stage_replica -> tp must not be forced.
-    apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
-    assert "tensor_parallel_size" not in stages[1].yaml_engine_args
-
-
-def test_conflict_on_explicit_tp():
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = 4
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_equal_explicit_value_is_noop():
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = 2
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_explicit_none_conflicts_with_derived_value():
-    # An explicit YAML ``null`` (``tensor_parallel_size: null``) is a *present*
-    # value, not a missing key, so a strategy deriving a non-None size must raise
-    # rather than silently clobber the explicit None.
-    stages = _qwen_stages()
-    stages[0].yaml_engine_args["tensor_parallel_size"] = None
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_missing_key_is_filled_without_conflict():
-    # A genuinely absent key (never set in the YAML) is filled by the strategy
-    # and must NOT raise — the contrast case to an explicit None.
-    stages = _qwen_stages()
-    assert "tensor_parallel_size" not in stages[0].yaml_engine_args
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_num_replicas_conflict():
-    stages = _qwen_stages()
-    stages[1].yaml_runtime["num_replicas"] = 3
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"talker": [_stage_replica(2)]})
-
-
-def test_device_count_ok_template():
-    stages = _qwen_stages()
-    stages[0].yaml_runtime["devices"] = "0,1"
-    apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_device_count_ok_pool():
-    # tp=2 -> world=2; 2 replicas -> pool of 4 device ids is valid.
-    stages = _qwen_stages()
-    stages[1].yaml_runtime["devices"] = "0,1,2,3"
-    apply_strategy_specs(stages, {"talker": [_tp(2), _stage_replica(2)]})
-    assert stages[1].yaml_runtime["num_replicas"] == 2
-
-
-def test_device_count_mismatch():
-    stages = _qwen_stages()
-    stages[0].yaml_runtime["devices"] = "0,1,2"
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_unknown_role_raises():
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"nonexistent": [_tp(2)]})
-
-
-def test_conflicting_lb_policy_across_roles():
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(
-            stages,
-            {
-                "talker": [_stage_replica(2, "round_robin")],
-                "code2wav": [_stage_replica(2, "least_queue")],
-            },
-        )

tests/config/composable_parallel/test_factory_integration.py

@@ -1,141 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam."""
-
-from __future__ import annotations
-
-import logging
-from pathlib import Path
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    RouteByStage,
-    StrategyApplyError,
-    StrategySpec,
-    TakeRank,
-    apply_strategy_specs,
-)
-from vllm_omni.config.config_factory import StageConfigFactory
-from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
-from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-_DEPLOY = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "qwen2_5_omni.yaml"
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def _qwen_stages():
-    pipeline = OMNI_PIPELINES["qwen2_5_omni"]
-    deploy = load_deploy_config(_DEPLOY)
-    return merge_pipeline_deploy(pipeline, deploy)
-
-
-def _stage(stages, role):
-    return next(s for s in stages if s.model_stage == role)
-
-
-def test_overlay_tp_on_thinker():
-    stages = _qwen_stages()
-    # The bundled deploy pins the thinker to one GPU; a TP=2 strategy needs a
-    # matching 2-GPU layout (mirrors what a TP2 deploy would declare).
-    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
-    result = apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-    assert _stage(result.stages, "thinker").yaml_engine_args["tensor_parallel_size"] == 2
-
-
-def test_device_guard_rejects_tp_on_single_gpu_deploy():
-    # The bundled deploy pins the thinker to a single GPU, so the pre-spawn
-    # device check must refuse a TP=2 strategy on it.
-    stages = _qwen_stages()
-    with pytest.raises(StrategyApplyError):
-        apply_strategy_specs(stages, {"thinker": [_tp(2)]})
-
-
-def test_overlay_stage_replica_on_talker():
-    stages = _qwen_stages()
-    result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]})
-    assert _stage(result.stages, "talker").yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def test_overlay_mixed_roles():
-    stages = _qwen_stages()
-    _stage(stages, "thinker").yaml_runtime["devices"] = "0,1"
-    result = apply_strategy_specs(
-        stages,
-        {
-            "thinker": [_tp(2)],
-            "talker": [_stage_replica(2, "round_robin")],
-            "code2wav": [_stage_replica(2, "round_robin")],
-        },
-    )
-    by_role = {s.model_stage: s for s in result.stages}
-    assert by_role["thinker"].yaml_engine_args["tensor_parallel_size"] == 2
-    assert by_role["talker"].yaml_runtime["num_replicas"] == 2
-    assert by_role["code2wav"].yaml_runtime["num_replicas"] == 2
-    assert result.omni_lb_policy == "round-robin"
-
-
-def _resolved(stages, role):
-    """Return a role's fully-resolved (post-CLI) OmegaConf stage config."""
-    stage = next(s for s in stages if s.model_stage == role)
-    return stage.to_omegaconf()
-
-
-def test_device_check_survives_cli_override():
-    # Strategy replicates the talker (1-GPU template -> valid at apply time),
-    # but a CLI --stage_1_devices with 3 ids must NOT slip past the device
-    # guard: effective world=1, replicas=2 admits only 1 or 2 device ids.
-    pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
-    with pytest.raises(StrategyApplyError):
-        StageConfigFactory._create_from_registry(
-            "qwen2_5_omni",
-            pipeline_cfg,
-            cli_overrides={"stage_1_devices": "0,1,2"},
-            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
-        )
-
-
-def test_cli_overrides_strategy_with_warning():
-    # Strategy derives num_replicas=2 for the talker; a CLI override to 3 wins
-    # (it is the most explicit user action) but must be surfaced loudly rather
-    # than silently. The resulting layout (1 template device) stays valid.
-    #
-    # vLLM's logger sets propagate=False, so attach a handler directly to it
-    # rather than relying on pytest's caplog (which listens on the root logger).
-    messages: list[str] = []
-
-    class _Capture(logging.Handler):
-        def emit(self, record: logging.LogRecord) -> None:
-            messages.append(record.getMessage())
-
-    log = logging.getLogger("vllm_omni.config.config_factory")
-    handler = _Capture(level=logging.WARNING)
-    log.addHandler(handler)
-    try:
-        pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"]
-        stages, _ = StageConfigFactory._create_from_registry(
-            "qwen2_5_omni",
-            pipeline_cfg,
-            cli_overrides={"stage_1_num_replicas": 3},
-            strategy_specs={"talker": [_stage_replica(2, "round_robin")]},
-        )
-    finally:
-        log.removeHandler(handler)
-
-    # CLI value wins in the resolved config.
-    assert _resolved(stages, "talker").runtime.num_replicas == 3
-    # ...and the override was warned about, naming the conflicting field.
-    assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages)

tests/config/composable_parallel/test_spec_types.py

@@ -1,42 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for the vendored spec/pattern data types."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel.aggregation import AggregationPattern, TakeRank
-from vllm_omni.config.composable_parallel.routing import RoutingPattern
-from vllm_omni.config.composable_parallel.spec import MeshAxisSpec
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_mesh_axis_rejects_bad_kind():
-    with pytest.raises(ValueError):
-        MeshAxisSpec(kind="not_a_kind", size=2)  # type: ignore[arg-type]
-
-
-def test_mesh_axis_rejects_nonpositive_size():
-    with pytest.raises(ValueError):
-        MeshAxisSpec(kind="tp", size=0)
-
-
-def test_routing_pattern_is_closed():
-    with pytest.raises(TypeError):
-
-        class RogueRouting(RoutingPattern):  # noqa: F811
-            pass
-
-
-def test_aggregation_pattern_is_closed():
-    with pytest.raises(TypeError):
-
-        class RogueAggregation(AggregationPattern):  # noqa: F811
-            pass
-
-
-def test_take_rank_type_guard():
-    with pytest.raises(TypeError):
-        TakeRank(rank=True)  # bool is not a valid rank

tests/config/composable_parallel/test_strategy_loader.py

@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for parsing/loading strategy files."""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.config.composable_parallel.routing import Broadcast, RouteByStage
-from vllm_omni.config.composable_parallel.strategy_loader import (
-    StrategyLoadError,
-    load_strategy_specs,
-    parse_strategy_specs,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_parse_basic():
-    data = {
-        "strategies": {
-            "thinker": [{"axis": "tp", "size": 2}],
-            "talker": [{"axis": "stage_replica", "size": 2, "routing": "round_robin"}],
-        }
-    }
-    specs = parse_strategy_specs(data)
-    assert set(specs) == {"thinker", "talker"}
-
-    (tp_spec,) = specs["thinker"]
-    assert tp_spec.mesh_axis.kind == "tp"
-    assert tp_spec.mesh_axis.size == 2
-    assert isinstance(tp_spec.routing, Broadcast)
-
-    (sr_spec,) = specs["talker"]
-    assert sr_spec.mesh_axis.kind == "stage_replica"
-    assert isinstance(sr_spec.routing, RouteByStage)
-    assert sr_spec.routing.routing_policy == "round_robin"
-
-
-def test_parse_without_strategies_key():
-    # A bare role mapping (no top-level "strategies") is also accepted.
-    specs = parse_strategy_specs({"thinker": [{"axis": "tp", "size": 4}]})
-    assert specs["thinker"][0].mesh_axis.size == 4
-
-
-def test_l1_owner_goes_to_shard_extension():
-    specs = parse_strategy_specs({"talker": [{"axis": "stage_replica", "size": 2, "l1_owner": "delegated"}]})
-    assert specs["talker"][0].shard_extension["l1_owner"] == "delegated"
-
-
-def test_missing_axis_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"size": 2}]})
-
-
-def test_missing_size_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp"}]})
-
-
-def test_routing_on_non_policy_axis_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp", "size": 2, "routing": "random"}]})
-
-
-def test_bad_size_raises():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": [{"axis": "tp", "size": "two"}]})
-
-
-def test_entries_must_be_list():
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": {"axis": "tp", "size": 2}})
-
-
-def test_non_mapping_entry_raises():
-    # A list whose elements are not mappings (e.g. a bare string) must raise a
-    # StrategyLoadError, not an opaque TypeError from dict(entry).
-    with pytest.raises(StrategyLoadError):
-        parse_strategy_specs({"thinker": ["tp"]})
-
-
-def test_load_from_file(tmp_path):
-    path = tmp_path / "strategy.yaml"
-    path.write_text(
-        "strategies:\n"
-        "  thinker:\n"
-        "    - axis: tp\n"
-        "      size: 2\n"
-        "  talker:\n"
-        "    - axis: stage_replica\n"
-        "      size: 2\n"
-        "      routing: least_queue\n"
-    )
-    specs = load_strategy_specs(str(path))
-    assert specs["thinker"][0].mesh_axis.size == 2
-    assert specs["talker"][0].routing.routing_policy == "least_queue"

tests/config/composable_parallel/test_translator.py

@@ -1,195 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for StrategySpec -> OmniParallelConfig translation."""
-
-from __future__ import annotations
-
-import dataclasses
-
-import pytest
-
-from vllm_omni.config.composable_parallel import (
-    AxisTranslationError,
-    Broadcast,
-    FanInByStage,
-    MeshAxisSpec,
-    PartitionByHash,
-    PipelineMicrobatch,
-    RouteByStage,
-    StrategySpec,
-    TakeRank,
-    Union,
-    translate_strategy_stack,
-)
-from vllm_omni.config.composable_parallel.aggregation import StitchPipeline
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _tp(size: int) -> StrategySpec:
-    return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank())
-
-
-def _dp(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("dp", MeshAxisSpec("dp", size), RouteByStage(policy), Union())
-
-
-def _pp(size: int) -> StrategySpec:
-    return StrategySpec("pp", MeshAxisSpec("pp", size), PipelineMicrobatch(), StitchPipeline())
-
-
-def _ep(size: int) -> StrategySpec:
-    return StrategySpec("ep", MeshAxisSpec("ep", size), Broadcast(), Union())
-
-
-def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec:
-    return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage())
-
-
-def test_tp_only():
-    cfg = translate_strategy_stack([_tp(2)])
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.data_parallel_size == 1
-    assert cfg.world_size == 2
-    assert cfg.stage_replica_size == 1
-    assert cfg.l1_owners == {"tp": "engine"}
-
-
-def test_dp_is_engine_owned():
-    cfg = translate_strategy_stack([_dp(2)])
-    assert cfg.data_parallel_size == 2
-    assert cfg.world_size == 2
-    # engine DP does not emit an omni LB policy (vLLM's internal DP LB is used).
-    assert cfg.omni_lb_policy is None
-    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
-    assert cfg.l1_owners["dp"] == "engine"
-
-
-def test_tp_times_dp_world_size():
-    cfg = translate_strategy_stack([_tp(2), _dp(2)])
-    assert cfg.world_size == 4
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.data_parallel_size == 2
-
-
-def test_stage_replica_is_delegated():
-    cfg = translate_strategy_stack([_stage_replica(3, "least_queue")])
-    # stage_replica is NOT a world dimension.
-    assert cfg.world_size == 1
-    assert cfg.stage_replica_size == 3
-    assert cfg.omni_lb_policy == "least-queue-length"
-    assert cfg.l1_owners["stage_replica"] == "delegated"
-    # stage_replica_size is not an engine kwarg; it's a per-stage deploy knob.
-    assert "stage_replica_size" not in cfg.as_engine_kwargs()
-
-
-def test_stage_replica_lb_policy_only_when_replicated():
-    cfg = translate_strategy_stack([_stage_replica(1)])
-    assert cfg.stage_replica_size == 1
-    assert cfg.as_engine_kwargs().get("omni_lb_policy") is None
-
-
-def test_tp_with_stage_replica():
-    cfg = translate_strategy_stack([_tp(2), _stage_replica(2)])
-    assert cfg.world_size == 2  # tp only
-    assert cfg.tensor_parallel_size == 2
-    assert cfg.stage_replica_size == 2
-    assert cfg.omni_lb_policy == "round-robin"
-
-
-def test_ep_must_match_tp_dp_product():
-    cfg = translate_strategy_stack([_tp(2), _ep(2)])
-    assert cfg.enable_expert_parallel is True
-    assert cfg.as_engine_kwargs()["enable_expert_parallel"] is True
-
-
-def test_ep_mismatch_raises():
-    with pytest.raises(Exception) as exc:
-        translate_strategy_stack([_tp(2), _ep(4)])
-    assert "must equal" in str(exc.value)
-
-
-def test_duplicate_kind_rejected():
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([_tp(2), _tp(2)])
-
-
-def test_unsupported_kind_rejected():
-    spec = StrategySpec("sp", MeshAxisSpec("sp_ulysses", 2), Broadcast(), TakeRank())
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_hash_routing_rejected():
-    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), PartitionByHash(), Union())
-    with pytest.raises(NotImplementedError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_invalid_routing_policy_rejected():
-    spec = StrategySpec("dp", MeshAxisSpec("dp", 2), RouteByStage("bogus"), Union())
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_lb_policy_never_in_engine_kwargs_even_when_replicated():
-    # omni_lb_policy is pipeline-wide (applied at engine construction), not a
-    # per-stage engine arg, so it must never leak into as_engine_kwargs().
-    cfg = translate_strategy_stack([_tp(2), _stage_replica(2, "least_queue")])
-    assert cfg.omni_lb_policy == "least-queue-length"
-    assert "omni_lb_policy" not in cfg.as_engine_kwargs()
-
-
-def test_stage_replica_hash_routing_rejected():
-    spec = StrategySpec("sr", MeshAxisSpec("stage_replica", 2), RouteByStage("hash"), FanInByStage())
-    with pytest.raises(NotImplementedError):
-        translate_strategy_stack([spec])
-
-
-def test_dp_wrong_owner_rejected():
-    spec = StrategySpec(
-        "dp",
-        MeshAxisSpec("dp", 2),
-        RouteByStage("round_robin"),
-        Union(),
-        shard_extension={"l1_owner": "delegated"},
-    )
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-def test_stage_replica_wrong_owner_rejected():
-    spec = StrategySpec(
-        "sr",
-        MeshAxisSpec("stage_replica", 2),
-        RouteByStage("round_robin"),
-        FanInByStage(),
-        shard_extension={"l1_owner": "engine"},
-    )
-    with pytest.raises(AxisTranslationError):
-        translate_strategy_stack([spec])
-
-
-@pytest.mark.parametrize(
-    "specs",
-    [
-        [_tp(2)],
-        [_dp(2)],
-        [_tp(2), _dp(2)],
-        [_tp(2), _ep(2)],
-        [_tp(2), _stage_replica(2)],
-        [_pp(2)],
-    ],
-)
-def test_as_engine_kwargs_are_valid_omni_engine_args(specs):
-    # Guard against desync: every key emitted by as_engine_kwargs() must be a
-    # real OmniEngineArgs field, otherwise constructing the per-stage engine args
-    # would explode at runtime. Keep this CPU-only by checking field names rather
-    # than instantiating the engine.
-    from vllm_omni.engine.arg_utils import OmniEngineArgs
-
-    valid_fields = {f.name for f in dataclasses.fields(OmniEngineArgs)}
-    cfg = translate_strategy_stack(specs)
-    kwargs = cfg.as_engine_kwargs()
-    unknown = set(kwargs) - valid_fields
-    assert not unknown, f"as_engine_kwargs emitted non-OmniEngineArgs fields: {unknown}"

tests/dfx/perf/tests/test_lingbot_video_vllm_omni.json

@@ -0,0 +1,34 @@
+[
+    {
+        "test_name": "test_lingbot_video_single_device",
+        "description": "Single-device LingBot dense T2V baseline at 320x192, 9 frames, 2 steps.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "robbyant/lingbot-video-dense-1.3b",
+            "serve_args": {
+                "model-class-name": "LingBotVideoPipeline",
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "320x192_frames9_steps2",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 320,
+                "height": 192,
+                "num-frames": 9,
+                "fps": 24,
+                "num-inference-steps": 2,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "warmup-requests": 1,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 2,
+                "enable-negative-prompt": true,
+                "skip-performance-assertion": true
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_voxcpm2.json

@@ -3,10 +3,7 @@
     "test_name": "test_voxcpm2",
     "server_params": {
       "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true,
-      "extra_cli_args": [
-        "--trust-remote-code"
-      ]
+      "trust_remote_code": true
     },
     "benchmark_params": [
       {

tests/diffusion/models/lingbot_video/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/lingbot_video/test_lingbot_video_transformer.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _tiny_transformer(**overrides):
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
+
+    config = {
+        "patch_size": (1, 1, 1),
+        "in_channels": 2,
+        "out_channels": 2,
+        "hidden_size": 16,
+        "num_attention_heads": 1,
+        "depth": 0,
+        "intermediate_size": 32,
+        "text_dim": 8,
+        "freq_dim": 8,
+        "axes_dims": (4, 4, 8),
+        "axes_lens": (32, 32, 32),
+    }
+    config.update(overrides)
+    return LingBotVideoTransformer3DModel(**config)
+
+
+def test_joint_position_ids_video_then_text_order():
+    from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import make_joint_position_ids
+
+    positions = make_joint_position_ids(text_len=3, grid_t=1, grid_h=2, grid_w=2, device=torch.device("cpu"))
+
+    assert positions.shape == (7, 3)
+    assert positions[:4, 0].tolist() == [4, 4, 4, 4]
+    assert positions[:4, 1:].tolist() == [[0, 0], [0, 1], [1, 0], [1, 1]]
+    assert positions[4:].tolist() == [[1, 0, 0], [2, 0, 0], [3, 0, 0]]
+
+
+def test_tiny_transformer_depth_zero_forward_shape():
+    model = _tiny_transformer()
+    hidden_states = torch.randn(1, 2, 1, 2, 2)
+    timestep = torch.tensor([300.0])
+    encoder_hidden_states = torch.randn(1, 3, 8)
+    encoder_attention_mask = torch.ones(1, 3, dtype=torch.long)
+
+    with torch.no_grad():
+        out = model(
+            hidden_states,
+            timestep,
+            encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            return_dict=False,
+        )[0]
+
+    assert out.shape == hidden_states.shape
+    assert torch.isfinite(out).all()
+
+
+def test_packed_attention_uses_sdpa_fallback_without_flash_varlen(monkeypatch):
+    from vllm_omni.diffusion.models.lingbot_video import lingbot_video_transformer as module
+
+    monkeypatch.setattr(module, "flash_attn_varlen_func_v3", None)
+    attn = module.LingBotVideoAttention(
+        hidden_size=8,
+        num_heads=2,
+        norm_eps=1e-6,
+        qkv_bias=False,
+        out_bias=False,
+    )
+    captured = {}
+
+    def fake_sdpa_forward(query, key, value, attn_metadata):
+        captured["mask"] = attn_metadata.attn_mask
+        return torch.zeros_like(query)
+
+    monkeypatch.setattr(attn.attn.sdpa_fallback, "forward", fake_sdpa_forward)
+    x = torch.randn(1, 5, 8)
+    rotary = torch.ones(1, 5, 2, dtype=torch.complex64)
+    packed_indices = {
+        "cu_seqlens_kv": torch.tensor([0, 2, 5], dtype=torch.int32),
+        "max_seqlen_in_batch_kv": 3,
+        "attention_mask": module._packed_block_attention_mask([2, 3], x.device),
+    }
+
+    out = attn(x, rotary, packed_indices=packed_indices)
+
+    assert out.shape == x.shape
+    mask = captured["mask"]
+    assert mask.shape == (1, 1, 5, 5)
+    assert mask[0, 0, :2, :2].all()
+    assert mask[0, 0, 2:, 2:].all()
+    assert not mask[0, 0, :2, 2:].any()
+    assert not mask[0, 0, 2:, :2].any()
+
+
+def test_tiny_transformer_rejects_invalid_rope_dims():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoTransformer3DModel
+
+    with pytest.raises(AssertionError, match="head_dim"):
+        LingBotVideoTransformer3DModel(
+            hidden_size=16,
+            num_attention_heads=1,
+            axes_dims=(4, 4, 4),
+            depth=0,
+        )
+
+
+def test_transformer_to_keeps_sensitive_modules_in_fp32():
+    model = _tiny_transformer()
+
+    model.to(dtype=torch.bfloat16)
+
+    assert model.patch_embedder.weight.dtype == torch.bfloat16
+    assert model.time_embedder.linear_1.weight.dtype == torch.float32
+    assert model.norm_out_modulation[1].weight.dtype == torch.float32

tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_pipeline():
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline
+
+    pipeline = object.__new__(LingBotVideoPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.default_negative_prompt = "default negative"
+    pipeline.od_config = SimpleNamespace(flow_shift=None)
+    return pipeline
+
+
+def _make_request_batch(prompt, **sampling_overrides):
+    sampling = OmniDiffusionSamplingParams(**sampling_overrides)
+    request = OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id="lingbot-test")
+    return DiffusionRequestBatch([request])
+
+
+def test_lingbot_video_pipeline_import_and_registry():
+    from vllm_omni.diffusion.models.lingbot_video import (
+        LingBotVideoPipeline,
+        LingBotVideoTransformer3DModel,
+        get_lingbot_video_post_process_func,
+    )
+    from vllm_omni.diffusion.registry import _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS
+
+    assert LingBotVideoPipeline is not None
+    assert LingBotVideoTransformer3DModel is not None
+    assert get_lingbot_video_post_process_func is not None
+    assert _DIFFUSION_MODELS["LingBotVideoPipeline"] == (
+        "lingbot_video",
+        "pipeline_lingbot_video",
+        "LingBotVideoPipeline",
+    )
+    assert _DIFFUSION_POST_PROCESS_FUNCS["LingBotVideoPipeline"] == "get_lingbot_video_post_process_func"
+
+
+def test_component_discovery_declarations():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
+
+    assert LingBotVideoPipeline._dit_modules == ["transformer"]
+    assert LingBotVideoPipeline._encoder_modules == ["text_encoder"]
+    assert LingBotVideoPipeline._vae_modules == ["vae"]
+    assert LingBotVideoPipeline.supports_step_execution is False
+
+
+def test_extra_body_params_include_video_flow_shift_alias():
+    from vllm_omni.model_extras import get_extra_body_params
+
+    params = get_extra_body_params("LingBotVideoPipeline")
+    assert "shift" in params
+    assert "flow_shift" in params
+    assert "batch_cfg" in params
+
+
+def test_extract_prompt_string_and_mapping():
+    from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import _extract_prompt
+
+    string_req = SimpleNamespace(prompt="a robot waves")
+    assert _extract_prompt(string_req) == ("a robot waves", None)
+
+    mapping_req = SimpleNamespace(prompt={"prompt": "a robot walks", "negative_prompt": "low quality"})
+    assert _extract_prompt(mapping_req) == ("a robot walks", "low quality")
+
+    with pytest.raises(TypeError, match="Unsupported LingBot prompt type"):
+        _extract_prompt(SimpleNamespace(prompt=["not", "supported"]))
+
+
+def test_check_inputs_accepts_t2v_shapes_and_rejects_invalid_values():
+    from vllm_omni.diffusion.models.lingbot_video import LingBotVideoPipeline
+
+    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=1)
+    LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=9)
+
+    with pytest.raises(ValueError, match="num_frames"):
+        LingBotVideoPipeline.check_inputs(height=192, width=320, num_frames=8)
+    with pytest.raises(ValueError, match="height"):
+        LingBotVideoPipeline.check_inputs(height=190, width=320, num_frames=9)
+
+
+def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested():
+    from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func
+
+    frames = torch.ones(2, 4, 4, 3)
+    post = get_lingbot_video_post_process_func(SimpleNamespace())
+
+    assert post(frames, SimpleNamespace(output_type="pt")) is frames
+    array = post(frames, SimpleNamespace(output_type="np"))
+    assert isinstance(array, np.ndarray)
+    assert array.shape == (2, 4, 4, 3)
+
+
+def test_forward_resolves_t2v_sampling_and_flow_shift_alias():
+    pipeline = _make_pipeline()
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(9, 192, 320, 3)
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        {"prompt": "a robot arm", "negative_prompt": "bad hands"},
+        height=192,
+        width=320,
+        num_frames=9,
+        num_inference_steps=2,
+        guidance_scale=3.0,
+        seed=42,
+        output_type="pt",
+        extra_args={"flow_shift": 4.0, "batch_cfg": True},
+    )
+
+    out = pipeline.forward(req)
+
+    assert torch.equal(out.output, torch.zeros(9, 192, 320, 3))
+    assert len(calls) == 1
+    call = calls[0]
+    assert call["prompt"] == "a robot arm"
+    assert call["negative_prompt"] == "bad hands"
+    assert call["height"] == 192
+    assert call["width"] == 320
+    assert call["num_frames"] == 9
+    assert call["num_inference_steps"] == 2
+    assert call["guidance_scale"] == 3.0
+    assert call["shift"] == 4.0
+    assert call["batch_cfg"] is True
+    assert call["output_type"] == "pt"
+
+
+def test_forward_prefers_shift_over_flow_shift_and_defaults_negative_prompt():
+    pipeline = _make_pipeline()
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        return torch.zeros(5, 192, 192, 3)
+
+    pipeline._generate = fake_generate
+    req = _make_request_batch(
+        "a robot arm",
+        height=192,
+        width=192,
+        num_frames=5,
+        num_inference_steps=2,
+        guidance_scale=1.0,
+        seed=42,
+        extra_args={"shift": 5.0, "flow_shift": 4.0},
+    )
+
+    pipeline.forward(req)
+
+    assert calls[0]["shift"] == 5.0
+    assert calls[0]["negative_prompt"] == "default negative"
+
+
+def test_forward_rejects_multi_request_batches():
+    pipeline = _make_pipeline()
+    first = OmniDiffusionRequest(
+        prompt="first",
+        sampling_params=OmniDiffusionSamplingParams(seed=1),
+        request_id="lingbot-first",
+    )
+    second = OmniDiffusionRequest(
+        prompt="second",
+        sampling_params=OmniDiffusionSamplingParams(seed=2),
+        request_id="lingbot-second",
+    )
+
+    with pytest.raises(ValueError, match="only supports one request"):
+        pipeline.forward(DiffusionRequestBatch([first, second]))
+
+
+def test_load_weights_rejects_external_weight_stream():
+    pipeline = _make_pipeline()
+
+    assert pipeline.load_weights([]) == set()
+    with pytest.raises(RuntimeError, match="components are loaded directly"):
+        pipeline.load_weights([("transformer.weight", torch.zeros(1))])

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -27,15 +27,8 @@ WIDTH = 1024
 PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
-# NOTE: re-baselined for the vLLM rebase (dev/vllm-align). HunyuanImage-3 is an
-# expert-parallel MoE model; the newer upstream vLLM (65a7b4628409, 585 commits
-# past the v0.24.0 tag that main tracks) shifted MoE kernel numerics, so the
-# structural-similarity metrics against the committed baseline drift to a
-# deterministic SSIM 0.9490 / PSNR 26.21 dB (the mean/p99 pixel-diff checks above
-# still pass — the image renders correctly). Thresholds relaxed to accept this
-# drift as the new bar. On main (older vLLM) the stricter 0.97 / 30.0 hold.
-SSIM_THRESHOLD = 0.94
-PSNR_THRESHOLD = 25.0
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
 
 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
 BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"

tests/e2e/offline_inference/test_qwen3_tts_customvoice.py

@@ -34,10 +34,12 @@ def get_cuda_graph_config():
             "stages": {
                 0: {
                     "max_num_seqs": 1,
+                    "gpu_memory_utilization": 0.2,
                     "enforce_eager": True,
                     "async_scheduling": False,
                 },
                 1: {
+                    "gpu_memory_utilization": 0.2,
                     "enforce_eager": True,
                     "async_scheduling": False,
                 },

tests/e2e/online_serving/test_lingbot_video.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Online serving smoke for ``robbyant/lingbot-video-dense-1.3b``.
+
+The LingBot dense PR supports text-to-video only, so this module exercises the
+native ``/v1/videos`` path with a small 9-frame request.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "robbyant/lingbot-video-dense-1.3b"
+PROMPT = "a robotic arm picks up a red block"
+NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--model-class-name", "LingBotVideoPipeline"],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_text_to_video_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "height": 192,
+            "width": 320,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 3.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_lingbot_video_expansion.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+L4 expansion coverage for ``robbyant/lingbot-video-dense-1.3b``.
+
+This is intentionally T2V-only. T2I, I2V, TI2V, MoE, and multi-GPU feature
+rows should be added in follow-up PRs when those paths are implemented.
+"""
+
+import json
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = "robbyant/lingbot-video-dense-1.3b"
+PROMPT = "a robotic arm picks up a red block"
+NEGATIVE_PROMPT = "low quality, blurry, watermark, text"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--model-class-name", "LingBotVideoPipeline"],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_cfg_off(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "height": 192,
+            "width": 192,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 1.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
+def test_batch_cfg_extra_params(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "model": omni_server.model,
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "height": 192,
+            "width": 320,
+            "num_frames": 9,
+            "fps": 24,
+            "num_inference_steps": 2,
+            "guidance_scale": 3.0,
+            "flow_shift": 3.0,
+            "seed": 42,
+            "extra_params": json.dumps({"batch_cfg": True}, separators=(",", ":")),
+        },
+    }
+    openai_client.send_video_diffusion_request(request_config)

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -410,44 +410,6 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
-    """
-    Regression: ``use_audio_in_video`` with default video loader sampling (no
-    ``media_io_kwargs.video.fps``) after a prior request that pinned fps.
-
-    Before the fix the second request could partial-hit mm-cache (audio cached,
-    video re-processed without audio) and fail with ``StopIteration`` in
-    ``replace_multimodal_special_tokens``.
-    """
-    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        video_data_url=video_data_url,
-        content_text="Reply with one word: OK.",
-    )
-
-    base_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": False,
-        "use_audio_in_video": True,
-        "modalities": ["text"],
-    }
-
-    # Prime mm-cache with explicit loader fps (historical failure mode).
-    openai_client.send_omni_request(
-        {
-            **base_config,
-            "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
-        }
-    )
-
-    # Default loader sampling: must succeed without explicit media_io fps.
-    openai_client.send_omni_request(dict(base_config))
-
-
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, openai_client) -> None:

tests/engine/test_async_omni_engine_stage_init.py

@@ -739,7 +739,7 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -772,7 +772,7 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -798,7 +798,7 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -874,54 +874,3 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key():
 
     assert "attention_config" not in engine_args_dict
     assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"
-
-
-def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch):
-    """Regression: two stages must never draw the same ZMQ port.
-
-    ``get_open_ports_list`` only dedups within a single call, so per-route
-    allocation used to let a later stage reuse an earlier stage's port. The
-    second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address
-    already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/
-    code2wav). ``OmniMasterServer`` now dedups every port it hands out.
-    """
-    import itertools
-
-    from vllm_omni.engine import stage_engine_startup as ses
-
-    # A colliding prefix (repeats + the master port 9000) followed by an endless
-    # fresh stream, so the only way to succeed is to redraw the collisions.
-    supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002))
-
-    def fake_get_open_ports_list(count):
-        return [next(supply) for _ in range(count)]
-
-    monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list)
-
-    server = ses.OmniMasterServer(
-        master_address="127.0.0.1",
-        master_port=9000,  # seed: a route must not reuse the registration port
-        stage_ids=[0, 1, 2],
-    )
-
-    ports = []
-    for sid in (0, 1, 2):
-        alloc = server.get_allocation(sid)
-        for addr in (
-            alloc.handshake_bind_address,
-            alloc.input_bind_address,
-            alloc.output_bind_address,
-        ):
-            ports.append(ses._port_from_zmq_address(addr))
-
-    assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}"
-    assert 9000 not in ports, "route reused the master registration port"
-
-
-def test_port_from_zmq_address_parsing():
-    from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address
-
-    assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277
-    assert _port_from_zmq_address(None) is None
-    assert _port_from_zmq_address("ipc:///tmp/sock") is None
-    assert _port_from_zmq_address("tcp://host:not-a-port") is None

tests/engine/test_strategy_lb_policy.py

@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine.
-
-``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy``
-and module-level logging, so we exercise it with a lightweight stub ``self`` to
-avoid standing up a full engine (which needs GPUs/model weights).
-"""
-
-from __future__ import annotations
-
-import pytest
-
-from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _Stub:
-    def __init__(self, current: str = "random") -> None:
-        self._omni_lb_policy = current
-
-
-_apply = AsyncOmniEngine._apply_strategy_lb_policy
-
-
-def test_no_derived_policy_is_noop():
-    stub = _Stub("random")
-    _apply(stub, None, {})
-    assert stub._omni_lb_policy == "random"
-
-
-def test_derived_overrides_default_random():
-    # "random" is the engine default == "unset", so the strategy value wins.
-    stub = _Stub("random")
-    _apply(stub, "round-robin", {})
-    assert stub._omni_lb_policy == "round-robin"
-
-
-def test_derived_overrides_when_flag_absent():
-    stub = _Stub("random")
-    _apply(stub, "least-queue-length", {"omni_lb_policy": None})
-    assert stub._omni_lb_policy == "least-queue-length"
-
-
-def test_explicit_matching_flag_is_fine():
-    stub = _Stub("round-robin")
-    _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"})
-    assert stub._omni_lb_policy == "round-robin"
-
-
-def test_explicit_conflicting_flag_raises():
-    stub = _Stub("round-robin")
-    with pytest.raises(ValueError, match="Conflicting load-balancer policy"):
-        _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"})
-
-
-def test_explicit_random_is_treated_as_unset():
-    # An explicit --omni-lb-policy=random is indistinguishable from the default,
-    # so a strategy value still overrides it (documented Phase-1 semantics).
-    stub = _Stub("random")
-    _apply(stub, "round-robin", {"omni_lb_policy": "random"})
-    assert stub._omni_lb_policy == "round-robin"

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -198,37 +198,16 @@ def _assert_realtime_smoke(result: dict) -> None:
     assert result["output_sample_rate"] > 0
 
 
-def _assert_realtime_accuracy(
-    result: dict,
-    threshold: float = 0.8,
-    whisper_model_size: str = "large-v3",
-) -> None:
-    """Assert that whisper transcription of audio output matches model text.
-
-    Args:
-        result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``.
-        threshold: Minimum cosine similarity between whisper text and model text.
-                   Default 0.8 is appropriate for sync mode (exact match expected).
-                   Use 0.65 for async_chunk mode where chunked upload with pacing
-                   delays introduces inherent variability in audio codec quality.
-        whisper_model_size: Whisper model used to transcribe the generated audio
-                   for the accuracy check. Defaults to ``large-v3``: the default
-                   ``small`` model mishears short Chinese TTS clips (observed:
-                   北京→韦京 and a dropped leading sentence, sim=0.443), which
-                   caused spurious sim<threshold failures under async_chunk codec
-                   variability even though audio generation was correct. large-v3
-                   transcribes these clips reliably, so a failure here now points
-                   at the model, not the ASR grader.
-    """
+def _assert_realtime_accuracy(result: dict) -> None:
     final_text = (result["transcription_text"] or "").strip()
     assert final_text, "Expected non-empty transcription (model text stream)"
 
     wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
-    whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip()
+    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
     assert whisper_text, "Whisper returned empty string for synthesized output audio"
 
     sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-    assert sim > threshold, (
+    assert sim > 0.8, (
         f"Output audio transcript should match model text (sim={sim:.3f}): "
         f"whisper={whisper_text!r}, model_text={final_text!r}"
     )
@@ -255,12 +234,7 @@ class TestQwen3OmniRealtimeWebSocket:
         )
 
         _assert_realtime_smoke(result)
-        # Use threshold=0.65 for async_chunk: chunked upload with send_delay_ms
-        # introduces inherent variability in audio codec quality. Whisper
-        # transcription is nondeterministic and may miss 1-2 words in async
-        # mode even when model text generation is correct.
-        # Sync mode uses the default 0.8 (exact match expected).
-        _assert_realtime_accuracy(result, threshold=0.65)
+        _assert_realtime_accuracy(result)
 
     @pytest.mark.advanced_model
     @pytest.mark.omni

tests/entrypoints/openai_api/test_serving_chat_speaker.py

@@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
     serving_chat.chat_template = None
     serving_chat.chat_template_content_format = "string"
     serving_chat.default_chat_template_kwargs = {}
-    serving_chat.online_renderer = mocker.MagicMock()
-    serving_chat.online_renderer.validate_chat_template.return_value = None
+    serving_chat.openai_serving_render = mocker.MagicMock()
+    serving_chat.openai_serving_render.validate_chat_template.return_value = None
     serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
     serving_chat._preprocess_chat = mocker.AsyncMock(
         side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")

tests/entrypoints/openai_api/test_serving_speech.py

@@ -39,8 +39,7 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
-from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext
-from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
     build_fish_voice_clone_prompt_ids,
@@ -974,29 +973,6 @@ class TestTTSMethods:
                 )
             )
 
-    def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
-        """Ming embedding uploads must be 192-dim before replacing stored voices."""
-        speech_server._tts_model_type = "ming_tts"
-        existing = {
-            "name": "bad_emb_voice",
-            "file_path": "/tmp/voice_samples/bad_emb_voice.safetensors",
-            "mime_type": "application/x-safetensors",
-            "embedding_source": "direct",
-            "embedding_dim": 192,
-        }
-        speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()}
-
-        with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"):
-            asyncio.run(
-                speech_server.upload_voice_embedding(
-                    embedding_json=json.dumps([0.0] * 191),
-                    consent="consent",
-                    name="bad_emb_voice",
-                )
-            )
-
-        assert speech_server.uploaded_speakers["bad_emb_voice"] == existing
-
     def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server):
         """Base task without ref_audio or speaker_embedding is rejected."""
         req = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
@@ -1581,86 +1557,6 @@ class TestTTSMethods:
         assert params["x_vector_only_mode"] == [True]
         assert "ref_audio" not in params
 
-    def test_ming_adapter_uploaded_direct_embedding_uses_embedding_not_audio(
-        self, speech_server, mocker: MockerFixture
-    ):
-        """Ming uploaded direct embeddings are loaded into prompt construction."""
-        speech_server.uploaded_speakers = {
-            "emb_voice": {
-                "name": "emb_voice",
-                "file_path": "/tmp/voice_samples/emb_voice.safetensors",
-                "mime_type": "application/x-safetensors",
-                "embedding_source": "direct",
-                "embedding_dim": 192,
-            }
-        }
-        fake_embedding = [0.1] * 192
-        mock_get_emb = mocker.patch.object(
-            speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding
-        )
-        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
-        mock_prompt = mocker.patch.object(
-            speech_server,
-            "_build_ming_dense_prompt",
-            return_value={"additional_information": {"speaker_count": 1}},
-        )
-
-        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
-        req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
-        prepared = asyncio.run(adapter.build(req, [], False))
-
-        mock_get_emb.assert_called_once_with("emb_voice")
-        mock_get_audio.assert_not_called()
-        mock_prompt.assert_called_once_with(req, ref_audio_data=None)
-        prompt_request = mock_prompt.call_args.args[0]
-        assert prompt_request.speaker_embedding == fake_embedding
-        assert len(prompt_request.speaker_embedding) == 192
-        assert prepared.prompt["additional_information"]["speaker_count"] == 1
-        assert prepared.model_type == "ming_tts"
-
-    def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture):
-        """Ming uploaded audio voices still resolve audio and stored ref_text."""
-        speech_server.uploaded_speakers = {
-            "audio_voice": {
-                "name": "audio_voice",
-                "file_path": "/tmp/voice_samples/audio_voice.wav",
-                "mime_type": "audio/wav",
-                "ref_text": "Reference transcript.",
-            }
-        }
-        ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
-        ref_audio_data = ([0.0, 0.1], 16000)
-        fake_embedding = [0.2] * 192
-        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
-        mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
-        mocker.patch.object(
-            speech_server,
-            "_resolve_ref_audio",
-            new=mocker.AsyncMock(return_value=ref_audio_data),
-        )
-        mock_extract = mocker.patch.object(
-            speech_server,
-            "_extract_ming_speaker_embeddings_from_ref_audio",
-            return_value=[fake_embedding],
-        )
-        mock_prompt = mocker.patch.object(
-            speech_server,
-            "_build_ming_dense_prompt",
-            return_value={"additional_information": {"speaker_count": 1}},
-        )
-
-        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
-        req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
-        prepared = asyncio.run(adapter.build(req, [], False))
-
-        mock_get_audio.assert_called_once_with("audio_voice")
-        mock_get_emb.assert_not_called()
-        mock_extract.assert_called_once_with([ref_audio_data])
-        mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
-        assert req.ref_text == "Reference transcript."
-        assert req.speaker_embedding == fake_embedding
-        assert prepared.model_type == "ming_tts"
-
     # ── regression: full flow from issue #1603 ──
 
     def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture):

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -343,7 +343,7 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     )
     mocker.patch(
         "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
+        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]),
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -32,7 +32,7 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
         lambda _model: True,
     )
 
-    stage_configs, _ = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
+    stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
     engine_args = stage_configs[0].engine_args
 
     assert engine_args.model_class_name == "DreamZeroPipeline"

tests/entrypoints/test_serve.py

@@ -9,7 +9,6 @@ import pytest
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
-from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,20 +31,6 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     assert not explicit["async_chunk"]
 
 
-def test_serve_parser_accepts_strategy_config() -> None:
-    """``--strategy-config`` must parse onto the ``strategy_config`` dest and be
-    forwarded as an explicit kwarg so the engine can overlay the strategy."""
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="subcommand")
-    cmd = OmniServeCommand()
-    cmd.subparser_init(subparsers)
-
-    argv = ["serve", "fake-model", "--omni", "--strategy-config", "/tmp/strategy.yaml"]
-    args = parser.parse_args(argv)
-    assert args.strategy_config == "/tmp/strategy.yaml"
-    assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml"
-
-
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",
@@ -95,97 +80,13 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
         run_headless(args)
 
 
-# ---------------------------------------------------------------------------
-# --stage-overrides parsing parity (headless vs standard path)
-# ---------------------------------------------------------------------------
-
-
-def test_parse_stage_overrides_valid_json() -> None:
-    """A valid JSON string is parsed into the nested per-stage dict."""
-    parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
-    assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
-
-
-def test_parse_stage_overrides_none_and_empty_return_none() -> None:
-    """No overrides (None / empty string) resolve to ``None``."""
-    assert parse_stage_overrides(None) is None
-    assert parse_stage_overrides("") is None
-
-
-def test_parse_stage_overrides_empty_dict_returns_none() -> None:
-    """An empty dict is falsy and must resolve to ``None``, locking in parity
-    with the original standard-path ``if stage_overrides_json:`` falsy check."""
-    assert parse_stage_overrides({}) is None
-
-
-def test_parse_stage_overrides_passes_through_non_str() -> None:
-    """An already-parsed mapping is returned unchanged (identity)."""
-    overrides = {"0": {"devices": "0"}}
-    assert parse_stage_overrides(overrides) is overrides
-
-
-def test_parse_stage_overrides_invalid_json_raises() -> None:
-    """Invalid JSON raises ValueError whose message matches the standard path
-    verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the
-    ``Got: <repr>`` suffix echoing the raw input."""
-    bad = "{not valid json}"
-    with pytest.raises(ValueError) as excinfo:
-        parse_stage_overrides(bad)
-    message = str(excinfo.value)
-    assert message.startswith("--stage-overrides is not valid JSON:")
-    assert f"Got: {bad!r}" in message
-
-
-def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
-    """Regression: the headless path must parse ``--stage-overrides`` (a JSON
-    string) and forward the parsed dict to ``load_and_resolve_stage_configs``,
-    mirroring the standard engine path. Previously it was dropped entirely,
-    silently producing a different per-stage device layout."""
-    captured: dict = {}
-
-    def _fake_resolve(*args, **kwargs):
-        captured.update(kwargs)
-        # Return a stage that does NOT match stage_id=0 so run_headless stops
-        # right after the resolver call (we only care about how it was called).
-        return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None)
-
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        side_effect=_fake_resolve,
-    )
-
-    args = _make_headless_args(
-        stage_id=0,
-        strategy_config="/tmp/strategy.yaml",
-        stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
-    )
-    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
-        run_headless(args)
-
-    assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
-    assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
-
-
-def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None:
-    """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the
-    shared ValueError instead of being silently ignored."""
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None),
-    )
-
-    args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}")
-    with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"):
-        run_headless(args)
-
-
 def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
     """Headless looks up its assigned stage_id in the loaded deploy YAML and
     fails fast when the launcher's --stage-id doesn't match any entry."""
     other_stage = SimpleNamespace(stage_id=99)
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [other_stage], None),
+        return_value=("/fake/stages.yaml", [other_stage]),
     )
 
     args = _make_headless_args(stage_id=0)
@@ -233,7 +134,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -309,7 +210,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -368,7 +269,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -448,7 +349,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg], None),
+        return_value=("/fake/stages.yaml", [stage_cfg]),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)

tests/entrypoints/test_stream_finish_reason.py

@@ -144,7 +144,7 @@ def _build_serving_chat():
         engine_client=mock_engine,
         models=models,
         response_role="assistant",
-        online_renderer=mock_render,
+        openai_serving_render=mock_render,
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",

tests/entrypoints/test_utils.py

@@ -318,7 +318,7 @@ class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self):
         """Ensure that dtype survives default stage creation."""
         kwargs = {"dtype": torch.float32}
-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
+        config_path, stage_configs = load_and_resolve_stage_configs(
             model="black-forest-labs/FLUX.2-klein-4B",
             stage_configs_path=None,
             kwargs=kwargs,
@@ -349,10 +349,10 @@ class TestLoadAndResolveStageConfigs:
         ]
         load_stage_configs = mocker.patch(
             "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-            return_value=(returned_stage_configs, None),
+            return_value=returned_stage_configs,
         )
 
-        config_path, stage_configs, _ = load_and_resolve_stage_configs(
+        config_path, stage_configs = load_and_resolve_stage_configs(
             model="dummy-model",
             stage_configs_path=str(deploy_path),
             kwargs={},
@@ -363,7 +363,6 @@ class TestLoadAndResolveStageConfigs:
             base_engine_args={},
             deploy_config_path=str(deploy_path),
             stage_overrides=None,
-            strategy_config_path=None,
         )
         assert config_path == str(deploy_path)
         assert len(stage_configs) == 2

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -18,11 +18,8 @@ import types
 
 import pytest
 import torch
-import torch.nn.functional as F
 from pytest_mock import MockerFixture
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 # Direct file import to avoid vllm_omni.__init__ patch dependencies.
 _MODELS = os.path.join(
     os.path.dirname(__file__),
@@ -545,92 +542,3 @@ class TestCodePredictorWrapperConfig:
             wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
         )
         assert graph_wrapper._prefix_graphs_enabled is True
-
-
-class TestGumbelMaxSampling:
-    """Sanity tests for the shared Gumbel-max sampling helper."""
-
-    @staticmethod
-    def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
-        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
-        return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator)
-
-    def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None:
-        """If one logit is dominant, Gumbel-max must always pick it."""
-        torch.manual_seed(0)
-        bsz, vocab = 4, 64
-        logits = torch.full((bsz, vocab), -1e4)
-        logits[:, 7] = 0.0  # only finite entry
-        out = self._sample(logits)
-        assert out.shape == (bsz, 1)
-        assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long))
-
-    def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None:
-        """Degenerate rows with one surviving finite entry should still return
-        a valid index instead of tripping multinomial-style asserts."""
-        bsz, vocab = 2, 8
-        logits = torch.full((bsz, vocab), float("-inf"))
-        logits[0, 3] = 0.0
-        logits[1, 5] = 0.0
-        out = self._sample(logits)
-        assert out[0].item() == 3
-        assert out[1].item() == 5
-
-    def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None:
-        """Same ``Generator.manual_seed`` should produce identical samples."""
-        torch.manual_seed(0)
-        logits = torch.randn(8, 32)
-        g1 = torch.Generator()
-        g1.manual_seed(123)
-        g2 = torch.Generator()
-        g2.manual_seed(123)
-        g3 = torch.Generator()
-        g3.manual_seed(456)
-        s1 = self._sample(logits, generator=g1)
-        s2 = self._sample(logits, generator=g2)
-        s3 = self._sample(logits, generator=g3)
-        assert torch.equal(s1, s2)
-        # Different seeds must (almost surely) yield a different sequence.
-        assert not torch.equal(s1, s3)
-
-    def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None:
-        """A per-row generator list should preserve row independence."""
-        torch.manual_seed(11)
-        logits = torch.randn(3, 16)
-        seeds = [101, 202, 303]
-        batched_generators = []
-        expected_rows = []
-        for row, seed in enumerate(seeds):
-            row_generator = torch.Generator()
-            row_generator.manual_seed(seed)
-            batched_generators.append(row_generator)
-
-            solo_generator = torch.Generator()
-            solo_generator.manual_seed(seed)
-            expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator))
-
-        batched = self._sample(logits, generator=batched_generators)
-        expected = torch.cat(expected_rows, dim=0)
-        assert torch.equal(batched, expected)
-
-    def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None:
-        logits = torch.randn(2, 8)
-        generators = [torch.Generator()]
-        with pytest.raises(ValueError, match="Expected 2 per-row generators"):
-            self._sample(logits, generator=generators)
-
-    def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None:
-        """Empirical frequency under Gumbel-max should be close to the
-        softmax probability mass.  Loose tolerance because we draw a
-        finite number of samples."""
-        torch.manual_seed(7)
-        vocab = 16
-        logits = torch.randn(vocab) * 2.0
-        target = F.softmax(logits, dim=-1, dtype=torch.float32)
-
-        n_draws = 20000
-        batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous()
-        samples = self._sample(batched_logits).squeeze(-1)
-        empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
-        # L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
-        assert (empirical - target).abs().sum().item() < 0.05

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

@@ -52,10 +52,6 @@ def _make_bare_talker():
     talker._cfm_buffers = None
     talker._last_audio_output_req_ids = []
     talker._batched_fsq_fusion_max_batch = 32
-    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
-    # default (False) to avoid AttributeError when test exercises forward
-    # path that checks this flag.
-    talker._enable_unified_decode_graph = False
     return talker
 
 
@@ -333,27 +329,23 @@ class TestDecodeBatchContract:
 
         for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
             torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
         assert [state.decode_step_count for state in batch_states] == [1, 1]
 
     def test_batch_decode_matches_per_request_state_updates(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
         seq_talker = self._make_decode_talker()
         batch_talker = self._make_decode_talker()
 
         seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 
         for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
 
         batch_talker._finish_decode_batch(
             [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -368,16 +360,12 @@ class TestDecodeBatchContract:
             torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
 
     def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
-        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
-            _DecodeResidualMeta,
-        )
-
         talker = self._make_decode_talker()
         talker._enable_batched_cfm = False
         states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
-            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -8,6 +8,7 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     talker2code2wav_async_chunk,
+    text2flow,
     text2flow_full_payload,
     text2flow_token_only,
 )
@@ -48,13 +49,13 @@ def _transfer_manager(
     )
 
 
-def test_text2flow_token_only_supports_batched_source_outputs():
+def test_text2flow_supports_batched_source_outputs():
     source_outputs = [
         _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
         _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}),
     ]
 
-    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)
 
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
@@ -63,6 +64,16 @@ def test_text2flow_token_only_supports_batched_source_outputs():
     assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
+def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
+    source_outputs = [
+        _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
+    ]
+
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)
+
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+
+
 def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
     source_outputs = [
         _source_output(

tests/model_executor/stage_input_processors/test_indextts2.py

@@ -7,6 +7,7 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.indextts2 import (
     _strip_stop_token,
+    talker2s2mel,
     talker2s2mel_full_payload,
     talker2s2mel_token_only,
 )
@@ -107,6 +108,118 @@ def test_strip_stop_token_all_stop():
     assert lat.shape[1] == 0
 
 
+# ── talker2s2mel ──
+
+
+def test_talker2s2mel_basic():
+    mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
+    latent = torch.randn(4, LATENT_DIM)
+    s_ref = torch.randn(1, 128)
+    ref_mel = torch.randn(1, 80, 200)
+    style = torch.randn(1, 256)
+
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=mel_codes,
+                latent=latent,
+                s_ref=s_ref,
+                ref_mel=ref_mel,
+                style=style,
+            )
+        ]
+    )
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert info["mel_codes"].shape == (1, 3)
+    assert info["latent"].shape == (1, 3, LATENT_DIM)
+    assert info["code_lens"].tolist() == [3]
+    assert "S_ref" in info
+    assert "ref_mel" in info
+    assert "style" in info
+    assert outputs[0]["prompt_token_ids"] == [0]
+
+
+def test_talker2s2mel_skips_unfinished():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_skips_empty_mel():
+    outputs = talker2s2mel(
+        source_outputs=[
+            _make_talker_output(
+                mel_codes=torch.tensor([]),
+                latent=torch.randn(0, LATENT_DIM),
+            )
+        ]
+    )
+
+    assert len(outputs) == 0
+
+
+def test_talker2s2mel_missing_optional_meta():
+    mel_codes = torch.tensor([10, 20])
+    latent = torch.randn(2, LATENT_DIM)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
+
+    assert len(outputs) == 1
+    info = outputs[0]["additional_information"]
+    assert "S_ref" not in info
+    assert "ref_mel" not in info
+    assert "style" not in info
+
+
+def test_talker2s2mel_batch():
+    out0 = _make_talker_output(
+        mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
+        latent=torch.randn(3, LATENT_DIM),
+        s_ref=torch.randn(1, 64),
+    )
+    out1 = _make_talker_output(
+        mel_codes=torch.tensor([5, 6, 7, 8]),
+        latent=torch.randn(4, LATENT_DIM),
+        ref_mel=torch.randn(1, 80, 100),
+    )
+
+    outputs = talker2s2mel(source_outputs=[out0, out1])
+
+    assert len(outputs) == 2
+    assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
+    assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
+
+
+def test_talker2s2mel_tensors_on_cpu():
+    device = torch.device("cpu")
+    mel_codes = torch.tensor([10, 20], device=device)
+    latent = torch.randn(2, LATENT_DIM, device=device)
+    s_ref = torch.randn(1, 64, device=device)
+
+    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
+
+    info = outputs[0]["additional_information"]
+    assert info["latent"].device.type == "cpu"
+    assert info["mel_codes"].device.type == "cpu"
+    assert info["S_ref"].device.type == "cpu"
+
+
+def test_talker2s2mel_no_multimodal_dict():
+    output = SimpleNamespace(
+        finished=True,
+        outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
+    )
+
+    results = talker2s2mel(source_outputs=[output])
+
+    assert len(results) == 0
+
+
 # ── full-payload connector path ──
 
 

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -11,40 +11,59 @@ import torch
 from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip
 from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
     MAX_CODE2WAV_TOKENS,
-    llm2code2wav_full_payload,
+    llm2code2wav,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
-    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
+def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
+    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
+    """Build minimal source_outputs with one talker output carrying codec codes."""
+    talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
+    return [talker_output]
+
+
+def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
+    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through."""
+    # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame:
+    # pad adds one row -> (9, 4) per frame, permuted and flattened.
+    # Pick enough frames to comfortably exceed the cap.
     frames = (MAX_CODE2WAV_TOKENS // 36) + 100
     codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
-    request = SimpleNamespace(request_id="req-long")
 
+    source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
+
+    # Attach caplog's handler directly to the module logger so the warning is
+    # captured regardless of propagation (vllm's logger configuration can
+    # interact badly with caplog.at_level's default root-handler path).
     target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio")
     target_logger.addHandler(caplog.handler)
     prev_level = target_logger.level
     target_logger.setLevel(logging.WARNING)
     try:
-        payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
+        prompts = llm2code2wav(source_outputs)
     finally:
         target_logger.removeHandler(caplog.handler)
         target_logger.setLevel(prev_level)
 
-    assert payload is not None
-    assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
+    assert len(prompts) == 1
+    assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS
+    assert any("truncating" in rec.getMessage() for rec in caplog.records), (
+        f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}"
+    )
 
 
-def test_llm2code2wav_full_payload_short_sequence_unchanged():
+def test_llm2code2wav_short_sequence_unchanged():
+    """Short codec sequences are returned without truncation."""
     codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    request = SimpleNamespace(request_id="req-short")
+    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
 
-    payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
+    prompts = llm2code2wav(source_outputs)
 
-    assert payload is not None
-    assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
+    assert len(prompts) == 1
+    # 4 frames + 1 pad row, flattened col-major → well below the cap
+    assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS
 
 
 def test_llm2code2wav_truncation_boundary_constant_matches_yaml():

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
 
 
 def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
@@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 0, 0]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
 
 
 def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
@@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
-    assert "code_predictor_codes" not in payload
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
@@ -247,6 +247,7 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
     assert payload["ids"]["all"] == [151644, 872, 3]
     assert payload["embed"]["prefill"].device.type == "cpu"
     assert payload["hidden_states"]["output"].device.type == "cpu"
+    assert payload["next_stage_prompt_len"] > 0
     assert payload["embed"]["prefill"].shape[0] == 2
     assert payload["hidden_states"]["output"].shape[0] == 2
 

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -12,9 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
     max_ic_for_chunk_size,
 )
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+    talker2code2wav,
     talker2code2wav_async_chunk,
     talker2code2wav_full_payload,
-    talker2code2wav_token_only,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -368,15 +368,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     assert rid in tm.request_payload
 
 
-def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
-    """``talker2code2wav_token_only`` only allocates placeholder prompt slots.
-
-    After the connector refactor, actual codec flattening (ref prepend +
-    codebook-major layout) is performed by ``talker2code2wav_full_payload`` on
-    the worker data plane.  The orchestrator hook still derives
-    ``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim
-    reference frames once the connector payload arrives.
-    """
+def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
     audio_codes = torch.tensor(
         [
@@ -395,38 +387,12 @@ def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav_token_only(stage.engine_outputs)
+    prompts = talker2code2wav(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
     assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
-    # 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers.
-    assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2))
-
-
-def test_full_payload_prepends_ref_code_and_flattens_codebook_major():
-    """Worker producer is authoritative for ref prepend + codec flatten."""
-    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
-    audio_codes = torch.tensor(
-        [
-            [0, 0, 0, 0],
-            [1, 2, 3, 4],
-            [5, 6, 7, 8],
-        ],
-        dtype=torch.long,
-    )
-    pooling_output = {
-        "codes.audio": audio_codes,
-        "codes.ref": ref_code,
-        "meta.ref_code_len": 2,
-    }
-    request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))
-
-    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
-
-    assert payload is not None
-    assert payload["meta"]["left_context_size"] == 2
-    assert payload["codes"]["audio"].tolist() == [
+    assert prompt["prompt_token_ids"] == [
         9,
         8,
         1,
@@ -467,7 +433,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav_token_only(stage.engine_outputs)
+    prompts = talker2code2wav(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -8,12 +8,61 @@ import pytest
 import torch
 
 from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
+    generator2tokenizer,
     generator2tokenizer_async_chunk,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+# ---- Helpers for generator2tokenizer (non-streaming) ----
+
+
+def _make_stage(engine_outputs):
+    return SimpleNamespace(engine_outputs=engine_outputs)
+
+
+def _make_output(audio_tensors):
+    """Build a single generator output with multimodal_output codes.audio."""
+    return SimpleNamespace(
+        outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
+    )
+
+
+# ---- Tests for generator2tokenizer ----
+
+
+class TestGenerator2Tokenizer:
+    def test_single_output_single_chunk(self):
+        """Single output with one audio tensor produces correct prompt_token_ids."""
+        audio = torch.tensor([10, 20, 30, 40])
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer(stage.engine_outputs)
+
+        assert len(result) == 1
+        assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
+        assert result[0]["multi_modal_data"] is None
+
+    def test_2d_audio_tensors_flattened(self):
+        """2D audio tensors (e.g., multi-codebook frames) are flattened."""
+        audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer(stage.engine_outputs)
+
+        assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
+
+    def test_no_outputs_yet_raises(self):
+        """Stage with engine_outputs=None raises RuntimeError."""
+        stage = _make_stage(engine_outputs=None)
+        with pytest.raises(TypeError):
+            generator2tokenizer(stage.engine_outputs)
+
+
+# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
+
+
 def _req(external_req_id: str, *, finished: bool):
     return SimpleNamespace(
         external_req_id=external_req_id,

tests/model_tests/diffusion/test_alignment.py

@@ -66,6 +66,7 @@ EXCLUDED_MODELS = [
     "AudioXPipeline",
     "HunyuanVideo15Pipeline",
     "HunyuanVideo15ImageToVideoPipeline",
+    "LingBotVideoPipeline",
     "MagiHumanPipeline",
     "OmniVoicePipeline",
     "OmniVoice",

tests/platforms/test_cuda_ir_op_priority.py

@@ -1,38 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
-
-from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
-    return VllmConfig(
-        device_config=DeviceConfig(device="cpu"),
-        compilation_config=CompilationConfig(backend=backend, mode=mode),
-    )
-
-
-@pytest.mark.parametrize(
-    "mode",
-    [CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
-)
-def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
-    """Regression for #4964: inductor-active configs must not switch to native-only."""
-    priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
-    assert priority.rms_norm == ["vllm_c", "native"]
-    assert priority.fused_add_rms_norm == ["vllm_c", "native"]
-
-
-def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
-    import vllm.envs as envs
-
-    monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
-    priority = CudaOmniPlatform.get_default_ir_op_priority(
-        _vllm_config(mode=CompilationMode.VLLM_COMPILE),
-    )
-    assert priority.rms_norm == ["oink", "vllm_c", "native"]
-    assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]

tests/test_config_factory.py

@@ -1166,10 +1166,7 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
-        # thinker2talker was removed: sync_process_input_func always wins.
-        assert s.custom_process_input_func is None
-        assert s.sync_process_input_func is not None
-        assert s.sync_process_input_func.endswith("thinker2talker_token_only")
+        assert s.custom_process_input_func is not None
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
@@ -1183,7 +1180,7 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
-        assert s.engine_output_type == "audio"
+        assert s.custom_process_input_func is not None
 
 
 class TestQwen2_5OmniPipeline:
@@ -1721,7 +1718,7 @@ class TestAuraOmniDeploy:
     def test_aura_omni_deploy_resolves_four_native_stages(self):
         pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni")
 
-        stages, _ = StageConfigFactory._create_from_registry(
+        stages = StageConfigFactory._create_from_registry(
             "qwen3_tts",
             pipeline_cfg,
             cli_overrides={},
@@ -1943,12 +1940,11 @@ class TestSentinelDefaultPrecedence:
             model_type,
             Q3_OMNI_ALL_STAGES_HF_CONFIG,
         )
-        stages, _ = StageConfigFactory._create_from_registry(
+        return StageConfigFactory._create_from_registry(
             "qwen3_omni_moe",
             pipeline_cfg,
             cli_overrides=cli_overrides,
         )
-        return stages
 
     def test_typed_kwarg_overrides_yaml(self):
         stages = self._stages({"max_num_seqs": 999})

tests/utils/test_tracking_parser.py

@@ -591,7 +591,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
     ns = p.parse_args(["--max-num-seqs", "999"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages, _ = StageConfigFactory._create_from_registry(
+    stages = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -608,7 +608,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages):
     ns = p.parse_args([])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages, _ = StageConfigFactory._create_from_registry(
+    stages = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -630,7 +630,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages):
         ns = p.parse_args(["--config", "fake.yaml"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages, _ = StageConfigFactory._create_from_registry(
+    stages = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -648,7 +648,7 @@ def test_per_stage_override_routes_correctly(mock_stages):
     ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages, _ = StageConfigFactory._create_from_registry(
+    stages = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -670,7 +670,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages):
     ns = p.parse_args(["--enforce-eager"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages, _ = StageConfigFactory._create_from_registry(
+    stages = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,

vllm_omni/config/composable_parallel/__init__.py

@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Composable parallel strategies for vLLM-Omni.
-
-A small, runtime-agnostic vocabulary for declaring a stage's parallel layout
-(:class:`StrategySpec` + mesh/routing/aggregation patterns), plus the logic to
-translate a per-role stack of specs into concrete engine sizing and overlay it
-onto merged stage configs.
-
-The spec/pattern *types* are deliberately pure data with no engine imports; the
-translation and apply logic lives alongside them and consumes vLLM-Omni's own
-config objects.
-"""
-
-from vllm_omni.config.composable_parallel.aggregation import (
-    AGGREGATION_PATTERN_VARIANTS,
-    AggregationConflictError,
-    AggregationPattern,
-    AllGather,
-    Combine,
-    FanInByStage,
-    GatherDim,
-    StitchPipeline,
-    StitchSpatial,
-    TakeRank,
-    Union,
-)
-from vllm_omni.config.composable_parallel.apply import (
-    StrategyApplyError,
-    StrategyApplyResult,
-    apply_strategy_specs,
-    check_device_layout,
-)
-from vllm_omni.config.composable_parallel.routing import (
-    ROUTING_PATTERN_VARIANTS,
-    Broadcast,
-    DuplicateWithCondTag,
-    PartitionByHash,
-    PipelineMicrobatch,
-    RouteByExpert,
-    RouteByStage,
-    RoutingKeyError,
-    RoutingPattern,
-    ShardSequence,
-    ShardSpatial,
-)
-from vllm_omni.config.composable_parallel.spec import (
-    MESH_AXIS_KINDS,
-    KernelSpec,
-    LayerHookSpec,
-    MeshAxisKind,
-    MeshAxisSpec,
-    SpecMergeConflictError,
-    StrategySpec,
-)
-from vllm_omni.config.composable_parallel.translator import (
-    AxisTranslationError,
-    OmniParallelConfig,
-    translate_strategy_stack,
-)
-
-__all__ = [
-    # spec types
-    "MeshAxisKind",
-    "MESH_AXIS_KINDS",
-    "MeshAxisSpec",
-    "LayerHookSpec",
-    "KernelSpec",
-    "StrategySpec",
-    "SpecMergeConflictError",
-    # routing patterns
-    "RoutingPattern",
-    "RoutingKeyError",
-    "Broadcast",
-    "PartitionByHash",
-    "ShardSequence",
-    "ShardSpatial",
-    "RouteByExpert",
-    "PipelineMicrobatch",
-    "DuplicateWithCondTag",
-    "RouteByStage",
-    "ROUTING_PATTERN_VARIANTS",
-    # aggregation patterns
-    "AggregationPattern",
-    "AggregationConflictError",
-    "TakeRank",
-    "Union",
-    "GatherDim",
-    "AllGather",
-    "StitchSpatial",
-    "StitchPipeline",
-    "Combine",
-    "FanInByStage",
-    "AGGREGATION_PATTERN_VARIANTS",
-    # translation
-    "OmniParallelConfig",
-    "translate_strategy_stack",
-    "AxisTranslationError",
-    # apply
-    "apply_strategy_specs",
-    "StrategyApplyResult",
-    "check_device_layout",
-    "StrategyApplyError",
-]

vllm_omni/config/composable_parallel/aggregation.py

@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Closed aggregation-pattern hierarchy for declarative parallel strategies.
-
-An aggregation pattern describes how the per-worker results of one mesh axis
-are recombined into a single logical result (e.g. take one tensor-parallel
-rank's output, union disjoint data-parallel shards). The hierarchy is *closed*:
-subclasses may only be defined in this module.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass
-
-
-@dataclass(frozen=True)
-class AggregationPattern:
-    """Base class for closed aggregation patterns. Subclass only within this module."""
-
-    def __init_subclass__(cls, **kwargs: object) -> None:
-        super().__init_subclass__(**kwargs)
-        if cls.__module__ != __name__:
-            raise TypeError("AggregationPattern is closed to external subclassing")
-
-    def pattern_kind(self) -> str:
-        raise NotImplementedError
-
-
-class AggregationConflictError(ValueError):
-    """Raised when two aggregation inputs conflict for the same request id."""
-
-
-@dataclass(frozen=True)
-class TakeRank(AggregationPattern):
-    """Select one rank's result (TP — all ranks agree)."""
-
-    rank: int = 0
-
-    def __post_init__(self) -> None:
-        if not isinstance(self.rank, int) or isinstance(self.rank, bool):
-            raise TypeError(f"TakeRank.rank must be int, got {type(self.rank).__name__}")
-
-    def pattern_kind(self) -> str:
-        return "take_rank"
-
-
-@dataclass(frozen=True)
-class Union(AggregationPattern):
-    """Union of disjoint replica results (DP, HSDP)."""
-
-    def pattern_kind(self) -> str:
-        return "union"
-
-
-@dataclass(frozen=True)
-class GatherDim(AggregationPattern):
-    """Gather along a dimension and materialize on one designated rank.
-
-    Default designated rank is axis-rank 0 unless adapter overrides.
-    """
-
-    dim: int = 1
-
-    def pattern_kind(self) -> str:
-        return "gather_dim"
-
-
-@dataclass(frozen=True)
-class AllGather(AggregationPattern):
-    """All ranks materialize the same aggregate result across the axis."""
-
-    dim: int | None = None
-
-    def pattern_kind(self) -> str:
-        return "all_gather"
-
-
-@dataclass(frozen=True)
-class StitchSpatial(AggregationPattern):
-    """Stitch spatial patches (VAE-PP)."""
-
-    grid: tuple[int, int] = (1, 1)
-
-    def pattern_kind(self) -> str:
-        return "stitch_spatial"
-
-
-@dataclass(frozen=True)
-class StitchPipeline(AggregationPattern):
-    """Stitch pipeline stage outputs (PP)."""
-
-    def pattern_kind(self) -> str:
-        return "stitch_pipeline"
-
-
-@dataclass(frozen=True)
-class Combine(AggregationPattern):
-    """Custom reducer combination (CFG). Adapter resolves reducer_id."""
-
-    reducer_id: str = "default"
-
-    def pattern_kind(self) -> str:
-        return "combine"
-
-
-@dataclass(frozen=True)
-class FanInByStage(AggregationPattern):
-    """Fan-in results routed by stage (stage replicas)."""
-
-    def pattern_kind(self) -> str:
-        return "fan_in_by_stage"
-
-
-AGGREGATION_PATTERN_VARIANTS: tuple[type[AggregationPattern], ...] = (
-    TakeRank,
-    Union,
-    GatherDim,
-    AllGather,
-    StitchSpatial,
-    StitchPipeline,
-    Combine,
-    FanInByStage,
-)
-
-ReducerFn = Callable[..., object]

vllm_omni/config/composable_parallel/apply.py

@@ -1,264 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Apply declarative parallel strategies onto merged stage configs.
-
-This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses
-the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this
-module then overlays the engine sizing derived from a per-role
-``StrategySpec`` stack onto those stages.
-
-Design rules (one writer per axis):
-
-* A strategy only writes the axes it actually declares. If a stack has no ``tp``
-  axis, this module never touches ``tensor_parallel_size`` — it does not force a
-  default of 1.
-* **Conflict-on-explicit.** If a knob was already set explicitly by the deploy
-  YAML and the strategy derives a different value, we raise rather than silently
-  override. Equal values are a no-op; unset values are filled.
-* **Pre-spawn device check.** When a stage declares ``devices``, the device
-  count must be consistent with the engine world size (``tp * dp * pp``) and the
-  replica count, so misconfigurations fail at config time instead of at spawn.
-
-``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at
-orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a
-stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult`
-for the caller to pass to the orchestrator; it is never silently written into a
-stage where it would be ignored.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping, Sequence
-from dataclasses import dataclass, field
-from typing import Any, NoReturn
-
-from vllm.logger import init_logger
-
-from vllm_omni.config.composable_parallel.spec import StrategySpec
-from vllm_omni.config.composable_parallel.translator import (
-    OmniParallelConfig,
-    translate_strategy_stack,
-)
-
-logger = init_logger(__name__)
-
-# Role key -> the StageConfig fields each declared axis writes.
-_ENGINE_FIELD_BY_KIND: dict[str, str] = {
-    "tp": "tensor_parallel_size",
-    "dp": "data_parallel_size",
-    "pp": "pipeline_parallel_size",
-}
-
-RoleKey = str
-
-
-class StrategyApplyError(ValueError):
-    """Error for applying a strategy onto stage configs.
-
-    A single error type (rather than per-cause subclasses) keeps the public
-    surface small and consistent with the rest of the codebase; the specific
-    cause — an unmatched role, a conflict with an explicit deploy value, or a
-    device-count mismatch — is in the message and is logged before the raise so
-    it stays visible even when the type is unavailable to a caller (e.g. across a
-    server boundary).
-    """
-
-
-def _fail(msg: str) -> NoReturn:
-    """Log and raise :class:`StrategyApplyError`."""
-    logger.error("[composable_parallel] %s", msg)
-    raise StrategyApplyError(msg)
-
-
-@dataclass
-class StrategyApplyResult:
-    """Outcome of applying strategies to a stage list.
-
-    ``stages`` is the same list passed in (mutated in place and returned for
-    convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy
-    derived from any stage_replica axes (``None`` if none declared one); the
-    caller wires it into the orchestrator, since omni does not read it from
-    stage configs.
-
-    ``per_role_config`` is keyed by the role key the caller supplied — a
-    ``model_stage`` name (str). ``per_stage_config`` is the same information
-    keyed by the resolved integer ``stage_id`` — handy for re-validating a
-    stage after later layers (e.g. CLI overrides) have run.
-    """
-
-    stages: list[Any]
-    omni_lb_policy: str | None = None
-    per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict)
-    per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict)
-
-
-def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any:
-    """Find the single stage matching ``key`` by ``model_stage`` name.
-
-    Role keys are ``model_stage`` names (str) — the same human-readable names
-    that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``).
-    """
-    matches = [s for s in stages if getattr(s, "model_stage", None) == key]
-    descriptor = f"model_stage={key!r}"
-
-    if not matches:
-        available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages)
-        _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}")
-    if len(matches) > 1:
-        _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages")
-    return matches[0]
-
-
-def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None:
-    """Write ``field_name`` with conflict-on-explicit semantics.
-
-    ``engine_args`` is the raw mapping loaded straight from the deploy YAML
-    (before the engine-args dataclass is constructed), so a key is *present* only
-    when the YAML set it explicitly. We therefore key off membership rather than
-    ``.get``: an explicit ``None`` in the YAML is a real, present value and must
-    not be conflated with a key miss (which ``.get`` would report as ``None``).
-    """
-    if field_name in engine_args and engine_args[field_name] != value:
-        existing = engine_args[field_name]
-        _fail(
-            f"role {role!r}: strategy derives {field_name}={value} but the deploy config already "
-            f"set {field_name}={existing}. Remove one of them (the strategy is the single writer "
-            "for the axes it declares)."
-        )
-    engine_args[field_name] = value
-
-
-def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None:
-    """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset)."""
-    existing = runtime.get("num_replicas", 1)
-    if existing not in (1, value):
-        _fail(
-            f"role {role!r}: strategy derives num_replicas={value} but the deploy config already "
-            f"set num_replicas={existing}. Remove one of them."
-        )
-    runtime["num_replicas"] = value
-
-
-def _parse_device_count(devices: Any) -> int | None:
-    """Return the number of device ids in a ``devices`` value, or None."""
-    if devices is None:
-        return None
-    if isinstance(devices, (list, tuple)):
-        return len([d for d in devices])
-    text = str(devices).strip()
-    if not text:
-        return None
-    return len([d for d in text.split(",") if d.strip()])
-
-
-def check_device_layout(
-    devices: Any,
-    *,
-    tensor_parallel_size: int,
-    data_parallel_size: int,
-    pipeline_parallel_size: int,
-    num_replicas: int,
-    role: RoleKey,
-) -> None:
-    """Validate a stage's device count against its world size (× replicas in pool mode).
-
-    Raises :class:`StrategyApplyError` when the declared ``devices`` count is
-    neither the per-replica world size nor the full ``num_replicas`` pool.
-    Exposed publicly so the resolved (post-CLI) layout can be re-checked after
-    later override layers, not just the strategy-derived snapshot.
-    """
-    count = _parse_device_count(devices)
-    if count is None:
-        return
-    world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size)
-    replicas = int(num_replicas or 1)
-    # Two valid shapes: a single per-replica template (== world) or the full
-    # pool across replicas (== replicas * world). This mirrors omni's own
-    # per-replica device splitter.
-    valid = {world, replicas * world}
-    if count not in valid:
-        _fail(
-            f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} "
-            f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). "
-            f"Provide either {world} (per-replica template) or {replicas * world} "
-            f"(num_replicas={replicas} pool)."
-        )
-
-
-def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None:
-    """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`)."""
-    check_device_layout(
-        runtime.get("devices"),
-        tensor_parallel_size=cfg.tensor_parallel_size,
-        data_parallel_size=cfg.data_parallel_size,
-        pipeline_parallel_size=cfg.pipeline_parallel_size,
-        num_replicas=int(runtime.get("num_replicas", 1) or 1),
-        role=role,
-    )
-
-
-def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None:
-    engine_args = stage.yaml_engine_args
-    runtime = stage.yaml_runtime
-
-    declared = set(cfg.l1_owners.keys())
-    for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
-        if kind in declared:
-            _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role)
-    if "ep" in declared and cfg.enable_expert_parallel:
-        _set_explicit(engine_args, "enable_expert_parallel", True, role=role)
-    if "stage_replica" in declared:
-        _set_num_replicas(runtime, cfg.stage_replica_size, role=role)
-
-    _check_devices(runtime, cfg, role=role)
-
-
-def apply_strategy_specs(
-    stages: list[Any],
-    strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]],
-) -> StrategyApplyResult:
-    """Overlay per-role strategy specs onto a merged stage list.
-
-    Args:
-        stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``.
-        strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``)
-            to that stage's stack of ``StrategySpec`` (one per declared mesh
-            axis). Role keys are ``model_stage`` names (str).
-
-    Returns:
-        A :class:`StrategyApplyResult` with the mutated stages and the derived
-        pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one).
-
-    Raises:
-        StrategyApplyError: a role matched zero or multiple stages, a derived
-            value conflicts with an explicit deploy value (including two roles
-            deriving different ``omni_lb_policy`` values), or a stage's device
-            count is inconsistent with its world size.
-        AxisTranslationError: the spec stack is invalid/unsupported.
-        NotImplementedError: the spec stack requests routing not built yet.
-    """
-    result = StrategyApplyResult(stages=stages)
-    lb_policy: str | None = None
-    lb_owner: RoleKey | None = None
-
-    for key, specs in strategy_specs.items():
-        stage = _resolve_stage(stages, key)
-        cfg = translate_strategy_stack(specs)
-        _apply_to_stage(stage, cfg, role=key)
-        result.per_role_config[key] = cfg
-        stage_id = getattr(stage, "stage_id", None)
-        if stage_id is not None:
-            result.per_stage_config[int(stage_id)] = cfg
-
-        if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None:
-            if lb_policy is not None and lb_policy != cfg.omni_lb_policy:
-                _fail(
-                    f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role "
-                    f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is "
-                    "pipeline-wide; only one value is allowed."
-                )
-            lb_policy = cfg.omni_lb_policy
-            lb_owner = key
-
-    result.omni_lb_policy = lb_policy
-    return result

vllm_omni/config/composable_parallel/routing.py

@@ -1,138 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Closed routing-pattern hierarchy for declarative parallel strategies.
-
-A routing pattern describes how a request batch is distributed across the
-workers of one mesh axis (e.g. broadcast to every tensor-parallel rank, hash
-partition across data-parallel ranks). The hierarchy is *closed*: subclasses
-may only be defined in this module, so the set of patterns is a fixed,
-exhaustively matchable contract.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Literal
-
-
-@dataclass(frozen=True)
-class RoutingPattern:
-    """Base class for closed routing patterns. Subclass only within this module."""
-
-    def __init_subclass__(cls, **kwargs: object) -> None:
-        super().__init_subclass__(**kwargs)
-        if cls.__module__ != __name__:
-            raise TypeError("RoutingPattern is closed to external subclassing")
-
-    def pattern_kind(self) -> str:
-        raise NotImplementedError
-
-
-class RoutingKeyError(KeyError):
-    """Raised when a routing key is missing from a request."""
-
-    def __init__(self, request_id: object, missing_key: str) -> None:
-        self.request_id = request_id
-        self.missing_key = missing_key
-        super().__init__(f"Missing routing key {missing_key!r} for request_id={request_id!r}")
-
-
-@dataclass(frozen=True)
-class Broadcast(RoutingPattern):
-    """Every worker in the group sees the whole batch (TP, dense EP)."""
-
-    def pattern_kind(self) -> str:
-        return "broadcast"
-
-
-@dataclass(frozen=True)
-class PartitionByHash(RoutingPattern):
-    """Stable hash partition by request key (DP, HSDP)."""
-
-    key: str = "request_id"
-
-    def pattern_kind(self) -> str:
-        return "partition_by_hash"
-
-
-@dataclass(frozen=True)
-class ShardSequence(RoutingPattern):
-    """Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP)."""
-
-    dim: int = 1
-
-    def pattern_kind(self) -> str:
-        return "shard_sequence"
-
-
-@dataclass(frozen=True)
-class ShardSpatial(RoutingPattern):
-    """Shard along a spatial grid (VAE-PP)."""
-
-    grid: tuple[int, int] = (1, 1)
-
-    def pattern_kind(self) -> str:
-        return "shard_spatial"
-
-
-@dataclass(frozen=True)
-class RouteByExpert(RoutingPattern):
-    """Sparse MoE expert routing."""
-
-    router_id: str = "default"
-
-    def pattern_kind(self) -> str:
-        return "route_by_expert"
-
-
-@dataclass(frozen=True)
-class PipelineMicrobatch(RoutingPattern):
-    """Pipeline parallelism microbatch routing (PP)."""
-
-    microbatch_size: int = 1
-
-    def pattern_kind(self) -> str:
-        return "pipeline_microbatch"
-
-
-@dataclass(frozen=True)
-class DuplicateWithCondTag(RoutingPattern):
-    """Duplicate batch with a conditional tag (CFG)."""
-
-    flag_name: str = "guidance_branch"
-
-    def pattern_kind(self) -> str:
-        return "duplicate_with_cond_tag"
-
-
-@dataclass(frozen=True)
-class RouteByStage(RoutingPattern):
-    """Route requests by pipeline stage policy (stage replicas).
-
-    The ``random``/``round_robin``/``least_queue`` policies describe stateless
-    load balancing that a runtime (e.g. an engine's built-in DP load balancer,
-    or the omni StagePool balancer) can realize. ``hash`` describes
-    deterministic, key-stable routing that a load balancer cannot honor on its
-    own (see ``PartitionByHash``).
-    """
-
-    routing_policy: Literal["random", "round_robin", "least_queue", "hash"] = "round_robin"
-
-    def pattern_kind(self) -> str:
-        return "route_by_stage"
-
-
-ROUTING_PATTERN_VARIANTS: tuple[type[RoutingPattern], ...] = (
-    Broadcast,
-    PartitionByHash,
-    ShardSequence,
-    ShardSpatial,
-    RouteByExpert,
-    PipelineMicrobatch,
-    DuplicateWithCondTag,
-    RouteByStage,
-)
-
-# Optional callable slot for Combine-style custom routing reducers in adapters.
-RouterFn = Callable[..., object]

vllm_omni/config/composable_parallel/spec.py

@@ -1,136 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Core declarative types for parallel-strategy specification.
-
-A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes
-(kind + degree), how a batch is routed across that axis, and how per-worker
-results are aggregated back. A *stack* of specs (one per axis) fully describes a
-stage's parallel layout in a runtime-agnostic, data-only form that can be
-translated into concrete engine sizing.
-
-The mesh-axis kinds below enumerate every parallelism dimension the contract can
-describe. Only a subset is wired end-to-end today (see ``translator.py``):
-
-* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``,
-  ``stage_replica``.
-* **Reserved** (declared in the type system but not yet translatable — the
-  translator raises ``AxisTranslationError`` for them): ``sp_ulysses``,
-  ``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``.
-
-Reserved kinds fail fast at translation time rather than silently doing nothing,
-so declaring one is an explicit "not yet" rather than a no-op.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Any, Literal
-
-from vllm_omni.config.composable_parallel.aggregation import AggregationPattern
-from vllm_omni.config.composable_parallel.routing import RoutingPattern
-
-
-class MeshAxisKind(str, Enum):
-    """Enumerates every parallelism dimension the strategy contract can describe.
-
-    Subclassing ``str`` keeps existing string comparisons and dict keys working
-    (``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type
-    refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than
-    a silently-wrong bare string.
-    """
-
-    TP = "tp"
-    DP = "dp"
-    PP = "pp"
-    EP = "ep"
-    SP_ULYSSES = "sp_ulysses"
-    SP_RING = "sp_ring"
-    CFG = "cfg"
-    VAE_PP = "vae_pp"
-    HSDP = "hsdp"
-    STAGE_PP = "stage_pp"
-    STAGE_REPLICA = "stage_replica"
-    CP = "cp"
-
-
-# Bare string values, kept as a plain tuple so membership checks and error
-# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum
-# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks.
-MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind)
-
-HookCategory = Literal[
-    "linear",
-    "attention_pre",
-    "attention_post",
-    "ffn_pre",
-    "ffn_post",
-    "moe_dispatch",
-    "other",
-]
-
-HOOK_CATEGORY_ORDER: dict[HookCategory, int] = {
-    "linear": 0,
-    "attention_pre": 1,
-    "attention_post": 2,
-    "ffn_pre": 3,
-    "ffn_post": 4,
-    "moe_dispatch": 5,
-    "other": 6,
-}
-
-
-class SpecMergeConflictError(ValueError):
-    """Raised when merged hook/kernel specs have conflicting declarations."""
-
-
-@dataclass(frozen=True)
-class MeshAxisSpec:
-    """Declares one axis of a process mesh (kind + size)."""
-
-    kind: MeshAxisKind
-    size: int
-
-    def __post_init__(self) -> None:
-        if self.kind not in MESH_AXIS_KINDS:
-            raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}")
-        if self.size <= 0:
-            raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}")
-
-
-@dataclass(frozen=True)
-class LayerHookSpec:
-    """L2 hook slot referenced by StrategySpec; resolved by the model walker."""
-
-    hook_id: str
-    target: str | None = None
-    category: HookCategory = "other"
-    priority: int = 0
-    axis_index: int = 0
-
-
-@dataclass(frozen=True)
-class KernelSpec:
-    """L3 kernel slot referenced by StrategySpec; resolved by the kernel registry."""
-
-    kernel_id: str
-    target: str | None = None
-    category: HookCategory = "other"
-    priority: int = 0
-    axis_index: int = 0
-    group_axis_kind: MeshAxisKind | None = None
-    requires_collective: bool = False
-
-
-@dataclass(frozen=True)
-class StrategySpec:
-    """Declarative contract for one parallelism scheme (data + hook/kernel slots)."""
-
-    name: str
-    mesh_axis: MeshAxisSpec
-    routing: RoutingPattern
-    aggregation: AggregationPattern
-    layer_hook_specs: tuple[LayerHookSpec, ...] = ()
-    kernel_specs: tuple[KernelSpec, ...] = ()
-    shard_extension: Mapping[str, Any] = field(default_factory=dict)

vllm_omni/config/composable_parallel/strategy_loader.py

@@ -1,149 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Load per-role parallel strategies from a strategy file.
-
-The strategy file is a small, optional companion to the deploy config. It maps a
-logical role (``model_stage``, e.g. ``thinker``) to a list of axis declarations,
-each of which becomes one :class:`StrategySpec`. YAML is the only format wired up
-today, but the module/symbol names stay format-agnostic so a different source can
-be added without renaming the public surface::
-
-    strategies:
-      thinker:
-        - axis: tp
-          size: 2
-      talker:
-        - axis: stage_replica
-          size: 2
-          routing: round_robin   # random | round_robin | least_queue
-      code2wav:
-        - axis: stage_replica
-          size: 2
-          routing: round_robin
-
-Per-axis fields:
-
-* ``axis`` (required): mesh-axis kind (``tp``, ``dp``, ``pp``, ``ep``,
-  ``stage_replica`` are translatable today).
-* ``size`` (required): axis degree (> 0).
-* ``routing`` (optional): stateless policy for ``dp`` / ``stage_replica`` axes
-  (``random`` / ``round_robin`` / ``least_queue``). Ignored for axes whose
-  routing is fixed (``tp`` broadcast, ``pp`` microbatch).
-* ``l1_owner`` (optional): ``engine`` or ``delegated``; overrides the default
-  owner for the axis (recorded in ``shard_extension``).
-* ``name`` (optional): a label for the spec; defaults to the axis kind.
-
-Routing/aggregation patterns are filled in from sensible per-kind defaults so
-the file stays compact; the translator validates the combination.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from typing import Any
-
-from vllm_omni.config.composable_parallel.aggregation import (
-    AggregationPattern,
-    FanInByStage,
-    StitchPipeline,
-    TakeRank,
-    Union,
-)
-from vllm_omni.config.composable_parallel.routing import (
-    Broadcast,
-    PipelineMicrobatch,
-    RouteByStage,
-    RoutingPattern,
-)
-from vllm_omni.config.composable_parallel.spec import MeshAxisSpec, StrategySpec
-from vllm_omni.config.yaml_util import load_yaml_config, to_dict
-
-# Per-kind defaults for routing + aggregation, so the file only needs axis+size.
-_ROUTING_POLICY_KINDS = frozenset({"dp", "stage_replica"})
-
-
-class StrategyLoadError(ValueError):
-    """Raised when a strategy file is malformed."""
-
-
-def _default_routing(kind: str, routing_policy: str | None) -> RoutingPattern:
-    if kind in _ROUTING_POLICY_KINDS:
-        return RouteByStage(routing_policy=routing_policy or "round_robin")  # type: ignore[arg-type]
-    if kind == "pp":
-        return PipelineMicrobatch()
-    # tp, ep, and any other kind default to broadcast (the translator rejects
-    # kinds it cannot handle before it ever inspects routing).
-    return Broadcast()
-
-
-def _default_aggregation(kind: str) -> AggregationPattern:
-    if kind == "stage_replica":
-        return FanInByStage()
-    if kind == "pp":
-        return StitchPipeline()
-    if kind in ("dp", "ep"):
-        return Union()
-    return TakeRank()
-
-
-def _build_spec(role: str, entry: Mapping[str, Any]) -> StrategySpec:
-    if "axis" not in entry:
-        raise StrategyLoadError(f"role {role!r}: each strategy entry needs an 'axis' field")
-    if "size" not in entry:
-        raise StrategyLoadError(f"role {role!r}: strategy entry for axis {entry['axis']!r} needs a 'size'")
-
-    kind = str(entry["axis"])
-    try:
-        size = int(entry["size"])
-    except (TypeError, ValueError) as exc:
-        raise StrategyLoadError(f"role {role!r}: axis {kind!r} size must be an integer, got {entry['size']!r}") from exc
-
-    routing_policy = entry.get("routing")
-    if routing_policy is not None and kind not in _ROUTING_POLICY_KINDS:
-        raise StrategyLoadError(
-            f"role {role!r}: axis {kind!r} does not accept a 'routing' policy (only {sorted(_ROUTING_POLICY_KINDS)} do)"
-        )
-
-    shard_extension: dict[str, Any] = {}
-    l1_owner = entry.get("l1_owner")
-    if l1_owner is not None:
-        shard_extension["l1_owner"] = str(l1_owner)
-
-    return StrategySpec(
-        name=str(entry.get("name", kind)),
-        mesh_axis=MeshAxisSpec(kind=kind, size=size),  # type: ignore[arg-type]
-        routing=_default_routing(kind, routing_policy),
-        aggregation=_default_aggregation(kind),
-        shard_extension=shard_extension,
-    )
-
-
-def parse_strategy_specs(data: Mapping[str, Any]) -> dict[str, list[StrategySpec]]:
-    """Parse an already-loaded strategy mapping into per-role spec stacks."""
-    strategies = data.get("strategies", data)
-    if not isinstance(strategies, Mapping):
-        raise StrategyLoadError("strategy file must contain a 'strategies' mapping of role -> list of axes")
-
-    result: dict[str, list[StrategySpec]] = {}
-    for role, entries in strategies.items():
-        if not isinstance(entries, (list, tuple)):
-            raise StrategyLoadError(
-                f"role {role!r}: expected a list of axis declarations, got {type(entries).__name__}"
-            )
-        specs: list[StrategySpec] = []
-        for entry in entries:
-            if not isinstance(entry, Mapping):
-                raise StrategyLoadError(
-                    f"role {role!r}: each axis declaration must be a mapping, got {type(entry).__name__}"
-                )
-            specs.append(_build_spec(str(role), dict(entry)))
-        result[str(role)] = specs
-    return result
-
-
-def load_strategy_specs(path: str) -> dict[str, list[StrategySpec]]:
-    """Load and parse a strategy file into per-role spec stacks."""
-    data = to_dict(load_yaml_config(path))
-    if not isinstance(data, Mapping):
-        raise StrategyLoadError(f"strategy file {path!r} did not parse to a mapping")
-    return parse_strategy_specs(data)

vllm_omni/config/composable_parallel/translator.py

@@ -1,378 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Translate a ``StrategySpec`` stack into vLLM-Omni parallel sizing.
-
-Read a stack of strategy specs (one per mesh axis) and work out the
-tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage
-replicas, plus who owns each axis's request routing. The result is a plain,
-CPU-computable sizing struct (:class:`OmniParallelConfig`) keyed by the real
-``OmniEngineArgs`` / ``EngineArgs`` field names, so the deploy layer can splat it
-onto a stage's engine args.
-
-The distinction this module enforces — engine data parallelism (a true vLLM
-intra-engine world dimension) vs. omni stage replicas (independent engines fanned
-out by omni's coordinator) — is documented on the axis validators that depend on
-it (see :func:`_validate_dp` and :func:`_stage_replica_lb_policy`). Routing we do
-not support yet (key-stable / affinity routing) raises ``NotImplementedError``;
-any other invalid spec raises :class:`AxisTranslationError`.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping, Sequence
-from dataclasses import dataclass, field
-from typing import Literal, NoReturn, cast, get_args
-
-from vllm.logger import init_logger
-
-from vllm_omni.config.composable_parallel.routing import (
-    Broadcast,
-    PartitionByHash,
-    PipelineMicrobatch,
-    RouteByStage,
-    RoutingPattern,
-)
-from vllm_omni.config.composable_parallel.spec import MeshAxisKind, StrategySpec
-
-logger = init_logger(__name__)
-
-# Who owns an axis's request routing. A closed string type so an unexpected raw
-# value is caught statically and at runtime (``_VALID_L1_OWNERS`` is derived from
-# it, keeping one source of truth) rather than silently flowing through.
-L1Owner = Literal["delegated", "engine"]
-
-# DP, TP, PP, (dense) EP and stage_replica translate today. DP/TP/PP are true
-# world dimensions vLLM realizes intra-engine; dense EP is a flag over the
-# existing TP*DP ranks; stage_replica is an omni-coordinator-level fan-out of
-# independent engine replicas (NOT a vLLM world dimension). The remaining kinds
-# (sequence/context parallel, CFG, VAE pipelines, ...) need per-layer hooks or
-# custom collectives and arrive in later stages.
-_SUPPORTED_KINDS: tuple[MeshAxisKind, ...] = ("dp", "tp", "pp", "ep", "stage_replica")
-
-# Which EngineArgs field each axis kind *sizes*. EP and stage_replica are
-# deliberately absent: EP is not an independent world dimension but an
-# ``enable_expert_parallel`` flag whose degree equals tp*dp; stage_replica is a
-# per-stage deploy ``num_replicas`` count, not an intra-engine world dimension.
-_AXIS_TO_ENGINE_FIELD: dict[MeshAxisKind, str] = {
-    "tp": "tensor_parallel_size",
-    "dp": "data_parallel_size",
-    "pp": "pipeline_parallel_size",
-}
-
-# Default L1 owner per axis kind. ``stage_replica`` is the only delegated axis:
-# omni's coordinator/StagePool owns routing across replicas. ``dp`` is engine
-# data parallelism, realized by vLLM's own intra-engine DP load balancer.
-_DEFAULT_L1_OWNER: dict[MeshAxisKind, L1Owner] = {
-    "dp": "engine",
-    "tp": "engine",
-    "pp": "engine",
-    "ep": "engine",
-    "stage_replica": "delegated",
-}
-
-# How a RouteByStage policy maps onto omni's load balancer policy string. These
-# are the stateless options omni can actually do; note there's no "hash" here,
-# because omni has no key-stable balancer.
-_STAGE_POLICY_TO_OMNI_LB: dict[str, str] = {
-    "random": "random",
-    "round_robin": "round-robin",
-    "least_queue": "least-queue-length",
-}
-
-_VALID_L1_OWNERS: frozenset[str] = frozenset(get_args(L1Owner))
-
-
-class AxisTranslationError(ValueError):
-    """Error for an invalid or unsupported strategy spec.
-
-    A single error type (rather than a tree of subclasses) keeps the public
-    surface small and consistent with the rest of the codebase; the specific
-    cause is in the message and is logged before the raise so it is visible even
-    when the type is unavailable to a caller (e.g. across a server boundary).
-
-    Strategies that are *valid but not built yet* — key-stable / affinity routing
-    today — raise ``NotImplementedError`` instead, to distinguish "we haven't
-    implemented this" from "your config is wrong".
-    """
-
-
-def _fail(msg: str) -> NoReturn:
-    """Log and raise :class:`AxisTranslationError` for an invalid/unsupported spec."""
-    logger.error("[composable_parallel] %s", msg)
-    raise AxisTranslationError(msg)
-
-
-def _not_implemented(msg: str) -> NoReturn:
-    """Log and raise ``NotImplementedError`` for a valid-but-unbuilt strategy."""
-    logger.error("[composable_parallel] %s", msg)
-    raise NotImplementedError(msg)
-
-
-@dataclass(frozen=True)
-class OmniParallelConfig:
-    """Result of translating a spec stack into omni parallel sizing."""
-
-    tensor_parallel_size: int = 1
-    data_parallel_size: int = 1
-    pipeline_parallel_size: int = 1
-    # Dense expert parallelism: a flag over the TP*DP ranks, not a world axis.
-    enable_expert_parallel: bool = False
-    # Number of independent omni stage replicas (the ``stage_replica`` axis).
-    # This is NOT a vLLM world dimension; it is the per-stage deploy
-    # ``num_replicas`` count. 1 means a single (un-replicated) engine.
-    stage_replica_size: int = 1
-    # The omni StagePool LB policy string for the (delegated) stage_replica
-    # axis, if any. Only meaningful when ``stage_replica_size > 1``.
-    omni_lb_policy: str | None = None
-    # axis kind -> resolved L1 owner ("delegated" | "engine").
-    l1_owners: Mapping[MeshAxisKind, L1Owner] = field(default_factory=dict)
-
-    @property
-    def world_size(self) -> int:
-        # EP and stage_replica are intentionally excluded. EP reuses the TP*DP
-        # ranks rather than adding a dimension; stage_replica spins up separate
-        # engines (each its own world), not extra ranks in this engine's group.
-        # This product is what vLLM calls ``world_size_across_dp`` (tp * pp * dp).
-        return self.tensor_parallel_size * self.data_parallel_size * self.pipeline_parallel_size
-
-    @property
-    def delegated_axes(self) -> tuple[MeshAxisKind, ...]:
-        return tuple(kind for kind, owner in self.l1_owners.items() if owner == "delegated")
-
-    def as_engine_kwargs(self) -> dict[str, object]:
-        """Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names.
-
-        Two derived values are intentionally *not* emitted here because they are
-        not per-stage engine args:
-
-        * ``stage_replica_size`` — a per-stage deploy ``num_replicas`` knob
-          (``StageDeployConfig``); the deploy layer consumes it separately.
-        * ``omni_lb_policy`` — a pipeline-wide load-balancer policy the engine
-          reads once at construction (see ``StrategyApplyResult.omni_lb_policy``);
-          it is applied at the orchestrator level, not folded into per-stage args.
-        """
-        kwargs: dict[str, object] = {
-            "tensor_parallel_size": self.tensor_parallel_size,
-            "data_parallel_size": self.data_parallel_size,
-            "pipeline_parallel_size": self.pipeline_parallel_size,
-        }
-        if self.enable_expert_parallel:
-            kwargs["enable_expert_parallel"] = True
-        return kwargs
-
-
-def _is_affinity_dp_routing(routing: RoutingPattern) -> bool:
-    """True when DP routing demands key-stable (hash) placement."""
-    if isinstance(routing, PartitionByHash):
-        return True
-    if isinstance(routing, RouteByStage) and routing.routing_policy == "hash":
-        return True
-    return False
-
-
-def _resolve_l1_owner(spec: StrategySpec) -> L1Owner:
-    kind = spec.mesh_axis.kind
-    declared_owner = spec.shard_extension.get("l1_owner")
-    owner = declared_owner
-    if owner is None:
-        owner = _DEFAULT_L1_OWNER.get(kind, "engine")
-        if kind not in _DEFAULT_L1_OWNER:
-            logger.debug(
-                "[composable_parallel] axis %r has unknown kind %r with no default l1_owner; falling back to %r",
-                spec.name,
-                kind,
-                owner,
-            )
-        else:
-            logger.debug(
-                "[composable_parallel] axis %r (kind %r) declared no l1_owner; using default %r",
-                spec.name,
-                kind,
-                owner,
-            )
-    if owner not in _VALID_L1_OWNERS:
-        # Distinguish a bad value supplied via shard_extension from a bad default
-        # indexed out of _DEFAULT_L1_OWNER, so the source of the invalid owner is
-        # debuggable rather than swallowed.
-        source = "shard_extension" if declared_owner is not None else "_DEFAULT_L1_OWNER"
-        logger.debug(
-            "[composable_parallel] axis %r (kind %r) resolved invalid l1_owner %r from %s; valid owners are %s",
-            spec.name,
-            kind,
-            owner,
-            source,
-            sorted(_VALID_L1_OWNERS),
-        )
-        _fail(f"axis {kind!r} has invalid l1_owner {owner!r}; expected one of {sorted(_VALID_L1_OWNERS)}")
-    return cast(L1Owner, owner)
-
-
-def _validate_dp(spec: StrategySpec, owner: L1Owner) -> None:
-    """Validate an engine data-parallel axis.
-
-    DP is realized intra-engine by vLLM's own DP load balancer, so it is
-    engine-owned and emits no ``omni_lb_policy`` (that string configures omni's
-    *replica* balancer, a different layer). Key-stable (hash) request affinity is
-    not something vLLM's DP LB guarantees, so it is rejected rather than silently
-    dropped.
-    """
-    if owner != "engine":
-        _fail(
-            f"dp axis {spec.name!r} is engine data parallelism realized intra-engine by "
-            f"vLLM's DP load balancer; l1_owner must be 'engine', got {owner!r}. For "
-            "omni-coordinator-level request fan-out across independent replicas, use a "
-            "'stage_replica' axis."
-        )
-    if _is_affinity_dp_routing(spec.routing):
-        _not_implemented(
-            f"dp axis {spec.name!r} requests key-stable (hash) routing, which vLLM's "
-            "intra-engine DP load balancer does not guarantee — not supported yet. Use "
-            "RouteByStage(random|round_robin|least_queue) for stateless DP balancing."
-        )
-    if not isinstance(spec.routing, RouteByStage):
-        _fail(
-            f"dp axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) routing, "
-            f"got {type(spec.routing).__name__}"
-        )
-    if spec.routing.routing_policy not in _STAGE_POLICY_TO_OMNI_LB:
-        # Recognized-but-unimplemented routing (key-stable/hash) is already
-        # handled above via _is_affinity_dp_routing -> _not_implemented. Anything
-        # left here is an unknown/invalid policy value (e.g. a typo from YAML),
-        # i.e. invalid input -> AxisTranslationError, not NotImplementedError.
-        _fail(
-            f"dp axis {spec.name!r} has invalid routing_policy "
-            f"{spec.routing.routing_policy!r}; expected one of "
-            f"{sorted(_STAGE_POLICY_TO_OMNI_LB)}."
-        )
-
-
-def _stage_replica_lb_policy(spec: StrategySpec, owner: L1Owner) -> str:
-    """Return the omni StagePool LB policy for a delegated stage_replica axis.
-
-    A ``stage_replica`` axis is *not* a vLLM world dimension: it stands up N
-    independent engine replicas of one pipeline stage, coordinated by the omni
-    coordinator and balanced over a StagePool with stateless policies only
-    (random / round-robin / least-queue-length). It maps to the per-stage
-    ``num_replicas`` count plus the pipeline-level ``omni_lb_policy`` string, so
-    its ``l1_owner`` must be ``"delegated"`` (omni owns the routing). Key-stable
-    (hash) routing is rejected because omni has no key-stable balancer yet.
-    """
-    if owner != "delegated":
-        _fail(
-            f"stage_replica axis {spec.name!r} must be 'delegated' to omni's StagePool load "
-            f"balancer; got owner {owner!r}. Replica routing is owned by omni's coordinator."
-        )
-
-    routing = spec.routing
-    if _is_affinity_dp_routing(routing):
-        _not_implemented(
-            f"stage_replica axis {spec.name!r} requests key-stable (hash) routing, which needs a "
-            "dedicated load balancer — not implemented yet. Use "
-            "RouteByStage(random|round_robin|least_queue) to delegate to omni's load balancer."
-        )
-    if not isinstance(routing, RouteByStage):
-        _fail(
-            f"stage_replica axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) "
-            f"routing, got {type(routing).__name__}"
-        )
-    policy = _STAGE_POLICY_TO_OMNI_LB.get(routing.routing_policy)
-    if policy is None:
-        _not_implemented(
-            f"stage_replica axis {spec.name!r} routing_policy {routing.routing_policy!r} has no omni LB policy"
-        )
-    return policy
-
-
-def _validate_tp(spec: StrategySpec, owner: L1Owner) -> None:
-    if not isinstance(spec.routing, Broadcast):
-        _fail(f"tp axis {spec.name!r} expects Broadcast routing, got {type(spec.routing).__name__}")
-    if owner != "engine":
-        _fail(f"tp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
-
-
-def _validate_pp(spec: StrategySpec, owner: L1Owner) -> None:
-    if not isinstance(spec.routing, PipelineMicrobatch):
-        _fail(f"pp axis {spec.name!r} expects PipelineMicrobatch routing, got {type(spec.routing).__name__}")
-    if owner != "engine":
-        _fail(f"pp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
-
-
-def _validate_ep(spec: StrategySpec, owner: L1Owner) -> None:
-    # Dense EP: every rank still sees the whole batch, experts are sharded
-    # across ranks inside the engine (sparse MoE all-to-all is a later stage).
-    if not isinstance(spec.routing, Broadcast):
-        _fail(
-            f"ep axis {spec.name!r} expects Broadcast routing (dense expert parallel), "
-            f"got {type(spec.routing).__name__}"
-        )
-    if owner != "engine":
-        _fail(f"ep axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
-
-
-def translate_strategy_stack(specs: Sequence[StrategySpec]) -> OmniParallelConfig:
-    """Translate a spec stack into an ``OmniParallelConfig``.
-
-    Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and
-    stage_replica (omni replicas). Raises ``NotImplementedError`` for deferred
-    (affinity / key-stable) routing, and :class:`AxisTranslationError` for any
-    other invalid spec (a kind not yet translatable, a repeated kind, an owner
-    incompatible with the axis kind, or unsupported routing). The EP degree must
-    equal tensor_parallel_size * data_parallel_size.
-    """
-    sizes: dict[str, int] = {"tensor_parallel_size": 1, "data_parallel_size": 1, "pipeline_parallel_size": 1}
-    owners: dict[MeshAxisKind, L1Owner] = {}
-    omni_lb_policy: str | None = None
-    stage_replica_size = 1
-    enable_expert_parallel = False
-    ep_size: int | None = None
-
-    for spec in specs:
-        kind = spec.mesh_axis.kind
-        if kind not in _SUPPORTED_KINDS:
-            _fail(
-                f"axis kind {kind!r} is not translatable yet (supported: {list(_SUPPORTED_KINDS)}); "
-                "it is designed-for and lands in a later stage"
-            )
-        if kind in owners:
-            _fail(f"axis kind {kind!r} appears more than once in the spec stack")
-
-        owner = _resolve_l1_owner(spec)
-        if kind == "dp":
-            _validate_dp(spec, owner)
-        elif kind == "tp":
-            _validate_tp(spec, owner)
-        elif kind == "pp":
-            _validate_pp(spec, owner)
-        elif kind == "ep":
-            _validate_ep(spec, owner)
-            enable_expert_parallel = True
-            ep_size = spec.mesh_axis.size
-        elif kind == "stage_replica":
-            omni_lb_policy = _stage_replica_lb_policy(spec, owner)
-            stage_replica_size = spec.mesh_axis.size
-
-        if kind in _AXIS_TO_ENGINE_FIELD:
-            sizes[_AXIS_TO_ENGINE_FIELD[kind]] = spec.mesh_axis.size
-        owners[kind] = owner
-
-    if enable_expert_parallel:
-        # Dense EP shards experts across exactly the TP*DP ranks, so the declared
-        # EP degree must match that product (it is not its own world dimension,
-        # and PP is excluded). An EP axis of size 1 is a degenerate no-op that
-        # still sets the flag; it is only valid when TP*DP == 1.
-        ep_ranks = sizes["tensor_parallel_size"] * sizes["data_parallel_size"]
-        if ep_size != ep_ranks:
-            _fail(
-                f"ep axis size {ep_size} must equal tensor_parallel_size*data_parallel_size "
-                f"(={ep_ranks}); dense expert parallelism shards experts across exactly those ranks"
-            )
-
-    return OmniParallelConfig(
-        tensor_parallel_size=sizes["tensor_parallel_size"],
-        data_parallel_size=sizes["data_parallel_size"],
-        pipeline_parallel_size=sizes["pipeline_parallel_size"],
-        enable_expert_parallel=enable_expert_parallel,
-        stage_replica_size=stage_replica_size,
-        omni_lb_policy=omni_lb_policy,
-        l1_owners=owners,
-    )

vllm_omni/config/config_factory.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 
 import dataclasses
 import functools
-from collections.abc import Mapping
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -34,16 +33,6 @@ from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 logger = init_logger(__name__)
 
 
-# Default degree for any parallel axis / replica count that isn't set anywhere
-# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank.
-# TODO(composable_parallel): this "1" default and the per-axis device-layout
-# fallbacks are currently re-derived in the merge, apply, and reconcile layers.
-# Centralize them in one schema so the pre-spawn device guard and the engine-args
-# defaults can't drift apart. This is the light slice; the full device-layout
-# centralization is tracked as a follow-up.
-_DEFAULT_PARALLEL_DEGREE = 1
-
-
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -278,27 +267,14 @@ class StageConfigFactory:
         trust_remote_code: bool = False,
         cli_overrides: dict[str, Any] | None = None,
         deploy_config_path: str | None = None,
-        strategy_specs: Mapping[Any, Any] | None = None,
         **deprecated_kwargs: Any,
-    ) -> tuple[list[StageConfig] | None, str | None]:
+    ) -> list[StageConfig] | None:
         """Load pipeline + deploy config, merge with CLI overrides.
 
         Checks OMNI_PIPELINES first, since supported models should be explicitly
         registered. If a model is not registered in OMNI_PIPELINES, tries to fall
         back to using the Transformers config & finding pipelines that have overlapping
         supported architectures.
-
-        When ``strategy_specs`` is provided (a mapping of role -> list of
-        ``StrategySpec``), the derived parallel sizing is overlaid onto the
-        merged stages (see ``vllm_omni.config.composable_parallel``). This is
-        opt-in: omitting it leaves the existing merge path untouched.
-
-        Returns ``(stages, omni_lb_policy)``: the merged stages (``None`` when the
-        model is not in the pipeline registry and the caller should fall back to
-        the legacy YAML path) and the strategy-derived, pipeline-wide
-        ``omni_lb_policy`` (``None`` when no stage_replica axis set one). The
-        policy is returned rather than threaded through a mutable out-param so the
-        engine can apply it without every intermediate call carrying the dict.
         """
         if cli_overrides is None:
             cli_overrides = {}
@@ -314,9 +290,8 @@ class StageConfigFactory:
                 pipeline_cfg,
                 cli_overrides,
                 deploy_config_path,
-                strategy_specs=strategy_specs,
             )
-        return None, None
+        return None
 
     @classmethod
     def _create_from_registry(
@@ -325,17 +300,12 @@ class StageConfigFactory:
         pipeline_cfg: PipelineConfig,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None = None,
-        strategy_specs: Mapping[Any, Any] | None = None,
         **deprecated_kwargs: Any,
-    ) -> tuple[list[StageConfig], str | None]:
+    ) -> list[StageConfig]:
         """Create StageConfigs from pipeline registry + deploy YAML.
 
         Precedence: caller-typed (non-None) value > deploy YAML >
         StageDeployConfig dataclass default.
-
-        Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide
-        load-balancer policy (``None`` when no strategy set one) travels with the
-        stages instead of through a mutable out-param.
         """
         # Resolve deploy config path
         if deploy_config_path is None:
@@ -368,127 +338,12 @@ class StageConfigFactory:
 
         stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)
 
-        # Overlay declarative parallel strategies (opt-in) before CLI overrides.
-        applied = cls._apply_strategy_specs(stages, strategy_specs)
-
         explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None}
 
         for stage in stages:
             stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides)
 
-        # Re-validate the resolved layout now that CLI overrides are on top.
-        cls._reconcile_strategy_with_cli(stages, applied)
-
-        omni_lb_policy = applied.omni_lb_policy if applied is not None else None
-        return stages, omni_lb_policy
-
-    @staticmethod
-    def _apply_strategy_specs(
-        stages: list[StageConfig],
-        strategy_specs: Mapping[Any, Any] | None,
-    ) -> Any:
-        """Overlay derived parallel sizing onto merged stages (opt-in).
-
-        ``omni_lb_policy`` cannot be set from stage configs (omni reads it once
-        at orchestrator construction), so a derived non-default policy is logged
-        here and carried on the returned ``StrategyApplyResult`` for the caller
-        to hand to the orchestrator.
-
-        Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was
-        supplied) so the caller can re-validate the resolved layout once CLI
-        overrides have been merged on top, and read its ``omni_lb_policy``.
-        """
-        if not strategy_specs:
-            return None
-        from vllm_omni.config.composable_parallel import apply_strategy_specs
-
-        applied = apply_strategy_specs(stages, strategy_specs)
-        if applied.omni_lb_policy is not None:
-            logger.info(
-                "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied "
-                "to the orchestrator unless an explicit --omni-lb-policy was given.",
-                applied.omni_lb_policy,
-            )
-        return applied
-
-    @staticmethod
-    def _reconcile_strategy_with_cli(
-        stages: list[StageConfig],
-        applied: Any,
-    ) -> None:
-        """Reconcile CLI overrides applied *after* a strategy overlay.
-
-        CLI overrides are applied last (in ``to_omegaconf``), so for any stage a
-        strategy touched we must (a) warn loudly when a CLI arg overrides a
-        strategy-declared axis — the strategy is meant to be the single writer
-        for the axes it declares, so a silent CLI win is surprising — and
-        (b) re-run the device-count check against the *effective* (post-CLI)
-        ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI
-        ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip
-        past the pre-spawn guard that exists to prevent silent OOMs.
-        """
-        if applied is None:
-            return
-        from vllm_omni.config.composable_parallel import check_device_layout
-
-        # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg).
-        axis_fields = {
-            "tp": ("tensor_parallel_size", "tensor_parallel_size"),
-            "dp": ("data_parallel_size", "data_parallel_size"),
-            "pp": ("pipeline_parallel_size", "pipeline_parallel_size"),
-        }
-
-        for stage in stages:
-            cfg = applied.per_stage_config.get(stage.stage_id)
-            if cfg is None:
-                continue
-            overrides = stage.runtime_overrides or {}
-            declared = set(cfg.l1_owners.keys())
-
-            for kind, (field_name, attr) in axis_fields.items():
-                if kind in declared and overrides.get(field_name) is not None:
-                    cli_val = overrides[field_name]
-                    derived = getattr(cfg, attr)
-                    if cli_val != derived:
-                        logger.warning(
-                            "[composable_parallel] stage %s: CLI %s=%s overrides the "
-                            "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.",
-                            stage.stage_id,
-                            field_name,
-                            cli_val,
-                            field_name,
-                            derived,
-                        )
-            if "stage_replica" in declared and overrides.get("num_replicas") is not None:
-                cli_val = overrides["num_replicas"]
-                if cli_val != cfg.stage_replica_size:
-                    logger.warning(
-                        "[composable_parallel] stage %s: CLI num_replicas=%s overrides the "
-                        "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.",
-                        stage.stage_id,
-                        cli_val,
-                        cfg.stage_replica_size,
-                    )
-
-            def _eff(field_name: str, fallback: Any) -> Any:
-                val = overrides.get(field_name)
-                return val if val is not None else fallback
-
-            def _eff_degree(field_name: str, source: dict[str, Any]) -> int:
-                # Single place the per-axis "default to 1" fallback is applied,
-                # for both the override and the YAML-default sides (see the
-                # _DEFAULT_PARALLEL_DEGREE TODO).
-                value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE))
-                return int(value or _DEFAULT_PARALLEL_DEGREE)
-
-            check_device_layout(
-                _eff("devices", stage.yaml_runtime.get("devices")),
-                tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args),
-                data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args),
-                pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args),
-                num_replicas=_eff_degree("num_replicas", stage.yaml_runtime),
-                role=stage.model_stage,
-            )
+        return stages
 
     @classmethod
     def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -314,6 +314,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and kv_connector_output.invalid_block_ids:
@@ -556,20 +563,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
         # collect KV cache events from KV cache manager
         events = self.kv_cache_manager.take_events()
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -431,6 +431,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        # Merge connector-side stats (align with v0.14.0)
+        if kv_connector_stats and self.connector:
+            kv_stats = self.connector.get_kv_connector_stats()
+            if kv_stats:
+                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
@@ -638,20 +646,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
         # Collect and publish KV cache events (align with v0.14.0)
         events = self.kv_cache_manager.take_events()
         if self.connector is not None:

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -47,6 +47,7 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
+      seed: 42
       repetition_penalty: 1.05
 
   - stage_id: 2
@@ -66,6 +67,7 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
+      seed: 42
       repetition_penalty: 1.1
 
 platforms:

vllm_omni/deploy/qwen3_tts.yaml

@@ -35,7 +35,7 @@ connectors:
 
 stages:
   - stage_id: 0
-    max_num_seqs: 64
+    max_num_seqs: 10
     gpu_memory_utilization: 0.3
     trust_remote_code: true
     # Prefix caching on by default. The HS opt-out (talker only consumes the
@@ -43,9 +43,9 @@ stages:
     # PR make per-step bookkeeping negligible, and any text repetition across
     # requests (same voice or shared ref_audio) is a net throughput win.
     # Disable explicitly if memory headroom is tight on a single-GPU deploy.
-    enable_prefix_caching: false
+    enable_prefix_caching: true
     async_scheduling: true
-    max_num_batched_tokens: 32768
+    max_num_batched_tokens: 512
     max_model_len: 4096
     devices: "0"
     output_connectors:
@@ -67,7 +67,7 @@ stages:
 
   - stage_id: 1
     # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
-    max_num_seqs: 64
+    max_num_seqs: 10
     gpu_memory_utilization: 0.3
     enforce_eager: false
     trust_remote_code: true
@@ -95,8 +95,3 @@ platforms:
     stages:
       - stage_id: 0
         enforce_eager: false
-        compilation_config:
-          cudagraph_mode: PIECEWISE
-      - stage_id: 1
-        compilation_config:
-          cudagraph_mode: PIECEWISE

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -397,77 +397,16 @@ class DiffusersPipelineLoader:
 
         return model.eval()
 
-    @staticmethod
-    def _has_online_quant(model: nn.Module) -> bool:
-        """Whether any layer uses an online-quant method that defers weight
-        materialization onto the ``meta`` device (upstream vLLM
-        ``uses_meta_device=True``, e.g. online FP8)."""
-        for module in model.modules():
-            quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "uses_meta_device", False):
-                return True
-        return False
-
     def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
         """Process weights after loading for quantization methods.
 
         This handles vLLM's quantization methods that need to process weights
         after loading (e.g., FP8 online quantization from BF16/FP16 weights).
         """
-        # Newer upstream vLLM online-quant methods (uses_meta_device=True) create
-        # weights on the ``meta`` device and materialize them just-in-time as each
-        # layer's weights finish loading (via the layerwise online-process loader).
-        # Any "straggler" layers whose weights were not fully materialized during
-        # load (padded / partially-loaded layers) remain on ``meta``. Upstream's
-        # base_loader calls finalize_layerwise_processing() to materialize them;
-        # the diffusion loader must mirror that, otherwise the module.to() below
-        # raises "Cannot copy out of meta tensor; no data!". This whole meta-device
-        # handling is gated on online quant actually being in use, so that the
-        # proven code path for everything else (in particular FSDP/HSDP-sharded
-        # params, whose per-parameter .data cannot be cross-device reassigned) is
-        # left untouched. Import lazily so older vLLM (no meta-device quant) is
-        # unaffected.
-        has_online_quant = self._has_online_quant(model)
-        if has_online_quant:
-            from vllm.model_executor.model_loader.reload.layerwise import (
-                finalize_layerwise_processing,
-            )
-
-            # model_config is only dereferenced by finalize for vLLM Attention /
-            # MLAAttention layers; diffusion DiT models use their own attention and
-            # have none, so passing None is safe here.
-            finalize_layerwise_processing(model, model_config=None)
-
         for _, module in model.named_modules():
             quant_method = getattr(module, "quant_method", None)
-            if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
-                continue
-
-            if has_online_quant:
-                # Online quant may leave straggler params on the ``meta`` device.
-                # Move only real (non-meta) params onto the target device for
-                # processing and restore them afterward, mirroring upstream vLLM's
-                # device_loading_context — a blanket module.to(target_device) would
-                # raise NotImplementedError on meta params. Online quant initializes
-                # on the accelerator, so params are normally already on the target
-                # device and this loop is a no-op move; the point is to skip meta.
-                original_devices: dict[str, torch.device] = {}
-                for name, param in module.named_parameters():
-                    if param.device.type != "meta" and param.device != target_device:
-                        original_devices[name] = param.device
-                        param.data = param.data.to(target_device)
-
-                quant_method.process_weights_after_loading(module)
-
-                # Restore pre-existing params to their original device; leave any
-                # newly created (e.g. quantized) params on the target device.
-                for name, param in module.named_parameters():
-                    if name in original_devices:
-                        param.data = param.data.to(original_devices[name])
-            else:
-                # No meta params possible here. Preserve the original FSDP/HSDP-aware
-                # whole-module move (module.to()), which correctly handles sharded
-                # DTensor params that per-parameter .data reassignment cannot.
+            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
+                # Move module to target device for processing if needed
                 module_device = next(module.parameters(), None)
                 if module_device is not None:
                     module_device = module_device.device

vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py

@@ -19,7 +19,6 @@ from transformers.utils import cached_file
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
-from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
 from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
 from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
 from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
@@ -553,9 +552,4 @@ class Gr00tN1d7Processor(ProcessorMixin):
         return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
 
 
-# transformers >= 5.5 requires the config *class* (not a model_type string) as the
-# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing
-# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute
-# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in
-# modeling/gr00t_n1d7.py.
-AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)
+AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)

vllm_omni/diffusion/models/lingbot_video/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import (
+    LingBotVideoTransformer3DModel,
+)
+from vllm_omni.diffusion.models.lingbot_video.pipeline_lingbot_video import (
+    LingBotVideoPipeline,
+    get_lingbot_video_post_process_func,
+)
+
+__all__ = [
+    "LingBotVideoPipeline",
+    "LingBotVideoTransformer3DModel",
+    "get_lingbot_video_post_process_func",
+]

vllm_omni/diffusion/models/lingbot_video/lingbot_video_transformer.py

@@ -0,0 +1,702 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
+
+import math
+import os
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    flash_attn_varlen_func as flash_attn_varlen_func_v3,
+)
+from vllm_omni.diffusion.attention.layer import Attention
+
+LINGBOT_VIDEO_FP32_MODULES = (
+    "time_embedder",
+    "time_modulation",
+    "scale_shift_table",
+    "norm",
+    "norm1",
+    "norm2",
+    "norm_q",
+    "norm_k",
+    "norm_post_attn",
+    "norm_post_ffn",
+    "norm_out",
+    "norm_out_modulation",
+    "router",
+)
+
+
+def should_keep_in_fp32(name: str) -> bool:
+    return any(module_name in name.split(".") for module_name in LINGBOT_VIDEO_FP32_MODULES)
+
+
+def _all_to_all_split_cat(
+    local_input: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+) -> torch.Tensor:
+    world_size = dist.get_world_size(group)
+    input_list = [tensor.contiguous() for tensor in torch.tensor_split(local_input, world_size, scatter_dim)]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class LingBotVideoRMSNorm(nn.Module):
+    """RMSNorm with fp32 accumulation."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply complex RoPE to `(B, S, H, D)` attention tensors."""
+    with torch.amp.autocast("cuda", enabled=False):
+        x_c = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        out = torch.view_as_real(x_c * freqs_cis.unsqueeze(2)).flatten(3)
+        return out.type_as(x)
+
+
+class LingBotVideoRotaryEmbedding(nn.Module):
+    """Complex64 RoPE table indexed by position ids."""
+
+    def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float):
+        super().__init__()
+        self.axes_dims = tuple(axes_dims)
+        self.axes_lens = list(axes_lens)
+        self.theta = theta
+        self.freqs_cis = None
+
+    @staticmethod
+    def precompute_freqs_cis(dim: tuple[int, ...], end: tuple[int, ...], theta: float):
+        freqs_cis = []
+        for d, e in zip(dim, end):
+            freqs = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float64, device="cpu") / d))
+            timestep = torch.arange(e, device=freqs.device, dtype=torch.float64)
+            freqs = torch.outer(timestep, freqs).float()
+            freqs_cis.append(torch.polar(torch.ones_like(freqs), freqs).to(torch.complex64))
+        return freqs_cis
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        # position_ids: (S, 3) int -> (S, head_dim/2) complex64
+        device = position_ids.device
+        max_vals = position_ids.max(dim=0).values.tolist()
+        needs_rebuild = self.freqs_cis is None or any(
+            max_val >= axis_len for max_val, axis_len in zip(max_vals, self.axes_lens)
+        )
+        if needs_rebuild:
+            for i in range(len(self.axes_lens)):
+                if max_vals[i] >= self.axes_lens[i]:
+                    self.axes_lens[i] = int(max_vals[i] * 1.5) + 1
+            self.freqs_cis = self.precompute_freqs_cis(self.axes_dims, tuple(self.axes_lens), theta=self.theta)
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+        elif self.freqs_cis[0].device != device:
+            self.freqs_cis = [freqs_cis.to(device) for freqs_cis in self.freqs_cis]
+
+        return torch.cat([self.freqs_cis[i][position_ids[:, i]] for i in range(len(self.axes_dims))], dim=-1)
+
+
+def make_joint_position_ids(text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device) -> torch.Tensor:
+    """3D positions in [video; text] order. Text t-axis is 1..text_len; video t-axis starts at text_len+1.
+
+    Matches patchify_and_embed: cap start (1,0,0); vision start (cap_len+1,0,0);
+    freqs ordered with x first and cap second (same order as cat_interleave).
+    """
+    tt = torch.arange(grid_t, device=device, dtype=torch.int32) + (text_len + 1)
+    hh = torch.arange(grid_h, device=device, dtype=torch.int32)
+    ww = torch.arange(grid_w, device=device, dtype=torch.int32)
+    grid = torch.stack(torch.meshgrid(tt, hh, ww, indexing="ij"), dim=-1).flatten(0, 2)
+    text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1
+    text_pos = torch.stack([text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)], dim=-1)
+    return torch.cat([grid, text_pos], dim=0)  # (Nx + L, 3)
+
+
+def _cat_interleave(
+    a: torch.Tensor,
+    len_a: list[int],
+    b: torch.Tensor,
+    len_b: list[int],
+) -> torch.Tensor:
+    a_split = torch.split(a, len_a, dim=1)
+    b_split = torch.split(b, len_b, dim=1)
+    blocks: list[torch.Tensor] = []
+    for x_part, text_part in zip(a_split, b_split):
+        blocks.extend([x_part, text_part])
+    return torch.cat(blocks, dim=1)
+
+
+def _packed_block_attention_mask(sample_seq_lens: list[int], device: torch.device) -> torch.Tensor:
+    total_seq_len = sum(sample_seq_lens)
+    mask = torch.zeros(
+        1,
+        1,
+        total_seq_len,
+        total_seq_len,
+        dtype=torch.bool,
+        device=device,
+    )
+    start = 0
+    for seq_len in sample_seq_lens:
+        end = start + seq_len
+        mask[:, :, start:end, start:end] = True
+        start = end
+    return mask
+
+
+class LingBotVideoTextEmbedder(nn.Module):
+    """Matches CondProjection: RMSNorm(text_dim, eps=1e-6 fixed) -> Linear-SiLU-Linear."""
+
+    def __init__(self, text_dim: int, hidden_size: int):
+        super().__init__()
+        self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6)
+        self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.norm(x)
+        return self.linear_2(F.silu(self.linear_1(x)))
+
+
+class LingBotVideoAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, norm_eps, qkv_bias, out_bias):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps)
+        self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias)
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / math.sqrt(self.head_dim),
+            causal=False,
+            num_kv_heads=num_heads,
+        )
+
+    def forward(
+        self,
+        x,
+        rotary_emb,
+        attention_mask=None,
+        packed_indices: dict[str, torch.Tensor] | None = None,
+        parallel_config=None,
+    ):
+        B, S, _ = x.shape
+        if os.environ.get("LINGBOT_FUSED_QKV_LINEAR") == "1":
+            weight = torch.cat(
+                (self.to_q.weight, self.to_k.weight, self.to_v.weight),
+                dim=0,
+            )
+            bias = None
+            if self.to_q.bias is not None:
+                bias = torch.cat(
+                    (self.to_q.bias, self.to_k.bias, self.to_v.bias),
+                    dim=0,
+                )
+            qkv = F.linear(x, weight, bias)
+            q, k, v = qkv.view(B, S, 3, self.num_heads, self.head_dim).unbind(2)
+        else:
+            q = self.to_q(x).unflatten(2, (self.num_heads, self.head_dim))
+            k = self.to_k(x).unflatten(2, (self.num_heads, self.head_dim))
+            v = self.to_v(x).unflatten(2, (self.num_heads, self.head_dim))
+        q = apply_rotary_emb(self.norm_q(q), rotary_emb)
+        k = apply_rotary_emb(self.norm_k(k), rotary_emb)
+        if packed_indices is None:
+            out = self.attn(q, k, v, attn_metadata=AttentionMetadata(attn_mask=attention_mask))
+        else:
+            packed_attention_mask = packed_indices.get("attention_mask")
+            if packed_attention_mask is not None and parallel_config is None:
+                out = self.attn.sdpa_fallback.forward(
+                    q,
+                    k,
+                    v,
+                    AttentionMetadata(attn_mask=packed_attention_mask),
+                )
+                return self.to_out(out.flatten(2, 3).type_as(x))
+            if flash_attn_varlen_func_v3 is None:
+                raise RuntimeError(
+                    "A flash attention varlen function is required for packed context parallel attention."
+                )
+            if parallel_config is None:
+                result = flash_attn_varlen_func_v3(
+                    q=q.reshape(-1, self.num_heads, self.head_dim),
+                    k=k.reshape(-1, self.num_heads, self.head_dim),
+                    v=v.reshape(-1, self.num_heads, self.head_dim),
+                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
+                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
+                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
+                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
+                    causal=False,
+                )
+                out = result[0] if isinstance(result, tuple) else result
+                out = out.reshape(B, S, self.num_heads, self.head_dim)
+            else:
+                group = parallel_config.context_parallel_config._ulysses_mesh.get_group()
+                world_size = dist.get_world_size(group)
+                local_heads = self.num_heads // world_size
+                q_global = _all_to_all_split_cat(
+                    q.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                k_global = _all_to_all_split_cat(
+                    k.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                v_global = _all_to_all_split_cat(
+                    v.reshape(B, S, self.num_heads * self.head_dim),
+                    scatter_dim=2,
+                    gather_dim=1,
+                    group=group,
+                ).view(B, S * world_size, local_heads, self.head_dim)
+                q_flat = q_global.reshape(-1, local_heads, self.head_dim)
+                k_flat = k_global.reshape(-1, local_heads, self.head_dim)
+                v_flat = v_global.reshape(-1, local_heads, self.head_dim)
+                result = flash_attn_varlen_func_v3(
+                    q=q_flat,
+                    k=k_flat,
+                    v=v_flat,
+                    cu_seqlens_q=packed_indices["cu_seqlens_kv"],
+                    cu_seqlens_k=packed_indices["cu_seqlens_kv"],
+                    max_seqlen_q=packed_indices["max_seqlen_in_batch_kv"],
+                    max_seqlen_k=packed_indices["max_seqlen_in_batch_kv"],
+                    causal=False,
+                )
+                out_global = result[0] if isinstance(result, tuple) else result
+                out_global = out_global.reshape(B, S * world_size, local_heads * self.head_dim)
+                out = _all_to_all_split_cat(
+                    out_global,
+                    scatter_dim=1,
+                    gather_dim=2,
+                    group=group,
+                ).view(B, S, self.num_heads, self.head_dim)
+        return self.to_out(out.flatten(2, 3).type_as(x))
+
+
+class LingBotVideoMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LingBotVideoBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_attention_heads,
+        intermediate_size,
+        norm_eps,
+        qkv_bias,
+        out_bias,
+        num_experts,
+        num_experts_per_tok,
+        moe_intermediate_size,
+        decoder_sparse_step,
+        mlp_only_layers,
+        n_shared_experts,
+        score_func,
+        norm_topk_prob,
+        n_group,
+        topk_group,
+        routed_scaling_factor,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        h = hidden_size
+        self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * h))
+        self.norm1 = LingBotVideoRMSNorm(h, norm_eps)
+        self.attn = LingBotVideoAttention(h, num_attention_heads, norm_eps, qkv_bias, out_bias)
+        self.norm_post_attn = LingBotVideoRMSNorm(h, norm_eps)
+        self.norm2 = LingBotVideoRMSNorm(h, norm_eps)
+        if num_experts > 0:
+            raise NotImplementedError(
+                "LingBotVideoTransformer3DModel in vLLM-Omni currently supports "
+                "the dense checkpoint only. MoE/fused-expert support belongs in "
+                "the follow-up MoE PR."
+            )
+        self.ffn = LingBotVideoMLP(h, intermediate_size)
+        self.norm_post_ffn = LingBotVideoRMSNorm(h, norm_eps)
+
+    def forward(
+        self,
+        x,
+        temb6,
+        rotary_emb,
+        attention_mask=None,
+        moe_padding_mask=None,
+        packed_indices: dict[str, torch.Tensor] | None = None,
+        parallel_config=None,
+    ):
+        expected_tokens = x.shape[0] * x.shape[1]
+        if temb6.ndim != 2 or temb6.shape[0] != expected_tokens:
+            raise ValueError(
+                "LingBotVideoBlock expects token-level temb6 with shape "
+                f"(B*S, 6D); got {tuple(temb6.shape)} for hidden states {tuple(x.shape)}."
+            )
+        # AdaLN modulation and normalization stay in fp32 for the sensitive path.
+        mod = temb6.view(x.shape[0], x.shape[1], -1) + self.scale_shift_table.unsqueeze(0)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=-1)
+        gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
+        scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
+
+        # AdaLN modulation / norms run in fp32 (sensitive path); cast to the bulk
+        # compute dtype only at the bf16 Linear boundary. This replaces the old
+        # ambient autocast, which rounded Linear inputs to bf16 at the same point.
+        bulk_dtype = self.attn.to_q.weight.dtype
+        attn_in = (self.norm1(x) * scale_msa + shift_msa).to(bulk_dtype)
+        attn_out = self.attn(
+            attn_in,
+            rotary_emb,
+            attention_mask,
+            packed_indices=packed_indices,
+            parallel_config=parallel_config,
+        )
+        x = x + (gate_msa * self.norm_post_attn(attn_out)).to(x.dtype)
+
+        ffn_in = (self.norm2(x) * scale_mlp + shift_mlp).to(bulk_dtype)
+        ffn_out = self.ffn(ffn_in)
+        ffn_normed = self.norm_post_ffn(ffn_out)
+        x = x + (gate_mlp * ffn_normed).to(x.dtype)
+        return x
+
+
+class LingBotVideoTransformer3DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = False
+    _repeated_blocks = ["LingBotVideoBlock"]
+    _layerwise_offload_blocks_attr = "blocks"
+    _no_split_modules = ["LingBotVideoBlock"]
+    _keep_in_fp32_modules = list(LINGBOT_VIDEO_FP32_MODULES)
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        if dtype is None or dtype == torch.float32:
+            return super().to(*args, **kwargs)
+
+        dtype_is_floating = torch.is_floating_point(torch.empty((), dtype=dtype))
+        if not dtype_is_floating:
+            return super().to(*args, **kwargs)
+
+        if device is not None:
+            super().to(device=device, non_blocking=non_blocking)
+
+        for name, param in self.named_parameters():
+            if not torch.is_floating_point(param):
+                continue
+            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
+            param.data = param.data.to(dtype=target_dtype, non_blocking=non_blocking)
+            if param.grad is not None:
+                param.grad.data = param.grad.data.to(dtype=target_dtype, non_blocking=non_blocking)
+
+        for name, buffer in self.named_buffers():
+            if not torch.is_floating_point(buffer):
+                continue
+            target_dtype = torch.float32 if should_keep_in_fp32(name) else dtype
+            buffer.data = buffer.data.to(dtype=target_dtype, non_blocking=non_blocking)
+
+        return self
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: tuple[int, int, int] = (1, 2, 2),
+        in_channels: int = 16,
+        out_channels: int = 16,
+        hidden_size: int = 2048,
+        num_attention_heads: int = 16,
+        depth: int = 24,
+        intermediate_size: int = 6144,
+        text_dim: int = 2560,
+        freq_dim: int = 256,
+        norm_eps: float = 1e-6,
+        rope_theta: float = 256.0,
+        axes_dims: tuple[int, int, int] = (32, 48, 48),
+        axes_lens: tuple[int, int, int] = (8192, 1024, 1024),
+        qkv_bias: bool = False,
+        out_bias: bool = True,
+        patch_embed_bias: bool = True,
+        timestep_mlp_bias: bool = True,
+        num_experts: int = 0,
+        num_experts_per_tok: int = 8,
+        moe_intermediate_size: int = 512,
+        decoder_sparse_step: int = 1,
+        mlp_only_layers: tuple[int, ...] = (),
+        n_shared_experts: int | None = None,
+        score_func: str = "sigmoid",
+        norm_topk_prob: bool = True,
+        n_group: int | None = None,
+        topk_group: int | None = None,
+        routed_scaling_factor: float = 1.0,
+    ):
+        super().__init__()
+        head_dim = hidden_size // num_attention_heads
+        assert head_dim == sum(axes_dims), f"head_dim {head_dim} != sum(axes_dims) {sum(axes_dims)}"
+        mlp_only_layers = tuple(mlp_only_layers)
+
+        self.patch_embedder = nn.Linear(in_channels * math.prod(patch_size), hidden_size, bias=patch_embed_bias)
+        self.time_proj = Timesteps(freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(freq_dim, hidden_size, act_fn="silu", sample_proj_bias=timestep_mlp_bias)
+        self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+        self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size)
+        self.rope = LingBotVideoRotaryEmbedding(axes_dims, axes_lens, rope_theta)
+        self.blocks = nn.ModuleList(
+            [
+                LingBotVideoBlock(
+                    hidden_size=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    intermediate_size=intermediate_size,
+                    norm_eps=norm_eps,
+                    qkv_bias=qkv_bias,
+                    out_bias=out_bias,
+                    num_experts=num_experts,
+                    num_experts_per_tok=num_experts_per_tok,
+                    moe_intermediate_size=moe_intermediate_size,
+                    decoder_sparse_step=decoder_sparse_step,
+                    mlp_only_layers=mlp_only_layers,
+                    n_shared_experts=n_shared_experts,
+                    score_func=score_func,
+                    norm_topk_prob=norm_topk_prob,
+                    n_group=n_group,
+                    topk_group=topk_group,
+                    routed_scaling_factor=routed_scaling_factor,
+                    layer_idx=i,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps)
+        self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size))
+        self.proj_out = nn.Linear(hidden_size, math.prod(patch_size) * out_channels)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # (B, C, T, H, W)
+        timestep: torch.Tensor,  # (B,) in [0, 1000](= sigma*1000)
+        encoder_hidden_states: torch.Tensor,  # (B, L, text_dim)
+        encoder_attention_mask: torch.Tensor | None = None,  # (B, L) 1=valid
+        return_dict: bool = True,
+    ):
+        B, C, T, H, W = hidden_states.shape
+        pF, pH, pW = self.config.patch_size
+        gt, gh, gw = T // pF, H // pH, W // pW
+        n_video = gt * gh * gw
+        L = encoder_hidden_states.shape[1]
+        device = hidden_states.device
+        if encoder_attention_mask is not None:
+            text_lens = encoder_attention_mask.sum(dim=-1).long()
+        else:
+            text_lens = torch.full((B,), L, dtype=torch.long, device=device)
+        text_lens_list = [int(v) for v in text_lens.detach().cpu().tolist()]
+        packed_batch = B > 1
+
+        # patchify: token order (f h w), feature order (pf ph pw c) -- matches patchify_and_embed
+        patch_tokens = hidden_states.reshape(B, C, gt, pF, gh, pH, gw, pW)
+        patch_tokens = patch_tokens.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(
+            B,
+            n_video,
+            pF * pH * pW * C,
+        )
+        if packed_batch:
+            x = torch.cat(
+                [self.patch_embedder(patch_tokens[i : i + 1]) for i in range(B)],
+                dim=1,
+            )
+        else:
+            x = self.patch_embedder(patch_tokens)
+
+        if packed_batch:
+            text_parts = [
+                self.text_embedder(encoder_hidden_states[i : i + 1, : text_lens_list[i], :]) for i in range(B)
+            ]
+            text = torch.cat(text_parts, dim=1)
+            joint = _cat_interleave(
+                x,
+                [n_video] * B,
+                text,
+                text_lens_list,
+            )
+        else:
+            text = self.text_embedder(encoder_hidden_states)
+            joint = torch.cat([x, text], dim=1)  # [video; text]
+        joint_seq_len = joint.shape[1]
+
+        # Per-sample RoPE: video t-axis start = real text length of this sample + 1
+        rotary_parts = [self.rope(make_joint_position_ids(text_lens_list[i], gt, gh, gw, device)) for i in range(B)]
+        if packed_batch:
+            rotary = torch.cat(rotary_parts, dim=0).unsqueeze(0)
+        else:
+            rotary = torch.stack(rotary_parts, dim=0)  # (B, S, head_dim/2) complex64
+
+        parallel_config = getattr(self, "_parallel_config", None)
+        use_packed_attention = parallel_config is not None
+
+        attention_mask = None
+        moe_padding_mask = None
+        packed_indices = None
+        has_padding = encoder_attention_mask is not None and bool((text_lens < L).any())
+        if packed_batch or use_packed_attention:
+            sample_seq_lens = [n_video + text_len for text_len in text_lens_list]
+            cu_seqlens = torch.zeros(B + 1, device=device, dtype=torch.int32)
+            cu_seqlens[1:] = torch.cumsum(
+                torch.tensor(sample_seq_lens, device=device, dtype=torch.int32),
+                dim=0,
+            )
+            packed_indices = {
+                "cu_seqlens_kv": cu_seqlens,
+                "max_seqlen_in_batch_kv": max(sample_seq_lens),
+            }
+            if packed_batch and not use_packed_attention:
+                packed_indices["attention_mask"] = _packed_block_attention_mask(sample_seq_lens, device)
+            has_padding = False
+        if has_padding:
+            key_mask = torch.cat(
+                [torch.ones(B, n_video, dtype=torch.bool, device=device), encoder_attention_mask.bool()],
+                dim=1,
+            )
+            attention_mask = key_mask[:, None, None, :]  # (B,1,1,S) -> SDPA broadcast
+            moe_padding_mask = key_mask.reshape(-1).float()  # (B*S,)
+        packed_cp = packed_indices is not None and parallel_config is not None
+        padding_size = 0
+        if packed_cp:
+            cp_config = parallel_config.context_parallel_config
+            cp_world_size = int(getattr(cp_config, "ulysses_degree", getattr(cp_config, "_world_size", 1)))
+            padding_size = (cp_world_size - (joint_seq_len % cp_world_size)) % cp_world_size
+            if padding_size:
+                joint = torch.cat(
+                    [
+                        joint,
+                        torch.zeros(
+                            joint.shape[0],
+                            padding_size,
+                            joint.shape[2],
+                            device=joint.device,
+                            dtype=joint.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+                rotary = torch.cat(
+                    [
+                        rotary,
+                        torch.zeros(
+                            rotary.shape[0],
+                            padding_size,
+                            rotary.shape[2],
+                            device=rotary.device,
+                            dtype=rotary.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+                if packed_indices is None:
+                    raise RuntimeError("packed_indices must be initialized for packed context parallel.")
+                packed_indices["cu_seqlens_kv"] = torch.cat(
+                    [
+                        packed_indices["cu_seqlens_kv"],
+                        packed_indices["cu_seqlens_kv"][-1:] + padding_size,
+                    ],
+                    dim=0,
+                )
+                packed_indices["max_seqlen_in_batch_kv"] = max(
+                    int(packed_indices["max_seqlen_in_batch_kv"]),
+                    int(padding_size),
+                )
+                joint_seq_len = joint.shape[1]
+
+        timestep_for_embed = timestep.float()
+        timestep_proj = self.time_proj(timestep_for_embed)
+        t_emb = self.time_embedder(timestep_proj)  # (B, D)
+        if packed_batch:
+            temb_input = torch.cat(
+                [t_emb[i : i + 1].unsqueeze(1).expand(1, n_video + text_lens_list[i], -1) for i in range(B)],
+                dim=1,
+            )
+            if padding_size:
+                temb_input = torch.cat(
+                    [
+                        temb_input,
+                        torch.zeros(
+                            temb_input.shape[0],
+                            padding_size,
+                            temb_input.shape[2],
+                            device=temb_input.device,
+                            dtype=temb_input.dtype,
+                        ),
+                    ],
+                    dim=1,
+                )
+            temb6 = self.time_modulation(temb_input.reshape(joint_seq_len, -1))
+            temb6 = temb6.reshape(1, joint_seq_len, -1)
+        else:
+            temb_input = t_emb.unsqueeze(1).expand(B, joint_seq_len, -1)  # (B, S, D)
+            temb6 = self.time_modulation(temb_input.reshape(B * joint_seq_len, -1))
+            temb6 = temb6.reshape(B, joint_seq_len, -1)  # (B, S, 6D)
+
+        temb6 = temb6.reshape(temb6.shape[0] * temb6.shape[1], -1)
+
+        for block in self.blocks:
+            joint = block(
+                joint,
+                temb6,
+                rotary,
+                attention_mask,
+                moe_padding_mask,
+                packed_indices=packed_indices,
+                parallel_config=parallel_config,
+            )
+        final_mod = self.norm_out_modulation(temb_input.reshape(joint.shape[0] * joint.shape[1], -1))
+        shift, scale = final_mod.reshape(joint.shape[0], joint.shape[1], -1).chunk(2, dim=-1)
+        final_hidden = self.norm_out(joint) * (1.0 + scale) + shift
+        projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype))
+        if packed_cp:
+            if padding_size:
+                projected = projected[:, :-padding_size, :]
+        if packed_batch:
+            split_lengths: list[int] = []
+            for text_len in text_lens_list:
+                split_lengths.extend([n_video, text_len])
+            parts = torch.split(projected, split_lengths, dim=1)
+            x = torch.cat(parts[::2], dim=1).reshape(B, n_video, -1)
+        else:
+            x = projected[:, :n_video]
+
+        # unpatchify (matches the rearrange in postprocess)
+        Cout = self.config.out_channels
+        x = x.reshape(B, gt, gh, gw, pF, pH, pW, Cout)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(B, Cout, T, H, W)
+
+        if not return_dict:
+            return (x,)
+        return Transformer2DModelOutput(sample=x)

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -0,0 +1,695 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video).
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable, Mapping
+from contextlib import nullcontext
+from typing import Any, ClassVar
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from diffusers import AutoencoderKLWan
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.models.lingbot_video.lingbot_video_transformer import LingBotVideoTransformer3DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+TOKEN_LENGTH = 37698
+HIDDEN_STATE_SKIP_LAYER = 0
+LOW_NOISE_TAIL_V1_DEFAULT_STEPS = 2
+
+PROMPT_TEMPLATE = (
+    "<|im_start|>system\nGiven a user input that may include a text prompt alone, "
+    "a text prompt with an image reference, or a text prompt with a video reference "
+    'or a video reference alone, generate an "Enhanced prompt" that provides detailed '
+    "visual descriptions suitable for video generation. Evaluate the level of detail "
+    "in the user's input: if it is simple, enrich it by adding specifics about colors, "
+    "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
+    "progression, and spatial relationships to create vivid, concrete, and temporally "
+    "coherent scenes to create vivid and concrete scenes. Please generate only the "
+    "enhanced description for the prompt below and avoid including any additional "
+    "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
+DEFAULT_NEGATIVE_PROMPT = (
+    '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", '
+    '"pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", '
+    '"underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], '
+    '"artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", '
+    '"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", '
+    '"signature", "logo", "subtitles", "pillarboxed", "side bars", '
+    '"portrait image in landscape frame"], "temporal_and_motion_stability": ["flickering", '
+    '"jittery", "motion blur", "temporal inconsistency", "warping", "morphing", '
+    '"incoherent motion", "unnatural movement", "static object with sudden jump", '
+    '"frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", '
+    '"unrealistic texture", "deformed bottle", "liquid freezing improperly", '
+    '"distorted reflections"]}}'
+)
+
+
+def _dtype_from_name(value: Any, default: torch.dtype) -> torch.dtype:
+    if isinstance(value, torch.dtype):
+        return value
+    if value is None:
+        return default
+    normalized = str(value).lower()
+    if normalized in {"bf16", "bfloat16", "torch.bfloat16"}:
+        return torch.bfloat16
+    if normalized in {"fp16", "float16", "torch.float16"}:
+        return torch.float16
+    if normalized in {"fp32", "float32", "torch.float32"}:
+        return torch.float32
+    raise ValueError(f"Unsupported LingBot dtype: {value!r}.")
+
+
+def _extract_prompt(req: OmniDiffusionRequest) -> tuple[str, str | None]:
+    prompt_obj = req.prompt
+    if isinstance(prompt_obj, str):
+        return prompt_obj, None
+    if isinstance(prompt_obj, Mapping):
+        prompt = prompt_obj.get("prompt", "")
+        negative_prompt = prompt_obj.get("negative_prompt")
+        return str(prompt), None if negative_prompt is None else str(negative_prompt)
+    raise TypeError(f"Unsupported LingBot prompt type: {type(prompt_obj)!r}.")
+
+
+def _module_dtype(module: torch.nn.Module) -> torch.dtype:
+    try:
+        return next(module.parameters()).dtype
+    except StopIteration:
+        return torch.float32
+
+
+def _module_device(module: torch.nn.Module) -> torch.device:
+    try:
+        return next(module.parameters()).device
+    except StopIteration:
+        return torch.device("cpu")
+
+
+def _transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor:
+    sigma = timestep.float() / 1000.0
+    if transformer_dtype in {torch.bfloat16, torch.float16}:
+        sigma = sigma.to(transformer_dtype)
+    return (sigma * 1000.0).float()
+
+
+def _transformer_autocast(device: torch.device, transformer_dtype: torch.dtype):
+    if device.type != "cuda" or transformer_dtype not in {torch.bfloat16, torch.float16}:
+        return nullcontext()
+    return torch.autocast(device_type="cuda", dtype=transformer_dtype)
+
+
+def _group_global_rank(group: Any | None, group_rank: int) -> int:
+    if group is None:
+        return group_rank
+    get_global_rank = getattr(dist, "get_global_rank", None)
+    if get_global_rank is None:
+        return group_rank
+    return int(get_global_rank(group, group_rank))
+
+
+def _validate_refiner_sigmas(sigmas: np.ndarray, t_thresh: float | None = None) -> np.ndarray:
+    arr = np.asarray(list(sigmas), dtype=np.float64)
+    if arr.ndim != 1 or arr.size == 0:
+        raise ValueError("refiner sigma schedule must be a non-empty 1D list")
+    if not np.all(np.isfinite(arr)):
+        raise ValueError("refiner sigma schedule contains non-finite values")
+    if np.any(arr < 0.0) or np.any(arr > 1.0):
+        raise ValueError(f"refiner sigma schedule values must be in [0, 1], got {arr.tolist()}")
+    if arr.size > 1 and not np.all(np.diff(arr) < 0.0):
+        raise ValueError(f"refiner sigma schedule must be strictly descending, got {arr.tolist()}")
+    if t_thresh is not None and abs(float(arr[0]) - float(t_thresh)) > 1e-6:
+        raise ValueError(f"refiner sigma schedule must start at t_thresh={float(t_thresh)}, got {float(arr[0])}")
+    return arr
+
+
+def _compute_refiner_sigmas(
+    *,
+    sigma_max: float,
+    sigma_min: float,
+    num_inference_steps: int,
+    shift: float,
+    t_thresh: float | None,
+    tail_steps: int = 0,
+) -> np.ndarray | None:
+    if t_thresh is None:
+        return None
+    t_value = float(t_thresh)
+    if not (0.0 < t_value <= 1.0):
+        raise ValueError(f"refiner t_thresh must lie in (0, 1], got {t_value}")
+    steps = int(num_inference_steps)
+    if steps < 1:
+        raise ValueError(f"num_inference_steps must be >= 1, got {steps}")
+    tail = int(tail_steps or 0)
+    if tail < 0:
+        raise ValueError(f"refiner_sigma_tail_steps must be >= 0, got {tail}")
+
+    base = np.linspace(float(sigma_max), float(sigma_min), steps + 1).copy()[:-1]
+    shift_value = float(shift)
+    shifted = shift_value * base / (1.0 + (shift_value - 1.0) * base)
+    eps = 1e-6
+    sigmas = shifted[shifted <= t_value + eps]
+    if sigmas.size == 0 or abs(float(sigmas[0]) - t_value) > eps:
+        sigmas = np.concatenate([[t_value], sigmas])
+    if tail > 0:
+        start = float(sigmas[-1])
+        stop = min(float(sigma_min), start)
+        extra = np.linspace(start, stop, tail + 2, dtype=np.float64)[1:-1]
+        sigmas = np.concatenate([sigmas, extra])
+    return _validate_refiner_sigmas(sigmas, t_value).astype(np.float32)
+
+
+def _pad_prompt_embeds(
+    embeds: torch.Tensor,
+    mask: torch.Tensor,
+    target_length: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if embeds.shape[0] != 1:
+        raise ValueError(f"batched CFG helper expects batch=1 inputs, got {embeds.shape[0]}")
+    if embeds.shape[1] > target_length:
+        raise ValueError(f"cannot pad length {embeds.shape[1]} down to {target_length}")
+    pad_len = target_length - embeds.shape[1]
+    if pad_len == 0:
+        return embeds, mask
+    embed_pad = torch.zeros(embeds.shape[0], pad_len, embeds.shape[2], dtype=embeds.dtype, device=embeds.device)
+    mask_pad = torch.zeros(mask.shape[0], pad_len, dtype=mask.dtype, device=mask.device)
+    return torch.cat([embeds, embed_pad], dim=1), torch.cat([mask, mask_pad], dim=1)
+
+
+def _batch_cfg_prompt_inputs(
+    prompt_embeds: torch.Tensor,
+    prompt_mask: torch.Tensor,
+    negative_embeds: torch.Tensor,
+    negative_mask: torch.Tensor,
+    *,
+    null_cond_clone_zero: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if null_cond_clone_zero:
+        zero_negative = torch.zeros_like(prompt_embeds)
+        return (
+            torch.cat([prompt_embeds, zero_negative], dim=0),
+            torch.cat([prompt_mask, prompt_mask.clone()], dim=0),
+        )
+
+    target_length = max(int(prompt_embeds.shape[1]), int(negative_embeds.shape[1]))
+    prompt_padded, prompt_mask_padded = _pad_prompt_embeds(prompt_embeds, prompt_mask, target_length)
+    negative_padded, negative_mask_padded = _pad_prompt_embeds(negative_embeds, negative_mask, target_length)
+    return (
+        torch.cat([prompt_padded, negative_padded], dim=0),
+        torch.cat([prompt_mask_padded, negative_mask_padded], dim=0),
+    )
+
+
+def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def post_process_func(frames: torch.Tensor, sampling_params=None):
+        output_type = getattr(sampling_params, "output_type", None) or "pt"
+        if output_type == "np" and isinstance(frames, torch.Tensor):
+            return frames.float().cpu().numpy()
+        return frames
+
+    return post_process_func
+
+
+class LingBotVideoPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery):
+    """Native vLLM-Omni entry for the dense LingBot-Video checkpoint.
+
+    This PR supports the dense checkpoint only. MoE and fused-expert kernels are
+    intentionally left for the follow-up MoE PR.
+    """
+
+    supports_step_execution: ClassVar[bool] = False
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+        self.od_config = od_config
+        self.device = get_local_device()
+        self.vae_scale_factor_temporal = 4
+        self.vae_scale_factor_spatial = 8
+        self.token_length = TOKEN_LENGTH
+        self.hidden_state_skip_layer = HIDDEN_STATE_SKIP_LAYER
+        self.prompt_template = PROMPT_TEMPLATE
+        self._crop_start: int | None = None
+
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+        dtype = getattr(od_config, "dtype", torch.bfloat16)
+        model_config = getattr(od_config, "model_config", None) or {}
+        transformer_dtype = _dtype_from_name(model_config.get("transformer_dtype"), dtype)
+        text_encoder_dtype = _dtype_from_name(model_config.get("text_encoder_dtype"), dtype)
+        vae_dtype = _dtype_from_name(model_config.get("vae_dtype"), torch.float32)
+
+        transformer_subfolder = str(model_config.get("transformer_subfolder", "transformer"))
+        text_encoder_subfolder = str(model_config.get("text_encoder_subfolder", "text_encoder"))
+        processor_subfolder = str(model_config.get("processor_subfolder", "processor"))
+        vae_subfolder = str(model_config.get("vae_subfolder", "vae"))
+        scheduler_subfolder = str(model_config.get("scheduler_subfolder", "scheduler"))
+
+        self.transformer = LingBotVideoTransformer3DModel.from_pretrained(
+            model,
+            subfolder=transformer_subfolder,
+            torch_dtype=transformer_dtype,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        text_encoder_kwargs: dict[str, Any] = {
+            "dtype": text_encoder_dtype,
+            "local_files_only": local_files_only,
+        }
+        self.text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(
+            model,
+            subfolder=text_encoder_subfolder,
+            **text_encoder_kwargs,
+        ).to(self.device)
+        self.processor = Qwen3VLProcessor.from_pretrained(
+            model,
+            subfolder=processor_subfolder,
+            local_files_only=local_files_only,
+        )
+        self.vae = AutoencoderKLWan.from_pretrained(
+            model,
+            subfolder=vae_subfolder,
+            torch_dtype=vae_dtype,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        self.scheduler = FlowUniPCMultistepScheduler.from_pretrained(
+            model,
+            subfolder=scheduler_subfolder,
+            local_files_only=local_files_only,
+        )
+        self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True)))
+        self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+        super().to(*args, **kwargs)
+        if device is not None:
+            self.device = torch.device(device)
+            self.transformer.to(device=self.device, non_blocking=non_blocking)
+            self.text_encoder.to(device=self.device, non_blocking=non_blocking)
+            self.vae.to(device=self.device, non_blocking=non_blocking)
+        if dtype is not None:
+            self.transformer.to(dtype=dtype, non_blocking=non_blocking)
+        return self
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        consumed = list(weights)
+        if consumed:
+            raise RuntimeError(
+                f"{self.__class__.__name__}.load_weights received {len(consumed)} weight tensors; "
+                "LingBot components are loaded directly from subfolders during __init__."
+            )
+        return {name for name, _ in self.named_parameters()}
+
+    @staticmethod
+    def check_inputs(height: int, width: int, num_frames: int) -> None:
+        if num_frames != 1 and (num_frames - 1) % 4 != 0:
+            raise ValueError(f"`num_frames` must be 1 or 4n+1, got {num_frames}.")
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"`height` and `width` must be multiples of 16, got {height}x{width}.")
+
+    @staticmethod
+    def apply_text_to_template(text: str, template: str = PROMPT_TEMPLATE) -> str:
+        return template.format(text)
+
+    def _compute_crop_start(self) -> int:
+        if self._crop_start is None:
+            marker = "<|USER_INPUT_MARKER|>"
+            marked = self.prompt_template.format(marker)
+            marker_pos = marked.find(marker)
+            if marker_pos < 0:
+                self._crop_start = 0
+            else:
+                prefix = self.processor(
+                    text=marked[:marker_pos],
+                    images=None,
+                    videos=None,
+                    return_tensors="pt",
+                )
+                self._crop_start = int(prefix["input_ids"].shape[1])
+        return self._crop_start
+
+    def _build_prompt_inputs(self, prompt: str | list[str]):
+        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
+        texts = [self.apply_text_to_template(text, self.prompt_template) for text in prompts]
+        return self.processor(
+            text=texts,
+            images=None,
+            videos=None,
+            do_resize=False,
+            truncation=True,
+            max_length=self.token_length,
+            padding="longest",
+            return_tensors="pt",
+        )
+
+    @torch.no_grad()
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        *,
+        device: str | torch.device | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = torch.device(device) if device is not None else self.device
+        inputs = self._build_prompt_inputs(prompt).to(device)
+        outputs = self.text_encoder(
+            **inputs,
+            output_hidden_states=self.hidden_state_skip_layer is not None,
+        )
+        if self.hidden_state_skip_layer is not None:
+            prompt_embeds = outputs.hidden_states[-(self.hidden_state_skip_layer + 1)]
+        else:
+            prompt_embeds = outputs.last_hidden_state
+
+        prompt_mask = inputs["attention_mask"]
+        crop_start = self._compute_crop_start()
+        if crop_start > 0:
+            prompt_embeds = prompt_embeds[:, crop_start:]
+            prompt_mask = prompt_mask[:, crop_start:]
+
+        if prompt_embeds.shape[0] == 1:
+            true_len = int(prompt_mask[0].sum().item())
+            prompt_embeds = prompt_embeds[:, :true_len]
+            prompt_mask = prompt_mask[:, :true_len]
+        return prompt_embeds, prompt_mask
+
+    def prepare_latents(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        generator: torch.Generator | None,
+        latents: torch.Tensor | None,
+        device: torch.device,
+    ) -> torch.Tensor:
+        latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        latent_height = height // self.vae_scale_factor_spatial
+        latent_width = width // self.vae_scale_factor_spatial
+        shape = (1, self.transformer.config.in_channels, latent_frames, latent_height, latent_width)
+        if latents is None:
+            return randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        if tuple(latents.shape) != shape:
+            raise ValueError(f"`latents` shape must be {shape}, got {tuple(latents.shape)}.")
+        return latents.to(device=device, dtype=torch.float32)
+
+    def _dit_latent_to_vae(self, latents: torch.Tensor) -> torch.Tensor:
+        mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32)
+        std_inv = 1.0 / torch.tensor(self.vae.config.latents_std, device=latents.device, dtype=torch.float32)
+        mean = mean.view(1, -1, 1, 1, 1)
+        std_inv = std_inv.view(1, -1, 1, 1, 1)
+        return latents.float() / std_inv + mean
+
+    @torch.no_grad()
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        vae_device = _module_device(self.vae)
+        vae_dtype = _module_dtype(self.vae)
+        vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32)
+        if vae_latents.ndim == 5:
+            vae_latents = vae_latents.contiguous(memory_format=torch.channels_last_3d)
+        autocast_dtype = (
+            vae_dtype if vae_device.type == "cuda" and vae_dtype in {torch.bfloat16, torch.float16} else None
+        )
+        with torch.autocast("cuda", dtype=autocast_dtype or torch.bfloat16, enabled=autocast_dtype is not None):
+            decoded = self.vae.decode(vae_latents)
+        frames = decoded[0] if isinstance(decoded, tuple) else decoded.sample
+        frames = frames.float().clamp_(-1, 1)
+        frames = (frames + 1.0) / 2.0
+        frames = frames.permute(0, 2, 3, 4, 1).cpu()
+        return frames[0]
+
+    @torch.no_grad()
+    def _generate(
+        self,
+        *,
+        prompt: str,
+        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
+        height: int = 480,
+        width: int = 480,
+        num_frames: int = 81,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 6.0,
+        shift: float = 3.0,
+        generator: torch.Generator | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_mask: torch.Tensor | None = None,
+        output_type: str = "pt",
+        cfg_parallel_group: Any | None = None,
+        batch_cfg: bool = False,
+        null_cond_clone_zero: bool = False,
+        t_thresh: float | None = None,
+        refiner_sigma_tail_steps: int = LOW_NOISE_TAIL_V1_DEFAULT_STEPS,
+        offload_vae_during_denoise: bool = False,
+        **extra_args,
+    ) -> torch.Tensor:
+        del extra_args
+        self.check_inputs(height, width, num_frames)
+        device = self.device
+        do_cfg = guidance_scale > 1.0
+        effective_batch_cfg = bool(batch_cfg)
+        cfg_parallel = cfg_parallel_group is not None
+        cfg_parallel_rank = 0
+        if cfg_parallel:
+            if not dist.is_available() or not dist.is_initialized():
+                raise ValueError("`cfg_parallel_group` requires an initialized process group.")
+            if effective_batch_cfg:
+                raise ValueError("`cfg_parallel_group` and `batch_cfg` are mutually exclusive.")
+            if not do_cfg:
+                raise ValueError("CFG parallel requires `guidance_scale > 1.0`.")
+            cfg_parallel_rank = dist.get_rank(cfg_parallel_group)
+            cfg_parallel_world_size = dist.get_world_size(cfg_parallel_group)
+            if cfg_parallel_world_size != 2:
+                raise ValueError(f"CFG parallel currently requires exactly 2 ranks, got {cfg_parallel_world_size}.")
+
+        if prompt_embeds is not None:
+            if prompt_mask is None:
+                raise ValueError("`prompt_mask` is required when `prompt_embeds` is provided.")
+            prompt_embeds = prompt_embeds.to(device=device)
+            prompt_mask = prompt_mask.to(device=device)
+        if negative_prompt_embeds is not None:
+            if negative_prompt_mask is None:
+                raise ValueError("`negative_prompt_mask` is required when `negative_prompt_embeds` is provided.")
+            negative_prompt_embeds = negative_prompt_embeds.to(device=device)
+            negative_prompt_mask = negative_prompt_mask.to(device=device)
+
+        negative_embeds = None
+        negative_mask = None
+        if cfg_parallel and cfg_parallel_rank == 1:
+            if negative_prompt_embeds is not None:
+                negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+            else:
+                negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
+            prompt_embeds = prompt_mask = None
+        else:
+            if prompt_embeds is None:
+                prompt_embeds, prompt_mask = self.encode_prompt(prompt, device=device)
+            if do_cfg and not cfg_parallel:
+                if null_cond_clone_zero:
+                    negative_embeds = torch.zeros_like(prompt_embeds)
+                    negative_mask = prompt_mask.clone()
+                elif negative_prompt_embeds is not None:
+                    negative_embeds, negative_mask = negative_prompt_embeds, negative_prompt_mask
+                else:
+                    negative_embeds, negative_mask = self.encode_prompt(negative_prompt, device=device)
+
+        latents = self.prepare_latents(num_frames, height, width, generator, latents, device)
+        sigmas = _compute_refiner_sigmas(
+            sigma_max=float(self.scheduler.sigma_max),
+            sigma_min=float(self.scheduler.sigma_min),
+            num_inference_steps=num_inference_steps,
+            shift=shift,
+            t_thresh=t_thresh,
+            tail_steps=refiner_sigma_tail_steps,
+        )
+        if sigmas is None:
+            self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
+        else:
+            self.scheduler.set_timesteps(int(sigmas.shape[0]), device=device, sigmas=sigmas, shift=1.0)
+
+        transformer_dtype = _module_dtype(self.transformer)
+        vae_restore_device: torch.device | None = None
+        vae_offloaded = False
+        if offload_vae_during_denoise and output_type != "latent":
+            vae_device = _module_device(self.vae)
+            if vae_device.type == "cuda":
+                self.vae.to("cpu")
+                torch.accelerator.empty_cache()
+                vae_restore_device = vae_device
+                vae_offloaded = True
+
+        cfg_latent_src = _group_global_rank(cfg_parallel_group, 0)
+        cfg_uncond_src = _group_global_rank(cfg_parallel_group, 1)
+        for timestep in self.progress_bar(self.scheduler.timesteps):
+            if cfg_parallel:
+                dist.broadcast(latents, src=cfg_latent_src, group=cfg_parallel_group)
+            timestep_batch = _transformer_timestep(timestep, transformer_dtype).expand(1).to(device)
+            latent_model_input = latents
+            if cfg_parallel:
+                if cfg_parallel_rank == 0:
+                    branch_embeds = prompt_embeds
+                    branch_mask = prompt_mask
+                else:
+                    branch_embeds = negative_embeds
+                    branch_mask = negative_mask
+                if branch_embeds is None:
+                    raise RuntimeError("CFG branch embeddings were not initialized.")
+                with _transformer_autocast(device, transformer_dtype):
+                    branch_noise_pred = self.transformer(
+                        latent_model_input,
+                        timestep_batch,
+                        branch_embeds.to(transformer_dtype),
+                        encoder_attention_mask=branch_mask,
+                        return_dict=False,
+                    )[0].float()
+                if cfg_parallel_rank == 0:
+                    noise_pred = branch_noise_pred
+                    noise_pred_uncond = torch.empty_like(noise_pred)
+                else:
+                    noise_pred_uncond = branch_noise_pred
+                dist.broadcast(noise_pred_uncond, src=cfg_uncond_src, group=cfg_parallel_group)
+                if cfg_parallel_rank != 0:
+                    continue
+                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+            else:
+                if prompt_embeds is None:
+                    raise RuntimeError("Prompt embeddings were not initialized.")
+                prompt_model_input = prompt_embeds.to(transformer_dtype)
+                if do_cfg and effective_batch_cfg:
+                    if negative_embeds is None or negative_mask is None:
+                        raise RuntimeError("Negative embeddings were not initialized for CFG.")
+                    cfg_embeds, cfg_mask = _batch_cfg_prompt_inputs(
+                        prompt_model_input,
+                        prompt_mask,
+                        negative_embeds.to(transformer_dtype),
+                        negative_mask,
+                        null_cond_clone_zero=False,
+                    )
+                    cfg_latents = torch.cat([latent_model_input, latent_model_input], dim=0)
+                    cfg_timesteps = torch.cat([timestep_batch, timestep_batch], dim=0)
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_batched = self.transformer(
+                            cfg_latents,
+                            cfg_timesteps,
+                            cfg_embeds,
+                            encoder_attention_mask=cfg_mask,
+                            return_dict=False,
+                        )[0].float()
+                    noise_pred, noise_pred_uncond = noise_batched.chunk(2, dim=0)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+                else:
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_pred = self.transformer(
+                            latent_model_input,
+                            timestep_batch,
+                            prompt_model_input,
+                            encoder_attention_mask=prompt_mask,
+                            return_dict=False,
+                        )[0].float()
+
+                if do_cfg and not effective_batch_cfg:
+                    if negative_embeds is None or negative_mask is None:
+                        raise RuntimeError("Negative embeddings were not initialized for CFG.")
+                    with _transformer_autocast(device, transformer_dtype):
+                        noise_pred_uncond = self.transformer(
+                            latent_model_input,
+                            timestep_batch,
+                            negative_embeds.to(transformer_dtype),
+                            encoder_attention_mask=negative_mask,
+                            return_dict=False,
+                        )[0].float()
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred - noise_pred_uncond)
+
+            latents = self.scheduler.step(noise_pred, timestep, latents, return_dict=False, generator=generator)[0]
+
+        if cfg_parallel:
+            dist.barrier(group=cfg_parallel_group)
+            if cfg_parallel_rank != 0:
+                return latents if output_type == "latent" else []
+
+        if output_type == "latent":
+            return latents
+        if output_type in {"pt", "np"}:
+            if vae_offloaded and vae_restore_device is not None:
+                self.vae.to(device=vae_restore_device)
+                torch.accelerator.empty_cache()
+            return self._decode_latents(latents)
+        raise ValueError(f"Unsupported output_type: {output_type}")
+
+    @torch.inference_mode()
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        if req.num_reqs != 1:
+            raise ValueError(f"LingBotVideoPipeline only supports one request per batch, got {req.num_reqs}.")
+        request = req.requests[0]
+        prompt, prompt_negative = _extract_prompt(request)
+        sampling = request.sampling_params
+        extra_args = dict(sampling.extra_args or {})
+
+        generator = sampling.generator
+        if isinstance(generator, list):
+            generator = generator[0] if generator else None
+        if generator is None and sampling.seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(int(sampling.seed))
+
+        height = sampling.height if sampling.height is not None else extra_args.pop("height", 480)
+        width = sampling.width if sampling.width is not None else extra_args.pop("width", 480)
+        num_frames = sampling.num_frames or extra_args.pop("num_frames", 81)
+        num_inference_steps = (
+            sampling.num_inference_steps
+            if sampling.num_inference_steps is not None
+            else extra_args.pop("num_inference_steps", 40)
+        )
+        guidance_scale = (
+            sampling.guidance_scale
+            if sampling.guidance_scale_provided or sampling.guidance_scale > 0
+            else extra_args.pop("guidance_scale", 6.0)
+        )
+        shift = extra_args.pop(
+            "shift",
+            extra_args.pop("flow_shift", getattr(self.od_config, "flow_shift", None) or 3.0),
+        )
+        negative_prompt = extra_args.pop("negative_prompt", prompt_negative or self.default_negative_prompt)
+        output_type = (
+            sampling.output_type or getattr(self.od_config, "output_type", None) or extra_args.pop("output_type", "pt")
+        )
+        if output_type not in {"pt", "np", "latent"}:
+            output_type = "pt"
+        sampling.output_type = output_type
+
+        frames = self._generate(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            shift=shift,
+            generator=generator,
+            latents=sampling.latents,
+            output_type=output_type,
+            batch_cfg=bool(extra_args.pop("batch_cfg", False)),
+            null_cond_clone_zero=bool(extra_args.pop("null_cond_clone_zero", False)),
+            offload_vae_during_denoise=bool(extra_args.pop("offload_vae_during_denoise", False)),
+            t_thresh=extra_args.pop("t_thresh", None),
+            refiner_sigma_tail_steps=int(extra_args.pop("refiner_sigma_tail_steps", LOW_NOISE_TAIL_V1_DEFAULT_STEPS)),
+        )
+        return DiffusionOutput(output=frames)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -17,6 +17,7 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -554,8 +555,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -592,8 +593,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -1006,7 +1007,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/diffusion/registry.py

@@ -256,6 +256,11 @@ _DIFFUSION_MODELS = {
         "pipeline_hunyuan_video_1_5_i2v",
         "HunyuanVideo15I2VPipeline",
     ),
+    "LingBotVideoPipeline": (
+        "lingbot_video",
+        "pipeline_lingbot_video",
+        "LingBotVideoPipeline",
+    ),
     "MagiHumanPipeline": (
         "magi_human",
         "pipeline_magi_human",
@@ -526,6 +531,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Flux2Pipeline": "get_flux2_post_process_func",
     "HunyuanVideo15Pipeline": "get_hunyuan_video_15_post_process_func",
     "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func",
+    "LingBotVideoPipeline": "get_lingbot_video_post_process_func",
     "MagiHumanPipeline": "get_magi_human_post_process_func",
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",

vllm_omni/engine/arg_utils.py

@@ -434,9 +434,6 @@ class OrchestratorArgs:
     stage_configs_path: str | None = None
     deploy_config: str | None = None
     stage_overrides: str | None = None  # raw JSON string; parsed downstream
-    # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
-    # derived sizing onto merged stages, then drops it before per-stage engine args.
-    strategy_config: str | None = None
 
     # === Mode Switches (orchestrator reads, DeployConfig redistributes) ===
     async_chunk: bool | None = None

vllm_omni/engine/async_omni_engine.py

@@ -58,10 +58,7 @@ from vllm_omni.engine.stage_runtime import (
     create_stage_runtime,
 )
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import (
-    load_and_resolve_stage_configs,
-    parse_stage_overrides,
-)
+from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 
@@ -1126,41 +1123,11 @@ class AsyncOmniEngine:
 
         return result
 
-    def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
-        """Apply a strategy-derived ``omni_lb_policy`` to the engine.
-
-        Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is
-        the engine default and is treated as "unset" (indistinguishable from no
-        flag), so a strategy value overrides it. If the user explicitly passed a
-        non-default policy that conflicts with the strategy-derived one, raise so
-        the mismatch is not silently ignored.
-        """
-        if not derived:
-            return
-        explicit = kwargs.get("omni_lb_policy")
-        user_set = explicit is not None and str(explicit) != "random"
-        if user_set:
-            if str(explicit) != str(derived):
-                raise ValueError(
-                    f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given "
-                    f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. "
-                    "Drop --omni-lb-policy to use the strategy value, or make them match."
-                )
-            return
-        if self._omni_lb_policy != str(derived):
-            logger.info(
-                "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).",
-                derived,
-                self._omni_lb_policy,
-            )
-            self._omni_lb_policy = str(derived)
-
     def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
         stage_configs_path = kwargs.get("stage_configs_path", None)
         deploy_config_path = kwargs.pop("deploy_config", None)
-        strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
         explicit_stage_configs = kwargs.pop("stage_configs", None)
         if explicit_stage_configs is not None:
@@ -1175,23 +1142,27 @@ class AsyncOmniEngine:
             base_kwargs = kwargs
 
         # Parse --stage-overrides JSON string if provided
-        stage_overrides = parse_stage_overrides(stage_overrides_json)
+        stage_overrides = None
+        if stage_overrides_json:
+            if isinstance(stage_overrides_json, str):
+                try:
+                    stage_overrides = json.loads(stage_overrides_json)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(
+                        f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}"
+                    ) from exc
+            else:
+                stage_overrides = stage_overrides_json
 
-        config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
+        config_path, stage_configs = load_and_resolve_stage_configs(
             model,
             stage_configs_path,
             base_kwargs,
             default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
         )
 
-        # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is
-        # an orchestrator-level knob (read once at construction), so apply it here
-        # rather than as a per-stage config field.
-        self._apply_strategy_lb_policy(strategy_lb_policy, kwargs)
-
         # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
         for cfg in stage_configs:
             try:
@@ -1567,18 +1538,16 @@ class AsyncOmniEngine:
         # and can cause CUDA OOM for subsequent engine instances (especially
         # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
         #
-        # CuMemAllocator.sleep() is NOT idempotent — calling it on already-
-        # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator
-        # cuMemRelease (double-free of the memory handle).  Use release_pools()
-        # instead, which is the designed cleanup path: it drops MemPool refs
-        # and lets the destructor/free path handle asleep entries correctly
-        # (returns a null handle so the C extension skips unmap/release).
+        # Discard mode (level=2) is correct at shutdown: there is no benefit to
+        # keeping a CPU backup when the engine is being torn down.
         try:
             from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
 
             if cumem_available:
                 allocator = CuMemAllocator.get_instance()
-                allocator.release_pools()
+                # Sleep at level 2 discards all pool memory from the GPU
+                # without creating CPU backups — cheapest and fastest.
+                allocator.sleep()
                 logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
         except Exception:
             pass

vllm_omni/engine/output_processor.py

@@ -28,12 +28,6 @@ from vllm_omni.outputs import OmniRequestOutput
 
 logger = init_logger(__name__)
 
-# Keys whose values may arrive as scalar tensors (e.g. sample rate) but
-# represent metadata, not concatenatable data.  When these appear in
-# ``MultimodalPayload.tensors`` they must be relocated to ``.metadata``
-# before consolidation so the tensor concatenation path does not attempt
-# a dimensioned cat on a 0-d tensor.
-_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"})
 
 # ---------------------------------------------------------------------------
 # Module-level helpers
@@ -218,18 +212,6 @@ class OmniRequestState(RequestState):
             modality = OutputModality.TEXT
         strategy = get_accumulation_strategy(modality)
 
-        # Move known metadata scalar tensors (e.g. sample rate) from
-        # tensors to metadata so they are handled by the REPLACE path
-        # below instead of failing tensor concatenation.
-        for key in _METADATA_TENSOR_KEYS:
-            v = self.mm_accumulated.tensors.pop(key, None)
-            if v is not None:
-                if isinstance(v, list):
-                    # Keep the last value (REPLACE semantics)
-                    self.mm_accumulated.metadata[key] = v[-1]
-                else:
-                    self.mm_accumulated.metadata[key] = v
-
         try:
             for k, v in list(self.mm_accumulated.tensors.items()):
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):

vllm_omni/engine/stage_engine_core_client.py

@@ -230,7 +230,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
-        logger.debug(
+        logger.info(
             "[%s] stage-%s [rep-%s] add request: %s",
             self.__class__.__name__,
             self.stage_id,

vllm_omni/engine/stage_engine_startup.py

@@ -138,21 +138,6 @@ class StageReplicaResources:
 # ---------------------------------------------------------------------------
 
 
-def _port_from_zmq_address(address: str | None) -> int | None:
-    """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``.
-
-    Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return
-    ``None`` so callers can skip seeding them into the port-dedup set.
-    """
-    if not address:
-        return None
-    tail = address.rsplit(":", 1)[-1]
-    try:
-        return int(tail)
-    except ValueError:
-        return None
-
-
 class OmniMasterServer:
     """Registration server for single-stage engine startup."""
 
@@ -173,25 +158,10 @@ class OmniMasterServer:
         self._stage_configs: dict[StageRoute, Any] = {}
         self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {}
         self._stage_config_events: dict[StageRoute, threading.Event] = {}
-        # Ports already handed out by *this* server. ``get_open_ports_list`` only
-        # guarantees uniqueness *within* a single call; a per-route call cannot
-        # see ports drawn for other routes, so two stages/replicas can draw the
-        # same ephemeral port and the second engine to ``bind()`` it dies with
-        # ``zmq.error.ZMQError: Address already in use``. Multi-stage models
-        # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit
-        # this intermittently. We dedup every allocated port against this set,
-        # seeded with the registration port (and coordinator ROUTER port) that
-        # are already bound on the same host.
-        self._allocated_ports: set[int] = set()
-        if master_port:
-            self._allocated_ports.add(int(master_port))
         # Coordinator ROUTER address echoed back in every registration reply
         # so OmniCoordClientForStage knows where to connect from inside the
         # engine subprocess.
         self._coordinator_router_address = coordinator_router_address
-        coord_port = _port_from_zmq_address(coordinator_router_address)
-        if coord_port is not None:
-            self._allocated_ports.add(coord_port)
         # Fires only for *newly assigned* (auto-assigned) replicas, not for
         # head-side pre-allocated slots that already have head-side clients.
         self._on_register = on_register
@@ -265,7 +235,7 @@ class OmniMasterServer:
 
         self._stage_config_events[route] = threading.Event()
         self._stage_coordinator_addresses[route] = StageCoordinatorAddresses()
-        hs_port, inp_port, out_port = self._alloc_unique_ports(3)
+        hs_port, inp_port, out_port = get_open_ports_list(count=3)
         alloc = StageAllocation(
             handshake_bind_address=f"tcp://{self._address}:{hs_port}",
             handshake_connect_address=f"tcp://{self._address}:{hs_port}",
@@ -277,32 +247,6 @@ class OmniMasterServer:
         self._stage_routes[route] = alloc
         return alloc
 
-    def _alloc_unique_ports(self, count: int) -> list[int]:
-        """Return ``count`` open ports unique across every route this server owns.
-
-        ``get_open_ports_list`` dedups only within its own call, so we redraw any
-        port that collides with one already recorded in ``self._allocated_ports``
-        (registration/coordinator ports plus every previously-allocated route
-        port) and register the winners before returning. Callers already hold the
-        relevant allocation context (``__init__`` is single-threaded;
-        registration holds ``self._alloc_lock``).
-        """
-        picked: list[int] = []
-        # Bounded retry budget so a pathological host (near port exhaustion)
-        # fails loudly instead of spinning forever.
-        for _ in range(64):
-            for port in get_open_ports_list(count=count - len(picked)):
-                if port in self._allocated_ports:
-                    continue
-                self._allocated_ports.add(port)
-                picked.append(port)
-                if len(picked) == count:
-                    return picked
-        raise RuntimeError(
-            f"[OmniMasterServer] Could not allocate {count} unique open ports "
-            f"(host={self._address}); {len(self._allocated_ports)} already in use."
-        )
-
     def _next_free_replica_id(self, stage_id: int) -> int:
         """Return the next replica id to assign for an auto-assign registration.
 

vllm_omni/entrypoints/cli/serve.py

@@ -254,15 +254,6 @@ class OmniServeCommand(CLISubcommand):
             help="Path to a deploy config YAML (new format with stages/engine_args). "
             "Mutually exclusive with --stage-configs-path.",
         )
-        omni_config_group.add_argument(
-            "--strategy-config",
-            type=str,
-            default=None,
-            help="Path to a composable-parallel strategy.yaml. Only applies to "
-            "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing "
-            "is overlaid onto the registry-merged stages before per-stage engine "
-            "args are built.",
-        )
         omni_config_group.add_argument(
             "--stage-overrides",
             type=str,
@@ -775,10 +766,7 @@ def run_headless(args: TrackingNamespace) -> None:
         load_omni_transfer_config_for_model,
         prepare_engine_environment,
     )
-    from vllm_omni.entrypoints.utils import (
-        load_and_resolve_stage_configs,
-        parse_stage_overrides,
-    )
+    from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
 
     model = args.model
     stage_id: int | None = args.stage_id
@@ -814,18 +802,11 @@ def run_headless(args: TrackingNamespace) -> None:
             args.replica_id,
         )
 
-    # Parse --stage-overrides (raw JSON string) exactly like the standard
-    # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and
-    # standard launches resolve to the same per-stage device layout.
-    stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides"))
-
-    config_path, stage_configs, _ = load_and_resolve_stage_configs(
+    config_path, stage_configs = load_and_resolve_stage_configs(
         model,
         stage_configs_path,
         args_dict,
         deploy_config_path=args_dict.get("deploy_config"),
-        stage_overrides=stage_overrides,
-        strategy_config_path=args_dict.get("strategy_config"),
     )
 
     # Locate the stage config that matches stage_id.

vllm_omni/entrypoints/openai/api_server.py

@@ -32,7 +32,7 @@ from starlette.routing import Route
 from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
-from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
+from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
@@ -60,11 +60,12 @@ from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
-from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
+from vllm.entrypoints.serve.disagg.serving import ServingTokens
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
 from vllm.entrypoints.serve.instrumentator.basic import base
+from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
@@ -84,7 +85,6 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
     OpenAIServingTranslation,
 )
 from vllm.logger import init_logger
-from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
@@ -451,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     # Add process-specific prefix to stdout and stderr.
     decorate_logs("APIServer", skip_if_decorated=True)
 
-    listen_address, sock = setup_openai_server(args, reuse_port=False)
+    listen_address, sock = setup_openai_server(args)
 
     # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -830,12 +830,6 @@ async def omni_init_app_state(
         if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
             resolved_chat_template = _load_model_chat_template_json(args.model)
 
-    chat_template_config = ChatTemplateConfig(
-        chat_template=resolved_chat_template,
-        chat_template_content_format=args.chat_template_content_format,
-        trust_request_chat_template=args.trust_request_chat_template,
-    )
-
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()
         assert isinstance(tool_server, DemoToolServer)
@@ -901,10 +895,11 @@ async def omni_init_app_state(
     await state.openai_serving_models.init_static_loras()
 
     # NOTE: kept aligned with upstream `init_app_state`:
-    # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream).
-    state.online_renderer = OnlineRenderer(
+    # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split).
+    state.openai_serving_render = OpenAIServingRender(
         model_config=engine_client.model_config,
         renderer=engine_client.renderer,
+        model_registry=state.openai_serving_models.registry,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -921,7 +916,7 @@ async def omni_init_app_state(
         OpenAIServingResponses(
             engine_client,
             state.openai_serving_models,
-            state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -942,7 +937,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -969,7 +964,7 @@ async def omni_init_app_state(
         OpenAIServingCompletion(
             engine_client,
             state.openai_serving_models,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -984,7 +979,9 @@ async def omni_init_app_state(
             state.openai_serving_models,
             supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if any(task in POOLING_TASKS for task in supported_tasks)
         else None
@@ -994,7 +991,9 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if "embed" in supported_tasks
         else None
@@ -1004,7 +1003,9 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            trust_request_chat_template=args.trust_request_chat_template,
         )
         if "classify" in supported_tasks
         else None
@@ -1013,9 +1014,8 @@ async def omni_init_app_state(
         ServingScores(
             engine_client,
             state.openai_serving_models,
-            supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            chat_template_config=chat_template_config,
+            score_template=resolved_chat_template,
             log_error_stack=args.log_error_stack,
         )
         if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
@@ -1023,7 +1023,7 @@ async def omni_init_app_state(
     )
     state.serving_tokenization = ServingTokenization(
         state.openai_serving_models,
-        state.online_renderer,
+        state.openai_serving_render,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -1055,7 +1055,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            online_renderer=state.online_renderer,
+            openai_serving_render=state.openai_serving_render,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -1074,7 +1074,7 @@ async def omni_init_app_state(
         ServingTokens(
             engine_client,
             state.openai_serving_models,
-            state.online_renderer,
+            state.openai_serving_render,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,

vllm_omni/entrypoints/openai/serving_audio_generate.py

@@ -3,7 +3,7 @@ import asyncio
 import torch
 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -41,7 +41,6 @@ except ImportError:
     soundfile = None
 
 
-from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -65,6 +64,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     ToolCall,
     UsageInfo,
 )
+from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
 )
@@ -150,21 +150,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     # Harmony flag (always False for vllm-omni models)
     use_harmony: bool = False
 
-    @property
-    def tool_call_id_type(self) -> str:
-        """Return the tool call ID type, delegating to model config.
-
-        Upstream vLLM removed the stored ``tool_call_id_type`` attribute
-        from ``OpenAIServingChat`` after the ParserManager refactor; the
-        field is now resolved on demand via ``get_tool_call_id_type``.
-        """
-        try:
-            from vllm.entrypoints.chat_utils import get_tool_call_id_type
-
-            return get_tool_call_id_type(self.model_config)
-        except Exception:
-            return "random"
-
     def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
         """Check if streamed tokens should go through the tool-call parser.
 
@@ -383,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tool_dicts = [tool.model_dump() for tool in request.tools]
 
             if not self.use_harmony:
-                error_check_ret = self.online_renderer.validate_chat_template(
+                error_check_ret = self.openai_serving_render.validate_chat_template(
                     request_chat_template=request.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs,
                     trust_request_chat_template=self.trust_request_chat_template,
@@ -411,7 +396,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
             else:
                 should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self.online_renderer._make_request_with_harmony(
+                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
                     request, should_include_tools
                 )
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -21,12 +21,12 @@ import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
-from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
     RequestResponseMetadata,
 )
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.inputs import tokens_input
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
@@ -1395,7 +1395,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         emb_dim = len(embedding)
         if self._tts_model_type == "ming_tts":
             if emb_dim != 192:
-                raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}")
+                logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim)
         else:
             dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
             if dim_err is not None:

vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py

@@ -25,18 +25,9 @@ class MingTTSAdapter(ARTTSAdapter):
         ref_audio_source = request.ref_audio
         voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
         if ref_audio_source is None and voice_lower in server.uploaded_speakers:
-            speaker_info = server.uploaded_speakers[voice_lower]
-            if speaker_info.get("embedding_source") == "direct":
-                if request.speaker_embedding is None:
-                    request.speaker_embedding = server._get_uploaded_speaker_embedding(request.voice)
-                if request.speaker_embedding is None:
-                    raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing")
-            else:
-                ref_audio_source = server._get_uploaded_audio_data(request.voice)
-                if not ref_audio_source:
-                    raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing")
-                if request.ref_text is None:
-                    request.ref_text = speaker_info.get("ref_text")
+            ref_audio_source = server._get_uploaded_audio_data(request.voice)
+            if request.ref_text is None:
+                request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text")
         ref_audio_data = None
         if isinstance(ref_audio_source, list):
             ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source)

vllm_omni/entrypoints/utils.py

@@ -1,4 +1,3 @@
-import json
 import os
 import types
 from collections import Counter
@@ -317,8 +316,7 @@ def load_stage_configs_from_model(
     base_engine_args: dict | None = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-    strategy_config_path: str | None = None,
-) -> tuple[list, str | None]:
+) -> list:
     """Load stage configurations from model's default config file.
 
     For models registered in the pipeline registry (new path), uses
@@ -332,15 +330,9 @@ def load_stage_configs_from_model(
         base_engine_args: Base engine args to merge as CLI overrides.
         deploy_config_path: Optional explicit deploy config path.
         stage_overrides: Per-stage overrides from --stage-overrides.
-        strategy_config_path: Optional path to a composable-parallel
-            ``strategy.yaml`` whose derived sizing is overlaid onto the
-            registry-merged stages (opt-in; ignored on the legacy YAML path).
 
     Returns:
-        ``(stage_configs, omni_lb_policy)``: the list of stage configuration
-        dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
-        (``None`` when no strategy set one). The policy is returned rather than
-        written into a caller-provided mutable dict.
+        List of stage configuration dictionaries
     """
     if base_engine_args is None:
         base_engine_args = {}
@@ -351,44 +343,26 @@ def load_stage_configs_from_model(
             for key, val in overrides.items():
                 cli_overrides[f"stage_{stage_id_str}_{key}"] = val
 
-    strategy_specs = None
-    if strategy_config_path is not None:
-        from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
-
-        strategy_specs = load_strategy_specs(strategy_config_path)
-
-    stages, omni_lb_policy = StageConfigFactory.create_from_model(
+    stages = StageConfigFactory.create_from_model(
         model,
         trust_remote_code=cli_overrides.get("trust_remote_code", False),
         cli_overrides=cli_overrides,
         deploy_config_path=deploy_config_path,
-        strategy_specs=strategy_specs,
     )
     if stages is not None:
         # Convert StageConfig objects to OmegaConf for backward compat
-        return [stage.to_omegaconf() for stage in stages], omni_lb_policy
+        return [stage.to_omegaconf() for stage in stages]
 
-    # Legacy fallback: load from YAML. A composable-parallel strategy cannot be
-    # applied here (it overlays onto registry-merged stages), so warn rather than
-    # silently dropping the operator's --strategy-config.
-    if strategy_config_path is not None:
-        logger.warning(
-            "--strategy-config (%s) was provided but model %r resolves via the "
-            "legacy stage_configs YAML path, which does not support "
-            "composable-parallel strategies; the strategy is ignored. Use a "
-            "registry-based model to apply it.",
-            strategy_config_path,
-            model,
-        )
+    # Legacy fallback: load from YAML
     stage_config_path = resolve_model_config_path(model)
     if stage_config_path is None:
-        return [], None
+        return []
     stage_configs = load_stage_configs_from_yaml(
         config_path=stage_config_path,
         base_engine_args=base_engine_args,
         prefer_stage_engine_args=True,
     )
-    return stage_configs, None
+    return stage_configs
 
 
 def load_stage_configs_from_yaml(
@@ -527,25 +501,6 @@ def filter_stages(
         return stage_configs
 
 
-def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
-    """Parse the ``--stage-overrides`` value into a per-stage override dict.
-
-    ``value`` may be a raw JSON string (as supplied on the CLI) or an
-    already-parsed mapping. Returns ``None`` when no overrides are given.
-
-    Raises:
-        ValueError: when ``value`` is a string that is not valid JSON.
-    """
-    if not value:
-        return None
-    if isinstance(value, str):
-        try:
-            return json.loads(value)
-        except json.JSONDecodeError as exc:
-            raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc
-    return value
-
-
 def load_and_resolve_stage_configs(
     model: str,
     stage_configs_path: str | None,
@@ -553,8 +508,7 @@ def load_and_resolve_stage_configs(
     default_stage_cfg_factory: Any = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-    strategy_config_path: str | None = None,
-) -> tuple[str, list, str | None]:
+) -> tuple[str, list]:
     """Load stage configurations from model or YAML file with fallback to defaults.
 
     Args:
@@ -567,13 +521,9 @@ def load_and_resolve_stage_configs(
             Mutually exclusive with ``stage_configs_path``.
         stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
             Keys are stage_id strings, values are dicts of overrides.
-        strategy_config_path: Optional path to a composable-parallel
-            ``strategy.yaml`` overlaid onto the registry-merged stages.
 
     Returns:
-        Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
-        the strategy-derived pipeline-wide load-balancer policy (``None`` when no
-        strategy set one), returned for the engine to apply.
+        Tuple of (config_path, stage_configs)
     """
     if stage_configs_path is not None and deploy_config_path is not None:
         raise ValueError(
@@ -601,15 +551,13 @@ def load_and_resolve_stage_configs(
                 stage_configs_path,
             )
 
-    omni_lb_policy: str | None = None
     if deploy_config_path is not None:
         config_path = deploy_config_path
-        stage_configs, omni_lb_policy = load_stage_configs_from_model(
+        stage_configs = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -619,11 +567,10 @@ def load_and_resolve_stage_configs(
                 stage_configs = []
     elif stage_configs_path is None:
         config_path = resolve_model_config_path(model)
-        stage_configs, omni_lb_policy = load_stage_configs_from_model(
+        stage_configs = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             stage_overrides=stage_overrides,
-            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -638,7 +585,7 @@ def load_and_resolve_stage_configs(
     stage_configs = filter_stages(config_path, stage_configs, kwargs)
     logger.debug(f"stage_configs: {stage_configs}")
 
-    return config_path, stage_configs, omni_lb_policy
+    return config_path, stage_configs
 
 
 def get_final_stage_id_for_e2e(

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -23,7 +23,22 @@ from dataclasses import dataclass
 import torch
 from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
 from vllm.v1.kv_cache_interface import SlidingWindowSpec
-from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+try:
+    from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+except ModuleNotFoundError:
+
+    def register_kv_cache_spec(*, manager_class, uniform_type_base_spec=None):
+        """Compatibility shim for vLLM versions without kv_cache_spec_registry."""
+
+        def decorator(spec_class):
+            from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map
+
+            spec_manager_map[spec_class] = manager_class
+            return spec_class
+
+        return decorator
+
 
 # ── Slot mapping ────────────────────────────────────────────────────────────
 

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -28,9 +28,6 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None
-_UNIFORM_EPS = 1e-20
-
 
 # ===================================================================
 # HF-numerics-compatible layers for code predictor
@@ -608,30 +605,6 @@ class CodePredictorWrapper(nn.Module):
                 values.add(parsed)
         return values
 
-    @staticmethod
-    def _normalize_generators(
-        generator: _GeneratorLike, batch_size: int
-    ) -> torch.Generator | list[torch.Generator | None] | None:
-        if generator is None or isinstance(generator, torch.Generator):
-            return generator
-
-        row_generators = list(generator)
-        if len(row_generators) != batch_size:
-            raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.")
-        return row_generators
-
-    @classmethod
-    def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
-        """Sample ``logits`` via Gumbel-max with optional per-row generators."""
-        row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
-        u = torch.empty_like(logits, dtype=torch.float32)
-        if isinstance(row_generators, list):
-            for row, row_generator in enumerate(row_generators):
-                u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
-        else:
-            u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
-        return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
-
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
         all_seq_lens = list(range(2, max_seq))
         if not self._prefix_graph_seq_lens:
@@ -807,7 +780,6 @@ class CodePredictorWrapper(nn.Module):
         bsz = int(layer0_code.shape[0])
         if generators is not None and len(generators) != bsz:
             raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
-        sample_generator: _GeneratorLike = generators if generators is not None else generator
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -884,20 +856,9 @@ class CodePredictorWrapper(nn.Module):
 
             logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
 
-            # Sample next code via Gumbel-max.
-            #
-            # ``argmax_i(logits_i + Gumbel_i)`` with
-            # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is
-            # distributionally identical to sampling from ``softmax(logits)``.
-            # In this file the motivations are practical rather than graph
-            # related: it is measurably cheaper than ``softmax + multinomial``
-            # on the B x 2048 shapes used here, it stays well-defined for
-            # degenerate masked rows with a surviving finite entry (and is more
-            # defensive than ``multinomial`` around fully-masked/NaN inputs),
-            # and the helper below can honor either one batch generator or one
-            # generator per seeded row.
+            # Sample next code
             if stored_mode:
-                # "stored" mode: top-k -> top-p -> Gumbel-max
+                # "stored" mode: top-k -> top-p -> softmax -> multinomial
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)
                     logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
@@ -908,15 +869,17 @@ class CodePredictorWrapper(nn.Module):
                     remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
-                code = self._sample_codes_gumbel(logits, generator=sample_generator)
+                probs = F.softmax(logits, dim=-1, dtype=torch.float32)
+                code = self._multinomial(probs, generator, generators)
             else:
-                # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
+                # "per_call" mode: temperature-scaled + top-k
                 if use_sampling:
                     scaled = logits * inv_temperature
                     if top_k > 0:
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                    code = self._sample_codes_gumbel(scaled, generator=sample_generator)
+                    probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
+                    code = self._multinomial(probs, generator, generators)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -5,8 +5,9 @@
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
   * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
-    bulk tensors arrive via ``text2flow_full_payload`` on the connector.
+    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
+    from the emitted speech tokens. ``text2flow`` is the
+    ``custom_process_input_func`` variant.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
@@ -46,6 +47,7 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.text2flow",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),

vllm_omni/model_executor/models/mimo_audio/pipeline.py

@@ -56,6 +56,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
+            custom_process_input_func=f"{_PROC}.llm2code2wav",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -17,6 +17,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 from vllm.model_executor.models.vision import (
@@ -38,7 +39,6 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
-    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.platforms import current_omni_platform

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py

@@ -4,7 +4,6 @@ from collections.abc import Iterable, Iterator, Mapping
 from typing import Any
 
 import numpy as np
-import PIL.Image as PILImage
 import torch
 from torch import nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -27,21 +26,16 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerProcessingInfo,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
 )
-from vllm.model_executor.models.qwen2_5_omni_thinker import (
-    Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase,
-)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
 )
-from vllm.model_executor.models.qwen2_5_omni_thinker import (
-    Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase,
-)
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -62,13 +56,12 @@ from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
+from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing.processor import (
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.utils.collection_utils import is_list_of
 
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
@@ -81,143 +74,11 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
-def _presampled_videos_hf_kwargs(
-    mm_data: Mapping[str, object],
-    mm_kwargs: Mapping[str, object],
-) -> Mapping[str, object]:
-    """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader.
-
-    When ``video_metadata`` is present (emitted by
-    ``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled
-    according to ``media_io_kwargs``, so the HF processor must not re-sample
-    them, and it needs the sampled fps (instead of its default) to compute
-    ``video_second_per_grid`` correctly. Otherwise the audio/video temporal
-    alignment breaks under ``use_audio_in_video=True``.
-    """
-    video_metadata = mm_data.get("video_metadata")
-    if not video_metadata:
-        return mm_kwargs
-
-    mm_kwargs = dict(mm_kwargs)
-    videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {})
-    videos_kwargs["do_sample_frames"] = False
-
-    def _compute_sampled_video_fps(metadata) -> float | None:
-        duration = getattr(metadata, "duration", None)
-        indices = getattr(metadata, "frames_indices", None)
-        if not duration or not indices or duration <= 0:
-            return None
-        return len(indices) / float(duration)
-
-    # An explicit user-provided fps takes precedence.
-    if "fps" not in videos_kwargs:
-        fps_values = [_compute_sampled_video_fps(m) for m in video_metadata]
-        # HF accepts a single fps per call and uses it for every video's
-        # video_second_per_grid.
-        known_fps = [fps for fps in fps_values if fps is not None]
-        unique_fps = set(known_fps)
-        if len(unique_fps) == 1:
-            videos_kwargs["fps"] = known_fps[0]
-        elif len(unique_fps) > 1:
-            logger.warning(
-                f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}."
-            )
-            videos_kwargs["fps"] = known_fps[0]
-
-    mm_kwargs["videos_kwargs"] = videos_kwargs
-    return mm_kwargs
-
-
-class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
-    """Video items that carry the loader's ``(frames, metadata)`` tuples.
-
-    The tuples are kept in ``data`` so that mm hashing and cache-miss
-    re-parsing retain the metadata; ``get_processor_data`` unpacks them into
-    the ``videos`` + ``video_metadata`` arguments of the HF processor.
-    """
-
-    def get_processor_data(self) -> Mapping[str, object]:
-        from transformers.video_utils import VideoMetadata
-
-        videos = [frames for frames, _ in self.data]
-        video_metadata = [
-            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
-        ]
-        return {"videos": videos, "video_metadata": video_metadata}
-
-
-class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase):
-    def _parse_video_data(self, data):
-        if data is None or isinstance(data, dict) or self.is_embeddings(data):
-            return super()._parse_video_data(data)
-
-        # Normalize to a list of per-video items (same as the base parser).
-        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
-            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
-        ):
-            data_items = [data]
-        elif isinstance(data, (np.ndarray, torch.Tensor)):
-            data_items = [elem for elem in data]
-        elif isinstance(data, tuple) and len(data) == 2:
-            data_items = [data]
-        else:
-            data_items = data
-
-        videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items]
-
-        # Metadata is optional: videos fetched by vLLM's media connector
-        # arrive as (frames, metadata) tuples, while plain arrays (e.g.
-        # offline `multi_modal_data` or dummy profiling inputs) carry no
-        # metadata and parse as before.
-        if any(metadata is None for _, metadata in videos_with_metadata):
-            return super()._parse_video_data(data)
-
-        return Qwen2_5OmniVideoProcessorItems(
-            videos_with_metadata,
-            metadata=[metadata for _, metadata in videos_with_metadata],
-        )
-
-
-class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase):
-    def get_data_parser(self):
-        feature_extractor = self.get_feature_extractor()
-
-        return Qwen2_5OmniThinkerMultiModalDataParser(
-            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
-            target_sr=feature_extractor.sampling_rate,
-            target_channels=self.get_target_channels(),
-            expected_hidden_size=self._get_expected_hidden_size(),
-        )
-
-
 class Qwen2_5OmniThinkerMultiModalProcessor(
     _Qwen2_5OmniThinkerMultiModalProcessorBase,
 ):
     """Override to fix use_audio_in_video detection when mm cache returns None."""
 
-    def _cached_apply_hf_processor(self, inputs, timing_ctx):
-        # If use_audio_in_video, process video and audio together; otherwise,
-        # skip cache to avoid errors.
-        # Prevents partial cache hits from causing processing failures.
-        if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"):
-            return self._apply_hf_processor(inputs, timing_ctx)
-        return super()._cached_apply_hf_processor(inputs, timing_ctx)
-
-    def _call_hf_processor(
-        self,
-        prompt: str,
-        mm_data: Mapping[str, object],
-        mm_kwargs: Mapping[str, object],
-        tok_kwargs: Mapping[str, object],
-    ):
-        mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
-        return super()._call_hf_processor(
-            prompt=prompt,
-            mm_data=mm_data,
-            mm_kwargs=mm_kwargs,
-            tok_kwargs=tok_kwargs,
-        )
-
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -51,6 +51,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             hf_config_name="talker_config",
             engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.thinker2talker",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
@@ -68,6 +69,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
+            custom_process_input_func=f"{_PROC}.talker2code2wav",
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -74,6 +74,7 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
+    Qwen2_5OmniThinkerMultiModalProcessor,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
@@ -119,9 +120,6 @@ from vllm.transformers_utils.processor import cached_processor_from_config
 
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerMultiModalDataParser,
-    Qwen2_5OmniThinkerMultiModalProcessor,
-    _presampled_videos_hf_kwargs,
 )
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
@@ -623,19 +621,6 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
         assert isinstance(feature_extractor, WhisperFeatureExtractor)
         return feature_extractor
 
-    def get_data_parser(self):
-        feature_extractor = self.get_feature_extractor()
-
-        # Install the Omni parser, which keeps video metadata from vLLM's
-        # video loader (see Qwen2_5OmniVideoProcessorItems); the parser
-        # inherited from Qwen2AudioProcessingInfo would drop it.
-        return Qwen2_5OmniThinkerMultiModalDataParser(
-            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
-            target_sr=feature_extractor.sampling_rate,
-            target_channels=1,
-            expected_hidden_size=self._get_expected_hidden_size(),
-        )
-
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
@@ -686,11 +671,8 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         mm_data = dict(mm_data)
-        mm_kwargs = dict(mm_kwargs)
         audios = mm_data.pop("audios", [])
 
-        mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs))
-
         def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
             length = x.shape[-1]
             if length % hop_length != 0:
@@ -945,13 +927,9 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
             audio_in_video_item_idx += 1
 
-            # Prefer the HF-computed value (temporal_patch_size / sampled fps)
-            # over request kwargs;
-            second_per_grid_ts = out_mm_data.get("second_per_grid_ts")
-            if second_per_grid_ts is None:
-                second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
-            if second_per_grid_ts is not None:
-                video_second_per_grid_t = float(second_per_grid_ts[item_idx])
+            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
+            if second_per_grid_ts:
+                video_second_per_grid_t = second_per_grid_ts[item_idx]
             else:
                 video_second_per_grid_t = 2.0
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -13,7 +13,6 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorWrapper,
     CodePredictorWrapperConfig,
 )
-from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
@@ -22,7 +21,7 @@ Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel
 
 
 class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper):
-    """Qwen3-TTS code predictor (per-call sampling, projection)."""
+    """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection)."""
 
     def __init__(
         self,
@@ -36,7 +35,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             vllm_config=vllm_config,
             cp_config=config,
             wrapper_config=CodePredictorWrapperConfig(
-                use_cuda_graphs=current_omni_platform.is_npu(),
+                use_cuda_graphs=True,
                 use_parallel_embedding=False,
                 use_projection=(config.hidden_size != talker_config.hidden_size),
                 return_proj_buf=False,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -324,33 +324,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
-        self.talker_mtp_graph_safe = True
         # talker_mtp samples with per-row generators, so explicitly-seeded
         # requests stay batched instead of one scalar forward per row (#4883).
-        # This is only valid while talker_mtp receives the *unpadded* active
-        # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp
-        # (padded batches + a single captured RNG stream) whenever full
-        # cudagraphs are enabled, so disable per-row generators in that case.
-        # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is
-        # not reproducible (matching Qwen3-Omni, which has no per-row seeding);
-        # eager mode keeps the per-row generator fast-path.
-        # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an
-        # upstream (runner) concern, and the runner already re-derives the same
-        # wrap decision in ``_init_talker_mtp``. Once all TTS models declare
-        # ``talker_mtp_graph_safe`` we should move this gating into the model
-        # runner (models only expose the static capability flag, and users can
-        # then configure ``cudagraph_mode`` freely). This also means mtp seeding
-        # is silently dropped when the decode batch > 1 under full cudagraphs;
-        # the follow-up should refactor the per-row generator in the runner to
-        # be cudagraph-safe so seeded batches stay reproducible.
-        cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None)
-        talker_mtp_graph_wrapped = (
-            self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
-        )
-        self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped
-        self.use_async_omni_output = True
-        self.eager_omni_postprocess_before_async_output = True
-        self.omni_pooler_payload_include_hidden = False
+        # Only valid while talker_mtp receives the unpadded active batch (this
+        # talker is not graph-wrapped); a padded batch would need the runner to
+        # pad the generators list as well.
+        self.talker_mtp_accepts_per_row_generators = True
+
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -119,6 +119,30 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
     return decoded
 
 
+def text2flow(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    _requires_multimodal_data: bool = True,
+):
+    """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs."""
+    engine_inputs: list[OmniTokensPrompt] = []
+    for source_output in source_outputs:
+        output = source_output.outputs[0]
+        multi_modal_data = output.multimodal_output
+        if multi_modal_data is None:
+            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
+
+        prefix_ids = _ensure_list(source_output.prompt_token_ids)
+        raw_output_ids = _ensure_list(output.cumulative_token_ids)
+        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
+        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
+        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
+        additional_info = dict(multi_modal_data)
+        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
+        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
+    return engine_inputs
+
+
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -114,6 +114,37 @@ def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | Non
     raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}")
 
 
+def slow_ar_to_dac_decoder(
+    source_outputs: list[Any],
+    _prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder."""
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    slow_ar_outputs = source_outputs
+    dac_inputs: list[OmniTokensPrompt] = []
+
+    for output in slow_ar_outputs:
+        out = output.outputs[0]
+        # audio_codes shape: [num_frames, num_codebooks]
+        audio_codes = out.multimodal_output["audio_codes"].to(torch.long)
+        # Filter zero-padded frames.
+        valid_mask = audio_codes.any(dim=1)
+        audio_codes = audio_codes[valid_mask]
+        # Codebook-major flat: [num_codebooks * num_frames]
+        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
+        dac_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=None,
+            )
+        )
+    return dac_inputs
+
+
 def slow_ar_to_dac_decoder_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/indextts2.py

@@ -143,6 +143,65 @@ def _request_id(request: Any) -> str:
     return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
 
 
+def talker2s2mel(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Legacy orchestrator path: collect all Stage 0 output, format Stage 1 input."""
+    _old_nt = torch.get_num_threads()
+    torch.set_num_threads(1)
+    try:
+        return _talker2s2mel_impl(source_outputs)
+    finally:
+        torch.set_num_threads(_old_nt)
+
+
+def _talker2s2mel_impl(source_outputs: list[Any]) -> list[Any]:
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    s2mel_inputs: list[OmniTokensPrompt] = []
+
+    for i, talker_output in enumerate(source_outputs):
+        if not talker_output.finished:
+            continue
+
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        if not isinstance(mm, dict):
+            logger.warning("Talker output %d has no multimodal_output dict", i)
+            continue
+
+        codes_dict = mm.get("codes", {})
+        mel_codes = codes_dict.get("mel")
+        if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
+            logger.warning("Talker output %d has empty mel_codes", i)
+            continue
+        mel_codes = mel_codes.to(torch.long)
+
+        hs_dict = mm.get("hidden_states", {})
+        latent = hs_dict.get("latent")
+        if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
+            logger.warning("Talker output %d has empty latent", i)
+            continue
+
+        meta = mm.get("meta", {})
+        if not isinstance(meta, dict):
+            meta = {}
+
+        additional_information = _build_s2mel_additional_information(mel_codes, latent, meta, context="talker2s2mel")
+        s2mel_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0],
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=additional_information,
+            )
+        )
+
+    return s2mel_inputs
+
+
 def talker2s2mel_token_only(
     source_outputs: list[Any],
     prompt: Any = None,

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 import torch
+from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -9,6 +10,7 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
+from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
 
 logger = init_logger(__name__)
@@ -228,6 +230,88 @@ def llm2code2wav_async_chunk(
     )
 
 
+def llm2code2wav(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """
+    Process talker outputs to create code2wav inputs.
+
+    Workflow:
+    1. Extract talker's codec code outputs (8-layer RVQ codes)
+    2. Flatten codes for code2wav input
+    3. Package for code2wav stage
+
+    Args:
+        prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required
+
+    Returns:
+        List of OmniTokensPrompt for code2wav stage
+    """
+    talker_outputs = source_outputs
+    code2wav_inputs = []
+
+    # Process each talker output
+    for i, talker_output in enumerate(talker_outputs):
+        output = talker_output.outputs[0]
+
+        # Extract codec codes from talker output
+        # Expected shape: [8, seq_len] (8-layer RVQ codes)
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
+        mm_hs = mm.get("hidden_states", {})
+        if "audio" in mm_codes:
+            codec_codes = mm_codes["audio"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
+            is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
+            non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
+            if len(non_zero_indices) == 0:
+                # All codec codes are zero - skip this request with a warning
+                request_id = getattr(talker_output, "request_id", f"unknown_{i}")
+                logger.warning(
+                    "Skipping request %s: all codec codes are zero (empty output from Stage-0). "
+                    "This may indicate the model failed to generate valid audio codes.",
+                    request_id,
+                )
+            else:
+                if len(non_zero_indices) < codec_codes.shape[0]:
+                    codec_codes = codec_codes[non_zero_indices]
+        elif "output" in mm_hs and "audio" not in mm_codes:
+            codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
+        else:
+            raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")
+
+        pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
+
+        code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec)
+        code_final = code_final.tolist()
+
+        # Guard against flattened sequences longer than code2wav's max_model_len.
+        # Without this, add_request raises ``could not broadcast input array
+        # from shape (N,) into shape (max_model_len,)`` and kills the engine
+        # core (see issue #2683). Mirrors the offline end2end.py safeguard.
+        if len(code_final) > MAX_CODE2WAV_TOKENS:
+            request_id = getattr(talker_output, "request_id", f"unknown_{i}")
+            logger.warning(
+                "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.",
+                request_id,
+                len(code_final),
+                MAX_CODE2WAV_TOKENS,
+            )
+            code_final = code_final[:MAX_CODE2WAV_TOKENS]
+
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=code_final,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return code2wav_inputs
+
+
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # AR runner's `flatten_payload` converts the model emit
@@ -245,9 +329,8 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
 def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor:
     """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor.
 
-    Shared by the sync placeholder builder (``llm2code2wav_token_only``) and
-    the full-payload producer (``llm2code2wav_full_payload``) so both size the
-    downstream codec sequence off the same non-zero frames.
+    Mirrors the zero-row filter in the orchestrator-path `llm2code2wav`
+    body (see this file's ``llm2code2wav`` around line 224).
     """
     if codec_codes.ndim != 4 or codec_codes.numel() == 0:
         return codec_codes

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -10,6 +10,7 @@ from typing import Any
 
 import torch
 from vllm.inputs import TextPrompt
+from vllm.platforms import current_platform
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -24,7 +25,9 @@ from vllm_omni.data_entry_keys import (
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
+    extract_language_from_prompt,
     extract_language_from_request,
+    extract_speaker_from_prompt,
     extract_speaker_from_request,
 )
 
@@ -566,31 +569,25 @@ def thinker2talker_full_payload(
         output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
         all_token_ids = list(prompt_token_ids) + list(output_token_ids)
 
-    # Drop the terminal stop-token row only when more than one row was
-    # accumulated; trimming a single row would ship 0 conditioning tensors
-    # while ids still has tokens and break talker prefill alignment.
-    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 1:
+    # Trim the trailing stop-token row from the accumulated thinker output.
+    # The accumulator captures one hidden-state row per executed thinker
+    # forward (prefill + every decode step including the one that emitted
+    # the stop_token), so for a finished request thinker_emb has exactly one
+    # row more than the rows the talker should consume. async_chunk's
+    # chunk-0 path naturally captures only the prefill / non-stop portion,
+    # which is why the [async_chunk] parametrization passes while [default]
+    # over-generates one codec frame on short outputs (e.g.
+    # test_one_word_prompt_001[default]: audio extends "London" with
+    # spurious phonemes).
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
         thinker_emb_prefill = thinker_emb[:-1]
     else:
         thinker_emb_prefill = thinker_emb
-    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 1:
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
         thinker_hid_prefill = thinker_hid[:-1]
     else:
         thinker_hid_prefill = thinker_hid
 
-    emb_rows = int(thinker_emb_prefill.shape[0]) if isinstance(thinker_emb_prefill, torch.Tensor) else 0
-    hid_rows = int(thinker_hid_prefill.shape[0]) if isinstance(thinker_hid_prefill, torch.Tensor) else 0
-    if len(all_token_ids) > 0 and (emb_rows == 0 or hid_rows == 0):
-        logger.warning(
-            "thinker2talker_full_payload: empty thinker conditioning for req=%s "
-            "(ids_len=%s embed_rows=%s hidden_rows=%s); withholding payload.",
-            rid,
-            len(all_token_ids),
-            emb_rows,
-            hid_rows,
-        )
-        return None
-
     payload: OmniPayload = {
         "embed": {
             "prefill": thinker_emb_prefill.detach().cpu(),
@@ -602,6 +599,7 @@ def thinker2talker_full_payload(
         "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
+    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
     speaker = extract_speaker_from_request(request)
     if speaker is not None:
         payload["speaker"] = speaker
@@ -611,32 +609,37 @@ def thinker2talker_full_payload(
     return payload
 
 
-def thinker2talker_token_only(
+def thinker2talker(
     source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
     streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """Orchestrator-side placeholder builder for Stage-1 (Talker) when
-    ``async_chunk=False``.
-
-    After the communication-layer refactor, this function only allocates a
-    placeholder ``prompt_token_ids`` of the correct length so the scheduler can
-    reserve KV-cache slots. It does **not** forward bulk tensors or voice
-    metadata.
-
-    Bulk talker conditioning (embed / hidden_states / ids) and per-request
-    voice metadata (``speaker`` / ``language``) are packed by Stage-0
-    ``thinker2talker_full_payload`` via ``extract_speaker_from_request`` /
-    ``extract_language_from_request``, sent over the connector, and merged into
-    the Talker worker's ``model_intermediate_buffer`` on recv. The orchestrator
-    path therefore sets ``additional_information=None``.
-
-    ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature
-    compatibility with other orchestrator input processors; they are unused.
     """
+    Process thinker outputs to create talker inputs.
+
+    Workflow:
+    1. Extract thinker's text generation outputs (token IDs + hidden states)
+    2. Split hidden states into: prompt embeddings + generated embeddings
+    3. Package for talker with additional information
+
+    In PD disaggregation mode, merges prefill-stage prompt embeddings with
+    decode-stage generated embeddings before handing off to the talker.
+
+    Args:
+        prompt: Original prompt data
+        requires_multimodal_data: Whether multimodal data is required
+
+    Returns:
+        List of OmniTokensPrompt for talker stage
+    """
+    thinker_outputs = source_outputs
     talker_inputs: list[OmniTokensPrompt] = []
-    for i, thinker_output in enumerate(source_outputs):
+
+    device = torch.device(current_platform.device_type)
+
+    # Process each thinker output
+    for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
         req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
         prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
@@ -653,12 +656,138 @@ def thinker2talker_token_only(
             )
         thinker_sequences = prompt_token_ids + output_ids
         thinker_input_ids = prompt_token_ids
-        info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
-        prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
+        new_seq_length = len(prompt_token_ids + output_ids) - 1
+        thinker_mm_raw = getattr(output, "multimodal_output", None)
+        if not isinstance(thinker_mm_raw, Mapping):
+            logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id)
+            continue
+        thinker_mm: OmniPayload = thinker_mm_raw
+        mm_hs = thinker_mm.get("hidden_states", {})
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
+        emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY)
+        hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY)
+        if emb_layer is None or hid_layer is None:
+            logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id)
+            continue
+        thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
+        thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
+
+        prefill_mm: dict[str, Any] | None = None
+        prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context)
+
+        if prefill_mm is not None:
+            expected_total = len(prompt_token_ids) + len(output_ids)
+            try:
+                thinker_emb, thinker_hid = _merge_pd_embeddings(
+                    thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total
+                )
+            except Exception as exc:
+                logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
+
+        payload = OmniPayloadStruct(
+            embed=EmbeddingsStruct(
+                prefill=thinker_emb,
+                tts_bos=_resolve_tts_token_embedding(
+                    "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+                tts_eos=_resolve_tts_token_embedding(
+                    "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+                tts_pad=_resolve_tts_token_embedding(
+                    "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
+                ),
+            ),
+            hidden_states=HiddenStatesStruct(output=thinker_hid),
+            ids=IdsStruct(all=thinker_sequences, prompt=thinker_input_ids),
+            speaker=extract_speaker_from_prompt(prompt, index=i),
+            language=extract_language_from_prompt(prompt, index=i),
+        )
+        info = to_dict(payload)
+        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
+
         talker_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=[0] * prompt_len,
-                additional_information=None,
+                additional_information=info,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return talker_inputs
+
+
+def thinker2talker_token_only(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
+) -> list[OmniTokensPrompt]:
+    """Non-async-chunk Stage-1 input builder for the connector data plane.
+
+    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
+    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
+    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
+    The orchestrator only needs to ship a placeholder prefill prompt of the
+    correct length so the scheduler can allocate KV-cache slots.
+
+    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
+    here from the user prompt so the worker's line-408 buffer seed picks it
+    up. The connector-side ``extract_speaker_from_request`` reads the
+    strongly-typed ``request.additional_information.entries["speaker"]`` which
+    currently does not always round-trip the user-supplied voice; until that
+    plumbing is normalized, providing the small fields directly preserves
+    voice selection (regression discovered on Buildkite 9668:
+    ``test_speaker_002[default]`` lost the preset voice).
+    """
+    talker_inputs: list[OmniTokensPrompt] = []
+    for i, thinker_output in enumerate(source_outputs):
+        output = thinker_output.outputs[0]
+        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
+        # Skip-on-missing parity with thinker2talker_full_payload: if the
+        # connector builder would drop this request (no MM dict or missing
+        # hidden-state layers), do the same here so the worker buffer
+        # presence agrees with the orchestrator's scheduling decision.
+        thinker_mm_raw = getattr(output, "multimodal_output", None)
+        if not isinstance(thinker_mm_raw, Mapping):
+            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
+            continue
+        mm_hs = thinker_mm_raw.get("hidden_states", {})
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
+        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
+            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
+            continue
+        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
+        output_ids = _ensure_list(output.cumulative_token_ids)
+        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
+        if is_streaming_session:
+            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
+                req_id,
+                prompt_token_ids,
+                output_ids,
+                getattr(streaming_context, "new_prompt_len_snapshot", None),
+                streaming_context,
+                clear_state=bool(getattr(thinker_output, "finished", False)),
+            )
+        thinker_sequences = prompt_token_ids + output_ids
+        thinker_input_ids = prompt_token_ids
+        info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
+        prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
+
+        # Forward only small voice metadata; bulk tensors come from the
+        # connector path via _sync_local_stage_payloads.
+        small_info: dict[str, Any] = {}
+        speaker = extract_speaker_from_prompt(prompt, index=i)
+        if speaker is not None:
+            small_info["speaker"] = speaker
+        language = extract_language_from_prompt(prompt, index=i)
+        if language is not None:
+            small_info["language"] = language
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=(small_info if small_info else None),
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
             )
@@ -820,5 +949,59 @@ def talker2code2wav_full_payload(
     )
     return {
         "codes": {"audio": codec_codes},
+        "code_predictor_codes": codec_codes,
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
+
+
+def talker2code2wav(
+    source_outputs: list[Any],
+    _prompt: OmniTokensPrompt | TextPrompt | None = None,
+    _requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
+) -> list[OmniTokensPrompt]:
+    """
+    Process talker outputs to create code2wav inputs.
+
+    Workflow:
+    1. Extract talker's codec code outputs (8-layer RVQ codes)
+    2. Flatten codes for code2wav input
+    3. Package for code2wav stage
+
+    Args:
+    Returns:
+        List of OmniTokensPrompt for code2wav stage
+    """
+    talker_outputs = source_outputs
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    # Process each talker output
+    for i, talker_output in enumerate(talker_outputs):
+        output = talker_output.outputs[0]
+        req_id = str(getattr(talker_output, "request_id", f"idx-{i}"))
+        cur_seq_len = len(output.cumulative_token_ids) - 1
+        seq_len = cur_seq_len
+        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
+        if is_streaming_session:
+            seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
+        mm_raw = getattr(output, "multimodal_output", None)
+        if not isinstance(mm_raw, Mapping):
+            logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id)
+            continue
+        mm: OmniPayload = mm_raw
+        if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]:
+            logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id)
+            continue
+        # Extract codec codes from talker output
+        # Expected shape: [8, seq_len] (8-layer RVQ codes)
+        codec_codes = (
+            mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
+        )  # 16, seq_len
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return code2wav_inputs

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -38,6 +38,103 @@ def _qwen3_tts_empty_finished_payload():
     }
 
 
+def talker2code2wav(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[Any]:
+    """Non-async: collect all talker codes, then pass to code2wav at once."""
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    talker_outputs = source_outputs
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    for i, talker_output in enumerate(talker_outputs):
+        if not talker_output.finished:
+            # Non-async decode should only run once, after talker has
+            # accumulated the final code sequence.
+            continue
+        output = talker_output.outputs[0]
+        mm = output.multimodal_output
+        mm_codes = mm.get("codes", {})
+
+        # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
+        audio_codes = mm_codes["audio"].to(torch.long)
+        token_ids = output.cumulative_token_ids
+
+        # token_ids provides an upper bound on the newly generated codec span.
+        # audio_codes may still contain zero-padded / invalid rows, so trim only
+        # after filtering valid frames instead of trying to align EOS indices.
+        seq_len = max(len(token_ids) - 1, 0)
+        # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g.
+        # stop_token_id=2150 exceeds codebook_size=2048), and negative
+        # sentinels (e.g. -1 padding).
+        _CODEBOOK_SIZE = 2048
+        valid_mask = (
+            (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
+        )
+        audio_codes = audio_codes[valid_mask]
+        if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
+            audio_codes = audio_codes[-seq_len:]
+        ref_code = mm_codes.get("ref")
+        ref_code_len = mm.get("meta", {}).get("ref_code_len")
+        if isinstance(ref_code_len, torch.Tensor):
+            ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
+        elif ref_code_len is None:
+            ref_code_len = 0
+        else:
+            ref_code_len = int(ref_code_len)
+        if isinstance(ref_code, list):
+            ref_code = ref_code[0] if ref_code else None
+        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+            ref_code = ref_code.to(torch.long).cpu().contiguous()
+            if ref_code.ndim == 1:
+                num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16
+                if ref_code.numel() % num_quantizers != 0:
+                    logger.warning(
+                        "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
+                        ref_code.numel(),
+                        num_quantizers,
+                    )
+                    ref_code = None
+                else:
+                    ref_code = ref_code.reshape(-1, num_quantizers)
+            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_len > 0 and int(ref_code.shape[0]) > ref_code_len:
+                logger.warning(
+                    "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.",
+                    int(ref_code.shape[0]),
+                    ref_code_len,
+                )
+                ref_code = ref_code[:ref_code_len]
+            if not isinstance(ref_code, torch.Tensor):
+                ref_code_len = 0
+            else:
+                ref_code_len = int(ref_code.shape[0])
+                audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
+        else:
+            ref_code_len = 0
+        # Code2Wav expects codebook-major flat: [Q*num_frames]
+        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
+        additional_information = to_dict(
+            OmniPayloadStruct(
+                meta=MetaStruct(left_context_size=ref_code_len) if ref_code_len > 0 else None,
+                speaker=extract_speaker_from_prompt(prompt, index=i),
+                language=extract_language_from_prompt(prompt, index=i),
+            )
+        )
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=codec_codes,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+                additional_information=additional_information if additional_information else None,
+            )
+        )
+    return code2wav_inputs
+
+
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_codes = multimodal_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
@@ -254,7 +351,10 @@ _NUM_QUANTIZERS_DEFAULT = 16
 
 
 def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
-    """Filter zero-padded, out-of-range, and negative-padded codec frames."""
+    """Filter zero-padded, out-of-range, and negative-padded codec frames.
+
+    Mirrors the orchestrator-path body in `talker2code2wav` above.
+    """
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return audio_codes
     if audio_codes.ndim != 2:
@@ -267,8 +367,10 @@ def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
 
 def _coerce_ref_code_len(raw) -> int:
     """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"]
-    raw value (Tensor | int | None) into a non-negative int; clamps any
-    negative input to 0 since downstream code treats this as a frame count."""
+    raw value (Tensor | int | None) into a non-negative int.  Mirrors the
+    extraction inlined in the legacy ``talker2code2wav`` path; clamps any
+    negative input to 0 since downstream code treats this as a non-negative
+    frame count."""
     if isinstance(raw, torch.Tensor):
         value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0
     elif raw is None:
@@ -279,7 +381,7 @@ def _coerce_ref_code_len(raw) -> int:
 
 
 def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int):
-    """Coerce ref_code into a [ref_len, Q] tensor or None."""
+    """Coerce ref_code into a [ref_len, Q] tensor or None.  Mirrors orchestrator path."""
     if isinstance(ref_code, list):
         ref_code = ref_code[0] if ref_code else None
     if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0:
@@ -305,7 +407,8 @@ def talker2code2wav_token_only(
 
     Sized to the expected codec token count (codebook-major flat:
     Q * (ref_frames + audio_frames)).  Speaker / language metadata are
-    extracted from `prompt` and threaded via `additional_information`.
+    extracted from `prompt` and threaded via `additional_information`
+    (orchestrator-style; same as the legacy `talker2code2wav` builder).
     Actual codec ids are delivered via the worker connector payload built
     by `talker2code2wav_full_payload`.
     """
@@ -370,8 +473,8 @@ def talker2code2wav_full_payload(
     Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT
     across steps via flatten_payload), latest `pooling_output["codes.ref"]`
     (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`.
-    Filters invalid frames, crops to seq_len, prepends ref, and flattens
-    codebook-major for code2wav consumption.
+    Replicates the orchestrator-path body of `talker2code2wav` (filter,
+    crop to seq_len, prepend ref, codebook-major flatten).
     """
     del transfer_manager
     rid = getattr(request, "request_id", "?")