[WIP] Optimize Helios accuracy-maintained path

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

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.20.0"
+          VLLM_VERSION: "v0.21.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

.buildkite/scripts/hardware_ci/run-xpu-test.sh

@@ -52,5 +52,6 @@ time timeout -k 30 30m docker run \
     pip install tblib==3.1.0
     cd /workspace/vllm-omni
     pytest -v -s -m "core_model and xpu and B60"
+    export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
     pytest -v -s -m "advanced_model and xpu and B60"
 '

.buildkite/test-ready.yml

@@ -457,6 +457,51 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "Diffusion · Helios Distilled Acc/Perf Compare"
+        commands:
+          - |
+            timeout 90m bash -c "
+              set -euxo pipefail
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \
+                -m 'full_model and benchmark' \
+                --run-level 'full_model'
+            "
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/result.json"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/artifacts/**/*.mp4"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "Diffusion · Qwen Image Test"
         commands:

.buildkite/test-weekly.yml

@@ -84,3 +84,65 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
+  - group: "Reliability Test - Invalid parameters Test"
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    steps:
+      - label: "Invalid parameters Test · H100"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      # Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``.
+      - label: "Invalid parameters Test · L4"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"

.claude/skills/diffusion-perf-opt/SKILL.md

@@ -0,0 +1,488 @@
+---
+name: diffusion-perf-opt
+description: Diagnose and optimize vLLM Omni diffusion workloads, especially Wan/Qwen/Flux-style image and video generation. Use when Codex is asked to analyze profiling traces, choose parallel strategies, inspect torch profiler trace.json or trace.json.gz timelines, estimate optimization ROI, investigate GPU idle/free bubbles, compare USP/CFG/HSDP/VAE parallelism, or design operator/host/quantization optimizations for vLLM Omni.
+---
+
+# vLLM Omni Performance Optimization
+
+Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.
+
+## First Questions
+
+Before proposing changes, ask for the optimization scene if it is not already known:
+
+- GPU model, card count, topology, and whether NVLink is present.
+- Model and pipeline, for example Wan2.2 I2V A14B.
+- User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
+- Three runnable commands for each model/workload, or scripts that generate them:
+  - **Server startup command**: exact `vllm serve` command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
+  - **User request command**: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
+  - **User benchmark command**: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
+- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
+- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
+- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
+
+## Workflow
+
+1. **Freeze the measurement protocol and commands.**
+   - Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
+   - Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
+   - Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
+   - Present all duration and latency values in **milliseconds (ms)** in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
+   - Keep profiler modes separate in the command set:
+     - Baseline/benchmark commands should avoid torch profiler and stack collection.
+     - Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
+   - If any of the three commands is missing, ask for it or create a proposed script before proceeding.
+   - Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.
+
+2. **Collect a real baseline.**
+   - Disable PyTorch profiler and stack collection.
+   - Disable or fix torch.compile state so A/B is fair.
+   - Avoid `--enforce-eager` for production-speed baselines unless eager is the target.
+   - Prefer `--log-stats` and `--enable-diffusion-pipeline-profiler` for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
+   - Run warmup requests and exclude first-request lazy init.
+
+3. **Model the parallel strategy before testing.**
+   - Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
+   - Select a small candidate matrix rather than exhaustively testing everything.
+   - Typical 2-card candidates: `USP=2`, `CFG=2`, `USP=1/HSDP on-off`, VAE parallel on-off if memory allows.
+   - Typical 4-card candidates: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE parallel on-off.
+   - Typical 8-card candidates for official CFG-enabled video diffusion models:
+     - Primary candidate: `CFG=2 x USP=4` with VAE patch parallel across all 8 ranks.
+     - Compare against: `CFG=1 x USP=8` to test whether larger sequence parallel groups beat CFG branch parallelism.
+     - Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
+     - If long-video `diffuse` remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as `CFG=2 x USP=2 x Ring=2`.
+     - Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
+   - Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
+   - Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.
+
+4. **Search for the best parallel configuration.**
+   - Start with a small matrix that answers one question per comparison:
+     - `CFG parallel vs larger USP`: compare `CFG=2 x USP=world/2` against `CFG=1 x USP=world` for CFG-enabled workloads.
+     - `VAE patch parallel value`: compare the best DiT strategy with VAE patch parallel enabled and disabled.
+     - `HSDP cost`: compare HSDP on/off only when both configurations fit memory.
+     - `Ulysses vs Ring`: test a Ring/Ulysses hybrid only after long-sequence `diffuse` is confirmed dominant.
+   - For each configuration, create or record a stable config id, for example `A_cfg2_usp4_vaepp8_hsdp_tiling`.
+   - For each config id, capture:
+     - exact server command and environment variables,
+     - observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
+     - exact request command for every scenario,
+     - final server response metrics and client-side elapsed time,
+     - output artifact paths and any failed/empty responses.
+   - Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
+   - Compare configurations by stage, not only end-to-end latency. A configuration can improve `diffuse` while hurting `vae.decode`; record both effects.
+   - Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.
+
+5. **Run targeted A/B tests.**
+   - Change one variable per test.
+   - Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
+   - Record latency, stage timings, memory, output quality, and logs.
+   - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
+
+6. **Collect diagnostic trace only after narrowing hypotheses.**
+   - Use torch profiler for a small number of requests.
+   - Run two separate diagnostic traces instead of mixing concerns:
+     - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
+     - **Host-stack trace**: enable `torch_profiler_with_stack=True` and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
+   - Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
+   - Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where `diffuse` dominates.
+   - If profiler endpoints are available, run one warmup request first, then call `/start_profile`, run one profiled request, and call `/stop_profile`. This keeps model initialization and warmup out of the diagnostic trace.
+   - Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
+   - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
+   - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
+
+7. **Analyze host, communication, and operators.**
+   - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
+   - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
+   - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
+   - Rank operator work by total CUDA time and by repeated small-kernel launch count.
+
+8. **Produce an optimization plan.**
+   - Classify candidates as P0/P1/P2.
+   - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
+   - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
+   - End the plan with a user-facing candidate selection table. The assistant
+     should not automatically choose a risky optimization just because it is
+     technically possible. Present the options clearly and let the user decide
+     which item is worth implementing based on latency target, engineering
+     budget, memory headroom, and quality tolerance.
+   - Organize the plan by optimization layer, not by a flat list of ideas:
+     - host/runtime optimization,
+     - measurement/benchmark reliability,
+     - parallelism and communication,
+     - VAE encode/decode and media pre/post processing,
+     - operator fusion and layout cleanup,
+     - attention main-path optimization,
+     - algorithmic, precision, or approximation changes.
+   - For each layer, tie every candidate back to evidence from baseline
+     metrics, diagnostic trace, source code, or output quality requirements.
+     Do not list generic optimizations without a trace or workload reason.
+
+## Priority Rules
+
+- **P0:** low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable `empty_cache`, scheduler coefficient caching.
+- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
+- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
+
+Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
+
+## Optimization Layers
+
+After baseline, parallel-search, and diagnostic traces, summarize optimization
+opportunities by layer. This is the core of the performance analysis: the goal
+is to connect evidence to a scoped implementation and a validation plan.
+
+### Host and Runtime Optimization
+
+Purpose: remove CPU/Python stalls, synchronization points, allocator overhead,
+and request-path overhead that leave GPU lanes empty.
+
+Evidence to look for:
+
+- High `idle_pct` or large `GAP` blocks in `trace_analyzer.py`.
+- Host-stack trace lines such as `torch.cuda.empty_cache`,
+  `cudaStreamSynchronize`, `cudaDeviceSynchronize`, Python locks, scheduler
+  waits, image/video preprocessing, or repeated small allocation paths.
+- Difference between client wall-clock time and server `inference_time_s`.
+
+Typical candidates:
+
+- Make avoidable `torch.cuda.empty_cache()` optional or guard it by memory
+  headroom.
+- Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU
+  computations when the request shape/steps are fixed.
+- Remove avoidable host-device synchronizations and blocking logging/stat calls.
+- Move expensive preprocessing out of the critical path or cache fixed prompt,
+  image, and transform work for benchmark scenarios.
+- Ensure benchmark scripts record client-side elapsed time, HTTP status, output
+  path, and server response metrics.
+
+Priority guidance:
+
+- Usually P0 when the change is measurement reliability or an obvious removable
+  synchronization.
+- Usually P1 when it changes scheduling, memory lifetime, or request execution
+  order.
+
+Validation:
+
+- Re-run non-profiler baseline with same workload and seed.
+- Confirm peak memory headroom if disabling cache cleanup.
+- Confirm generated output exists and quality/seed behavior is unchanged.
+
+### Parallelism and Communication Optimization
+
+Purpose: choose the right decomposition for CFG branches, sequence tokens, model
+weights, VAE tiles, and rank topology.
+
+Evidence to look for:
+
+- Baseline A/B across `CFG`, `USP/SP`, `Ring`, `HSDP/FSDP`, and VAE patch
+  parallelism.
+- Stage timing shifts: `diffuse`, `vae.encode`, `vae.decode`, and server
+  end-to-end.
+- NCCL kernel time from trace, not only `user_annotation` time.
+- Rank imbalance across SP group ranks or CFG branch ranks.
+- Memory headroom and OOM risk.
+
+Typical candidates:
+
+- `CFG=2 x USP=world/2` versus `CFG=1 x USP=world` for CFG-enabled models.
+- VAE patch parallel on/off or patch size tuning.
+- HSDP/FSDP on/off only if both configurations fit memory.
+- Ulysses versus Ulysses+Ring only after long-sequence `diffuse` is confirmed
+  dominant and all-to-all is suspected.
+- Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
+- Buffer reuse or preallocation for FSDP/HSDP all-gather paths.
+
+Priority guidance:
+
+- P0/P1 for configuration-only changes with strong measured wins.
+- P1 for buffer reuse or rank mapping changes.
+- P2 for invasive distributed algorithm changes.
+
+Validation:
+
+- Measure by stage and memory, not only end-to-end.
+- Use one-variable A/B with identical prompt/media/seed/shape/steps.
+- When communication is suspected, compare rank0-3 in one USP group and rank0
+  versus rank4 across CFG branches for `CFG=2 x USP=4`.
+
+### VAE and Media Pipeline Optimization
+
+Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.
+
+Evidence to look for:
+
+- Large `vae.encode` or `vae.decode` in low-overhead stage timings.
+- Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video
+  transforms.
+- VAE kernels or cuDNN convolution in operator trace.
+- Whether every rank needs the final decoded tensor.
+
+Typical candidates:
+
+- Keep VAE patch parallel enabled when it has clear measured benefit.
+- Reduce VAE gather/broadcast to only ranks that need the final media output.
+- Reuse tile metadata, split buffers, or gather buffers.
+- Evaluate bf16/autocast behavior for VAE only with visual quality checks.
+- Avoid redundant image conversion, resize, or tensor construction in repeated
+  benchmark runs.
+
+Priority guidance:
+
+- P0/P1 if VAE is a large share of the target workload or if a host gap is
+  obvious and low risk.
+- Lower priority when `diffuse` dominates and VAE is already patch-parallelized.
+
+Validation:
+
+- Compare `vae.encode`, `vae.decode`, server end-to-end, and peak memory.
+- Check output video integrity, artifacts, flicker, and seed stability.
+
+### Operator Fusion and Layout Cleanup
+
+Purpose: reduce high-frequency small kernels, memory bandwidth pressure, layout
+conversions, and launch overhead in transformer and VAE blocks.
+
+Evidence to look for:
+
+- Top operator tables showing many `aten::copy_`, `aten::cat`,
+  `split_with_sizes_copy`, `aten::add`, `aten::mul`, `aten::div`, norm,
+  activation, RoPE, or reshape/layout kernels.
+- `ops_rankN.xlsx` `by_shape` sheet showing repeated small shapes inside the
+  same block path.
+- Trace lanes showing many short kernels between larger GEMM/attention kernels.
+- Source code patterns with repeated elementwise chains or layout conversions.
+
+Typical fusion targets:
+
+- AdaLayerNorm / RMSNorm / LayerNorm plus scale/shift fusion.
+- RoPE fusion with Q/K layout preparation when shapes are stable.
+- Residual add, scale, gate, and elementwise chains.
+- MLP gate/up/down cleanup, such as fusing activation and multiply around
+  GEMM outputs when feasible.
+- QKV projection and reshape/split/cat path cleanup.
+- Attention pre/post layout cleanup to avoid unnecessary copies, cats, and
+  splits around sequence parallel all-to-all.
+
+Priority guidance:
+
+- P1 when implemented with existing PyTorch/Triton/local helper patterns and
+  validated against exact outputs or tolerances.
+- P2 when it requires custom CUDA/Triton kernels, changes numerics, or touches
+  attention math directly.
+
+Validation:
+
+- First prove the operator family is material for the target shape.
+- Use non-profiler A/B for latency and stage timing.
+- Use quality regression checks for generated video stability.
+- Check compile behavior and graph breaks if using `torch.compile`.
+
+### Attention Main-Path Optimization
+
+Purpose: address the dominant self-attention cost when FlashAttention or other
+attention kernels dominate CUDA time.
+
+Evidence to look for:
+
+- Operator trace where FlashAttention/SDPA kernels dominate total CUDA time.
+- Attention shape from `ops_rankN.xlsx` `by_shape`, model code, or trace
+  metadata.
+- Whether attention cost scales with latent frames, latent H/W, patch size, or
+  CFG duplication.
+- Layout/copy/all-to-all work around attention.
+
+Typical candidates:
+
+- Verify the attention backend and shape are on the intended fast kernel path.
+- Compare supported attention backends only with identical workload and quality
+  settings.
+- Reduce attention input size by safe model/config choices when allowed:
+  latent resolution, frame count, patching, boundary ratio, or windowing.
+- Remove avoidable layout conversions before/after attention.
+- Reuse condition-side KV or other static inputs if the model structure allows.
+- Consider custom kernels, sparse/window/linear attention, or approximation only
+  after quality risk is accepted.
+
+Priority guidance:
+
+- P1 for backend/config/layout changes with preserved math.
+- P2 for approximate attention, sparsity, custom kernels, or any change that can
+  alter quality/temporal consistency.
+
+Validation:
+
+- Always include output quality and seed behavior checks.
+- Compare `diffuse`, server end-to-end, peak memory, and attention kernel time
+  in diagnostic traces if needed.
+
+### Algorithmic, Precision, and Approximation Optimization
+
+Purpose: reduce mathematical work or precision cost beyond local code cleanup.
+
+Evidence to look for:
+
+- A single operator family dominates even after low-risk runtime, parallel, and
+  fusion work.
+- Memory bandwidth or compute utilization suggests precision or quantization
+  could matter.
+- The user explicitly allows quality-preserving or approximate methods.
+
+Typical candidates:
+
+- FP8/quantization for transformer or selected projections.
+- Sparsity or Rainfusion-style acceleration.
+- Reduced steps, scheduler changes, distillation, or caching across frames.
+- Approximate attention or linear attention.
+
+Priority guidance:
+
+- Usually P2 because quality, numerics, and implementation risk are high.
+
+Validation:
+
+- Requires strict A/B, visual quality review, temporal flicker checks, seed
+  stability, and possibly human evaluation.
+
+### Interpolation, Super-Resolution, and E2E Pipeline Optimization
+
+Purpose: optimize the whole user-visible video product, not only the base
+diffusion invocation. Some deployments trade base-model latency against
+post-processing, interpolation, or super-resolution stages.
+
+Evidence to look for:
+
+- E2E latency breakdown across base generation, interpolation, super-resolution,
+  encoding, storage, and response streaming.
+- Fast/slow GPU or fast/slow stage analysis across multiple cards and pipeline
+  stages.
+- User quality target: resolution, FPS, temporal smoothness, and acceptable
+  post-processing artifacts.
+
+Typical candidates:
+
+- Add or optimize a frame interpolation stage when it reduces required base
+  model frames for the same perceived FPS.
+- Add or optimize a super-resolution model when generating lower base
+  resolution plus SR is faster for the target quality.
+- Analyze E2E2 pipeline behavior: client request, service scheduling, diffusion,
+  VAE/media, post-process, file write, and response.
+- Identify fast/slow cards or stages and rebalance pipeline placement.
+
+Priority guidance:
+
+- P1 when using proven interpolation/SR components without changing diffusion
+  math.
+- P2 when quality risk is high or the pipeline adds significant operational
+  complexity.
+
+Validation:
+
+- Measure E2E wall-clock, per-stage server timings, output FPS/resolution,
+  artifacts, flicker, and user-visible quality.
+
+### Optimization Candidate Library
+
+Use this table as a compact menu, not as automatic recommendations. Pick items
+only when baseline metrics, trace evidence, source inspection, or quality
+tolerance supports them.
+
+| Layer | Candidate | Evidence | Priority | Validation focus |
+|---|---|---|---|---|
+| Measurement | Freeze server/request/benchmark commands | Missing or drifting commands | P0 | Repeatable non-profiler A/B |
+| Measurement | Separate baseline and diagnostic profiler runs | Profiler used for latency claims | P0 | Low-overhead stage timings |
+| Host/runtime | Guard or remove avoidable `empty_cache` | Host-stack gaps or sync stalls | P0/P1 | Latency, peak memory, OOM safety |
+| Host/runtime | Cache scheduler coefficients/timesteps | Repeated tiny CPU/GPU work | P0/P1 | Same seed/output, stage timing |
+| Host/runtime | Reduce framework scheduling overhead | Client time exceeds server time | P1 | E2E latency, throughput |
+| Parallel | `CFG=2 x USP=world/2` vs `USP=world` | CFG doubles forward work | P0/P1 | `diffuse`, NCCL, memory |
+| Parallel | Tune VAE patch parallelism | VAE encode/decode is material | P0/P1 | VAE time, output correctness |
+| Parallel | HSDP on/off or buffer reuse | HSDP affects memory/all-gather | P1 | Memory, latency, OOM risk |
+| Parallel | Ulysses vs Ulysses+Ring | Long sequence all-to-all suspected | P1/P2 | Rank balance, NCCL kernels |
+| Cross-attn | Disable SP for short condition tokens | Cross-attn comm exceeds compute | P1 | `diffuse`, correctness |
+| VAE/media | Reduce VAE gather/broadcast | Rank traces show VAE wait | P1 | Rank balance, output file |
+| VAE/media | Reuse tile metadata/buffers | Tile split/merge host gaps | P1 | `vae.encode/decode`, memory |
+| VAE/media | VAE bf16/autocast | VAE float kernels are slow | P1/P2 | Artifacts, flicker, seed stability |
+| Operator fusion | AdaLayerNorm/LayerNorm fusion | Norm plus scale/shift kernels | P1 | Numeric tolerance, latency |
+| Operator fusion | RMSNorm fusion | Many small RMSNorm kernels | P1 | Numeric tolerance, latency |
+| Operator fusion | RoPE cache/fuse/layout cleanup | RoPE copy/reshape kernels | P1/P2 | Kernel count, correctness |
+| Operator layout | QKV or attention layout cleanup | Copy/cat/split around attention | P1 | Copy kernels, compile behavior |
+| Attention | Verify backend fast path | FA/SDPA dominates trace | P1 | `diffuse`, attention kernels |
+| Attention | FA to LA or selected-head LA | Attention remains dominant | P2 | Quality, temporal stability |
+| Precision | Transformer FP8/quantization | Compute/bandwidth bound and allowed | P2 | Quality, speed, stability |
+| Sparsity | Rainfusion-style acceleration | DiT compute remains dominant | P2 | Prompt diversity, quality |
+| Pipeline | Frame interpolation | Fewer base frames can meet FPS | P1/P2 | E2E latency, motion artifacts |
+| Pipeline | Super-resolution | Lower base res plus SR may win | P1/P2 | Detail quality, artifacts |
+| E2E | Fast/slow-card analysis | Multi-card stragglers | P0/P1 | Per-rank/stage wall-clock |
+
+### Optimization Plan Template
+
+Use this table shape when reporting the next work items:
+
+| Priority | Layer | Candidate | Evidence | Expected benefit | Implementation path | Validation | Quality risk |
+|---|---|---|---|---|---|---|---|
+| P0 | Host/runtime | Guard `empty_cache` | Host-stack gap points to `torch.cuda.empty_cache` | Small latency reduction, less idle | Add config/env guard | Non-profiler A/B, memory check | Low |
+| P1 | Operator fusion | RMSNorm/AdaLayerNorm fusion | High-frequency norm/elementwise kernels | Lower launch/bandwidth overhead | Use existing fusion helper or targeted Triton | A/B + output check | Medium |
+| P1/P2 | Attention | Attention layout/backend investigation | FA kernel dominates CUDA time | Potentially large | Inspect shapes/backend and remove layout copies | A/B + trace + quality | Medium/high |
+
+Then present a short selection prompt using the same rows:
+
+```text
+Which candidate should we implement next?
+
+1. P0 Host/runtime: guard empty_cache
+   - Expected benefit: small but low-risk latency reduction.
+   - Risk: possible memory increase/OOM if memory headroom is insufficient.
+
+2. P1 Operator fusion: inspect by_shape and implement first norm/RoPE/layout fusion
+   - Expected benefit: medium if high-frequency small kernels are confirmed.
+   - Risk: numerical/compile/quality validation needed.
+
+3. P1/P2 Attention: FA/LA/backend/layout investigation
+   - Expected benefit: potentially large.
+   - Risk: high quality and implementation risk.
+```
+
+If the user has not chosen an item, default to explaining tradeoffs and asking
+which candidate to execute. Only proceed autonomously on low-risk P0 measurement
+or instrumentation fixes.
+
+## Analysis Helpers
+
+PyTorch profiler traces are Chrome/Perfetto-compatible JSON files, usually
+`trace_rankN.json` or `trace_rankN.json.gz`. They normally contain a top-level
+`traceEvents` list, though some exporters emit the raw event list directly.
+
+Use the checked-in analyzer from the repository root:
+
+```bash
+python3 .claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py \
+  vllm_profile/.../trace_rank0.json.gz \
+  --min-gap-ms 5 \
+  --topn 20
+```
+
+For rank imbalance or communication questions, pass all relevant rank traces in
+one command. For host gaps, lower `--min-gap-ms` to `1` and use a host-stack
+trace. Read `gpu_span_s`, `busy_union_s`, `idle_union_s`, `idle_pct`, `GAP`
+blocks, `Top GPU/operator events by total duration`, and `Top NCCL-like events
+by category`. Treat `cat=user_annotation` NCCL ranges as enclosing annotations;
+prefer `cat=kernel` or `cat=gpu_user_annotation` for real device work.
+
+The analyzer summarizes timing only. It does not parse tensor shapes, attribute
+overlap to individual streams, prove quality, or provide final latency claims.
+Use `ops_rankN.xlsx` or PyTorch key averages for shape analysis, and re-test any
+optimization with non-profiler baseline commands.
+
+Read `references/optimization-playbook.md` when drafting the optimization table or comparing candidate techniques.
+
+## vLLM Omni Heuristics
+
+- Cross-attention usually should not use USP/SP when text/image condition token count is much smaller than latent video tokens. Confirm via trace; in Wan2.2 I2V, self-attention dominates cross-attention.
+- VAE bf16/autocast is often worthwhile but requires visual quality checks.
+- VAE patch parallel can help decode/encode but may add gather/merge/broadcast overhead. Check whether all ranks need the final decoded tensor.
+- HSDP/FSDP is primarily a memory strategy. If the model fits without it, run an on/off latency comparison.
+- Scheduler work can create small host/device gaps; cache tiny solve coefficients when timesteps/order are known.
+- `torch.cuda.empty_cache()` can prevent OOM but creates synchronization/idle. Make it optional if memory headroom is sufficient.
+- `Command Buffer Full` in profiler output is profiler overhead, not a model optimization target.

.claude/skills/diffusion-perf-opt/agents/openai.yaml

@@ -0,0 +1,4 @@
+interface:
+  display_name: "vLLM Omni Perf Opt"
+  short_description: "Profile and optimize vLLM Omni diffusion performance."
+  default_prompt: "Analyze and optimize a vLLM Omni diffusion workload using baselines, traces, parallel strategy modeling, and prioritized optimization plans."

.claude/skills/diffusion-perf-opt/references/optimization-playbook.md

@@ -0,0 +1,82 @@
+# Optimization Playbook
+
+Use this reference when turning trace evidence into prioritized vLLM Omni optimizations.
+
+## Parallel Strategy
+
+Start from workload shape:
+
+- Long video tokens, high resolution, many frames: self-attention and FFN dominate. USP/SP is usually valuable.
+- CFG enabled with positive and negative prompts: transformer forward count doubles. CFG parallel may beat USP on small/moderate sequence lengths.
+- Cross-attention: skip SP/USP when condition tokens are much smaller than latent tokens; extra all-to-all can exceed compute saved.
+- HSDP/FSDP: use to fit memory first. If memory allows, compare without HSDP because all-gathers and lazy init add overhead.
+- VAE parallel: useful when encode/decode is material. Check gather/merge/broadcast and whether every rank needs final decoded output.
+
+Candidate matrices:
+
+- 2 cards: `USP=2`, `CFG=2`, HSDP on/off, VAE PP on/off.
+- 4 cards: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE PP on/off.
+- 8 cards: add pipeline/stage parallel only after measuring orchestration and inter-stage transfer overhead.
+
+## Host And Free Bubble Analysis
+
+Look for device intervals where no kernel, memcpy, or memset is active.
+
+Common causes:
+
+- First request FSDP/HSDP lazy init.
+- Python preprocessing: image load, resize, `to_tensor`, tokenizer, request marshaling.
+- `torch.cuda.empty_cache()` or explicit synchronizations.
+- Scheduler small linear algebra such as tiny `torch.linalg.solve`.
+- VAE tile split/merge Python loops.
+- Profiler overhead such as CUPTI command buffer flush/full.
+
+Also inspect trace-external time:
+
+- API request parsing and multipart upload.
+- Engine/orchestrator dispatch.
+- Worker RPC and queue wait.
+- Output serialization, video write, base64, HTTP response.
+
+## Operator Optimization Candidates
+
+Evaluate necessity before implementation:
+
+- **FA to LA / LA for head:** high potential only if attention is the main CUDA hotspot and target shapes fit LA well.
+- **Cross-attention BSND/no transpose:** do only if transpose/copy around cross-attn is measurable.
+- **AdaLayerNorm fusion:** fuse norm, scale/shift, gate, residual, and dtype conversion if repeated elementwise kernels are significant.
+- **RMSNorm fusion:** useful for Q/K norms and VAE norms when small-kernel launch count is high.
+- **FP32 LayerNorm to bf16:** quality-sensitive; run visual regression.
+- **LA preprocess:** integrate only if preprocess overhead is cached or amortized.
+- **RoPE move/cache/fuse:** useful when RoPE small kernels are visible around every self-attn block.
+- **VAE RMSNorm replacement:** only after VAE internal trace shows norm as a real component of decode/encode.
+- **Scheduler fusion/cache:** cache coefficient solves and fuse CFG combine with scheduler arithmetic when possible.
+
+Before writing a custom kernel, check:
+
+- Local vLLM Omni implementations.
+- Upstream vLLM kernels and attention backends.
+- SGLang, FlashInfer, FlashAttention, and Triton examples.
+- Shape constraints and layout expectations.
+
+## Quantization And Sparsity
+
+Treat these as high-risk optimizations:
+
+- Rainfusion/sparse diffusion can be high reward but needs model-specific quality validation.
+- VAE decode quantization can reduce decode time but usually has smaller end-to-end upside than transformer attention/FFN work.
+- FP8 or approximate kernels require seed-stable visual regression, temporal flicker checks, and prompt diversity.
+
+## Output Format For Recommendations
+
+Return a table with:
+
+- Priority: P0/P1/P2.
+- Optimization.
+- Evidence from trace or code.
+- Expected benefit.
+- Risk/quality impact.
+- Implementation path.
+- Validation command or A/B test.
+
+State clearly when a recommendation is a hypothesis requiring new trace with shapes.

.claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+"""Summarize torch profiler Chrome traces for vLLM Omni diffusion optimization."""
+
+from __future__ import annotations
+
+import argparse
+import collections
+import gzip
+import json
+from pathlib import Path
+from typing import Any
+
+GPU_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
+CPU_CATS = {"python_function", "user_annotation", "cpu_op", "cuda_runtime", "cuda_driver"}
+
+
+def open_trace(path: Path):
+    if path.suffix == ".gz":
+        return gzip.open(path, "rt")
+    return path.open("rt")
+
+
+def event_name(event: dict[str, Any]) -> str:
+    return str(event.get("name", ""))
+
+
+def summarize_trace(path: Path, min_gap_us: float, topn: int) -> None:
+    with open_trace(path) as f:
+        data = json.load(f)
+
+    events = data if isinstance(data, list) else data.get("traceEvents", [])
+    gpu: list[tuple[float, float, float, str, str, Any, Any]] = []
+    cpu: list[tuple[float, float, float, str, str, Any, Any]] = []
+    by_gpu_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
+    by_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
+    nccl: dict[tuple[str, str], list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
+
+    for event in events:
+        dur = event.get("dur")
+        ts = event.get("ts")
+        if dur is None or ts is None or dur <= 0:
+            continue
+        cat = str(event.get("cat", ""))
+        name = event_name(event)
+        row = (float(ts), float(ts + dur), float(dur), name, cat, event.get("pid"), event.get("tid"))
+        stat = by_name[name]
+        stat[0] += 1
+        stat[1] += dur
+        stat[2] = max(stat[2], dur)
+        if cat in GPU_CATS:
+            gpu.append(row)
+            gstat = by_gpu_name[name]
+            gstat[0] += 1
+            gstat[1] += dur
+            gstat[2] = max(gstat[2], dur)
+        elif cat in CPU_CATS:
+            cpu.append(row)
+        if "nccl" in name.lower():
+            nstat = nccl[(cat, name)]
+            nstat[0] += 1
+            nstat[1] += dur
+            nstat[2] = max(nstat[2], dur)
+
+    print(f"\n== {path}")
+    print(f"events={len(events)} gpu_events={len(gpu)} cpu_events={len(cpu)}")
+    if not gpu:
+        print("No GPU events found.")
+        return
+
+    gpu.sort()
+    merged: list[list[Any]] = []
+    for start, end, dur, name, cat, pid, tid in gpu:
+        if not merged or start > merged[-1][1]:
+            merged.append([start, end, [(start, end, dur, name, cat, pid, tid)]])
+        else:
+            merged[-1][1] = max(merged[-1][1], end)
+            merged[-1][2].append((start, end, dur, name, cat, pid, tid))
+
+    span = max(end for _, end, *_ in gpu) - min(start for start, *_ in gpu)
+    busy = sum(end - start for start, end, _ in merged)
+    idle = span - busy
+    print(
+        f"gpu_span_s={span / 1e6:.3f} "
+        f"busy_union_s={busy / 1e6:.3f} "
+        f"idle_union_s={idle / 1e6:.3f} "
+        f"idle_pct={idle / span * 100:.2f}"
+    )
+
+    interesting_cpu = [
+        row
+        for row in cpu
+        if row[2] >= 1000
+        and (
+            row[4] in {"python_function", "user_annotation"}
+            or "cudaStreamSynchronize" in row[3]
+            or "cudaDeviceSynchronize" in row[3]
+            or "cudaLaunch" in row[3]
+            or "cudaMemcpy" in row[3]
+        )
+    ]
+
+    gaps = []
+    for idx in range(1, len(merged)):
+        gap_start = merged[idx - 1][1]
+        gap_end = merged[idx][0]
+        gap_dur = gap_end - gap_start
+        if gap_dur >= min_gap_us:
+            prev_event = max(merged[idx - 1][2], key=lambda x: x[1])
+            next_event = min(merged[idx][2], key=lambda x: x[0])
+            mid = (gap_start + gap_end) / 2
+            containers = [row for row in interesting_cpu if row[0] <= mid <= row[1]]
+            containers = sorted(containers, key=lambda x: x[2])[:8]
+            gaps.append((gap_dur, gap_start, gap_end, prev_event, next_event, containers))
+
+    print(f"gaps_ge_{min_gap_us / 1000:.3f}ms count={len(gaps)} sum_s={sum(g[0] for g in gaps) / 1e6:.3f}")
+    for gap_dur, gap_start, gap_end, prev_event, next_event, containers in sorted(gaps, reverse=True)[:topn]:
+        print(f"\nGAP {gap_dur / 1000:.3f} ms ts={gap_start:.0f}->{gap_end:.0f}")
+        print(f"  prev {prev_event[4]} {prev_event[2] / 1000:.3f} ms {prev_event[3][:160]}")
+        print(f"  next {next_event[4]} {next_event[2] / 1000:.3f} ms {next_event[3][:160]}")
+        for row in containers:
+            print(f"  in   {row[4]} {row[2] / 1000:.3f} ms {row[3][:180]}")
+
+    print("\nTop GPU/operator events by total duration:")
+    for name, (count, total, max_dur) in sorted(by_gpu_name.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
+        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms {name[:180]}")
+
+    print("\nTop NCCL-like events by category:")
+    for (cat, name), (count, total, max_dur) in sorted(nccl.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
+        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms cat={cat} {name[:160]}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("traces", nargs="+", type=Path, help="trace.json or trace.json.gz files")
+    parser.add_argument("--min-gap-ms", type=float, default=5.0, help="minimum GPU idle gap to print")
+    parser.add_argument("--topn", type=int, default=20, help="number of gaps/hotspots to print")
+    args = parser.parse_args()
+
+    for trace in args.traces:
+        summarize_trace(trace, args.min_gap_ms * 1000.0, args.topn)
+
+
+if __name__ == "__main__":
+    main()

.claude/skills/readme.md

@@ -18,6 +18,9 @@ include:
 
 - `add-diffusion-model`: guides integration of a new diffusion model into
   `vllm-omni`
+- `diffusion-perf-opt`: guides diffusion model performance optimization,
+  including profiling traces, parallel strategies, stage timing analysis, and
+  benchmark-driven tuning
 - `add-omni-model`: covers addition of new omni-modality model support
 - `add-tts-model`: covers integration of new TTS models and related serving
   workflows

README.md

@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
 

benchmarks/diffusion/README.md

@@ -30,7 +30,7 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 **Notes**
 
-- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
+- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
 - If you run the server on another host or port, pass `--base-url` accordingly.
 
 ## 2. Supported Datasets
@@ -95,8 +95,9 @@ You can point to your own trace using `--dataset-path`.
 
 ### Basic flags
 
-- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
+- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
 - `--model`: The OpenAI-compatible `model` field.
+- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
 - `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
 - `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
 - `--num-prompts`: Number of requests to send.

benchmarks/diffusion/backends.py

@@ -151,7 +151,7 @@ async def async_request_chat_completions(
     return output
 
 
-async def async_request_openai_images(
+async def async_request_openai_image_generations(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
@@ -278,7 +278,7 @@ async def async_request_v1_videos(
                 return output
 
         # invoke a poll request (GET /v1/videos/{video_id})
-        poll_interval = 2.0  # Unit(s)
+        poll_interval = 0.2  # Unit(s)
         timeout_seconds = 600.0
         deadline = time.perf_counter() + timeout_seconds
         job_url = f"{input.api_url}/{job_id}"
@@ -348,12 +348,38 @@ async def async_request_v1_videos(
     return output
 
 
+LEGACY_BACKEND_ENDPOINT_ALIASES = {
+    "vllm-omni": "/v1/chat/completions",
+    "openai": "/v1/images/generations",
+}
+
+
+def normalize_endpoint(value: str) -> str:
+    endpoint = str(value).strip()
+    if not endpoint:
+        raise ValueError("endpoint must not be empty.")
+    endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
+        endpoint,
+        LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
+    )
+    if not endpoint.startswith("/"):
+        endpoint = f"/{endpoint}"
+    return endpoint
+
+
+def endpoint_filename_token(value: str) -> str:
+    token = normalize_endpoint(value).lstrip("/")
+    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
+        token = token.replace(bad, "_")
+    return token or "endpoint"
+
+
 backends_function_mapping = {
     "2i": {
-        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
-        "openai": (async_request_openai_images, "/v1/images/generations"),
+        "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
+        "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
     },
     "2v": {
-        "v1/videos": (async_request_v1_videos, "/v1/videos"),
+        "/v1/videos": (async_request_v1_videos, "/v1/videos"),
     },
 }

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -6,31 +6,31 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 
-Supports multiple backends:
-    - vllm-omni: Uses /v1/chat/completions endpoint (default)
-    - openai: Uses /v1/images/generations endpoint
-    - v1/videos: Use /v1/videos endpoint
+Supports multiple endpoints:
+    - /v1/chat/completions: OpenAI chat-compatible image requests
+    - /v1/images/generations: OpenAI image generation requests
+    - /v1/videos: Async video jobs
 
 Usage:
-    # Video (v1/videos backend)
+    # Video (/v1/videos endpoint)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
+        --endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset vbench --task i2v --num-prompts 10
+        --endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10
 
 
-    # Image (vllm-omni backend)
+    # Image (/v1/chat/completions endpoint)
     t2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
+        --endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset random --task t2i --num-prompts 1 \
+        --endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
         --max-concurrency 1 --enable-negative-prompt \
         --random-request-config '[
             {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
@@ -41,18 +41,18 @@ Usage:
 
     i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
+        --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
 
-    # Image (openai backend)
+    # Image (/v1/images/generations endpoint)
     t2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend openai --dataset vbench --task t2i --num-prompts 10 \
+        --endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
     # Video (v1/videos)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset random --task t2v --num-prompts 1 \
+        --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
         --max-concurrency 1 --enable-negative-prompt \
         --random-request-config '[
             {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
@@ -68,6 +68,7 @@ import glob
 import json
 import logging
 import os
+import pathlib
 import random
 import tempfile
 import time
@@ -80,7 +81,12 @@ from typing import Any
 import aiohttp
 import numpy as np
 import requests
-from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
+from backends import (
+    RequestFuncInput,
+    RequestFuncOutput,
+    backends_function_mapping,
+    normalize_endpoint,
+)
 from PIL import Image
 from tqdm.asyncio import tqdm
 
@@ -593,6 +599,13 @@ class RandomDataset(BaseDataset):
         if self.enable_negative_prompt:
             extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
 
+        request_fields = {
+            "width",
+            "height",
+            "num_frames",
+            "num_inference_steps",
+            "fps",
+        }
         params = {
             "width": self.args.width,
             "height": self.args.height,
@@ -602,9 +615,11 @@ class RandomDataset(BaseDataset):
         }
         if self._sampled_requests:
             profile = self._sampled_requests[idx]
-            params.update(profile)
+            params.update({key: value for key, value in profile.items() if key in request_fields})
+            extra_body.update({key: value for key, value in profile.items() if key not in request_fields})
+        prompt = self.args.random_prompt or f"Random prompt {idx} for benchmarking diffusion models"
         return RequestFuncInput(
-            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            prompt=prompt,
             api_url=self.api_url,
             model=self.model,
             seed=self.args.seed,
@@ -864,6 +879,35 @@ def calculate_metrics(
     return metrics
 
 
+def _response_suffix(endpoint: str) -> str:
+    if endpoint == "/v1/videos":
+        return ".mp4"
+    return ".json"
+
+
+def _save_success_responses(
+    outputs: list[RequestFuncOutput],
+    requests_list: list[RequestFuncInput],
+    output_dir: str,
+    endpoint: str,
+) -> list[str]:
+    path = pathlib.Path(output_dir).expanduser().resolve()
+    path.mkdir(parents=True, exist_ok=True)
+    saved_paths: list[str] = []
+    suffix = _response_suffix(endpoint)
+
+    for index, (request, output) in enumerate(zip(requests_list, outputs, strict=False)):
+        if not output.success:
+            continue
+        item_path = path / f"{index:05d}_{request.request_id}{suffix}"
+        if isinstance(output.response_body, bytes):
+            item_path.write_bytes(output.response_body)
+        else:
+            item_path.write_text(json.dumps(output.response_body, indent=2, sort_keys=True), encoding="utf-8")
+        saved_paths.append(str(item_path))
+    return saved_paths
+
+
 def wait_for_service(base_url: str, timeout: int = 120) -> None:
     print(f"Waiting for service at {base_url}...")
     start_time = time.time()
@@ -883,6 +927,14 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
         time.sleep(1)
 
 
+def _default_endpoint_for_task(task: str) -> str:
+    if task in {"t2v", "i2v", "ti2v"}:
+        return "/v1/videos"
+    if task in {"t2i", "i2i", "ti2i"}:
+        return "/v1/chat/completions"
+    raise ValueError(f"Unsupported task for endpoint resolution: {task}")
+
+
 async def benchmark(args):
     # Construct base_url if not provided
     if args.base_url is None:
@@ -902,18 +954,23 @@ async def benchmark(args):
             f"Valid image tasks: {sorted(IMAGE_TASKS)}"
         )
 
-    valid_backends = sorted(backends_function_mapping[task_type].keys())
+    raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
+    if raw_endpoint is None:
+        raw_endpoint = _default_endpoint_for_task(args.task)
+    args.endpoint = normalize_endpoint(raw_endpoint)
 
-    if args.backend not in valid_backends:
+    valid_endpoints = sorted(backends_function_mapping[task_type].keys())
+
+    if args.endpoint not in valid_endpoints:
         logger.error(
-            f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
-            f"Valid backends for this task type: {valid_backends}\n"
-            f"Example usage: --task {args.task} --backend {valid_backends[0]}"
+            f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
+            f"Valid endpoints for this task type: {valid_endpoints}\n"
+            f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
         )
-        raise ValueError("Backend validation failed. See log above for valid options.")
+        raise ValueError("Endpoint validation failed. See log above for valid options.")
 
-    # Setup API URL and request function based on backend
-    request_func, api_url = backends_function_mapping[task_type][args.backend]
+    # Setup API URL and request function based on endpoint.
+    request_func, api_url = backends_function_mapping[task_type][args.endpoint]
     api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
@@ -976,15 +1033,22 @@ async def benchmark(args):
     metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
 
     # Add configuration info to metrics for JSON output
-    metrics["backend"] = args.backend
+    metrics["endpoint"] = args.endpoint
     metrics["model"] = args.model
     metrics["dataset"] = args.dataset
     metrics["task"] = args.task
+    if args.save_response_dir:
+        metrics["saved_responses"] = _save_success_responses(
+            outputs,
+            requests_list,
+            args.save_response_dir,
+            args.endpoint,
+        )
 
-    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
+    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
 
     # Section 1: Configuration
-    print("{:<40} {:<15}".format("Backend:", args.backend))
+    print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
     print("{:<40} {:<15}".format("Model:", args.model))
     print("{:<40} {:<15}".format("Dataset:", args.dataset))
     print("{:<40} {:<15}".format("Task:", args.task))
@@ -1047,12 +1111,17 @@ if __name__ == "__main__":
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
     parser.add_argument("--port", type=int, default=8091, help="Server port.")
     parser.add_argument("--model", type=str, default="default", help="Model name.")
+    parser.add_argument(
+        "--endpoint",
+        type=str,
+        default=None,
+        help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
+    )
     parser.add_argument(
         "--backend",
         type=str,
-        default="vllm-omni",
-        choices=["vllm-omni", "openai", "v1/videos"],
-        help="Backend to target the benchmark to.",
+        default=None,
+        help=argparse.SUPPRESS,
     )
     parser.add_argument(
         "--dataset",
@@ -1131,6 +1200,12 @@ if __name__ == "__main__":
     )
     parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
     parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.")
+    parser.add_argument(
+        "--save-response-dir",
+        type=str,
+        default=None,
+        help="Directory to save successful raw responses. Video responses are saved as MP4 files.",
+    )
     parser.add_argument(
         "--slo",
         action="store_true",
@@ -1166,6 +1241,12 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--random-prompt",
+        type=str,
+        default=None,
+        help="Prompt to use for every synthetic random request. Defaults to generated placeholder prompts.",
+    )
     parser.add_argument(
         "--num-input-images",
         type=int,

benchmarks/diffusion/performance_dashboard/helios_distilled_serving_performance.md

@@ -0,0 +1,222 @@
+# Helios-Distilled Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **BestWishYsh/Helios-Distilled** using vLLM-Omni. It follows the same serving benchmark structure as the Qwen-Image and Wan2.2 dashboards.
+
+---
+
+# 1. Overview
+
+Helios-Distilled is a text-to-video diffusion model. Helios uses an autoregressive chunking pattern that generates 33 frames per chunk. For stable comparisons, choose `num_frames` as a multiple of 33, such as 99.
+
+This document covers:
+
+* Service launch configuration
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance and accuracy comparison methodology
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA H100 for CI accuracy/perf gate; local bring-up also used L20X |
+| Diffusion Attention Backend | FlashAttention |
+| Model | `BestWishYsh/Helios-Distilled` |
+
+---
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve BestWishYsh/Helios-Distilled --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter | Description |
+| --- | --- |
+| `--port` | Serving port |
+| `--max-num-seqs` | Maximum number of active sequences |
+| `--enable-diffusion-pipeline-profiler` | Optional stage-level profiling |
+| `--profiler-config` | Optional torch profiler configuration |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Scripts
+
+## 4.1 Online Serving Benchmark
+
+Use the shared diffusion serving benchmark, as with Qwen-Image and Wan2.2:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --endpoint /v1/videos \
+    --model BestWishYsh/Helios-Distilled \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --disable-tqdm \
+    --random-request-config '[
+        {
+            "width": 640,
+            "height": 384,
+            "num_inference_steps": 50,
+            "num_frames": 99,
+            "fps": 16,
+            "guidance_scale": 1.0,
+            "is_enable_stage2": true,
+            "pyramid_num_stages": 3,
+            "pyramid_num_inference_steps_list": [1, 1, 1],
+            "is_amplify_first_chunk": false,
+            "weight": 1
+        }
+    ]'
+```
+
+`--random-request-config` forwards standard fields such as `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` as first-class request fields. Model-specific fields are forwarded in the request body, which is required for Helios stage-2 settings.
+
+## 4.2 Warmed Serving Benchmark
+
+For local bring-up, use the same shared serving benchmark with one warmup request and one measured request:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --endpoint /v1/videos \
+    --model BestWishYsh/Helios-Distilled \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --warmup-requests 1 \
+    --warmup-num-inference-steps 1 \
+    --random-prompt "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting." \
+    --seed 42 \
+    --disable-tqdm \
+    --output-file metrics.json \
+    --save-response-dir responses \
+    --random-request-config '[
+        {
+            "width": 640,
+            "height": 384,
+            "num_inference_steps": 50,
+            "num_frames": 99,
+            "fps": 16,
+            "guidance_scale": 1.0,
+            "is_enable_stage2": true,
+            "pyramid_num_stages": 3,
+            "pyramid_num_inference_steps_list": [1, 1, 1],
+            "is_amplify_first_chunk": false,
+            "weight": 1
+        }
+    ]'
+```
+
+## 4.3 PR-vs-Main Accuracy and Performance Comparison
+
+The PR-vs-main accuracy/performance gate lives in the Helios accuracy pytest. It checks out `origin/main`, starts a server for each checkout, and drives both with `diffusion_benchmark_serving.py`:
+
+```bash
+pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \
+    -m 'full_model and benchmark' \
+    --run-level 'full_model'
+```
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A (99 frames, 384x640)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Frame count: 99 (`33 x 3`)
+* FPS: 16
+* Stage-2 steps: `[1, 1, 1]`
+
+```json
+[
+  {
+    "width": 640,
+    "height": 384,
+    "num_inference_steps": 50,
+    "num_frames": 99,
+    "fps": 16,
+    "guidance_scale": 1.0,
+    "is_enable_stage2": true,
+    "pyramid_num_stages": 3,
+    "pyramid_num_inference_steps_list": [1, 1, 1],
+    "is_amplify_first_chunk": false,
+    "weight": 1
+  }
+]
+```
+
+---
+
+# 6. Performance Metrics
+
+| Metric | Description | Unit |
+| --- | --- | --- |
+| Mean Latency | Mean request latency | seconds |
+| P99 Latency | P99 request latency | seconds |
+| Generated FPS | `num_frames / measured_time_s` | frames/second |
+| Candidate/reference latency ratio | PR latency divided by reference latency | ratio |
+
+---
+
+# 7. Accuracy Metrics
+
+The PR-vs-main gate compares encoded video frames extracted from MP4 outputs:
+
+| Metric | Threshold | Direction |
+| --- | --- | --- |
+| PSNR | `>= 35 dB` | Higher is better |
+| MAE | `<= 3.0 / 255` | Lower is better |
+| Candidate/reference latency ratio | `<= 1.03` | Lower is better |
+
+Local L20X bring-up for this PR produced:
+
+| Variant | Measured Latency (ms) | Generated FPS |
+| --- | --- | --- |
+| `origin/main` | 5317.95 | 18.62 |
+| Candidate PR | 5283.27 | 18.74 |
+
+Encoded-video similarity:
+
+| Metric | Value |
+| --- | --- |
+| PSNR | 38.66 dB |
+| MAE | 1.89 / 255 |
+| Cosine similarity | 0.99978 |
+
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record vLLM and vLLM-Omni versions
+* Record benchmark parameters: resolution, frame count, FPS, seed, concurrency, and stage-2 steps
+* Use frame counts that are multiples of 33
+* Keep prompt, seed, resolution, and sampling parameters identical when comparing PR vs main
+* Ensure no background workload on GPUs during testing
+* Use the same MP4 export and frame extraction path for both variants
+
+---
+
+This document serves as the Helios-Distilled serving performance reference under vLLM-Omni.

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -58,7 +58,7 @@ Record these parameters when reporting performance results.
 
 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend vllm-omni \
+    --endpoint /v1/chat/completions \
     --dataset <DATASET_NAME> \
     --task t2i \
     --num-prompts <N> \
@@ -71,7 +71,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 | Parameter              | Description                       |
 | ---------------------- | --------------------------------- |
-| `--backend`            | Serving backend (use `vllm-omni`) |
+| `--endpoint`           | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
 | `--dataset`            | Dataset name (`random` or custom) |
 | `--task`               | Task type (e.g., `t2i`)           |
 | `--num-prompts`        | Total number of requests          |
@@ -129,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend vllm-omni \
+    --endpoint /v1/chat/completions \
     --dataset random \
     --task t2i \
     --num-prompts 1 \

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

@@ -53,7 +53,7 @@ Record these parameters when reporting performance results.
 
 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend v1/videos \
+    --endpoint /v1/videos \
     --dataset <DATASET_NAME> \
     --task t2v \
     --num-prompts <N> \
@@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 | Parameter              | Description                       |
 | ---------------------- | --------------------------------- |
-| `--backend`            | Serving backend (use `v1/videos`) |
+| `--endpoint`           | API endpoint (use `/v1/videos`; leading `/` is optional) |
 | `--dataset`            | Dataset name (`random` or custom) |
 | `--task`               | Task type (e.g., `t2v`)           |
 | `--num-prompts`        | Total number of requests          |
@@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend v1/videos \
+    --endpoint /v1/videos \
     --dataset random \
     --task t2v \
     --num-prompts 1 \

benchmarks/diffusion/quantization_quality.py

@@ -148,6 +148,7 @@ def _build_omni_kwargs(args, quantization=None):
 def _generate_image(omni, args, prompt, seed):
     """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
     from vllm_omni.platforms import current_omni_platform
 
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
@@ -165,8 +166,9 @@ def _generate_image(omni, args, prompt, seed):
     elapsed = time.perf_counter() - start
     peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
 
-    first = outputs[0]
-    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    req_out = OmniRequestOutput.unwrap_result(outputs)
+    if not req_out.images:
+        raise ValueError("Could not extract image output from result.")
     img = req_out.images[0]
     return img, elapsed, peak_mem
 

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.20.0
+ARG VLLM_VERSION=v0.21.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -6,10 +6,11 @@ nav:
     - getting_started/installation/*
   - Serving:
     - OpenAI-Compatible API:
-      - Diffusion Chat API: serving/diffusion_chat_api.md
+      - Chat Completions API: serving/chat_completions_api.md
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
+      - Videos: serving/videos_api.md
       - Audio Generation: serving/audio_generate_api.md
       - Streaming Video Input: serving/video_stream_api.md
   - Examples:

docs/design/feature/pipeline_parallel.md

@@ -0,0 +1,356 @@
+# Pipeline Parallel
+
+This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
+image-to-video pipelines as the reference implementations.
+
+## Implementation Checklist
+
+Adding Pipeline Parallel support requires:
+
+1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
+2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
+3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
+4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
+5. ✅ **Test parity** - compare PP results against the single-GPU baseline
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is Pipeline Parallelism?
+
+Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
+different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
+
+For each denoising step:
+
+1. Rank 0 starts the forward pass using the current latents.
+2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
+   intermediate tensors downstream.
+3. The last PP rank produces the final noise prediction.
+4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
+
+This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
+combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
+
+### Architecture
+
+vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
+
+| Method                            | Purpose                       | Automatic Behavior                                                              |
+|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
+| `diffuse()`                       | Denoising loop boundary       | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit      |
+| `predict_noise_maybe_with_cfg()`  | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
+| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync   | Runs scheduler on the last PP rank, returns updated latents to rank 0           |
+| `_sync_pp_send()`                 | Flush pending async sends     | Waits for outstanding `isend` handles before later collectives or decode        |
+
+`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
+defines how a local stage executes.
+
+### How It Works
+
+`predict_noise_maybe_with_cfg()` automatically switches between these modes:
+
+- **PP disabled** (`pipeline_parallel_size == 1`):
+    - Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
+- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
+    - Rank 0 starts with the input latents
+    - Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
+    - The last rank returns the final noise prediction
+    - When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
+      PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
+      `cfg_parallel_size > 1` is available to avoid this overhead.
+- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
+    - Each PP pipeline carries one CFG branch
+    - The last PP rank all-gathers across the CFG group
+    - CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
+
+`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
+
+- **PP disabled**:
+    - Falls back to `scheduler_step_maybe_with_cfg()`
+- **PP enabled**:
+    - Only the last PP rank has `noise_pred` and runs the scheduler step
+    - The resulting latents are sent back to rank 0
+    - Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
+
+This asynchronous design avoids unnecessary blocking between denoising steps.
+
+When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
+`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
+ensuring that the final non-blocking PP send completes before decode or later collectives.
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
+
+`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
+class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
+`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
+
+`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
+supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
+`combine_cfg_noise()`.
+
+**Example:**
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
+import torch.nn as nn
+
+
+class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
+    ...
+```
+
+The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
+`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
+`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
+
+### Step 2: Make model forward and `predict_noise()` PP-aware
+
+The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
+two inputs:
+
+- Normal input from rank 0, typically passed as `hidden_states` or `x`
+- `intermediate_tensors` from upstream PP ranks
+
+The standard model forward pattern is:
+
+1. If `intermediate_tensors` is present, read the local hidden state from it.
+2. Run only this rank's layer slice.
+3. Return `IntermediateTensors(...)` on non-last PP ranks.
+4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
+   for common pipelines.
+
+**Minimal example:**
+
+```python
+from vllm.sequence import IntermediateTensors
+from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
+
+
+def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
+    if intermediate_tensors is not None:
+        hidden_states = intermediate_tensors["hidden_states"]
+
+    for i in range(self.start_layer, self.end_layer):
+        hidden_states = self.layers[i](hidden_states)
+
+    pp_group = get_pp_group()
+    if not pp_group.is_last_rank:
+        return IntermediateTensors({"hidden_states": hidden_states})
+    return (hidden_states,)
+```
+
+### Step 3: Partition the transformer layers
+
+The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
+implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
+`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
+`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
+PP-aware weight loading.
+
+To prepare a transformer for PP, implement the following pieces in order.
+
+#### 3.1 Split the transformer layers across PP ranks
+
+Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
+usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
+
+The goal is:
+
+- every PP rank knows its `[start_layer, end_layer)` range
+- non-local layers do not execute on this rank
+- the forward can resume from incoming `intermediate_tensors`
+
+By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
+divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
+override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
+
+```bash
+# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
+export VLLM_PP_LAYER_PARTITION=8,12,12,8
+```
+
+The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
+the total number of transformer layers.
+
+#### 3.2 Expose `make_empty_intermediate_tensors`
+
+Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
+`make_empty_intermediate_tensors_factory(...)`.
+
+This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
+keys and hidden dimension.
+
+**Example:**
+
+```python
+from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
+
+self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+    ["hidden_states"],
+    inner_dim,
+)
+```
+
+For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
+factory is created with that key and the transformer hidden size.
+
+#### 3.3 Return `IntermediateTensors` on non-last PP ranks
+
+Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
+ranks and return `IntermediateTensors(...)` on non-last ranks.
+
+This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
+
+#### 3.4 Skip non-local weights during `load_weights()`
+
+When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
+rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
+
+If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
+placeholders.
+
+The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
+and fused parameters.
+
+If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
+contract.
+
+### Step 4: Use the standard denoising contract
+
+vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
+`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
+`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
+synchronization.
+
+In PP mode:
+
+- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
+- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
+  combined result.
+- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
+  consumed.
+- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
+
+---
+
+## Testing
+
+Use an offline inference script with `pipeline_parallel_size > 1`.
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2.mp4 \
+--pipeline-parallel-size=2
+```
+
+For PP + CFG-Parallel together:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2_cfg2.mp4 \
+--pipeline-parallel-size=2 \
+--cfg-parallel-size=2
+```
+
+**Verify:**
+
+1. The run completes without hangs at the PP stage boundary
+2. Output quality matches the non-PP baseline within normal numerical variance
+3. Peak memory per GPU drops compared with the single-rank model
+4. There are no outstanding communication errors before decode
+
+---
+
+## Troubleshooting
+
+### Issue: `TypeError` on import — `CFGParallelMixin` missing
+
+**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
+
+```
+TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
+```
+
+or:
+
+```
+TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
+```
+
+**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
+and lists `PipelineParallelMixin` first in the MRO.
+
+**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
+
+
+class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
+    ...
+```
+
+### Issue: Non-last PP ranks crash when calling `predict_noise`
+
+**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
+
+**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
+
+**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
+`intermediate_tensors` when present.
+
+### Issue: PP outputs differ from single-GPU baseline
+
+**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
+
+**Causes & Solutions:**
+
+- **Local layer partitioning is wrong**
+    - Verify each rank runs only its `[start_layer, end_layer)` slice
+- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
+    - Return `IntermediateTensors({...})` until the last PP stage
+- **CFG branch wiring is incorrect**
+    - When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Component               | Path                                                       | Notes                                                                                                                   |
+|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
+| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py`     | Core PP communication and scheduler helpers                                                                             |
+| `CFGParallelMixin`      | `vllm_omni/diffusion/distributed/cfg_parallel.py`          | Default `predict_noise()` tuple normalization and CFG helper fallback                                                   |
+| Wan2.2 transformer      | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py`  | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
+| Wan2.2 T2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py`     | Reference PP + CFG integration for text-to-video                                                                        |
+| Wan2.2 I2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video                                                                       |
+| PP tests                | `tests/diffusion/distributed/test_pipeline_parallel.py`    | Baseline parity and async communication tests                                                                           |

docs/serving/chat_completions_api.md

@@ -1,4 +1,4 @@
-# Diffusion Chat Completions API
+# Chat Completions API
 
 vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
 endpoint using diffusion models. This page explains how to pass generation

docs/serving/videos_api.md

@@ -0,0 +1,152 @@
+# Videos API
+
+vLLM-Omni provides an OpenAI-compatible video generation API for diffusion
+video models. The API supports asynchronous video jobs through `/v1/videos` and
+a synchronous benchmark-oriented endpoint through `/v1/videos/sync`.
+
+Each server instance runs a single model specified at startup with
+`vllm serve <model> --omni`.
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+### Create a Video Job
+
+```bash
+create_response=$(curl -s http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16" \
+  -F "num_inference_steps=40")
+
+video_id=$(echo "${create_response}" | jq -r '.id')
+```
+
+### Poll and Download
+
+```bash
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
+```
+
+## API Reference
+
+### Endpoints
+
+| Endpoint | Method | Description |
+|----------|--------|-------------|
+| `/v1/videos` | `POST` | Create an asynchronous video generation job |
+| `/v1/videos/sync` | `POST` | Generate a video synchronously and return raw video bytes |
+| `/v1/videos/{video_id}` | `GET` | Retrieve job status and metadata |
+| `/v1/videos` | `GET` | List stored video jobs |
+| `/v1/videos/{video_id}/content` | `GET` | Download generated video content |
+| `/v1/videos/{video_id}` | `DELETE` | Delete a video job and stored output |
+
+### Request Parameters
+
+`POST /v1/videos` and `POST /v1/videos/sync` accept `multipart/form-data`.
+
+#### OpenAI-style fields
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | string | **required** | Text prompt for video generation |
+| `model` | string | server's model | Optional model name |
+| `seconds` | string | null | Requested clip duration in seconds |
+| `size` | string | null | Requested output size in `WIDTHxHEIGHT` format |
+| `user` | string | null | Optional user identifier |
+
+#### vLLM-Omni extension fields
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
+| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
+| `width` | integer | model default | Output video width |
+| `height` | integer | model default | Output video height |
+| `num_frames` | integer | model default | Number of generated frames |
+| `fps` | integer | model default | Output frames per second |
+| `num_inference_steps` | integer | model default | Number of diffusion steps |
+| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
+| `guidance_scale_2` | number | null | CFG guidance scale for the high-noise stage |
+| `boundary_ratio` | number | null | Boundary split ratio for multi-stage denoising |
+| `flow_shift` | number | null | Scheduler flow-shift value |
+| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
+| `seed` | integer | null | Random seed for reproducibility |
+| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
+| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
+| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
+| `frame_interpolation_scale` | number | null | RIFE inference scale |
+| `frame_interpolation_model_path` | string | null | Local path or Hugging Face repo for the interpolation model |
+| `lora` | string | null | JSON-encoded LoRA configuration object |
+| `extra_params` | string | null | JSON-encoded object for additional model-specific parameters |
+
+### Create Response
+
+`POST /v1/videos` returns a job record:
+
+```json
+{
+  "id": "video-123",
+  "status": "queued",
+  "created_at": 1701234567
+}
+```
+
+The final content is available from `/v1/videos/{video_id}/content` after the
+job status becomes `completed`.
+
+### Synchronous Response
+
+`POST /v1/videos/sync` blocks until generation finishes and returns raw video
+bytes. It is useful for benchmarks and simple scripts that do not need job
+storage or polling.
+
+## Examples
+
+### Image-to-Video
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=animate this image with subtle camera movement" \
+  -F "input_reference=@input.png" \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16"
+```
+
+### Synchronous Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/videos/sync \
+  -F "prompt=A small robot walking through a neon city" \
+  -F "width=854" \
+  -F "height=480" \
+  -F "num_frames=80" \
+  -F "fps=16" \
+  -o output.mp4
+```
+
+## Storage
+
+Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
+stored:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+```
+
+## Model-Specific Examples
+
+For complete text-to-video and image-to-video walkthroughs, see:
+
+- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
+- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)

docs/user_guide/diffusion/frame_interpolation.md

@@ -28,7 +28,6 @@ Frame interpolation is currently supported for:
 
 - `WanPipeline` (Wan2.2 text-to-video)
 - `WanImageToVideoPipeline`
-- `Wan22TI2VPipeline`
 
 ## Request Parameters
 

docs/user_guide/diffusion/parallelism/overview.md

@@ -4,13 +4,14 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
 
 ## Supported Methods
 
-| Method | Description |
-|--------|-------------|
-| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
-| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
-| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
-| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
-| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
-| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
+| Method                                             | Description                                                                                                         |
+|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
+| **[Tensor Parallelism](tensor_parallel.md)**       | Shards DiT weights across GPUs to reduce per-GPU memory                                                             |
+| **[Sequence Parallelism](sequence_parallel.md)**   | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
+| **[CFG-Parallel](cfg_parallel.md)**                | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation                         |
+| **[Pipeline Parallelism](pipeline_parallel.md)**   | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory             |
+| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
+| **[HSDP](hsdp.md)**                                | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs              |
+| **[Expert Parallelism](expert_parallel.md)**       | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0)                                          |
 
 See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.

docs/user_guide/diffusion/parallelism/pipeline_parallel.md

@@ -0,0 +1,205 @@
+# Pipeline Parallelism Guide
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
+part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
+devices.
+
+It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
+Parallelism.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        pipeline_parallel_size=2,
+    ),
+)
+
+outputs = omni.generate(
+    {"prompt": "A cinematic drone shot over snowy mountains"},
+    OmniDiffusionSamplingParams(
+        num_inference_steps=40,
+        num_frames=81,
+        height=704,
+        width=1280,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python scripts under:
+
+- `examples/offline_inference/text_to_video/text_to_video.py`
+- `examples/offline_inference/image_to_video/image_to_video.py`
+
+Text-to-video example:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2.mp4 \
+--pipeline-parallel-size=2
+```
+
+Pipeline Parallelism can also be combined with CFG-Parallel:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+--width=1280 \
+--height=704 \
+--guidance-scale=5.0 \
+--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
+--output=t2v_5B_pp2_cfg2.mp4 \
+--pipeline-parallel-size=2 \
+--cfg-parallel-size=2
+```
+
+### Online Serving
+
+Enable Pipeline Parallelism in online serving:
+
+```bash
+# Default PP configuration
+vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
+
+# PP + CFG-Parallel
+vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
+  --pipeline-parallel-size 2 \
+  --cfg-parallel-size 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`
+
+| Parameter                | Type | Default | Description                                                                                                      |
+|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
+| `pipeline_parallel_size` | int  | 1       | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
+
+> [!NOTE]
+> Total GPU count is the product of all enabled distributed dimensions, for example
+`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
+
+### Manual Layer Partitioning
+
+By default, transformer layers are distributed evenly across PP ranks. You can override this with the
+`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
+
+```bash
+# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
+export VLLM_PP_LAYER_PARTITION=8,12,12,8
+```
+
+The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
+the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
+ranks.
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large diffusion transformers that do not fit comfortably on one GPU
+- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
+- Combining with CFG-Parallel or other distributed methods on supported models
+
+**Not for:**
+
+- Single GPU setups
+- Models that do not support Pipeline Parallelism
+  (check [supported models](../../diffusion_features.md#supported-models))
+- Very small models where inter-stage communication overhead may outweigh the benefit
+
+### Expected Behavior
+
+Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
+Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
+primarily a latency optimization.
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: No benefit from Pipeline Parallelism
+
+**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
+
+**Solutions**:
+
+1. **Check your goal:**
+
+```python
+# PP is mainly for memory scaling, not guaranteed latency speedup
+parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
+```
+
+2. **Check model support:**
+    - Verify your model in [supported models](../../diffusion_features.md#supported-models)
+    - PP is currently validated only on selected pipelines
+
+3. **Combine with other methods when appropriate:**
+    - PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
+
+### Common Issue 2: PP pipeline fails at import
+
+**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
+
+**Solutions**:
+
+1. Inherit both mixins.
+2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
+3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
+
+---
+
+## Summary
+
+1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
+2. ✅ **Use Supported Models** - Verify your model supports PP in
+   [supported models](../../diffusion_features.md#supported-models)
+3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
+4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers

docs/user_guide/diffusion_features.md

@@ -35,14 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
 
 Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
 
-| Method | Description | Best For |
-|--------|-------------|----------|
-| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
-| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
-| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
-| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
-| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
-| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+| Method                                                                 | Description                                                              | Best For                                                          |
+|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
+| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)**           | Sequence parallelism via all-to-all communication                        | High-resolution images (>1536px) or long videos with 2-8 GPUs     |
+| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)**       | Sequence parallelism via ring-based communication                        | Videos, very long sequences, memory-constrained, with 2-8 GPUs    |
+| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)**              | Splits CFG positive/negative branches across devices                     | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs    |
+| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)**     | Shards model weights across devices                                      | Large models that don't fit in single GPU, with 2+ GPUs           |
+| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
+| **[HSDP](diffusion/parallelism/hsdp.md)**                              | Weight sharding via FSDP2, redistributed on-demand at runtime            | Very large models (14B+) on limited VRAM, combinable with SP      |
+| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)**     | Shards MoE expert MLP blocks across devices                              | MoE diffusion models (e.g., HunyuanImage3.0)                      |
 
 #### Startup Optimization
 
@@ -105,32 +106,32 @@ The following tables show which models support each feature:
 
 ### ImageGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅  | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
-| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
+| **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
+| **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
+| **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
+| **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -138,15 +139,15 @@ The following tables show which models support each feature:
 
 ### VideoGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
-| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
-| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
-| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
+| **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
+| **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 
 **Frame Interpolation Support**
 
@@ -155,9 +156,9 @@ The following tables show which models support each feature:
 
 ### AudioGen
 
-| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
+| Model                 | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Stable-Audio-Open** |     ✅     |     ❌      |           ❓           |       ❓        |         ❌         |          ❌          |   ✅    |             ✅             |          ❌           |       ✅        |        ❌         |
 
 
 ## Feature Compatibility
@@ -257,7 +258,7 @@ Measured on NVIDIA H800:
 
 **Parallelism Methods:**
 
-- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
+- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
 
 **Memory Optimization:**
 

docs/user_guide/examples/online_serving/glm_image.md

@@ -101,7 +101,7 @@ guides.
 
 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
+[Chat Completions API guide](../../../serving/chat_completions_api.md)).
 When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
 or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
 the supported generation controls as top-level fields directly. For image

docs/user_guide/examples/online_serving/image_to_image.md

@@ -208,7 +208,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
     ```
 
     For details on how generation parameters are handled across different clients, see the
-    [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
+    [Chat Completions API guide](../../../../serving/chat_completions_api.md).
 
 ### Layered Image Generation (Qwen-Image-Layered)
 
@@ -361,7 +361,7 @@ Provide multiple images in `content` (order matters):
 
 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
+[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
 When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
 endpoint, pass the supported generation controls as top-level form fields
 directly. For image dimensions and count, use `size` and `n` rather than

docs/user_guide/examples/online_serving/text_to_image.md

@@ -216,7 +216,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
     ```
 
     For details on how generation parameters are handled across different clients, see the
-    [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
+    [Chat Completions API guide](../../../../serving/chat_completions_api.md).
 
 ### Multimodal Input (Text + Structured Content)
 
@@ -237,7 +237,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
 
 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
+[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
 When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
 endpoint, pass the supported generation controls as top-level JSON fields
 directly. For image dimensions and count, use `size` and `n` rather than

docs/user_guide/quantization/overview.md

@@ -138,3 +138,98 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
 build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
 build_quant_config(None)
 ```
+
+## Output Similarity Comparison Tool
+
+Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
+to compare a reference diffusion run with a quantized candidate run using the
+same prompt, seed, resolution, scheduler settings, and inference steps. The
+tool compares final decoded images or video frames, and also reports generation
+latency and worker-reported peak memory when available.
+
+This is useful when validating whether online quantization, an offline
+pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation
+quality close to the BF16 reference.
+
+### Online Quantization Example
+
+```bash
+python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
+  --task t2i \
+  --model Qwen/Qwen-Image \
+  --candidate-quantization fp8 \
+  --ignored-layers img_mlp \
+  --prompt "a cup of coffee on the table" \
+  --height 512 --width 512 \
+  --num-inference-steps 20 \
+  --seed 142 \
+  --output-json /tmp/qwen_image_fp8_similarity/result.json \
+  --save-output-dir /tmp/qwen_image_fp8_similarity/images \
+  --enforce-eager
+```
+
+### Offline Checkpoint Example
+
+Use `--candidate-model` when the candidate is already quantized or lives at a
+different model path:
+
+```bash
+python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
+  --task t2i \
+  --reference-model Qwen/Qwen-Image \
+  --candidate-model /path/to/qwen-image-fp8-checkpoint \
+  --prompt "a cup of coffee on the table" \
+  --height 512 --width 512 \
+  --num-inference-steps 20 \
+  --seed 142 \
+  --output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json
+```
+
+If the checkpoint does not include a loadable quantization config, pass one
+explicitly:
+
+```bash
+--candidate-quantization-config-json '{"method":"fp8"}'
+```
+
+### Output Metrics
+
+The output JSON includes `output_metrics`, `reference_generation`, and
+`candidate_generation`.
+
+| Metric | Direction | Meaning |
+|--------|-----------|---------|
+| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. |
+| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. |
+| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. |
+| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. |
+| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. |
+| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. |
+| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. |
+| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. |
+
+Recommended starting thresholds for same-seed diffusion comparisons:
+
+| Metric | Smoke threshold | Stricter target | Notes |
+|--------|-----------------|-----------------|-------|
+| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. |
+| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. |
+| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. |
+| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. |
+
+These thresholds are heuristics. Tune them by model family, task, resolution,
+quantization method, and deployment tolerance. For release gating, pair the
+numeric report with visual inspection of saved reference and candidate outputs.
+
+The tool intentionally reports separate quality, latency, and memory metrics
+instead of a single consolidated similarity score. A single score can hide
+important tradeoffs, for example a candidate with good PSNR but a meaningful
+memory regression, or a candidate with low average error but localized visual
+artifacts. If you need a project-specific pass/fail gate, define it as an
+explicit policy over the individual metrics.
+
+Pixel-level metrics do not measure semantic consistency. For higher-cost
+evaluation, you can complement this report with a vision-language judge that
+describes the reference and candidate outputs and compares those descriptions.
+Keep that semantic check separate from this lightweight tool so users can
+choose whether the additional model cost and latency are appropriate.

examples/offline_inference/hunyuan_image3/end2end.py

@@ -77,6 +77,24 @@ def parse_args():
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
     parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument(
+        "--diffusion-kv-cache-dtype",
+        type=str,
+        default=None,
+        help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.",
+    )
+    parser.add_argument(
+        "--diffusion-kv-cache-skip-steps",
+        type=str,
+        default=None,
+        help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.",
+    )
+    parser.add_argument(
+        "--diffusion-kv-cache-skip-layers",
+        type=str,
+        default=None,
+        help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.",
+    )
     parser.add_argument(
         "--additional-config",
         type=str,
@@ -140,6 +158,9 @@ def main():
         "init_timeout": args.init_timeout,
         "enforce_eager": args.enforce_eager,
         "mode": _MODALITY_MODE[args.modality],
+        "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype,
+        "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps,
+        "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers,
     }
 
     if additional_config is not None:
@@ -237,6 +258,9 @@ def main():
         print(f"  Inference steps: {args.steps}")
         print(f"  Guidance scale: {args.guidance_scale}")
         print(f"  Seed: {args.seed}")
+        print(f"  diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}")
+        print(f"  diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}")
+        print(f"  diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}")
     if args.modality == "text2img":
         print(f"  Output size: {args.width}x{args.height}")
     if args.image_path:

examples/offline_inference/image_to_video/image_to_video.py

@@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "mxfp8", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+    )
+
+    # Distributed and parallel execution
     parser.add_argument(
         "--ulysses-degree",
         type=int,
@@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
-    parser.add_argument(
-        "--cache-backend",
-        type=str,
-        default=None,
-        choices=["cache_dit", "tea_cache"],
-        help=(
-            "Cache backend to use for acceleration. "
-            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
-            "Default: None (no cache acceleration)."
-        ),
-    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
+        "--pipeline-parallel-size",
+        type=int,
+        default=1,
+        help="Number of pipeline parallel stages.",
     )
     parser.add_argument(
         "--profiler-config",
@@ -326,6 +334,7 @@ def main():
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
     )
     omni_kwargs = dict(
         model=args.model,
@@ -366,7 +375,8 @@ def main():
     print(f"  diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
     print(
         f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
-        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
+        f" pipeline_parallel_size={args.pipeline_parallel_size}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")

examples/offline_inference/text_to_video/text_to_video.py

@@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "mxfp8", "int8", "gguf"],
+        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+    )
+
+    # Distributed and parallel execution
     parser.add_argument(
         "--ulysses-mode",
         type=str,
@@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
+    parser.add_argument(
+        "--pipeline-parallel-size",
+        type=int,
+        default=1,
+        help="Number of pipeline parallel stages.",
+    )
     parser.add_argument(
         "--enable-expert-parallel",
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--profiler-config",
-        type=parse_profiler_config,
-        default=None,
-        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -262,6 +270,7 @@ def main():
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
         vae_patch_parallel_size=args.vae_patch_parallel_size,
+        pipeline_parallel_size=args.pipeline_parallel_size,
         enable_expert_parallel=args.enable_expert_parallel,
     )
 
@@ -307,7 +316,8 @@ def main():
     print(
         f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
         f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
-        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
+        f" enable_expert_parallel={args.enable_expert_parallel}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")

pyproject.toml

@@ -252,3 +252,4 @@ fullset = "fullset"
 Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
+ND = "ND"

recipes/Baidu/ERNIE-Image.md

@@ -1,4 +1,6 @@
-# ERNIE-Image for text-to-image generation
+# ERNIE-Image
+
+> Text-to-image online serving (ERNIE-Image 8B)
 
 ## Summary
 

recipes/LTX/LTX-2.3.md

@@ -1,4 +1,4 @@
-# LTX-2.3 Text-to-Video with Audio on 1x GPU (96GB VRAM)
+# LTX-2.3
 
 > 22B parameter text-to-video + audio generation model served via vLLM-Omni
 

recipes/LTX/LTX-2.md

@@ -1,6 +1,6 @@
-# LTX-2 T2V and I2V
+# LTX-2
 
-> LTX-2 for text-to-video & image-to-video on 1x H200 141GB
+> Text-to-video and image-to-video serving
 
 ## Summary
 

recipes/Qwen/Qwen-Image-Edit.md

@@ -0,0 +1,171 @@
+# Qwen-Image-Edit
+
+> Image editing serving
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen-Image-Edit`
+- Task: Text-guided single-image editing
+- Mode: Online serving with the OpenAI-compatible multimodal endpoint
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe as a baseline for serving `Qwen/Qwen-Image-Edit` for single-image editing on H200 141GB. Two configurations are validated below, sharing one software environment:
+
+| Configuration | Parallelism | Notes |
+| --- | --- | --- |
+| `1x H200` | none | Single-card baseline; ~60 GiB peak VRAM on a 141 GB card |
+| `2x H200` | TP=2 | Lower single-request latency (≈ 1.60× speedup) and lower per-GPU VRAM |
+
+## References
+
+- Model: <https://huggingface.co/Qwen/Qwen-Image-Edit>
+- Example guide: [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md)
+- User guide: [`docs/user_guide/examples/online_serving/image_to_image.md`](../../docs/user_guide/examples/online_serving/image_to_image.md)
+
+## Hardware Support
+
+### GPU
+
+The two hardware configurations below share the same software environment.
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12+
+- Driver / runtime: CUDA-capable runtime matching the repository build
+- vLLM: use versions from your current checkout, >=0.20.0
+- vLLM-Omni: use versions from your current checkout
+
+#### 1x H200 141GB
+
+##### Command
+
+```bash
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen-Image-Edit --omni \
+  --port 8092
+```
+
+The bundled launcher is equivalent:
+
+```bash
+bash examples/online_serving/image_to_image/run_server.sh
+```
+
+The server is ready when the log shows `INFO:     Application startup complete.` (preceded by `Started server process [PID]` and `Waiting for application startup.`).
+
+##### Verification
+
+Run a single-image edit against the public example asset:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png
+
+PROMPT="Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'"
+IMG_B64=$(base64 -w0 qwen-bear.png)
+
+curl -s http://localhost:8092/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d "$(jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
+    messages: [{
+      role: "user",
+      content: [
+        {type: "text", text: $prompt},
+        {type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}}
+      ]
+    }],
+    extra_body: {
+      height: 1024,
+      width: 1024,
+      num_inference_steps: 50,
+      guidance_scale: 1.0,
+      seed: 42
+    }
+  }')" \
+  | jq -r '.choices[0].message.content[0].image_url.url' \
+  | cut -d',' -f2- | base64 -d > qwen_image_edit_1card.png
+```
+
+The decoded PNG is a 1024×1024 RGB image of the bear mascot dancing under stylized moonlight, surrounded by floating stars and a "Be Kind" text bubble.
+
+You can also run the offline CLI without standing up a server:
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --image qwen-bear.png \
+  --prompt "Let this mascot dance under the moon, ..." \
+  --output qwen_image_edit_offline.png \
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
+```
+
+Validated reference numbers for 1024×1024, 50 denoising steps, batch size 1, `guidance_scale=1.0`, seed 42 (`Cold` = first request after startup, `Warm` = identical subsequent request):
+
+| Cold latency | Warm latency | Peak GPU memory |
+| ---:         | ---:         | ---:            |
+| 17.80 s      | 17.78 s      | 61,091 MiB (~59.66 GiB) |
+
+##### Notes
+
+- Generation parameters live under `extra_body` for `/v1/chat/completions`, or as top-level form fields when using `/v1/images/edits`. See [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md) for the full request / response schema.
+- Memory: a 1024×1024 / 50-step edit uses ~60 GiB on a single H200, well within the 141 GB headroom.
+- Multi-image variants of Qwen-Image-Edit (e.g. `Qwen-Image-Edit-2509`) are intentionally out of scope for this recipe and are not validated here.
+
+#### 2x H200 141GB (TP=2)
+
+##### Command
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen-Image-Edit --omni \
+  --port 8092 \
+  --tensor-parallel-size 2
+```
+
+In addition to the `Application startup complete.` line, the TP=2 server's `non-default args` log line echoes `'tensor_parallel_size': 2`, and a separate `Worker N: Initialized device and distributed environment` line is emitted for each rank — one per visible GPU.
+
+##### Verification
+
+Run the same `curl` command as in the **1x H200** Verification section above against the TP=2 server (same prompt, same seed, same image). Save the result to `qwen_image_edit_2card_tp2.png` so the two outputs sit side by side for comparison. During TP=2 inference, `nvidia-smi` shows sustained `~99 %` utilization on both GPUs, which is the expected signal that TP wiring is active.
+
+Same workload as the 1x table above (1024×1024, 50 steps, `guidance_scale=1.0`, seed 42):
+
+| Configuration | Cold latency | Warm latency | Peak per-GPU memory |
+| ---           | ---:         | ---:         | ---:                |
+| 1x H200       | 17.80 s      | 17.78 s      | 61,091 MiB (~59.66 GiB) |
+| 2x H200 TP=2  | 11.16 s      | 11.09 s      | 49,377 MiB (~48.22 GiB) |
+
+TP=2 reduces single-request warm latency from `17.78 s` to `11.09 s` (≈ **1.60×**) and drops per-GPU peak VRAM by ~11.4 GiB.
+
+##### Notes
+
+- TP=2 shards model weights across both cards, which is what produces both the latency reduction and the per-GPU VRAM drop above.
+- CFG-Parallel (`--cfg-parallel-size 2`) is supported on the Qwen-Image family but only when `cfg_scale > 1`. The verification command above uses `guidance_scale=1.0` to match the public example image, so it does not exercise CFG-parallel; a separate run with `cfg_scale > 1` is needed if you want to benchmark that path.
+- For deeper acceleration knobs (Cache-DiT, sequence parallel, HSDP) see [`docs/user_guide/diffusion/parallelism/overview.md`](../../docs/user_guide/diffusion/parallelism/overview.md); this recipe intentionally documents only the validated TP=2 baseline.
+
+#### Profiling
+
+The numbers below are from a warm request at 1024×1024, 50 steps, and `guidance_scale=1.0`, measured with `--enable-diffusion-pipeline-profiler` and a 250 ms `nvidia-smi --query-gpu=memory.used` sampler. For TP=2, latency uses rank max, and VRAM is per GPU.
+
+Per-stage latency:
+
+| Stage | TP=1 | TP=2 | Notes |
+| --- | ---: | ---: | --- |
+| `text_encoder.forward` | 0.132 s | 0.135 s | ~1.0× |
+| `vae.encode` | 0.040 s | 0.041 s | ~1.0× |
+| `diffuse` | 17.105 s | 10.382 s | **1.65×** speedup |
+| `vae.decode` | 0.064 s | 0.066 s | ~1.0× |
+| `forward` | 17.372 s | 10.652 s | 1.63× speedup |
+
+Per-GPU VRAM:
+
+| Snapshot | TP=1 | TP=2 | Notes |
+| --- | ---: | ---: | --- |
+| Static VRAM | 58,777 MiB | 47,065 MiB | after model load, before request |
+| Peak VRAM | 61,091 MiB | 49,377 MiB | during inference |
+| Request-time increment | 2,314 MiB | 2,312 MiB | peak minus static |
+
+The latency gain mainly comes from the DiT denoising stage (`diffuse` in the table above). The other measured stages, such as text encoder and VAE encode/decode, stay essentially unchanged between TP=1 and TP=2, which is consistent with them not benefiting meaningfully from TP in this pipeline. The `forward` row's 1.63× speedup tracks `diffuse`, since `diffuse` accounts for about 98% of total pipeline time.
+
+For memory, TP=2 saves about 11.4 GiB per GPU, mostly from static model-state memory, likely dominated by tensor-parallel sharding of the DiT weights. The request-time increment is almost unchanged, so TP=2 does not materially reduce transient inference-time memory for this single-image, batch-1 workload.

recipes/Qwen/Qwen-Image.md

@@ -1,4 +1,6 @@
-# Qwen-Image for text-to-image serving with optional continuous batching on 1x A100 80GB
+# Qwen-Image
+
+> Text-to-image serving with optional step-wise continuous batching
 
 ## Summary
 
@@ -102,7 +104,7 @@ To replay the batching benefit with matched warmup:
 
 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-  --backend vllm-omni \
+  --endpoint /v1/chat/completions \
   --dataset vbench \
   --task t2i \
   --model Qwen/Qwen-Image \

recipes/Qwen/Qwen3-Omni.md

@@ -1,5 +1,7 @@
 # Qwen3-Omni
 
+> Online serving for multimodal chat
+
 ## Summary
 
 - Vendor: Qwen

recipes/Qwen/Qwen3-TTS.md

@@ -1,4 +1,6 @@
-# Qwen3-TTS for text-to-speech
+# Qwen3-TTS
+
+> Text-to-speech serving (CustomVoice / VoiceDesign / Base)
 
 ## Summary
 

recipes/README.md

@@ -19,34 +19,29 @@ recipes/
   Qwen/
     Qwen3-Omni.md
     Qwen3-TTS.md
-  Tencent-Hunyuan/
-    HunyuanVideo.md
+  Tencent/
+    Covo-Audio-Chat.md
 ```
 
 ## Available Recipes
 
-- [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md): text-to-image serving recipe for
-  Qwen-Image on `1x A100 80GB`, including optional step-wise continuous batching replay
-- [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
-  multimodal chat on `1x A100 80GB`
-- [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md): online
-  serving recipe for audio chat on `1x A100 80GB`
-- [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md): text-to-speech serving recipe
-  for Qwen3-TTS (CustomVoice / VoiceDesign / Base) on `1x H100/A100 80GB`
-- [`LTX/LTX-2.md`](./LTX/LTX-2.md): text-to-video and image-to-video serving
-  recipe for LTX-2 on `1x H200 141GB`
-- [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving
-  recipe for Wan2.2 14B on `8x Ascend NPU (A2/A3)`
-- [`Tencent-Hunyuan/HunyuanImage-3.0-Instruct.md`](./Tencent-Hunyuan/HunyuanImage-3.0-Instruct.md):
-  DiT-only text-to-image serving and benchmark recipe for HunyuanImage-3.0-Instruct
-  on `4x H100/H800 80GB`
-- [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md):
-  online serving recipe for multimodal chat (`4x H100 80GB`) and standalone TTS (`1x H100 80GB`)
-- [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md): text-to-image serving
-  online serving recipe for ERNIE-Image 8B on `1x RTX 4090 24GB` or `2x RTX 4090 24GB`
-- [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md): online serving recipe for TTS on `1x A800 80GB`
-- [`audiox/AudioX.md`](./audiox/AudioX.md): offline + online recipe for AudioX
-  unified text/video→audio diffusion on `1x L4 24GB`
+| Recipe | Task | Hardware |
+|--------|------|----------|
+| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
+| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
+| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
+| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
+| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
+| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
+| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
+| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
+| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
+| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
+| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
+| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
+| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark | 4x H100/H800 80GB |
+| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
+| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like

recipes/SenseNova/SenseNova-U1.md

@@ -1,4 +1,6 @@
-# SenseNova-U1 for unified image generation and understanding
+# SenseNova-U1
+
+> Unified image generation and understanding
 
 ## Summary
 

recipes/TEMPLATE.md

@@ -1,6 +1,6 @@
 # Recipe Title
 
-> Example: Qwen3-Omni for speech chat on 1x A100 80GB
+> Example: Qwen3-Omni for speech chat
 
 ## Summary
 

recipes/Tencent/Covo-Audio-Chat.md

@@ -1,4 +1,6 @@
-# Covo-Audio-Chat for audio chat on 1x A100 80GB
+# Covo-Audio-Chat
+
+> Online serving for audio chat
 
 ## Summary
 

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -1,7 +1,6 @@
-# HunyuanImage-3.0-Instruct DiT Image Generation on 4x GPU
+# HunyuanImage-3.0-Instruct
 
-> DiT-only text-to-image recipe for HunyuanImage-3.0-Instruct with FP8,
-> tensor parallelism, sequence parallelism, and CFG parallelism.
+> DiT-only text-to-image serving and benchmark
 
 ## Summary
 

recipes/Wan-AI/Wan2.2-I2V.md

@@ -1,4 +1,6 @@
-# Wan2.2 Image To Video
+# Wan2.2 Image-to-Video
+
+> Image-to-video serving (Wan2.2 14B)
 
 ## Summary
 

recipes/Wan-AI/Wan2.2-S2V.md

@@ -1,4 +1,6 @@
-# Wan2.2 Speech To Video
+# Wan2.2 Speech-to-Video
+
+> Speech-to-video serving (Wan2.2 14B)
 
 ## Summary
 

recipes/fishaudio/Fish-Speech-S2-Pro.md

@@ -1,5 +1,7 @@
 # Fish Speech S2 Pro
 
+> Online serving for TTS
+
 ## Summary
 
 - Vendor: FishAudio

recipes/inclusionAI/Ming-flash-omni-2.0.md

@@ -1,4 +1,6 @@
-# Ming-flash-omni 2.0 for omni-speech chat and standalone TTS
+# Ming-flash-omni 2.0
+
+> Online serving for multimodal chat + standalone TTS
 
 ## Summary
 

tests/benchmarks/metrics/test_metrics.py

@@ -63,5 +63,96 @@ def test_total_input_aggregated_from_output_prompt_len():
     )
 
 
+# ============================================================================
+# TTFT suppression for pure-audio (TTS) benchmarks
+# ============================================================================
+
+
+class _EmptyAwareTokenizer:
+    """Minimal tokenizer stub: token count == len(text), so '' -> 0 tokens.
+
+    Mirrors production where a TTS speech endpoint returns empty generated_text,
+    making total_output == 0 (the real CI path uses a real tokenizer, not None).
+    """
+
+    def __call__(self, text, add_special_tokens=False):
+        class _R:
+            pass
+
+        r = _R()
+        r.input_ids = [0] * len(text)
+        return r
+
+
+def _make_tts_output(prompt_len: int) -> MixRequestFuncOutput:
+    """Pure-TTS output: no text tokens, only audio. ttft is left unset (0.0)."""
+    output = MixRequestFuncOutput()
+    output.success = True
+    output.prompt_len = prompt_len
+    output.output_tokens = 0
+    output.generated_text = ""
+    output.ttft = 0.0
+    output.text_latency = 1.0
+    output.latency = 1.0
+    output.start_time = 0.0
+    output.itl = []
+    output.audio_ttfp = 0.05
+    output.audio_rtf = 0.2
+    output.audio_duration = 5.0
+    output.audio_frames = 120000
+    output.input_audio_duration = 0.0
+    output.error = ""
+    return output
+
+
+_TTS_PERCENTILE_METRICS = ["ttft", "e2el", "audio_rtf", "audio_ttfp", "audio_duration"]
+
+
+def test_tts_benchmark_omits_ttft(capsys):
+    """Pure-TTS run (total_output == 0) must not print a Time to First Token section."""
+    outputs = [_make_tts_output(100), _make_tts_output(120)]
+
+    calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=10.0,
+        tokenizer=_EmptyAwareTokenizer(),
+        selected_percentiles=[99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=_TTS_PERCENTILE_METRICS,
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    out = capsys.readouterr().out
+    assert "Time to First Token" not in out, "TTS bench must not surface a meaningless TTFT"
+    assert "Time to First Packet" in out, "audio TTFP must still be reported"
+    assert "End-to-end Latency" in out, "e2el must still be reported"
+
+
+def test_text_benchmark_still_reports_ttft(capsys):
+    """Regression guard: real text generation (total_output > 0) keeps TTFT."""
+    outputs = [_make_output(100, output_tokens=10), _make_output(120, output_tokens=10)]
+
+    calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=10.0,
+        tokenizer=_EmptyAwareTokenizer(),
+        selected_percentiles=[99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=_TTS_PERCENTILE_METRICS,
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    out = capsys.readouterr().out
+    assert "Time to First Token" in out, "text benchmarks must keep TTFT"
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/benchmarks/test_diffusion_backends_metrics.py

@@ -1,6 +1,11 @@
 import pytest
 
-from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions
+from benchmarks.diffusion.backends import (
+    RequestFuncInput,
+    async_request_chat_completions,
+    endpoint_filename_token,
+    normalize_endpoint,
+)
 
 
 class _MockResponse:
@@ -29,6 +34,32 @@ class _MockSession:
         return _MockResponse(self._payload)
 
 
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_normalization_accepts_optional_leading_slash():
+    assert normalize_endpoint("v1/videos") == "/v1/videos"
+    assert normalize_endpoint("/v1/videos") == "/v1/videos"
+    assert normalize_endpoint("v1/chat/completions") == "/v1/chat/completions"
+    assert normalize_endpoint("v1/images/generations") == "/v1/images/generations"
+
+
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_normalization_accepts_legacy_backend_aliases():
+    assert normalize_endpoint("vllm-omni") == "/v1/chat/completions"
+    assert normalize_endpoint("openai") == "/v1/images/generations"
+
+
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@pytest.mark.cpu
+def test_endpoint_filename_token_drops_leading_slash():
+    assert endpoint_filename_token("/v1/videos") == "v1_videos"
+    assert endpoint_filename_token("v1/chat/completions") == "v1_chat_completions"
+
+
 @pytest.mark.core_model
 @pytest.mark.benchmark
 @pytest.mark.cpu

tests/conftest.py

@@ -36,6 +36,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config):
 _ASSERTION_EXPORT_NAMES = (
     "assert_audio_speech_response",
     "assert_diffusion_response",
+    "assert_http_error",
     "assert_image_diffusion_response",
     "assert_image_valid",
     "assert_omni_response",
@@ -54,6 +55,7 @@ _MEDIA_EXPORT_NAMES = (
 _STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",)
 _RUNTIME_EXPORT_NAMES = (
     "DiffusionResponse",
+    "HttpResponse",
     "OmniResponse",
     "OmniRunner",
     "OmniRunnerHandler",

tests/dfx/perf/scripts/diffusion_result_template.json

@@ -1,7 +1,7 @@
 [
   {
     "test_name": null,
-    "backend": null,
+    "endpoint": null,
     "timestamp": null,
     "server_params": {
       "model": null,
@@ -43,7 +43,7 @@
       "stage_durations_mean": {},
       "stage_durations_p50": {},
       "stage_durations_p99": {},
-      "backend": null,
+      "endpoint": null,
       "model": null,
       "dataset": null,
       "task": null
@@ -51,6 +51,7 @@
     "log_file": null,
     "Model": null,
     "Framework": null,
+    "API Endpoint": null,
     "Hardware": null,
     "Deployment": null,
     "Task": null,

tests/dfx/perf/scripts/run_benchmark.py

@@ -1,8 +1,6 @@
 import json
 import os
-import subprocess
 import threading
-from datetime import datetime
 from pathlib import Path
 from typing import Any
 
@@ -14,6 +12,7 @@ from tests.dfx.conftest import (
     create_unique_server_params,
     get_benchmark_params_for_server,
     load_configs,
+    run_benchmark,
 )
 from tests.helpers.runtime import OmniServer
 
@@ -98,91 +97,6 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
-def run_benchmark(
-    args: list,
-    test_name: str,
-    flow,
-    dataset_name: str,
-    num_prompt,
-    *,
-    baseline_config: dict[str, Any] | None = None,
-    sweep_index: int | None = None,
-    request_rate: Any | None = None,
-    max_concurrency: Any | None = None,
-    random_input_len: Any | None = None,
-    random_output_len: Any | None = None,
-) -> Any:
-    """Run a single benchmark iteration and return the parsed result JSON.
-
-    After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the same
-    per-metric resolved thresholds as ``assert_result`` (via ``_baseline_thresholds_for_step``).
-    When ``random_input_len`` / ``random_output_len`` are set, they are also written into the result JSON;
-    omitted keys when not configured.
-    """
-    current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
-    ri = _safe_filename_token(random_input_len)
-    ro = _safe_filename_token(random_output_len)
-    result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
-    if "--result-filename" in args:
-        print(f"The result file will be overwritten by {result_filename}")
-    command = (
-        ["vllm", "bench", "serve", "--omni"]
-        + args
-        + [
-            "--save-result",
-            "--result-dir",
-            os.environ.get("BENCHMARK_DIR", "tests"),
-            "--result-filename",
-            result_filename,
-        ]
-    )
-    process = subprocess.Popen(
-        command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True
-    )
-
-    for line in iter(process.stdout.readline, ""):
-        print(line, end=" ")
-
-    for line in iter(process.stderr.readline, ""):
-        print(line, end=" ")
-
-    if "--result-dir" in command:
-        index = command.index("--result-dir")
-        result_dir = command[index + 1]
-    else:
-        result_dir = "./"
-
-    result_path = os.path.join(result_dir, result_filename)
-    if not os.path.exists(result_path):
-        with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
-            template_result: dict[str, Any] = json.load(f)
-        Path(result_path).parent.mkdir(parents=True, exist_ok=True)
-        with open(result_path, "w", encoding="utf-8") as f:
-            json.dump(template_result, f, ensure_ascii=False, indent=2)
-        print(f"Benchmark result file not generated, fallback to template: {result_path}")
-        result = template_result
-    else:
-        with open(result_path, encoding="utf-8") as f:
-            result = json.load(f)
-
-    if baseline_config:
-        result["baseline"] = _baseline_thresholds_for_step(
-            baseline_config,
-            sweep_index=sweep_index,
-            request_rate=request_rate,
-            max_concurrency=max_concurrency,
-        )
-    else:
-        result["baseline"] = {}
-    if random_input_len is not None:
-        result["random_input_len"] = random_input_len
-    if random_output_len is not None:
-        result["random_output_len"] = random_output_len
-    with open(result_path, "w", encoding="utf-8") as f:
-        json.dump(result, f, ensure_ascii=False, indent=2)
-    return result
-
-
 benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
 
 

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -5,16 +5,15 @@ This runner separates two concepts:
 
 1. ``server_type``: how the serving process is started.
    Currently only ``vllm-omni`` is supported here.
-2. ``benchmark_backend``: which serving API the benchmark client calls.
-   Examples: ``vllm-omni`` for ``/v1/chat/completions`` and ``v1/videos``
-   for async video jobs.
+2. ``benchmark_endpoint``: which serving API the benchmark client calls.
+   Examples: ``/v1/chat/completions`` and ``/v1/videos``.
 
 A config JSON file is REQUIRED via --test-config-file:
   pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
 All benchmark results for a session are consolidated into a single JSON file under
 BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
-Each entry in the file contains the test metadata (test_name, backend, benchmark_params,
+Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
 timestamp) together with the raw metrics returned by the benchmark script.
 """
 
@@ -33,6 +32,8 @@ from typing import Any, cast
 import psutil
 import pytest
 
+from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
+
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -400,7 +401,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
                 "model": cfg["server_params"]["model"],
                 "serve_args_dict": serve_args_dict,
                 "serve_args": _build_serve_args(serve_args_dict),
-                "benchmark_backend": cfg.get("benchmark_backend"),
+                "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                 "server_params": cfg["server_params"],
             }
         )
@@ -482,7 +483,7 @@ def run_benchmark(
     model: str,
     params: dict[str, Any],
     test_name: str,
-    backend: str = "vllm-omni",
+    endpoint: str = "/v1/chat/completions",
     server_cfg: dict[str, Any] | None = None,
     source_file: str = "",
 ) -> dict[str, Any]:
@@ -490,18 +491,19 @@ def run_benchmark(
 
     The raw metrics are written to a temporary file by the subprocess.  After
     the run completes the metrics are merged with full metadata (test_name,
-    backend, benchmark_params, timestamp, flat reporting fields) and appended
+    endpoint, benchmark_params, timestamp, flat reporting fields) and appended
     to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE).  The
     temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
     to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
     the record.
     """
     timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    endpoint = normalize_endpoint(endpoint)
 
     log_dir = BENCHMARK_RESULT_DIR / "logs"
     log_dir.mkdir(parents=True, exist_ok=True)
-    backend_label = backend.replace("/", "_")
-    log_file = log_dir / f"{test_name}_{backend_label}_{timestamp}.log"
+    endpoint_label = endpoint_filename_token(endpoint)
+    log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
 
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
         tmp_result_file = Path(tmp.name)
@@ -517,8 +519,8 @@ def run_benchmark(
         str(port),
         "--model",
         model,
-        "--backend",
-        backend,
+        "--endpoint",
+        endpoint,
         "--dataset",
         params.get("dataset", "random"),
         "--task",
@@ -545,7 +547,7 @@ def run_benchmark(
     # cause truncated or out-of-order log output when stdout is piped).
     cmd = [cmd[0], "-u"] + cmd[1:]
 
-    print(f"\nRunning benchmark (backend={backend}): {' '.join(cmd)}")
+    print(f"\nRunning benchmark (endpoint={endpoint}): {' '.join(cmd)}")
     print(f"  Log file: {log_file}")
 
     # Redirect stdout + stderr directly to the log file at the OS level
@@ -597,7 +599,7 @@ def run_benchmark(
 
     record: dict[str, Any] = {
         "test_name": test_name,
-        "backend": backend,
+        "endpoint": endpoint,
         "timestamp": timestamp,
         "server_params": server_cfg.get("server_params"),
         "benchmark_params": params,
@@ -605,7 +607,7 @@ def run_benchmark(
         "log_file": str(log_file),
         "Model": model,
         "Framework": server_type,
-        "API Backend": backend,
+        "API Endpoint": endpoint,
         "Hardware": "",
         "Deployment": "",
         "Task": params.get("task", "t2i"),
@@ -669,21 +671,21 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
             assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
 
 
-def _default_benchmark_backend_for_task(task: str) -> str:
-    """Return the default client-side benchmark backend for a diffusion task."""
+def _default_benchmark_endpoint_for_task(task: str) -> str:
+    """Return the default client-side benchmark endpoint for a diffusion task."""
     if task in {"t2v", "i2v", "ti2v"}:
-        return "v1/videos"
+        return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
-        return "vllm-omni"
-    raise ValueError(f"Unsupported task for benchmark backend resolution: {task}")
+        return "/v1/chat/completions"
+    raise ValueError(f"Unsupported task for benchmark endpoint resolution: {task}")
 
 
-def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
+def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
     """Resolve which serving API the benchmark client should call."""
-    configured = server_cfg.get("benchmark_backend")
+    configured = server_cfg.get("benchmark_endpoint")
     if configured:
-        return cast(str, configured)
-    return _default_benchmark_backend_for_task(cast(str, params.get("task", "t2i")))
+        return normalize_endpoint(cast(str, configured))
+    return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
 
 
 # ---------------------------------------------------------------------------
@@ -711,7 +713,7 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     server_cfg = getattr(diffusion_server, "server_cfg", {})
-    backend = _resolve_benchmark_backend(server_cfg, params)
+    endpoint = _resolve_benchmark_endpoint(server_cfg, params)
 
     result = run_benchmark(
         host=diffusion_server.host,
@@ -719,13 +721,13 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
         model=diffusion_server.model,
         params=params,
         test_name=test_name,
-        backend=backend,
+        endpoint=endpoint,
         server_cfg=server_cfg,
         source_file=cast(str, CONFIG_FILE_PATH),
     )
 
     print(f"\n{'=' * 60}")
-    print(f"Results for {test_name} (server={diffusion_server.server_type}, backend={backend}):")
+    print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
     for key in (
         "throughput_qps",
         "latency_mean",

tests/dfx/reliability/invalid_param_test/conftest.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared env for ``http_invalid`` (real :func:`omni_server`)."""
+
+from __future__ import annotations
+
+import io
+import os
+
+import pytest
+from PIL import Image
+
+# Match ``tests/e2e/online_serving/*`` module-level env for subprocess serve.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_TEST_CLEAN_GPU_MEMORY", "0")
+
+
+@pytest.fixture
+def tiny_png_bytes() -> bytes:
+    img = Image.new("RGB", (32, 32), color="gray")
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()

tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/audio/generate`` validation (live Stable Audio Open server)."""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+# https://github.com/vllm-project/vllm-omni/issues/3649 — invalid inputs still HTTP 200 / timeout / missing validation.
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model="stabilityai/stable-audio-open-1.0",
+            server_args=[
+                "--trust-remote-code",
+                "--enforce-eager",
+                "--model-class-name",
+                "StableAudioPipeline",
+            ],
+        ),
+        id="stable_audio_open",
+        marks=hardware_marks(res={"cuda": "L4"}),
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"input": "   "}, ("cannot be empty", "input"), id="input_whitespace_only", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param({"input": 123}, ("input", "a valid string"), id="input_wrong_type"),
+        pytest.param(
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "Input should be"),
+            id="response_format_invalid",
+        ),
+        pytest.param({"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="speed_too_low"),
+        pytest.param({"speed": 99.0}, ("speed", "less_than_equal", "4"), id="speed_too_high"),
+        pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
+        pytest.param(
+            {"stream_format": "sse"},
+            ("sse", "stream_format", "value_error", "not a supported stream_format"),
+            id="stream_format_sse_blocked",
+        ),
+        pytest.param(
+            {"audio_length": -1.0}, ("audio_end_in_s", "higher than", "audio_start_in_s"), id="audio_length_negative"
+        ),
+        pytest.param({"audio_length": 0}, "audio_length", id="audio_length_zero", marks=_SKIP_ISSUE_3649),
+        # ``OmniOpenAIServingAudioGenerate`` only forwards ``audio_start`` when ``audio_length`` is set.
+        pytest.param(
+            {"audio_start": -0.5, "audio_length": 5.0},
+            ("audio_start_in_s", "between", "0", "512"),
+            id="audio_start_negative",
+        ),
+        pytest.param(
+            {"negative_prompt": ["noise"]}, ("negative_prompt", "a valid string"), id="negative_prompt_wrong_type"
+        ),
+        pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
+        pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param(
+            {"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
+        ),
+        pytest.param(
+            {"num_inference_steps": 6000},
+            "num_inference_steps",
+            id="num_inference_steps_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"guidance_scale": 1001.0}, "guidance_scale", id="guidance_scale_above_max", marks=_SKIP_ISSUE_3649
+        ),
+        pytest.param(
+            {"audio_length": 86401.0}, ("Requested audio length", "exceeds maximum"), id="audio_length_above_max"
+        ),
+        pytest.param(
+            {"audio_start": 86401.0, "audio_length": 1.0},
+            ("audio_start_in_s", "between", "0", "512"),
+            id="audio_start_above_max",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_audio_generate_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body = {"model": omni_server.model, "input": "ambient electronic pad"}
+    body.update(overrides)
+    openai_client.send_audio_generate_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_generate_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("missing", "input", "field required"),
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -0,0 +1,862 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HTTP + WebSocket validation for Qwen3-TTS Base: ``/v1/audio/speech``, ``/v1/audio/speech/stream``, batch, voices."""
+
+from __future__ import annotations
+
+import io
+import json
+import uuid
+import wave
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.stage_config import get_deploy_config_path
+
+# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni).
+# Align where possible with serving_speech defaults: _TTS_MAX_INSTRUCTIONS_LENGTH (500),
+# _TTS_MAX_NEW_TOKENS_MAX (4096). Voice form max lengths are not public constants in prod;
+# values below keep tests self-contained and deterministic.
+_SPEECH_API_MAX_INSTRUCTIONS_CHARS = 500
+_SPEECH_API_MAX_NEW_TOKENS = 4096
+_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024
+_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192
+_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size).
+_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024
+# Pop-only marker for ``test_speech_invalid_field_values``: omit ``ref_audio`` / ``ref_text`` so Base
+# ``speaker_embedding`` cases mirror standalone embedding-only JSON (no ref_audio clash).
+_SPEECH_INVALID_EMBEDDING_ONLY_BODY = "__speech_invalid_embedding_only_body__"
+# Batch-only: merged items inherit batch ``ref_audio`` / ``ref_text`` because ``_pick`` treats item ``None``
+# as "unset". Strip clone fields from the batch envelope before asserting CustomVoice preset errors.
+_SPEECH_BATCH_NO_CLONE_FIELDS = "__speech_batch_no_clone_fields__"
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+_L4_SPEECH_HW = hardware_marks(res={"cuda": "L4"})
+_SPEECH_SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+
+_QWEN3_TTS_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
+            stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="qwen3_tts",
+        marks=_L4_SPEECH_HW,
+    ),
+]
+
+
+def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
+    if loc == "batch":
+        body.update(overrides)
+        return
+    items = body["items"]
+    assert isinstance(items, list) and len(items) >= 1
+    first = items[0]
+    assert isinstance(first, dict)
+    body["items"] = [{**first, **overrides}]
+
+
+def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) -> bytes:
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        nframes = max(1, int(sample_rate * duration_s))
+        wf.writeframes(b"\x00\x00" * nframes)
+    return buf.getvalue()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/speech (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_http_request(
+        {
+            "raw_body": "{",
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("json", "decode", "invalid"),
+        }
+    )
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": 400,
+            "err_message": ("input", "field required", "missing"),
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"input": ""}, ("input", "empty"), id="input_empty"),
+        pytest.param({"input": "   "}, ("input", "empty"), id="input_whitespace_only"),
+        pytest.param({"voice": ""}, "Invalid voice", id="voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"instructions": 123}, ("instructions", "string_type", "valid string"), id="instructions_wrong_type"
+        ),
+        pytest.param(
+            {"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid"
+        ),
+        pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
+        pytest.param(
+            {"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked"
+        ),
+        pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"),
+        pytest.param(
+            {
+                "voice": "vivian",
+                "task_type": "CustomVoice",
+                "ref_audio": None,
+                "ref_text": None,
+            },
+            ("CustomVoice", "does not support"),
+            id="customvoice_task_not_supported",
+        ),
+        pytest.param(
+            {"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
+        ),
+        pytest.param({"language": ""}, ("language", "invalid language", "chinese"), id="language_empty"),
+        pytest.param({"ref_audio": "ftp://example.com/a.wav"}, ("ref_audio", "URL"), id="ref_audio_bad_scheme"),
+        pytest.param(
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="ref_audio_invalid_uri",
+        ),
+        pytest.param({"ref_text": ["x"]}, ("ref_text", "string_type", "valid string"), id="ref_text_wrong_type"),
+        pytest.param(
+            {"x_vector_only_mode": "wrong"},
+            ("x_vector_only_mode", "valid boolean"),
+            id="x_vector_only_mode_string_not_boolean",
+        ),
+        pytest.param(
+            {"speaker_embedding": [0.01] * _QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM},
+            ("mutually exclusive", "speaker_embedding", "ref_audio"),
+            id="speaker_embedding_with_ref_audio_mutually_exclusive",
+        ),
+        pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
+        pytest.param(
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="max_new_tokens_above_max",
+        ),
+        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+        pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
+        pytest.param(
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="initial_codec_chunk_frames_negative",
+        ),
+        pytest.param(
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="instructions_exceed_api_limit",
+        ),
+        # Base clone: whitespace-only ref_text fails serving validation (not pydantic).
+        pytest.param(
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="ref_text_whitespace_only_base_clone",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "speaker_embedding": [],
+            },
+            ("speaker_embedding", "non-empty"),
+            id="base_speaker_embedding_empty",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "x_vector_only_mode": True,
+                "speaker_embedding": [0.01] * 512,
+            },
+            ("speaker_embedding", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
+            id="base_speaker_embedding_wrong_dim_short",
+        ),
+        pytest.param(
+            {
+                _SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
+                "task_type": "Base",
+                "x_vector_only_mode": True,
+                "speaker_embedding": [0.01] * 2048,
+            },
+            ("speaker_embedding", "2048", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
+            id="base_speaker_embedding_wrong_dim_long",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    ov = dict(overrides)
+    embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False))
+    body = {
+        "model": omni_server.model,
+        "input": "Hello.",
+        "voice": "clone",
+        "language": "English",
+        "response_format": "wav",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    if embedding_only:
+        body.pop("ref_audio", None)
+        body.pop("ref_text", None)
+    body.update(ov)
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/speech/batch (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_batch_http_request(
+        {"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")}
+    )
+
+
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_audio_speech_batch_http_request(
+        {
+            "json": {"model": omni_server.model},
+            "timeout": 120,
+            "err_code": (400, 422),
+            "err_message": ("items", "field required", "missing"),
+        }
+    )
+
+
+# Align field-level checks with ``POST /v1/audio/speech`` where ``BatchSpeechRequest`` /
+# ``SpeechBatchItem`` expose the same knobs. Not on batch schema (single-speech only):
+# ``stream``, ``stream_format``, ``seed``, ``speaker_embedding`` — see dedicated speech tests.
+@pytest.mark.parametrize(
+    "loc, overrides, err_message",
+    [
+        pytest.param("item", {"input": ""}, ("input", "empty"), id="item_input_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "item",
+            {"input": "   "},
+            ("input", "empty"),
+            id="item_input_whitespace_only",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param("batch", {"voice": ""}, "Invalid voice", id="batch_voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param("item", {"voice": ""}, "Invalid voice", id="item_voice_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "batch",
+            {
+                "voice": "nonexistent_voice_xyz",
+                "task_type": "CustomVoice",
+                "ref_audio": None,
+                "ref_text": None,
+            },
+            "Invalid voice",
+            id="batch_voice_unknown_preset",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {
+                _SPEECH_BATCH_NO_CLONE_FIELDS: True,
+                "voice": "nonexistent_voice_xyz",
+                "task_type": "CustomVoice",
+            },
+            "Invalid voice",
+            id="item_voice_unknown_preset",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"instructions": 123},
+            ("instructions", "string_type", "valid string"),
+            id="batch_instructions_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"instructions": 123},
+            ("instructions", "string_type", "valid string"),
+            id="item_instructions_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="batch_instructions_exceed_api_limit",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
+            ("instructions", "too long"),
+            id="item_instructions_exceed_api_limit",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "wav"),
+            id="batch_response_format_invalid",
+        ),
+        pytest.param(
+            "item",
+            {"response_format": "mpeg"},
+            ("response_format", "literal_error", "wav"),
+            id="item_response_format_invalid",
+        ),
+        pytest.param("batch", {"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="batch_speed_too_low"),
+        pytest.param("item", {"speed": 10.0}, ("speed", "less_than_equal", "4"), id="item_speed_too_high"),
+        pytest.param(
+            "batch",
+            {"task_type": "InvalidEnum"},
+            ("task_type", "literal_error", "CustomVoice"),
+            id="batch_task_type_invalid",
+        ),
+        pytest.param(
+            "item",
+            {"task_type": "InvalidEnum"},
+            ("task_type", "literal_error", "CustomVoice"),
+            id="item_task_type_invalid",
+        ),
+        pytest.param(
+            "batch",
+            {"language": ""},
+            ("language", "invalid language", "chinese"),
+            id="batch_language_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"language": ""},
+            ("language", "invalid language", "chinese"),
+            id="item_language_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_audio": "ftp://example.com/a.wav"},
+            ("ref_audio", "URL"),
+            id="batch_ref_audio_bad_scheme",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"ref_audio": "ftp://example.com/a.wav"},
+            ("ref_audio", "URL"),
+            id="item_ref_audio_bad_scheme",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="batch_ref_audio_invalid_uri",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"ref_audio": "not_a_valid_uri"},
+            ("ref_audio", "url"),
+            id="item_ref_audio_invalid_uri",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"ref_text": ["x"]},
+            ("ref_text", "string_type", "valid string"),
+            id="batch_ref_text_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"ref_text": ["x"]},
+            ("ref_text", "string_type", "valid string"),
+            id="item_ref_text_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"x_vector_only_mode": "wrong_type"},
+            ("x_vector_only_mode", "bool_parsing", "validation error"),
+            id="batch_x_vector_only_mode_wrong_type",
+        ),
+        pytest.param(
+            "item",
+            {"x_vector_only_mode": "wrong_type"},
+            ("x_vector_only_mode", "bool_parsing", "validation error"),
+            id="item_x_vector_only_mode_wrong_type",
+        ),
+        pytest.param(
+            "batch",
+            {"max_new_tokens": 0},
+            ("max_new_tokens", "least 1"),
+            id="batch_max_new_tokens_below_min",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"max_new_tokens": 0},
+            ("max_new_tokens", "least 1"),
+            id="item_max_new_tokens_below_min",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="batch_max_new_tokens_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
+            ("max_new_tokens", "exceed 4096"),
+            id="item_max_new_tokens_above_max",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "batch",
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="batch_initial_codec_chunk_frames_negative",
+        ),
+        pytest.param(
+            "item",
+            {"initial_codec_chunk_frames": -1},
+            ("initial_codec_chunk_frames", "greater_than_equal", "0"),
+            id="item_initial_codec_chunk_frames_negative",
+        ),
+        # Base clone: whitespace-only ref_text (same serving check as single speech).
+        pytest.param(
+            "batch",
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="batch_ref_text_whitespace_only_base_clone",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "item",
+            {
+                "task_type": "Base",
+                "ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
+                "ref_text": "   ",
+            },
+            ("ref_text", "base task", "non-empty"),
+            id="item_ref_text_whitespace_only_base_clone",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_batch_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    loc: str,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    ov = dict(overrides)
+    strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False))
+    body = {
+        "model": omni_server.model,
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+        "language": "English",
+        "response_format": "wav",
+        "items": [{"input": "Hello."}],
+    }
+    if strip_clone:
+        body.pop("ref_audio", None)
+        body.pop("ref_text", None)
+    _apply_batch_overrides(body, loc=loc, overrides=ov)
+    openai_client.send_audio_speech_batch_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WebSocket /v1/audio/speech/stream
+# ─────────────────────────────────────────────────────────────────────────────
+
+# ``test_speech_stream_invalid_requests``: build ``session.config`` JSON with ``omni_server.model``.
+_SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object()
+
+
+@pytest.mark.parametrize(
+    "send_frames_spec, err_message",
+    [
+        pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
+        pytest.param(
+            json.dumps({"type": "input.text", "text": "hello"}),
+            "Expected session.config",
+            id="wrong_first_message_type",
+        ),
+        pytest.param(
+            _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES,
+            ("invalid session config", "response_format", "pcm"),
+            id="session_config_stream_audio_requires_pcm",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_speech_stream_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    send_frames_spec: Any,
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Bad first WebSocket text frames: non-JSON, wrong ``type``, or invalid ``session.config``."""
+    if send_frames_spec is _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES:
+        send_frames = json.dumps(
+            {
+                "type": "session.config",
+                "model": omni_server.model,
+                "response_format": "wav",
+                "stream_audio": True,
+            }
+        )
+    else:
+        send_frames = send_frames_spec
+    assert isinstance(send_frames, str)
+    openai_client.send_audio_speech_stream_ws_request(
+        {
+            "send_frames": send_frames,
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "err_message": err_message,
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/audio/voices (multipart + embedding-only) and DELETE /v1/audio/voices/{name}
+# ─────────────────────────────────────────────────────────────────────────────
+
+# Templates for ``test_voices_create_invalid_requests``: resolved in ``_finalize_voices_form_data``.
+_VF_LONG_CONSENT = object()
+_VF_LONG_REF_TEXT = object()
+_VF_LONG_SPEAKER_DESC = object()
+_VF_EMBEDDING_NAN_JSON = object()
+
+
+def _finalize_voices_form_data(template: dict[str, Any], uuid_hex: str) -> dict[str, Any]:
+    out: dict[str, Any] = {}
+    for key, val in template.items():
+        if val is _VF_LONG_CONSENT:
+            out[key] = "x" * (_VOICE_UPLOAD_MAX_CONSENT_LEN + 1)
+        elif val is _VF_LONG_REF_TEXT:
+            out[key] = "y" * (_VOICE_UPLOAD_MAX_REF_TEXT_CHARS + 1)
+        elif val is _VF_LONG_SPEAKER_DESC:
+            out[key] = "z" * (_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS + 1)
+        elif val is _VF_EMBEDDING_NAN_JSON:
+            out[key] = json.dumps([0.1] * 1023 + [float("nan")])
+        elif isinstance(val, str):
+            out[key] = val.replace("{uuid}", uuid_hex)
+        else:
+            out[key] = val
+    return out
+
+
+def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
+    """``None`` => no multipart file (embedding-only or neither)."""
+    if kind is None:
+        return None
+    if kind == "wav_ok":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "audio/wav")}
+    if kind == "wav_short":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(duration_s=0.3), "audio/wav")}
+    if kind == "wav_bad_body":
+        return {"audio_sample": ("clip.wav", b"not valid wav content", "audio/wav")}
+    if kind == "wav_pdf_type":
+        return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "application/pdf")}
+    raise AssertionError(f"unknown multipart kind {kind!r}")
+
+
+@pytest.mark.parametrize(
+    "multipart_kind, form_template, err_message",
+    [
+        pytest.param(
+            None,
+            {"consent": "c1", "name": "v1"},
+            ("audio_sample", "speaker_embedding", "provided"),
+            id="neither_audio_nor_embedding",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"name": "v_miss_consent_{uuid}"},
+            ("missing", "consent", "field required"),
+            id="missing_consent",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok"},
+            ("missing", "name", "field required"),
+            id="missing_name",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "", "name": "v_empty_consent_{uuid}"},
+            ("missing", "consent", "field required"),
+            id="empty_consent",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "   ", "name": "v_ws_consent_{uuid}"},
+            "consent",
+            id="whitespace_consent",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "bad/consent", "name": "v_bad_cons_{uuid}"},
+            "consent",
+            id="consent_path_sep",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": _VF_LONG_CONSENT, "name": "v_long_cons_{uuid}"},
+            ("consent", "too long", "Failed to save"),
+            id="consent_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": ""},
+            ("missing", "name", "field required"),
+            id="empty_name",
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": "evil/name"},
+            ("voice name", "invalid voice name"),
+            id="invalid_name_path_sep",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {"consent": "consent_ok", "name": "v_long_ref_{uuid}", "ref_text": _VF_LONG_REF_TEXT},
+            "ref_text",
+            id="ref_text_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {
+                "consent": "consent_ok",
+                "name": "v_long_desc_{uuid}",
+                "speaker_description": _VF_LONG_SPEAKER_DESC,
+            },
+            "speaker_description",
+            id="speaker_description_too_long",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_pdf_type",
+            {"consent": "consent_ok", "name": "v_bad_mime_{uuid}"},
+            ("mime", "unsupported", "audio/mp4"),
+            id="audio_unsupported_mime",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_short",
+            {"consent": "consent_ok", "name": "v_short_{uuid}"},
+            ("too short", "reference audio"),
+            id="audio_too_short",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_bad_body",
+            {"consent": "consent_ok", "name": "v_bad_wav_{uuid}"},
+            ("decode", "audio", "Format not recognised"),
+            id="audio_decode_error",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": "not valid json [[[",
+                "consent": "consent_emb",
+                "name": "v_bad_json_{uuid}",
+            },
+            ("speaker_embedding", "valid JSON"),
+            id="speaker_embedding_invalid_json",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": json.dumps([]),
+                "consent": "consent_emb",
+                "name": "v_empty_emb_{uuid}",
+            },
+            ("speaker_embedding", "non-empty"),
+            id="speaker_embedding_empty",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": _VF_EMBEDDING_NAN_JSON,
+                "consent": "consent_emb",
+                "name": "v_nan_emb_{uuid}",
+            },
+            ("speaker_embedding", "finite"),
+            id="speaker_embedding_nan",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            None,
+            {
+                "speaker_embedding": json.dumps([0.01] * 512),
+                "consent": "consent_emb",
+                "name": "v_wrong_dim_{uuid}",
+            },
+            ("expected", "dimensions", "1024"),
+            id="speaker_embedding_wrong_dim",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "wav_ok",
+            {
+                "consent": "consent_emb",
+                "name": "v_both_{uuid}",
+                "speaker_embedding": json.dumps([0.01] * 1024),
+            },
+            ("mutually exclusive", "audio_sample", "speaker_embedding"),
+            id="audio_and_embedding_mutually_exclusive",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_voices_create_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    multipart_kind: str | None,
+    form_template: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Invalid ``POST /v1/audio/voices`` bodies (multipart or embedding-only)."""
+    uid = uuid.uuid4().hex[:8]
+    cfg: dict[str, Any] = {
+        "data": _finalize_voices_form_data(form_template, uid),
+        "timeout": 120,
+        "err_code": 400,
+        "err_message": err_message,
+    }
+    files = _voices_upload_multipart_files(multipart_kind)
+    if files is not None:
+        cfg["files"] = files
+    openai_client.send_audio_voices_create_http_request(cfg)
+
+
+@pytest.mark.parametrize(
+    "voice_name, err_code, err_message",
+    [
+        pytest.param(
+            "missing-voice-404",
+            404,
+            "not found",
+            id="not_found",
+        ),
+        pytest.param(
+            "   ",
+            404,
+            "not found",
+            id="whitespace_only",
+        ),
+        pytest.param(
+            "evil/name",
+            404,
+            "not found",
+            id="path_separator",
+        ),
+        pytest.param(
+            ".",
+            405,
+            ("Method Not Allowed", "detail"),
+            id="dot_single",
+        ),
+        pytest.param(
+            "..",
+            404,
+            "not found",
+            id="dot_dot",
+        ),
+        pytest.param(
+            "bad\x00voice",
+            404,
+            "not found",
+            id="nul_byte",
+        ),
+        pytest.param(
+            "",
+            405,
+            "Method Not Allowed",
+            id="empty_name_path",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
+def test_voices_delete_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    voice_name: str,
+    err_code: int | tuple[int, ...],
+    err_message: str | tuple[str, ...],
+) -> None:
+    """DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405)."""
+    openai_client.send_audio_voices_delete_http_request(
+        {
+            "name": voice_name,
+            "timeout": 120,
+            "err_code": err_code,
+            "err_message": err_message,
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image-Edit"),
+        id="qwen_image_edit",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+# Populated in ``test_images_edits_invalid_requests`` from ``omni_server.model``.
+_IMAGE_EDITS_SERVER_MODEL = object()
+
+
+def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> dict[str, Any]:
+    return {k: (server_model if v is _IMAGE_EDITS_SERVER_MODEL else v) for k, v in template.items()}
+
+
+@pytest.mark.parametrize(
+    "include_image, data_template, err_message, err_code",
+    [
+        pytest.param(
+            False,
+            {"prompt": "make it brighter", "model": _IMAGE_EDITS_SERVER_MODEL},
+            ("image", "required", "Unprocessable Entity"),
+            422,
+            id="missing_image",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": "wrong-model-id"},
+            ("model", "mismatch"),
+            400,
+            id="model_mismatch",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "response_format": "url"},
+            ("response_format", "b64_json", "supported"),
+            400,
+            id="unsupported_response_format",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "output_compression": "101"},
+            ("less_than_equal", "output_compression", "100"),
+            400,
+            id="output_compression_above_max",
+        ),
+        pytest.param(
+            True,
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
+            ("Invalid layers", "3", "1"),
+            400,
+            id="invalid_layers",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_images_edits_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    tiny_png_bytes: bytes,
+    include_image: bool,
+    data_template: dict[str, Any],
+    err_message: str | tuple[str, ...],
+    err_code: int | tuple[int, ...],
+) -> None:
+    cfg: dict[str, Any] = {
+        "data": _finalize_edits_form_data(data_template, omni_server.model),
+        "timeout": 300,
+        "err_code": err_code,
+        "err_message": err_message,
+    }
+    if include_image:
+        cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")}
+    openai_client.send_images_edits_http_request(cfg)

tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_PARAMS = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image"),
+        id="qwen_image",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+# Full replacement body for ``test_images_generations_invalid_requests`` (no minimal prompt/model merge).
+_IMAGES_GEN_MISSING_PROMPT_JSON = object()
+
+
+def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
+    return {"prompt": "a simple red apple icon", "model": omni_server.model}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/images/generations (JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "body_spec, err_message",
+    [
+        pytest.param(_IMAGES_GEN_MISSING_PROMPT_JSON, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param(
+            {"prompt": "cat", "model": "not-the-loaded-checkpoint"},
+            ("model", "mismatch"),
+            id="model_mismatch",
+        ),
+        pytest.param({"prompt": ""}, ("prompt", "empty"), id="prompt_empty", marks=_SKIP_ISSUE_3649),
+        pytest.param({"prompt": "   "}, ("prompt", "empty"), id="prompt_whitespace_only", marks=_SKIP_ISSUE_3649),
+        pytest.param({"prompt": 123}, ("prompt", "string_type", "valid string"), id="prompt_wrong_type"),
+        pytest.param({"n": 0}, ("n", "greater_than_equal", "1"), id="n_below_min"),
+        pytest.param({"n": 11}, ("n", "less_than_equal", "10"), id="n_above_max"),
+        pytest.param({"size": "1024"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_height_token"),
+        pytest.param({"size": "not-a-size"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_separator"),
+        pytest.param({"size": "1024xabc"}, ("Invalid size format", "integers"), id="size_non_integer_height"),
+        pytest.param({"size": "1024x"}, ("Invalid size format", "integers"), id="size_incomplete_height"),
+        pytest.param({"size": "0x1024"}, ("Invalid size", "positive integers"), id="size_non_positive_width"),
+        pytest.param(
+            {"response_format": "url"},
+            ("response_format", "value_error", "b64_json"),
+            id="response_format_not_b64_json",
+        ),
+        pytest.param(
+            {"use_system_prompt": "bogus"},
+            ("use_system_prompt", "value_error", "dynamic"),
+            id="use_system_prompt_invalid",
+        ),
+        pytest.param(
+            {"num_inference_steps": 0},
+            ("num_inference_steps", "'greater_than_equal", "1"),
+            id="num_inference_steps_below_min",
+        ),
+        pytest.param(
+            {"num_inference_steps": 201},
+            ("num_inference_steps", "less_than_equal", "200"),
+            id="num_inference_steps_above_max",
+        ),
+        pytest.param(
+            {"guidance_scale": -1.0}, ("guidance_scale", "greater_than_equal", "0"), id="guidance_scale_negative"
+        ),
+        pytest.param(
+            {"guidance_scale": 21.0}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
+        ),
+        pytest.param(
+            {"true_cfg_scale": -1.0}, ("true_cfg_scale", "greater_than_equal", "0"), id="true_cfg_scale_negative"
+        ),
+        pytest.param(
+            {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
+        ),
+        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
+        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            {"seed": 2**32},
+            ("seed", "less_than_equal", "4294967295"),
+            id="seed_above_uint32",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"output_format": "gif"},
+            ("output_format", "value_error", "b64_json"),
+            id="output_format_invalid",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            {"vae_use_slicing": "wrong_type"},
+            ("vae_use_slicing", "bool_parsing", "validation error"),
+            id="vae_use_slicing_wrong_type",
+        ),
+        pytest.param({"user": 123}, ("user", "string_type", "valid string"), id="user_wrong_type"),
+        pytest.param(
+            {"negative_prompt": ["noise"]},
+            ("negative_prompt", "string_type", "valid string"),
+            id="negative_prompt_wrong_type",
+        ),
+        pytest.param(
+            {"generator_device": 1},
+            ("generator_device", "string_type", "valid string"),
+            id="generator_device_wrong_type",
+        ),
+        pytest.param(
+            {"lora": {"foo": "bar"}}, ("lora", "both name and path", "required"), id="lora_missing_required_fields"
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_images_generations_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    body_spec: Any,
+    err_message: str | tuple[str, ...],
+) -> None:
+    """Invalid JSON bodies for ``POST /v1/images/generations`` (missing prompt, model mismatch, bad fields)."""
+    if body_spec is _IMAGES_GEN_MISSING_PROMPT_JSON:
+        body: dict[str, object] = {"size": "1024x1024"}
+    else:
+        assert isinstance(body_spec, dict)
+        body = _minimal_images_gen_json(omni_server)
+        body.update(body_spec)
+    openai_client.send_images_generations_http_request(
+        {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.omni]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+
+def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]:
+    """Minimal valid chat body; individual tests override one offending field."""
+    return {
+        "model": omni_server.model,
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream": False,
+    }
+
+
+_QWEN3_OMNI_SERVER = [
+    pytest.param(
+        OmniServerParams(
+            model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            stage_config_path=get_deploy_config_path("qwen3_omni_moe.yaml"),
+            server_args=["--async-chunk"],
+        ),
+        id="qwen3_omni",
+    ),
+]
+
+
+def _chat_completions_request_without_expectations(omni_server: OmniServer, case_id: str) -> dict[str, Any]:
+    """Build ``send_chat_completions_http_request`` kwargs excluding ``err_code`` / ``err_message``."""
+    if case_id == "malformed_json":
+        return {"raw_body": "{not-json", "timeout": 120}
+    if case_id == "missing_messages":
+        return {"json": {"model": omni_server.model, "stream": False}, "timeout": 120}
+    if case_id == "model_mismatch":
+        return {
+            "json": {
+                "model": "this-model-is-not-served-on-this-instance",
+                "messages": [{"role": "user", "content": "hello"}],
+                "stream": False,
+            },
+            "timeout": 120,
+        }
+    body = _minimal_chat_json(omni_server)
+    if case_id == "stream_wrong_type":
+        body["stream"] = "not-a-boolean"  # type: ignore[assignment]
+    elif case_id == "temperature_wrong_type":
+        body["temperature"] = "hot"  # type: ignore[assignment]
+    elif case_id == "max_tokens_wrong_type":
+        body["max_tokens"] = "not-an-int"  # type: ignore[assignment] # type: ignore[assignment]
+    elif case_id == "modalities_list_bad":
+        body["modalities"] = [123]  # type: ignore[assignment]
+    elif case_id == "response_format_json_schema_incomplete":
+        body["response_format"] = {"type": "json_schema"}
+    elif case_id == "logprobs_wrong_type":
+        body["logprobs"] = "yes"  # type: ignore[assignment]
+    elif case_id == "logprobs_top_without_enabled":
+        body["logprobs"] = False
+        body["top_logprobs"] = 5
+    elif case_id == "speaker_unknown":
+        body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
+    else:
+        raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
+    return {"json": body, "timeout": 120}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/chat/completions
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize(
+    "case_id, err_code, err_message",
+    [
+        pytest.param(
+            "malformed_json",
+            400,
+            ("json_invalid", "decode error"),
+            id="malformed_json",
+        ),
+        pytest.param("missing_messages", 400, ("messages", "missing", "Field required"), id="missing_messages"),
+        pytest.param(
+            "model_mismatch",
+            404,
+            ("model", "does not exist", "NotFoundError"),
+            id="model_mismatch",
+        ),
+        pytest.param("stream_wrong_type", 400, ("stream", "bool_parsing", "not-a-boolean"), id="invalid_stream_type"),
+        pytest.param(
+            "temperature_wrong_type",
+            400,
+            ("temperature", "float_parsing", "a valid number"),
+            id="invalid_temperature_type",
+        ),
+        pytest.param(
+            "max_tokens_wrong_type", 400, ("max_tokens", "int_parsing", "integer"), id="invalid_max_tokens_type"
+        ),
+        pytest.param(
+            "modalities_list_bad",
+            400,
+            ("modalities", "value_error", ""),
+            id="modalities_list_bad_element",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "response_format_json_schema_incomplete",
+            400,
+            ("response_format", "value_error", "json_schema"),
+            id="invalid_response_format_json_schema",
+        ),
+        pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
+        pytest.param(
+            "logprobs_top_without_enabled",
+            400,
+            ("top_logprobs", "logprobs", "True"),
+            id="logprobs_top_without_logprobs_true",
+        ),
+        pytest.param(
+            "speaker_unknown",
+            400,
+            ("Invalid speaker", "Supported"),
+            id="speaker_unknown_preset",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_chat_completions_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    case_id: str,
+    err_code: int | tuple[int, ...],
+    err_message: str | tuple[str, ...],
+) -> None:
+    cfg = _chat_completions_request_without_expectations(omni_server, case_id)
+    cfg["err_code"] = err_code
+    cfg["err_message"] = err_message
+    openai_client.send_chat_completions_http_request(cfg)[0]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WS /v1/video/chat/stream
+# ─────────────────────────────────────────────────────────────────────────────
+
+# First-frame JSON with ``omni_server.model`` for ``test_video_chat_stream_invalid_requests``.
+_VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object()
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize(
+    "send_frames_spec, err_message",
+    [
+        pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
+        pytest.param(
+            _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR,
+            "Invalid session config",
+            id="session_modalities_scalar",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_video_chat_stream_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    send_frames_spec: Any,
+    err_message: str,
+) -> None:
+    if send_frames_spec is _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR:
+        send_frames = json.dumps({"type": "session.config", "model": omni_server.model, "modalities": "text"})
+    else:
+        send_frames = send_frames_spec
+    assert isinstance(send_frames, str)
+    openai_client.send_video_chat_stream_ws_request(
+        {
+            "send_frames": send_frames,
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "err_message": err_message,
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WS /v1/realtime
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
+def test_realtime_rejected_when_async_chunk_enabled(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+) -> None:
+    openai_client.send_realtime_ws_request(
+        {
+            "timeout": 120,
+            "ws_max_size": None,
+            "ws_json_type": "error",
+            "ws_error_code": "unsupported",
+            "err_message": "async_chunk",
+        }
+    )

tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
+
+_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
+
+_QWEN_IMAGE = [
+    pytest.param(
+        OmniServerParams(model="Qwen/Qwen-Image"),
+        id="qwen_image",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST omni sleep / omni wakeup (invalid JSON)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "endpoint, json_body, err_message",
+    [
+        pytest.param("sleep", {"level": 2}, ("stage_ids", "Field required", "Missing"), id="sleep_missing_stage_ids"),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [], "level": 2},
+            ("stage_ids", "Field required", "Missing"),
+            id="sleep_empty_stage_ids",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": "not-a-list", "level": 2},
+            ("stage_ids", "list_type", "a valid list"),
+            id="sleep_invalid_stage_ids_type",
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [0], "level": "high"},
+            ("level", "int_parsing", "a valid integer"),
+            id="sleep_invalid_level_type",
+        ),
+        pytest.param(
+            "sleep",
+            {"stage_ids": [0], "level": -1},
+            ("level", "greater"),
+            id="sleep_negative_level",
+            marks=_SKIP_ISSUE_3649,
+        ),
+        pytest.param("wakeup", {}, ("stage_ids", "missing", "Field required"), id="wakeup_missing_stage_ids"),
+        pytest.param(
+            "wakeup",
+            {"stage_ids": None},
+            ("stage_ids", "list_type", "a valid list"),
+            id="wakeup_invalid_stage_ids_type",
+        ),
+        pytest.param(
+            "wakeup",
+            {"stage_ids": []},
+            ("stage_ids", "length"),
+            id="wakeup_empty_stage_ids",
+            marks=_SKIP_ISSUE_3649,
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
+def test_omni_sleep_wakeup_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    endpoint: Literal["sleep", "wakeup"],
+    json_body: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    cfg: dict[str, Any] = {
+        "json": json_body,
+        "timeout": 120,
+        "err_code": 400,
+        "err_message": err_message,
+    }
+    if endpoint == "sleep":
+        openai_client.send_omni_sleep_http_request(cfg)
+    else:
+        openai_client.send_omni_wakeup_http_request(cfg)

tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``).
+
+WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in
+``test_invalid_omni_chat.py`` with the same Qwen3-Omni server module fixture (Wan cannot serve those paths).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
+
+_WAN_T2V = [
+    pytest.param(
+        OmniServerParams(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
+        id="wan22_t2v",
+        marks=hardware_marks(res={"cuda": "H100"}),
+    ),
+]
+
+
+def _minimal_video_form_data() -> dict[str, str]:
+    return {"prompt": "a cat walking on grass"}
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/videos (async)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "form_extra, err_message",
+    [
+        pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
+        pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
+        pytest.param(
+            {"num_inference_steps": "0"},
+            ("num_inference_steps", "greater_than_equal", "1"),
+            id="num_inference_steps_zero",
+        ),
+        pytest.param(
+            {"guidance_scale": "25"}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_post_videos_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    form_extra: dict[str, str] | None,
+    err_message: str | tuple[str, ...],
+) -> None:
+    data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
+    openai_client.send_videos_create_http_request(
+        {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# POST /v1/videos/sync
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "form_extra, err_message",
+    [
+        pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
+        pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
+        pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
+        pytest.param(
+            {"num_inference_steps": "201"},
+            ("num_inference_steps", "less_than_equal", "200"),
+            id="num_inference_steps_above_max",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_post_videos_sync_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    form_extra: dict[str, str] | None,
+    err_message: str | tuple[str, ...],
+) -> None:
+    data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
+    openai_client.send_videos_sync_http_request(
+        {"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos (list)
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "params, err_message",
+    [
+        pytest.param({"limit": -1}, ("limit", "greater_than_equal", "0"), id="invalid_limit_negative"),
+        pytest.param({"limit": 101}, ("limit", "less_than_equal", "100"), id="limit_above_max"),
+        pytest.param({"order": "newest"}, ("order", "literal_error", "asc", "desc"), id="invalid_order"),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_videos_list_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    params: dict[str, Any],
+    err_message: str | tuple[str, ...],
+) -> None:
+    openai_client.send_videos_list_http_request(
+        {"params": params, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos/{video_id}
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_retrieve_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# DELETE /v1/videos/{video_id}
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_delete_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# GET /v1/videos/{video_id}/content
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
+def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    openai_client.send_video_content_http_request(
+        {
+            "video_id": "does-not-exist",
+            "timeout": 120,
+            "err_code": 404,
+            "err_message": ("not found", "video"),
+        }
+    )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -0,0 +1,802 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import socket
+
+import pytest
+import torch
+from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
+from vllm.sequence import IntermediateTensors
+from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
+
+import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_classifier_free_guidance_rank,
+    get_pp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.parallel]
+
+
+def _find_free_port() -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return str(s.getsockname()[1])
+
+
+def update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for k, v in envs_dict.items():
+        os.environ[k] = v
+
+
+# ---------------------------------------------------------------------------
+# Shared stubs used by both unit and distributed tests
+# ---------------------------------------------------------------------------
+
+
+class FakeWork:
+    """Drop-in for torch.distributed.Work that records whether wait() was called."""
+
+    def __init__(self):
+        self.waited = False
+
+    def wait(self):
+        self.waited = True
+
+
+class SimpleScheduler:
+    """Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
+
+    def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
+        return (latents - 0.1 * noise_pred,)
+
+
+class FakeVAE:
+    def __init__(self, distributed_enabled: bool = False):
+        self.calls = 0
+        self.distributed_enabled = distributed_enabled
+
+    def decode(self, z: torch.Tensor):
+        """Original decode docstring."""
+        self.calls += 1
+        return (z + 1,)
+
+    def is_distributed_enabled(self) -> bool:
+        return self.distributed_enabled
+
+
+class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
+    """Minimal pipeline used to exercise PipelineParallelMixin.
+
+    Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
+    by real DiT models — so the PP layer-split logic is exercised faithfully.
+
+    Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
+    is initialized identically on every rank regardless of which ranks are active,
+    allowing the distributed output to be compared against the single-GPU baseline.
+
+    Args:
+        num_layers: Total number of Linear layers.
+        dim:        Input / hidden dimension.
+        seed:       Base RNG seed; layer ``i`` uses ``seed + i``.
+        device:     Target device for layer weights (default: CPU).
+        dtype:      Target dtype for layer weights (default: float32).
+    """
+
+    def __init__(
+        self,
+        num_layers: int = 4,
+        dim: int = 64,
+        seed: int = 42,
+        device: torch.device | None = None,
+        dtype: torch.dtype = torch.float32,
+    ):
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            num_layers,
+            lambda prefix: torch.nn.Linear(dim, dim, bias=False),
+            prefix="layers",
+        )
+
+        for i, layer in enumerate(self.layers):
+            if not isinstance(layer, PPMissingLayer):
+                torch.manual_seed(seed + i)
+                torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
+
+        self.layers.to(device=device, dtype=dtype)
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
+        self.scheduler = SimpleScheduler()
+
+    def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
+        """Layer-split forward pass.
+
+        * First PP rank: uses ``x`` from caller kwargs.
+        * Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
+          (which transparently waits for the async receive).
+        * Non-last PP ranks return ``IntermediateTensors``; the last rank
+          returns the plain noise-prediction tensor.
+        """
+        if intermediate_tensors is not None:
+            x = intermediate_tensors["hidden_states"]
+
+        for i in range(self.start_layer, self.end_layer):
+            x = self.layers[i](x)
+
+        pp_group = get_pp_group()
+        if not pp_group.is_last_rank:
+            return IntermediateTensors({"hidden_states": x})
+        return x
+
+
+# ---------------------------------------------------------------------------
+# 1.  AsyncLatents – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncLatents:
+    """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
+        return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
+
+    def test_resolve_returns_wrapped_tensor(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al._resolve() is t
+
+    def test_attribute_access_resolves(self):
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        assert al.shape == t.shape
+        assert al.dtype == t.dtype
+
+    def test_torch_function_protocol(self):
+        """torch ops that receive an AsyncLatents should see the underlying tensor."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        mask = torch.ones_like(t)
+        result = mask * al  # triggers __torch_function__
+        torch.testing.assert_close(result, mask * t)
+
+    def test_torch_function_with_list_arg(self):
+        """__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.cat([al, al], dim=0)
+        torch.testing.assert_close(result, torch.cat([t, t], dim=0))
+
+    def test_torch_tensor_conversion(self):
+        """torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
+        t = torch.randn(2, 4)
+        al = self._make(t)
+        result = torch.as_tensor(al)
+        assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
+
+    def test_handles_are_waited_before_resolve(self):
+        t = torch.randn(2, 4)
+        h1, h2 = FakeWork(), FakeWork()
+        al = self._make(t, handles=[h1, h2])
+        _ = al.shape  # trigger resolution
+        assert h1.waited and h2.waited, "Not all handles were waited on"
+
+    def test_postproc_callbacks_invoked_on_resolve(self):
+        t = torch.randn(2, 4)
+        log: list[int] = []
+        al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
+        _ = al.shape
+        assert log == [1, 2], f"postproc not called in order: {log}"
+
+    def test_idempotent_resolve(self):
+        """handle.wait() must not be called twice if _resolve() is called twice."""
+        t = torch.randn(2, 4)
+        h = FakeWork()
+        al = self._make(t, handles=[h])
+        _ = al.shape  # first resolve
+        h.waited = False  # reset sentinel
+        _ = al.dtype  # second resolve
+        assert not h.waited, "handle.wait() was called a second time"
+
+
+# ---------------------------------------------------------------------------
+# 2.  _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
+# ---------------------------------------------------------------------------
+
+
+class TestSyncPPSend:
+    """Verifies PipelineParallelMixin's internal PP-send flush."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline() -> PipelineParallelMixin:
+        # Instantiate a bare mixin — no layers, no distributed env needed.
+        # _sync_pp_send only touches _pp_send_work, so this is sufficient.
+        class _BarePP(PipelineParallelMixin, CFGParallelMixin):
+            pass
+
+        return _BarePP()
+
+    def test_noop_when_work_list_empty(self):
+        pipeline = self._make_pipeline()
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+    def test_waits_all_pending_handles(self):
+        pipeline = self._make_pipeline()
+        works = [FakeWork(), FakeWork(), FakeWork()]
+        pipeline._pp_send_work = works
+        pipeline._sync_pp_send()
+        assert all(w.waited for w in works), "Some handles were not waited on"
+
+    def test_clears_work_list_after_sync(self):
+        pipeline = self._make_pipeline()
+        pipeline._pp_send_work = [FakeWork()]
+        pipeline._sync_pp_send()
+        assert pipeline._pp_send_work == []
+
+
+class TestDiffuseWrapper:
+    """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
+
+    pytestmark = [pytest.mark.cpu]
+
+    def test_diffuse_flushes_pending_sends_on_success(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                return "done"
+
+        pipeline = _DiffusePP()
+
+        assert pipeline.diffuse() == "done"
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_flushes_pending_sends_on_exception(self):
+        work = FakeWork()
+
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                self._pp_send_work = [work]
+                raise RuntimeError("boom")
+
+        pipeline = _DiffusePP()
+
+        with pytest.raises(RuntimeError, match="boom"):
+            pipeline.diffuse()
+        assert work.waited
+        assert pipeline._pp_send_work == []
+
+    def test_diffuse_wrapper_preserves_metadata(self):
+        class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
+            def diffuse(self):
+                """Original diffuse docstring."""
+                return "done"
+
+        assert _DiffusePP.diffuse.__name__ == "diffuse"
+        assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
+
+
+class TestVaeDecodeGuard:
+    pytestmark = [pytest.mark.cpu]
+
+    @staticmethod
+    def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
+        class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
+            def __init__(self):
+                self.vae = FakeVAE(distributed_enabled=distributed_enabled)
+
+        return _DecodePP()
+
+    @staticmethod
+    def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
+        monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
+        monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
+
+    def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=1, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_calls_original_decode_on_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=True)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_skips_decode_on_non_first_stage(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline()
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)
+
+        assert pipeline.vae.calls == 0
+        assert output == (None,)
+
+    def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
+        self._set_rank(monkeypatch, world_size=2, first_stage=False)
+        pipeline = self._make_pipeline(distributed_enabled=True)
+        z = torch.ones(2, 3)
+
+        output = pipeline.vae.decode(z)[0]
+
+        assert pipeline.vae.calls == 1
+        torch.testing.assert_close(output, z + 1)
+
+    def test_decode_wrapper_preserves_metadata(self):
+        pipeline = self._make_pipeline()
+
+        assert pipeline.vae.decode.__name__ == "decode"
+        assert pipeline.vae.decode.__doc__ == "Original decode docstring."
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_cfg_mixin():
+    with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
+
+        class _MissingCFG(PipelineParallelMixin):
+            pass
+
+
+@pytest.mark.cpu
+def test_pipeline_parallel_requires_mro_before_cfg_mixin():
+    with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
+
+        class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Distributed test helpers
+# ---------------------------------------------------------------------------
+
+
+def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+    """Initialise the distributed environment for a spawned worker."""
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    init_distributed_environment()
+    return device
+
+
+def make_pipeline_and_inputs(
+    test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
+) -> tuple["MockPipelineParallel", dict, dict | None]:
+    """Create a MockPipelineParallel and seeded inputs from a test_config dict.
+
+    Must be called after ``initialize_model_parallel`` so that ``make_layers``
+    can read the PP group to determine this rank's layer slice.
+
+    Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
+    ``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
+    """
+    pipeline = MockPipelineParallel(
+        num_layers=test_config["num_layers"],
+        dim=test_config["dim"],
+        seed=test_config["model_seed"],
+        device=device,
+        dtype=dtype,
+    )
+
+    torch.manual_seed(test_config["input_seed"])
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(test_config["input_seed"])
+    pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
+
+    negative_kwargs = None
+    if do_true_cfg:
+        torch.manual_seed(test_config["input_seed"] + 1)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
+        negative_kwargs = {"x": neg_x}
+
+    return pipeline, pos_x, negative_kwargs
+
+
+# ---------------------------------------------------------------------------
+# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# ---------------------------------------------------------------------------
+
+
+def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=world_size)
+    pp_group = get_pp_group()
+
+    if pp_group.is_first_rank:
+        torch.manual_seed(77)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(77)
+        tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
+        handles = pp_group.isend_tensor_dict({"t": tensor})
+        for h in handles:
+            h.wait()
+        result_queue.put(("sent", tensor.cpu()))
+    else:
+        received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
+        result_queue.put(("received", received["t"].cpu()))
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
+
+    results = {label: tensor for label, tensor in [q.get(), q.get()]}
+    torch.testing.assert_close(
+        results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 4.  predict_noise_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+_baseline_cache: dict[tuple, torch.Tensor] = {}
+
+
+def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
+    """Compute expected single-GPU output using the same MockPipelineParallel.
+
+    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
+    Results are cached so identical configs are only computed once.
+    """
+    key = (
+        test_config["num_layers"],
+        test_config["dim"],
+        test_config["batch_size"],
+        test_config["model_seed"],
+        test_config["input_seed"],
+        test_config["cfg_scale"],
+        dtype,
+        do_true_cfg,
+    )
+    if key in _baseline_cache:
+        return _baseline_cache[key]
+
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    destroy_distributed_env()
+
+    _baseline_cache[key] = noise_pred.cpu()
+    return _baseline_cache[key]
+
+
+def predict_noise_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    test_config: dict,
+    result_queue,
+):
+    """Generic predict-noise worker parameterized by PP and CFG topology."""
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, dtype, device, do_true_cfg=do_true_cfg
+    )
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the non-last PP rank's async send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_last_rank:
+        assert noise_pred is not None
+        if cfg_rank == 0:
+            result_queue.put(noise_pred.cpu())
+    else:
+        assert noise_pred is None
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.float32,
+            4,
+            100,
+            1e-5,
+            1e-5,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-float32",
+        ),
+        pytest.param(
+            2,
+            1,
+            False,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            1,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-seq_cfg-bfloat16",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            torch.bfloat16,
+            4,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-bfloat16",
+        ),
+        pytest.param(
+            3,
+            1,
+            False,
+            torch.bfloat16,
+            6,
+            100,
+            1e-2,
+            1e-2,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
+            id="pp3-no_cfg-bfloat16",
+        ),
+    ],
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": num_layers,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    pp_q = manager.Queue()
+
+    world_size = pp_size * cfg_size
+    port = _find_free_port()
+    torch.multiprocessing.spawn(
+        predict_noise_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        nprocs=world_size,
+    )
+
+    pp_out = pp_q.get()
+
+    assert baseline_out.shape == pp_out.shape
+    torch.testing.assert_close(
+        pp_out,
+        baseline_out,
+        rtol=rtol,
+        atol=atol,
+        msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
+    )
+
+
+# ---------------------------------------------------------------------------
+# 5.  scheduler_step_maybe_with_cfg
+# ---------------------------------------------------------------------------
+
+
+def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
+    """Single-GPU reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        result = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+
+    destroy_distributed_env()
+    return result.cpu()
+
+
+def scheduler_step_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    test_config: dict,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+
+    pp_group = get_pp_group()
+    cfg_rank = get_classifier_free_guidance_rank()
+
+    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+        test_config, torch.float32, device, do_true_cfg=do_true_cfg
+    )
+    latents = positive_kwargs["x"]
+    t = torch.tensor(500, device=device)
+
+    with torch.inference_mode():
+        noise_pred = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=test_config["cfg_scale"],
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        latents = pipeline.scheduler_step_maybe_with_cfg(
+            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        )
+    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
+    # Flush the last PP rank's async latent send before barrier / process teardown.
+    pipeline._sync_pp_send()
+
+    if pp_group.is_first_rank and cfg_rank == 0:
+        resolved = latents.contiguous()
+        result_queue.put(resolved.cpu())
+
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    destroy_distributed_env()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(
+            2,
+            1,
+            False,
+            300,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
+            id="pp2-no_cfg",
+        ),
+        pytest.param(
+            2,
+            2,
+            True,
+            600,
+            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
+            id="pp2-cfg2-true_cfg",
+        ),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+    test_config = {
+        "num_layers": 4,
+        "dim": 64,
+        "batch_size": 2,
+        "cfg_scale": 7.5,
+        "model_seed": 42,
+        "input_seed": input_seed,
+    }
+
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+
+    mp_context = torch.multiprocessing.get_context("spawn")
+    manager = mp_context.Manager()
+    q = manager.Queue()
+
+    port = _find_free_port()
+    world_size = pp_size * cfg_size
+    torch.multiprocessing.spawn(
+        scheduler_step_worker,
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        nprocs=world_size,
+    )
+
+    result = q.get()
+    torch.testing.assert_close(
+        result,
+        baseline,
+        rtol=0,
+        atol=0,
+        msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
+    )

tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
+    Flux2KleinPipeline,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _StopAfterCheckInputsError(Exception):
+    pass
+
+
+def _make_pipeline():
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.vae_scale_factor = 8
+    pipeline.is_distilled = True
+    pipeline._guidance_scale = 0.0
+    return pipeline
+
+
+def _make_minimal_request(
+    prompt="valid prompt",
+    *,
+    num_inference_steps=None,
+):
+    """Build a minimal OmniDiffusionRequest-like object that forward() reads from."""
+    params = OmniDiffusionSamplingParams(
+        height=512,
+        width=512,
+        num_inference_steps=num_inference_steps,
+        seed=42,
+    )
+    req = MagicMock()
+    req.sampling_params = params
+    req.prompts = [prompt]
+    req.multi_modal_data = {}
+    return req
+
+
+def _capture_num_inference_steps(pipe, steps):
+    """Run pipe.forward() and capture the num_inference_steps passed to check_inputs."""
+    captured = {}
+
+    def fake_check_inputs(**kwargs):
+        captured["num_inference_steps"] = kwargs.get("num_inference_steps")
+        raise _StopAfterCheckInputsError
+
+    original = pipe.check_inputs
+    pipe.check_inputs = fake_check_inputs
+    try:
+        req = _make_minimal_request(num_inference_steps=steps)
+        pipe.forward(req)
+    except _StopAfterCheckInputsError:
+        pass
+    finally:
+        pipe.check_inputs = original
+    return captured.get("num_inference_steps")
+
+
+# --- Forward resolution tests (the actual regression) ---
+
+
+def test_forward_preserves_zero_not_default():
+    """#3703: num_inference_steps=0 must reach validation as 0, not as default 50."""
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, 0)
+    assert captured == 0, f"Expected 0, got {captured}"
+
+
+def test_forward_preserves_negative():
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, -1)
+    assert captured == -1
+
+
+def test_forward_none_uses_default():
+    """None means unset — forward() should substitute the pipeline default."""
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, None)
+    assert captured == 50
+
+
+def test_forward_preserves_positive():
+    pipe = _make_pipeline()
+    captured = _capture_num_inference_steps(pipe, 9)
+    assert captured == 9
+
+
+# --- check_inputs validation tests ---
+
+
+@pytest.mark.parametrize("steps", [0, -1])
+def test_check_inputs_rejects_non_positive(steps):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError):
+        pipe.check_inputs(
+            prompt="valid prompt",
+            height=512,
+            width=512,
+            num_inference_steps=steps,
+        )
+
+
+def test_check_inputs_accepts_none():
+    pipe = _make_pipeline()
+    pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=None)
+
+
+@pytest.mark.parametrize("steps", [1, 2])
+def test_check_inputs_accepts_positive(steps):
+    pipe = _make_pipeline()
+    pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=steps)

tests/diffusion/models/flux2/test_flux2_klein_validation.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
+    Flux2KleinPipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_pipeline():
+    pipeline = object.__new__(Flux2KleinPipeline)
+    pipeline.vae_scale_factor = 8
+    pipeline.is_distilled = True
+    pipeline._guidance_scale = 0.0
+    return pipeline
+
+
+def _check(pipeline, prompt, prompt_embeds=None):
+    pipeline.check_inputs(
+        prompt=prompt,
+        height=512,
+        width=512,
+        prompt_embeds=prompt_embeds,
+        num_inference_steps=4,
+        guidance_scale=0.0,
+    )
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "",
+        "   ",
+    ],
+)
+def test_rejects_empty_or_whitespace_prompt(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        ["valid prompt", ""],
+        ["valid prompt", "   "],
+        ["   "],
+    ],
+)
+def test_rejects_list_with_empty_or_whitespace_element(prompt):
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError, match="empty or whitespace-only"):
+        _check(pipe, prompt)
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "valid prompt",
+        ["valid prompt", "another valid prompt"],
+    ],
+)
+def test_accepts_valid_prompts(prompt):
+    pipe = _make_pipeline()
+    _check(pipe, prompt)
+
+
+def test_rejects_none_without_embeds():
+    pipe = _make_pipeline()
+    with pytest.raises(ValueError):
+        _check(pipe, None)
+
+
+def test_accepts_none_with_prompt_embeds():
+    pipe = _make_pipeline()
+    _check(pipe, None, prompt_embeds=torch.randn(1, 256, 4096))

tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py

@@ -1,97 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from PIL import Image
-from torch import nn
-
-from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
-    Wan22TI2VPipeline,
-    get_wan22_ti2v_pre_process_func,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
-    pipeline = object.__new__(Wan22TI2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
-    pipeline.vae = StubVAE(z_dim=4)
-    pipeline.vae_scale_factor_temporal = 4
-    pipeline.vae_scale_factor_spatial = 8
-    pipeline.progress_bar = noop_progress_bar
-    return pipeline
-
-
-def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
-    preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
-    request = SimpleNamespace(
-        prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
-        sampling_params=SimpleNamespace(height=None, width=None),
-    )
-
-    result = preprocess(request)
-
-    assert result.sampling_params.height == 672
-    assert result.sampling_params.width == 1344
-    assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
-
-
-def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents = torch.zeros(1, 4, 2, 4, 4)
-    calls = []
-
-    def fake_predict_noise_maybe_with_cfg(**kwargs):
-        calls.append(kwargs)
-        return torch.ones_like(latents)
-
-    pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg  # type: ignore[method-assign]
-    pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise  # type: ignore[method-assign]
-
-    result = pipeline.diffuse(
-        latents=latents,
-        timesteps=torch.tensor([7]),
-        prompt_embeds=torch.zeros(1, 2, 3),
-        negative_prompt_embeds=torch.zeros(1, 2, 3),
-        guidance_scale=3.0,
-        dtype=torch.float32,
-        attention_kwargs={"a": "b"},
-        num_latent_frames=2,
-        latent_height=4,
-        latent_width=4,
-    )
-
-    positive = calls[0]["positive_kwargs"]
-    assert calls[0]["do_true_cfg"] is True
-    assert positive["timestep"].shape == (1, 8)
-    torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
-    torch.testing.assert_close(positive["hidden_states"], latents)
-    torch.testing.assert_close(result, torch.ones_like(latents))
-
-
-def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
-    pipeline = _make_ti2v_pipeline()
-    latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
-        image=torch.zeros(1, 3, 16, 16),
-        batch_size=1,
-        num_channels_latents=4,
-        height=16,
-        width=16,
-        num_frames=5,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        generator=None,
-        latents=torch.zeros(1, 4, 2, 2, 2),
-    )
-
-    torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
-    assert latent_condition.shape == (1, 4, 1, 2, 2)
-    assert first_frame_mask[:, :, 0].sum() == 0
-    assert first_frame_mask[:, :, 1].sum() == 4

tests/diffusion/quantization/test_quantization_quality.py

@@ -32,17 +32,31 @@ Optional artifact dump:
 from __future__ import annotations
 
 import gc
+import importlib.util
 import json
 import os
+import sys
 from dataclasses import dataclass
 from pathlib import Path
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
 import torch
+from PIL import Image
 
 from tests.helpers.mark import hardware_marks
 
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_BENCH_MODULE_PATH = _REPO_ROOT / "benchmarks" / "diffusion" / "quantization_quality.py"
+_BENCH_MODULE_NAME = "benchmarks.diffusion.quantization_quality"
+
+if _BENCH_MODULE_NAME not in sys.modules:
+    _spec = importlib.util.spec_from_file_location(_BENCH_MODULE_NAME, _BENCH_MODULE_PATH)
+    _mod = importlib.util.module_from_spec(_spec)
+    sys.modules[_BENCH_MODULE_NAME] = _mod
+    _spec.loader.exec_module(_mod)
+
 # ---------------------------------------------------------------------------
 # Configuration — add new quantization methods / models here
 # ---------------------------------------------------------------------------
@@ -316,6 +330,37 @@ def _unload(omni):
 # Tests
 # ---------------------------------------------------------------------------
 
+
+def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch):
+    from benchmarks.diffusion.quantization_quality import _generate_image as benchmark_generate_image
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
+    monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
+    monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
+
+    image = Image.new("RGB", (2, 2))
+    inner = OmniRequestOutput.from_diffusion(request_id="req", images=[image])
+    outer = OmniRequestOutput(
+        request_id="req",
+        stage_id=0,
+        final_output_type="image",
+        request_output=inner,
+        finished=True,
+    )
+
+    class DummyOmni:
+        def generate(self, *_args, **_kwargs):
+            return [outer]
+
+    args = SimpleNamespace(height=2, width=2, num_inference_steps=1)
+    output, _elapsed, peak_mem = benchmark_generate_image(DummyOmni(), args, "prompt", 42)
+
+    assert output is image
+    assert peak_mem == 0.0
+
+
 _marks = hardware_marks(res={"cuda": "H100"})
 _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
 

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import (
+    VariantRun,
+    _build_variant_config,
+    _request_peak_memory_mb,
+    _run_summary,
+    compute_tensor_metrics,
+    compute_uint8_image_metrics,
+    metric_guidance,
+    parse_args,
+    summarize_output_image_metrics,
+)
+
+
+def test_compute_tensor_metrics_identical_tensors():
+    metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))
+
+    assert metrics["cosine_similarity"] == pytest.approx(1.0)
+    assert metrics["mae"] == 0.0
+    assert metrics["mse"] == 0.0
+    assert metrics["rmse"] == 0.0
+    assert metrics["max_abs"] == 0.0
+    assert metrics["l2"] == 0.0
+    assert metrics["relative_l2"] == 0.0
+
+
+def test_compute_uint8_image_metrics_adds_psnr():
+    lhs = np.zeros((2, 2, 3), dtype=np.uint8)
+    rhs = np.zeros((2, 2, 3), dtype=np.uint8)
+
+    metrics = compute_uint8_image_metrics(lhs, rhs)
+
+    assert math.isinf(metrics["psnr_db"])
+
+
+def test_summarize_output_image_metrics_stacks_pil_images():
+    reference = [Image.fromarray(np.zeros((2, 2, 3), dtype=np.uint8))]
+    candidate = [Image.fromarray(np.ones((2, 2, 3), dtype=np.uint8))]
+
+    summary = summarize_output_image_metrics(reference, candidate)
+
+    assert summary["num_images"] == 1
+    assert summary["image0_metrics"]["mae"] == 1.0
+    assert summary["all_images_metrics"]["mse"] == 1.0
+
+
+def test_run_summary_reports_worker_peak_memory():
+    summary = _run_summary(
+        VariantRun(
+            label="candidate",
+            result=object(),
+            generation_times_s=[1.0, 3.0],
+            peak_memory_mb=[100.0, 150.0],
+        )
+    )
+
+    assert summary["peak_memory_mb"] == 150.0
+    assert summary["avg_peak_memory_mb"] == 125.0
+    assert summary["max_peak_memory_mb"] == 150.0
+    assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
+
+
+def test_request_peak_memory_prefers_inner_diffusion_output():
+    inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
+    outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
+
+    assert _request_peak_memory_mb(outer) == 321.0
+
+
+def test_metric_guidance_describes_thresholds():
+    guidance = metric_guidance()
+
+    assert "cosine_similarity" in guidance["descriptions"]
+    assert guidance["recommended_thresholds"]["output_images_or_frames_uint8"]["psnr_db"]["recommended_min"] == 20.0
+    assert (
+        guidance["recommended_thresholds"]["performance"]["max_peak_memory_ratio_candidate_over_reference"][
+            "recommended_max"
+        ]
+        == 1.00
+    )
+
+
+def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantization():
+    args = SimpleNamespace(
+        model="Qwen/Qwen-Image",
+        candidate_model="Qwen/Qwen-Image-FP8",
+        candidate_quantization=None,
+        candidate_quantization_config_json=None,
+        candidate_ignored_layers=None,
+        ignored_layers=None,
+        candidate_gguf_model=None,
+        candidate_load_format="default",
+    )
+
+    config = _build_variant_config(args, "candidate")
+
+    assert config.model == "Qwen/Qwen-Image-FP8"
+    assert config.quantization is None
+    assert config.quantization_config is None
+
+
+def test_step_execution_defaults_to_false(monkeypatch):
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        [
+            "compare_diffusion_trajectory_similarity.py",
+            "--output-json",
+            "result.json",
+        ],
+    )
+
+    args = parse_args()
+
+    assert args.step_execution is False

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import queue
+import threading
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _make_request(
+    request_id: str,
+    *,
+    request_ids: list[str] | None = None,
+) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt_{request_id}"],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_ids=request_ids or [request_id],
+        request_id=request_id,
+    )
+
+
+def _make_engine() -> DiffusionEngine:
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.scheduler = RequestScheduler()
+    engine.scheduler.initialize(SimpleNamespace())
+    engine.executor = SimpleNamespace(shutdown=Mock())
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._out_queue = {}
+    engine._closed = False
+    engine._shutdown_complete = False
+    engine.abort_queue = queue.Queue()
+    engine._loop_started = False
+    engine.stop_event = None
+    engine.worker_thread = None
+    return engine
+
+
+def test_close_completes_pending_async_waiters() -> None:
+    engine = _make_engine()
+    event_loop = asyncio.new_event_loop()
+    try:
+        future = event_loop.create_future()
+        engine._out_queue["pending-req"] = future
+
+        engine.close()
+
+        assert future.done()
+        assert future.result().error == "DiffusionEngine is closed."
+    finally:
+        event_loop.close()
+
+
+def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
+    class RacingOutQueue(dict):
+        def __contains__(self, key):
+            return True
+
+        def pop(self, key, default=None):
+            return default
+
+    engine = _make_engine()
+    engine._out_queue = RacingOutQueue()
+    engine._finalize_finished_request = Mock(side_effect=AssertionError("should not finalize drained waiters"))
+
+    engine._handle_finished_requests({"pending-req"})
+
+    engine._finalize_finished_request.assert_not_called()
+
+
+def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
+    engine = _make_engine()
+    request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
+    sched_req_id = engine.scheduler.add_request(request)
+
+    engine.abort("batch-parent")
+    engine._process_aborts_queue()
+
+    state = engine.scheduler.get_request_state(sched_req_id)
+    assert state is not None
+    assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
+
+
+def test_close_rejects_late_async_requests() -> None:
+    engine = _make_engine()
+    event_loop = asyncio.new_event_loop()
+    try:
+        engine.main_loop = event_loop
+        engine.close()
+
+        with pytest.raises(RuntimeError, match="closed"):
+            engine.add_request(_make_request("late-req"))
+    finally:
+        event_loop.close()
+
+
+def test_close_resets_loop_started_for_dead_worker_thread() -> None:
+    engine = _make_engine()
+    engine._loop_started = True
+    engine.worker_thread = SimpleNamespace(is_alive=Mock(return_value=False))
+
+    engine.close()
+
+    assert engine._loop_started is False
+
+
+def test_close_defers_resource_shutdown_until_worker_thread_stops() -> None:
+    engine = _make_engine()
+    engine.scheduler.close = Mock()
+    engine._loop_started = True
+    worker_thread = SimpleNamespace(
+        is_alive=Mock(side_effect=[True, True, False, False]),
+        join=Mock(),
+    )
+    engine.worker_thread = worker_thread
+
+    engine.close()
+
+    worker_thread.join.assert_called_once_with(timeout=10)
+    engine.scheduler.close.assert_not_called()
+    engine.executor.shutdown.assert_not_called()
+    assert engine._shutdown_complete is False
+    assert engine._loop_started is True
+
+    engine.close()
+
+    engine.scheduler.close.assert_called_once()
+    engine.executor.shutdown.assert_called_once()
+    assert engine._shutdown_complete is True
+    assert engine._loop_started is False

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -136,6 +136,7 @@ def _make_engine_with_loop(
     calls and a real ``RequestScheduler``.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
 
     sched = RequestScheduler()
@@ -371,6 +372,7 @@ def test_collective_rpc_before_loop_starts_calls_executor_directly():
     on the caller's thread without enqueueing.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine._loop_started = False
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
@@ -391,6 +393,7 @@ def test_collective_rpc_before_loop_starts_serializes_concurrent_callers():
     serialize them so they cannot race on the shared executor MQ pair.
     """
     engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._closed = False
     engine._loop_started = False
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)

tests/diffusion/test_diffusion_scheduler.py

@@ -305,17 +305,55 @@ class TestRequestScheduler:
             prompts=["prompt_map_a", "prompt_map_b"],
             sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
             request_ids=["map-a", "map-b"],
+            request_id="map-parent",
         )
 
         sched_req_id = self.scheduler.add_request(request)
 
         assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
         assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
+        assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
 
         self.scheduler.pop_request_state(sched_req_id)
 
         assert self.scheduler.get_sched_req_id("map-a") is None
         assert self.scheduler.get_sched_req_id("map-b") is None
+        assert self.scheduler.get_sched_req_id("map-parent") is None
+
+    def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
+        self.scheduler.add_request(
+            OmniDiffusionRequest(
+                prompts=["prompt_existing"],
+                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+                request_ids=["existing-child"],
+                request_id="duplicate-parent",
+            )
+        )
+
+        colliding_request = OmniDiffusionRequest(
+            prompts=["prompt_unique"],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_ids=["unique-child"],
+            request_id="duplicate-parent",
+        )
+
+        with pytest.raises(ValueError, match="duplicate-parent"):
+            self.scheduler.add_request(colliding_request)
+
+        assert self.scheduler.get_request_state("unique-child") is None
+        assert self.scheduler.get_sched_req_id("unique-child") is None
+
+        valid_request = OmniDiffusionRequest(
+            prompts=["prompt_unique"],
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+            request_ids=["unique-child"],
+            request_id="unique-parent",
+        )
+        sched_req_id = self.scheduler.add_request(valid_request)
+
+        assert sched_req_id == "unique-child"
+        assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
+        assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
 
 
 class TestDiffusionEngine:
@@ -324,6 +362,8 @@ class TestDiffusionEngine:
         engine.scheduler = RequestScheduler()
         engine.scheduler.initialize(SimpleNamespace())
         engine._rpc_lock = threading.RLock()
+        engine._cv = threading.Condition(engine._rpc_lock)
+        engine._closed = False
         engine.abort_queue = queue.Queue()
 
         request = _make_request("engine")
@@ -343,6 +383,8 @@ class TestDiffusionEngine:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.scheduler = scheduler
         engine._rpc_lock = threading.RLock()
+        engine._cv = threading.Condition(engine._rpc_lock)
+        engine._closed = False
         engine.abort_queue = queue.Queue()
         engine.execute_fn = mocker.Mock(return_value=runner_output)
 
@@ -394,7 +436,8 @@ class TestDiffusionEngine:
     @pytest.mark.asyncio
     async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
-        engine._loop_started = False
+        engine._closed = False
+        engine._loop_started = True
         engine._init_lock = asyncio.Lock()
         engine.main_loop = asyncio.get_running_loop()
         engine.stop_event = threading.Event()
@@ -410,6 +453,7 @@ class TestDiffusionEngine:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._rpc_lock = threading.RLock()
         engine._cv = threading.Condition(engine._rpc_lock)
+        engine._closed = False
         engine.scheduler = RequestScheduler()
         engine.scheduler.initialize(SimpleNamespace())
         engine.abort_queue = queue.Queue()

tests/diffusion/test_diffusion_step_pipeline.py

@@ -305,6 +305,7 @@ def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
     engine.execute_fn = execute_fn
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
+    engine._closed = False
     engine.abort_queue = queue.Queue()
     return engine
 

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -79,6 +79,7 @@ def _make_engine(num_gpus: int = 1):
     engine.executor = executor
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
+    engine._closed = False
     engine._loop_started = False
     engine._rpc_queue = queue.Queue()
     engine.abort_queue = queue.Queue()

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -6,18 +6,18 @@ from time import time
 import pytest
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceInfo,
     LeastQueueLengthBalancer,
     RandomBalancer,
+    ReplicaInfo,
+    ReplicaStatus,
     RoundRobinBalancer,
-    StageStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_load_balancer_select_returns_valid_index():
-    """Verify RandomBalancer.select() returns a valid index for instances."""
+    """Verify RandomBalancer.select() returns a valid index for replicas."""
     # Task structure mirrors async_omni; RandomBalancer ignores task contents.
     task: dict = {
         "request_id": "test",
@@ -26,30 +26,30 @@ def test_load_balancer_select_returns_valid_index():
     }
 
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
@@ -58,38 +58,38 @@ def test_load_balancer_select_returns_valid_index():
 
     balancer = RandomBalancer()
 
-    index = balancer.select(task, instances)
+    index = balancer.select(task, replicas)
 
     assert isinstance(index, int)
-    assert 0 <= index < len(instances)
+    assert 0 <= index < len(replicas)
 
 
-def test_round_robin_balancer_cycles_instances():
+def test_round_robin_balancer_cycles_replicas():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=1,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -97,35 +97,35 @@ def test_round_robin_balancer_cycles_instances():
     ]
 
     balancer = RoundRobinBalancer()
-    results = [balancer.select({}, instances) for _ in range(5)]
+    results = [balancer.select({}, replicas) for _ in range(5)]
 
     # Default start_index=0 => 0,1,2,0,1
     assert results == [0, 1, 2, 0, 1]
 
 
-def test_round_robin_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_round_robin_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         RoundRobinBalancer().select({}, [])
 
 
 def test_round_robin_balancer_after_large_index_and_shorter_list():
-    """Large start_index % len(instances) then counter wraps with shorter list."""
+    """Large start_index % len(replicas) then counter wraps with shorter list."""
     now = time()
     two = [
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
@@ -138,30 +138,30 @@ def test_round_robin_balancer_after_large_index_and_shorter_list():
 
 def test_least_queue_length_balancer_picks_min_queue():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=2,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=0,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=5,
             last_heartbeat=now,
             registered_at=now,
@@ -169,41 +169,41 @@ def test_least_queue_length_balancer_picks_min_queue():
     ]
 
     balancer = LeastQueueLengthBalancer()
-    index = balancer.select({}, instances)
+    index = balancer.select({}, replicas)
     assert index == 1
 
 
-def test_least_queue_length_balancer_empty_instances_raises():
-    with pytest.raises(ValueError, match="instances must not be empty"):
+def test_least_queue_length_balancer_empty_replicas_raises():
+    with pytest.raises(ValueError, match="replicas must not be empty"):
         LeastQueueLengthBalancer().select({}, [])
 
 
 def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10002",
             output_addr="tcp://host:10002-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
         ),
-        InstanceInfo(
+        ReplicaInfo(
             input_addr="tcp://host:10003",
             output_addr="tcp://host:10003-out",
             stage_id=1,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=3,
             last_heartbeat=now,
             registered_at=now,
@@ -214,21 +214,21 @@ def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
         "vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
         return_value=2,
     )
-    assert balancer.select({}, instances) == 2
+    assert balancer.select({}, replicas) == 2
 
 
 def test_least_queue_length_balancer_negative_queue_raises():
     now = time()
-    instances = [
-        InstanceInfo(
+    replicas = [
+        ReplicaInfo(
             input_addr="tcp://host:10001",
             output_addr="tcp://host:10001-out",
             stage_id=0,
-            status=StageStatus.UP,
+            status=ReplicaStatus.UP,
             queue_length=-1,
             last_heartbeat=now,
             registered_at=now,
         ),
     ]
     with pytest.raises(ValueError, match="queue_length must be non-negative"):
-        LeastQueueLengthBalancer().select({}, instances)
+        LeastQueueLengthBalancer().select({}, replicas)

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -8,8 +8,8 @@ import pytest
 import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
-    InstanceList,
     OmniCoordClientForHub,
+    ReplicaList,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,8 +32,8 @@ def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> b
     return False
 
 
-def test_hub_client_caches_instance_list_from_pub():
-    """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list()."""
+def test_hub_client_caches_replica_list_from_pub():
+    """Verify OmniCoordClientForHub receives replica list updates from OmniCoordinator and caches for get_replica_list()."""
     ctx, pub, endpoint = _bind_pub()
 
     client = OmniCoordClientForHub(endpoint)
@@ -41,7 +41,7 @@ def test_hub_client_caches_instance_list_from_pub():
     time.sleep(0.2)
 
     now = time.time()
-    instances_payload = [
+    replicas_payload = [
         {
             "input_addr": "tcp://stage:10001",
             "output_addr": "tcp://stage:10001-out",
@@ -71,37 +71,37 @@ def test_hub_client_caches_instance_list_from_pub():
         },
     ]
 
-    payload = {"instances": instances_payload, "timestamp": now}
+    payload = {"replicas": replicas_payload, "timestamp": now}
     pub.send(json.dumps(payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3)
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 3)
 
-    inst_list = client.get_instance_list()
-    assert isinstance(inst_list, InstanceList)
-    assert len(inst_list.instances) == 3
+    rep_list = client.get_replica_list()
+    assert isinstance(rep_list, ReplicaList)
+    assert len(rep_list.replicas) == 3
 
-    for src, inst in zip(instances_payload, inst_list.instances, strict=True):
-        assert inst.input_addr == src["input_addr"]
-        assert inst.output_addr == src["output_addr"]
-        assert inst.stage_id == src["stage_id"]
-        assert inst.status.value == src["status"]
+    for src, rep in zip(replicas_payload, rep_list.replicas, strict=True):
+        assert rep.input_addr == src["input_addr"]
+        assert rep.output_addr == src["output_addr"]
+        assert rep.stage_id == src["stage_id"]
+        assert rep.status.value == src["status"]
 
-    stage0 = client.get_instances_for_stage(0)
-    stage1 = client.get_instances_for_stage(1)
+    stage0 = client.get_replicas_for_stage(0)
+    stage1 = client.get_replicas_for_stage(1)
 
-    assert all(inst.stage_id == 0 for inst in stage0.instances)
-    assert all(inst.stage_id == 1 for inst in stage1.instances)
+    assert all(rep.stage_id == 0 for rep in stage0.replicas)
+    assert all(rep.stage_id == 1 for rep in stage1.replicas)
 
-    # Send an updated list with fewer instances and verify cache refresh.
+    # Send an updated list with fewer replicas and verify cache refresh.
     updated_payload = {
-        "instances": instances_payload[:2],
+        "replicas": replicas_payload[:2],
         "timestamp": now + 1.0,
     }
     pub.send(json.dumps(updated_payload).encode("utf-8"))
 
-    assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2)
-    updated_list = client.get_instance_list()
-    assert len(updated_list.instances) == 2
+    assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 2)
+    updated_list = client.get_replica_list()
+    assert len(updated_list.replicas) == 2
 
     client.close()
     pub.close(0)

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -9,7 +9,7 @@ import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
-    StageStatus,
+    ReplicaStatus,
 )
 from vllm_omni.distributed.omni_coordinator import (
     omni_coord_client_for_stage as stage_client_module,
@@ -47,7 +47,7 @@ def test_stage_client_auto_register_on_init():
     event = _recv_event(router)
 
     assert event["event_type"] == "update"
-    assert event["status"] == StageStatus.UP.value
+    assert event["status"] == ReplicaStatus.UP.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -70,13 +70,13 @@ def test_stage_client_update_info_sends_correct_event():
     # Discard initial registration event.
     _recv_event(router)
 
-    client.update_info(status=StageStatus.ERROR)
+    client.update_info(status=ReplicaStatus.ERROR)
     client.update_info(queue_length=10)
 
     first = _recv_event(router)
     second = _recv_event(router)
 
-    assert first["status"] == StageStatus.ERROR.value
+    assert first["status"] == ReplicaStatus.ERROR.value
     assert first["stage_id"] == stage_id
     assert first["input_addr"] == input_addr
     assert first["output_addr"] == output_addr
@@ -107,7 +107,7 @@ def test_stage_client_close_sends_down_status():
     client.close()
 
     event = _recv_event(router)
-    assert event["status"] == StageStatus.DOWN.value
+    assert event["status"] == ReplicaStatus.DOWN.value
     assert event["stage_id"] == stage_id
     assert event["input_addr"] == input_addr
     assert event["output_addr"] == output_addr
@@ -274,6 +274,7 @@ def test_heartbeat_loop_retries_after_transient_send_failure():
     client._closed = False
     client._heartbeat_interval = 0.0
     client._stop_event = _FakeStopEvent()
+    client._on_heartbeat = None
 
     calls = {"count": 0}
 

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -11,14 +11,14 @@ from vllm.v1.utils import get_engine_client_zmq_addr
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
     OmniCoordinator,
-    StageStatus,
+    ReplicaStatus,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
-    """Receive InstanceList JSON from SUB socket. Returns None on timeout."""
+def _recv_replica_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
+    """Receive ReplicaList JSON from SUB socket. Returns None on timeout."""
     sub.setsockopt(zmq.RCVTIMEO, timeout_ms)
     try:
         data = sub.recv()
@@ -27,16 +27,16 @@ def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
         return None
 
 
-def _wait_for_instance_list(
+def _wait_for_replica_list(
     sub: zmq.Socket,
     expected_count: int,
     timeout: float = 3.0,
 ) -> dict | None:
-    """Wait until received InstanceList with expected_count active instances."""
+    """Wait until received ReplicaList with expected_count active replicas."""
     start = time.time()
     while time.time() - start < timeout:
-        msg = _recv_instance_list(sub, timeout_ms=500)
-        if msg is not None and len(msg.get("instances", [])) == expected_count:
+        msg = _recv_replica_list(sub, timeout_ms=500)
+        if msg is not None and len(msg.get("replicas", [])) == expected_count:
             return msg
     return None
 
@@ -45,7 +45,7 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None:
     """Drain queued SUB messages for a short window."""
     deadline = time.time() + max_seconds
     while time.time() < deadline:
-        _recv_instance_list(sub, timeout_ms=50)
+        _recv_replica_list(sub, timeout_ms=50)
 
 
 def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
@@ -81,7 +81,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     )
 
     # Wait for initial registration broadcast and clear any queued messages.
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
     _drain_sub_messages(sub)
 
@@ -96,7 +96,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
     deadline = time.time() + window_s
     recv_count = 0
     while time.time() < deadline:
-        if _recv_instance_list(sub, timeout_ms=100) is not None:
+        if _recv_replica_list(sub, timeout_ms=100) is not None:
             recv_count += 1
 
     assert recv_count < update_count // 2, (
@@ -110,8 +110,8 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
 
 
 def test_omni_coordinator_registration_broadcast():
-    """Verify that after multiple OmniCoordClientForStage instances register,
-    OmniCoordinator publishes an InstanceList containing all registered instances.
+    """Verify that after multiple OmniCoordClientForStage replicas register,
+    OmniCoordinator publishes a ReplicaList containing all registered replicas.
     """
     router_addr = get_engine_client_zmq_addr(
         local_only=False,
@@ -144,12 +144,12 @@ def test_omni_coordinator_registration_broadcast():
         OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
     ]
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected InstanceList with 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected ReplicaList with 3 replicas"
+    assert len(msg["replicas"]) == 3
     assert isinstance(msg["timestamp"], (int, float))
 
-    input_addrs = {inst["input_addr"] for inst in msg["instances"]}
+    input_addrs = {rep["input_addr"] for rep in msg["replicas"]}
     assert "tcp://stage:10001" in input_addrs
     assert "tcp://stage:10002" in input_addrs
     assert "tcp://stage:10003" in input_addrs
@@ -162,7 +162,7 @@ def test_omni_coordinator_registration_broadcast():
 
 
 def test_omni_coordinator_heartbeat_timeout_handling():
-    """Verify that when a stage instance stops sending heartbeats,
+    """Verify that when a stage replica stops sending heartbeats,
     OmniCoordinator marks it as unhealthy and excludes it from the active list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -201,23 +201,23 @@ def test_omni_coordinator_heartbeat_timeout_handling():
         "output_addr": "tcp://stage:c-out",
         "stage_id": 0,
         "event_type": "update",
-        "status": StageStatus.UP.value,
+        "status": ReplicaStatus.UP.value,
         "queue_length": 0,
     }
     dealer_c.send(json.dumps(reg_event).encode("utf-8"))
 
-    msg = _wait_for_instance_list(sub, expected_count=3)
-    assert msg is not None, "Expected initial 3 instances"
-    assert len(msg["instances"]) == 3
+    msg = _wait_for_replica_list(sub, expected_count=3)
+    assert msg is not None, "Expected initial 3 replicas"
+    assert len(msg["replicas"]) == 3
 
     # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s).
     time.sleep(8.0)
 
     # Receive the update (C should be ERROR and excluded from active list).
-    msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0)
-    assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout"
-    instances = msg_after_timeout.get("instances", [])
-    input_addrs = {inst["input_addr"] for inst in instances}
+    msg_after_timeout = _wait_for_replica_list(sub, expected_count=2, timeout=5.0)
+    assert msg_after_timeout is not None, "Expected ReplicaList with 2 replicas after timeout"
+    replicas = msg_after_timeout.get("replicas", [])
+    input_addrs = {rep["input_addr"] for rep in replicas}
 
     assert "tcp://stage:a" in input_addrs
     assert "tcp://stage:b" in input_addrs
@@ -232,8 +232,8 @@ def test_omni_coordinator_heartbeat_timeout_handling():
     sub_ctx.term()
 
 
-def test_omni_coordinator_instance_shutdown_handling():
-    """Verify that when a stage instance sends status='down',
+def test_omni_coordinator_replica_shutdown_handling():
+    """Verify that when a stage replica sends status='down',
     OmniCoordinator removes it from the active list and broadcasts an updated list.
     """
     router_addr = get_engine_client_zmq_addr(
@@ -261,18 +261,18 @@ def test_omni_coordinator_instance_shutdown_handling():
 
     client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
 
-    msg = _wait_for_instance_list(sub, expected_count=1)
+    msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None
-    assert len(msg["instances"]) == 1
-    assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown"
+    assert len(msg["replicas"]) == 1
+    assert msg["replicas"][0]["input_addr"] == "tcp://stage:shutdown"
 
     # Send down status (simulating graceful shutdown).
-    client.update_info(status=StageStatus.DOWN)
+    client.update_info(status=ReplicaStatus.DOWN)
 
-    # Receive updated list (should have 0 active instances).
-    msg = _wait_for_instance_list(sub, expected_count=0)
+    # Receive updated list (should have 0 active replicas).
+    msg = _wait_for_replica_list(sub, expected_count=0)
     assert msg is not None
-    assert len(msg["instances"]) == 0
+    assert len(msg["replicas"]) == 0
 
     client.close()
     coordinator.close()

tests/e2e/accuracy/test_helios_distilled.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import signal
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import reset_artifact_dir
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+MODEL_NAME = "BestWishYsh/Helios-Distilled"
+PROMPT = "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting."
+HEIGHT = 384
+WIDTH = 640
+NUM_FRAMES = 99
+FPS = 16
+SEED = 42
+GUIDANCE_SCALE = 1.0
+PYRAMID_NUM_INFERENCE_STEPS = (1, 1, 1)
+MIN_PSNR_DB = 35.0
+MAX_MAE = 3.0
+MAX_TIME_RATIO = 1.03
+WARMUP_REQUESTS = 1
+WARMUP_NUM_INFERENCE_STEPS = 1
+
+
+def _run(command: list[str], *, cwd: Path, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
+    print("+ " + " ".join(command), flush=True)
+    return subprocess.run(
+        command,
+        cwd=str(cwd),
+        env=env,
+        check=True,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _open_port(host: str = "127.0.0.1") -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind((host, 0))
+        return int(sock.getsockname()[1])
+
+
+def _wait_for_service(base_url: str, process: subprocess.Popen, *, timeout_s: float = 1200.0) -> None:
+    import requests
+
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        if process.poll() is not None:
+            raise RuntimeError(f"Server exited with code {process.returncode} before becoming ready.")
+        try:
+            response = requests.get(f"{base_url}/health", timeout=1)
+            if response.status_code == 200:
+                return
+        except requests.RequestException:
+            pass
+        time.sleep(2)
+    raise TimeoutError(f"Server at {base_url} did not become ready within {timeout_s:.0f}s.")
+
+
+def _stop_process(process: subprocess.Popen) -> None:
+    if process.poll() is not None:
+        return
+    try:
+        os.killpg(process.pid, signal.SIGTERM)
+        process.wait(timeout=30)
+    except Exception:
+        os.killpg(process.pid, signal.SIGKILL)
+        process.wait(timeout=30)
+
+
+def _extract_frames(video_path: Path, output_dir: Path) -> list[Image.Image]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    frame_pattern = output_dir / "frame_%05d.png"
+    _run(
+        [
+            "ffmpeg",
+            "-hide_banner",
+            "-loglevel",
+            "error",
+            "-i",
+            str(video_path),
+            "-vsync",
+            "0",
+            str(frame_pattern),
+        ],
+        cwd=output_dir,
+    )
+    frame_paths = sorted(output_dir.glob("frame_*.png"))
+    if not frame_paths:
+        raise RuntimeError(f"No frames extracted from {video_path}")
+    return [Image.open(path).convert("RGB").copy() for path in frame_paths]
+
+
+def _image_metrics(reference: np.ndarray, candidate: np.ndarray) -> dict[str, float]:
+    diff = reference.astype(np.float64) - candidate.astype(np.float64)
+    mse = float(np.mean(diff * diff))
+    mae = float(np.mean(np.abs(diff)))
+    max_abs = float(np.max(np.abs(diff)))
+    psnr_db = float("inf") if mse == 0.0 else 20.0 * math.log10(255.0) - 10.0 * math.log10(mse)
+    ref_flat = reference.astype(np.float64).reshape(-1)
+    cand_flat = candidate.astype(np.float64).reshape(-1)
+    ref_norm = float(np.linalg.norm(ref_flat))
+    cand_norm = float(np.linalg.norm(cand_flat))
+    if ref_norm == 0.0 and cand_norm == 0.0:
+        cosine = 1.0
+    else:
+        cosine = float(np.dot(ref_flat, cand_flat) / (ref_norm * cand_norm))
+    return {
+        "mae": mae,
+        "mse": mse,
+        "max_abs": max_abs,
+        "psnr_db": psnr_db,
+        "cosine_similarity": cosine,
+    }
+
+
+def _video_metrics(reference_frames: list[Image.Image], candidate_frames: list[Image.Image]) -> dict[str, Any]:
+    if len(reference_frames) != len(candidate_frames):
+        raise ValueError(f"Frame count mismatch: reference={len(reference_frames)} candidate={len(candidate_frames)}")
+
+    ref_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in reference_frames], axis=0)
+    cand_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in candidate_frames], axis=0)
+    mid = len(reference_frames) // 2
+    return {
+        "num_frames": len(reference_frames),
+        "frame0_metrics": _image_metrics(ref_stack[0], cand_stack[0]),
+        "mid_frame_index": mid,
+        "mid_frame_metrics": _image_metrics(ref_stack[mid], cand_stack[mid]),
+        "all_frames_metrics": _image_metrics(ref_stack, cand_stack),
+    }
+
+
+def _reference_repo(output_root: Path) -> Path:
+    configured = os.environ.get("VLLM_HELIOS_REFERENCE_REPO")
+    if configured:
+        reference = Path(configured).expanduser().resolve()
+        if not reference.exists():
+            raise FileNotFoundError(f"VLLM_HELIOS_REFERENCE_REPO does not exist: {reference}")
+        return reference
+
+    reference = output_root / "origin_main"
+    if reference.exists():
+        subprocess.run(["git", "worktree", "remove", "--force", str(reference)], cwd=str(REPO_ROOT), check=False)
+        if reference.exists():
+            shutil.rmtree(reference)
+    _run(["git", "fetch", "origin", "main"], cwd=REPO_ROOT)
+    _run(["git", "worktree", "add", "--detach", str(reference), "origin/main"], cwd=REPO_ROOT)
+    return reference
+
+
+def _server_env(repo: Path) -> dict[str, str]:
+    env = dict(os.environ)
+    env["PYTHONPATH"] = str(repo) + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _benchmark_config() -> str:
+    return json.dumps(
+        [
+            {
+                "width": WIDTH,
+                "height": HEIGHT,
+                "num_inference_steps": 50,
+                "num_frames": NUM_FRAMES,
+                "fps": FPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "is_enable_stage2": True,
+                "pyramid_num_stages": 3,
+                "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS),
+                "is_amplify_first_chunk": False,
+                "weight": 1,
+            }
+        ]
+    )
+
+
+def _start_server(*, repo: Path, variant_dir: Path) -> tuple[subprocess.Popen, str]:
+    port = _open_port()
+    base_url = f"http://127.0.0.1:{port}"
+    log_file = (variant_dir / "server.log").open("w", encoding="utf-8")
+    command = [
+        sys.executable,
+        "-m",
+        "vllm_omni.entrypoints.cli.main",
+        "serve",
+        MODEL_NAME,
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(port),
+        "--omni",
+    ]
+    print("+ " + " ".join(command), flush=True)
+    process = subprocess.Popen(
+        command,
+        cwd=str(repo),
+        env=_server_env(repo),
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        preexec_fn=os.setsid,
+        text=True,
+    )
+    try:
+        _wait_for_service(base_url, process)
+    except Exception:
+        log_file.close()
+        _stop_process(process)
+        raise
+    log_file.close()
+    return process, base_url
+
+
+def _run_serving_benchmark(*, repo: Path, variant_dir: Path, base_url: str) -> dict[str, Any]:
+    metrics_path = variant_dir / "metrics.json"
+    response_dir = variant_dir / "responses"
+    command = [
+        sys.executable,
+        str(REPO_ROOT / "benchmarks/diffusion/diffusion_benchmark_serving.py"),
+        "--base-url",
+        base_url,
+        "--endpoint",
+        "/v1/videos",
+        "--model",
+        MODEL_NAME,
+        "--dataset",
+        "random",
+        "--task",
+        "t2v",
+        "--num-prompts",
+        "1",
+        "--max-concurrency",
+        "1",
+        "--warmup-requests",
+        str(WARMUP_REQUESTS),
+        "--warmup-num-inference-steps",
+        str(WARMUP_NUM_INFERENCE_STEPS),
+        "--warmup-concurrency",
+        "1",
+        "--random-prompt",
+        PROMPT,
+        "--seed",
+        str(SEED),
+        "--disable-tqdm",
+        "--random-request-config",
+        _benchmark_config(),
+        "--output-file",
+        str(metrics_path),
+        "--save-response-dir",
+        str(response_dir),
+    ]
+    completed = _run(command, cwd=repo, env=_server_env(repo))
+    (variant_dir / "benchmark.log").write_text(completed.stdout, encoding="utf-8")
+    metrics = _read_json(metrics_path)
+    saved_responses = metrics.get("saved_responses") or []
+    if len(saved_responses) != 1:
+        raise RuntimeError(f"Expected one saved video response, got {saved_responses}")
+    return {
+        "repo": str(repo),
+        "metrics": metrics,
+        "video": saved_responses[0],
+        "log": str(variant_dir / "benchmark.log"),
+        "server_log": str(variant_dir / "server.log"),
+    }
+
+
+def _run_variant(*, label: str, repo: Path, output_root: Path) -> dict[str, Any]:
+    variant_dir = output_root / label
+    variant_dir.mkdir(parents=True, exist_ok=True)
+    process, base_url = _start_server(repo=repo, variant_dir=variant_dir)
+    try:
+        return _run_serving_benchmark(repo=repo, variant_dir=variant_dir, base_url=base_url)
+    finally:
+        _stop_process(process)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_helios_distilled_pr_matches_main_acc_perf(accuracy_artifact_root: Path) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("Helios-Distilled accuracy/perf comparison requires CUDA.")
+
+    output_root = reset_artifact_dir(accuracy_artifact_root / "helios_distilled_acc_perf")
+    reference = _reference_repo(output_root)
+    output_json = output_root / "result.json"
+
+    reference_result = _run_variant(label="reference", repo=reference, output_root=output_root)
+    candidate_result = _run_variant(label="candidate", repo=REPO_ROOT, output_root=output_root)
+
+    with TemporaryDirectory() as tmp:
+        tmp_dir = Path(tmp)
+        reference_frames = _extract_frames(Path(reference_result["video"]), tmp_dir / "reference_frames")
+        candidate_frames = _extract_frames(Path(candidate_result["video"]), tmp_dir / "candidate_frames")
+        output_metrics = _video_metrics(reference_frames, candidate_frames)
+
+    reference_time = float(reference_result["metrics"]["latency_mean"])
+    candidate_time = float(candidate_result["metrics"]["latency_mean"])
+    time_ratio = candidate_time / reference_time
+    result = {
+        "reference": reference_result,
+        "candidate": candidate_result,
+        "sampling": {
+            "model": MODEL_NAME,
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "num_frames": NUM_FRAMES,
+            "fps": FPS,
+            "seed": SEED,
+            "guidance_scale": GUIDANCE_SCALE,
+            "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS),
+            "is_amplify_first_chunk": False,
+        },
+        "performance": {
+            "reference_measured_time_s": reference_time,
+            "candidate_measured_time_s": candidate_time,
+            "candidate_over_reference_time_ratio": time_ratio,
+        },
+        "output_metrics": output_metrics,
+        "thresholds": {
+            "min_psnr_db": MIN_PSNR_DB,
+            "max_mae": MAX_MAE,
+            "max_time_ratio": MAX_TIME_RATIO,
+        },
+    }
+    output_json.write_text(json.dumps(result, indent=2, sort_keys=True, allow_nan=True), encoding="utf-8")
+    all_metrics = output_metrics["all_frames_metrics"]
+    print(
+        json.dumps(
+            {
+                "output_json": str(output_json),
+                "candidate_over_reference_time_ratio": time_ratio,
+                "all_frames_psnr_db": all_metrics["psnr_db"],
+                "all_frames_mae": all_metrics["mae"],
+                "all_frames_cosine_similarity": all_metrics["cosine_similarity"],
+            },
+            indent=2,
+            allow_nan=True,
+        )
+    )
+
+    assert all_metrics["psnr_db"] >= MIN_PSNR_DB
+    assert all_metrics["mae"] <= MAX_MAE
+    assert time_ratio <= MAX_TIME_RATIO
+    assert output_json.exists(), f"Expected Helios comparison result at {output_json}"

tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py

@@ -70,7 +70,7 @@ def test_text_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -93,7 +93,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -114,7 +114,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -135,7 +135,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -162,4 +162,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_omni_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"

tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py

@@ -87,7 +87,7 @@ def test_text_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -111,7 +111,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -133,7 +133,7 @@ def test_image_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -155,7 +155,7 @@ def test_video_to_text(omni_runner, omni_runner_handler):
         "modalities": ["text"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
     assert response.text_content and len(response.text_content.strip()) > 0
 
 
@@ -177,7 +177,7 @@ def test_video_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"
 
 
 # ------------------------------------------------------------------
@@ -204,4 +204,4 @@ def test_mix_to_audio(omni_runner, omni_runner_handler):
         "modalities": ["audio"],
     }
     response = omni_runner_handler.send_request(request_config)
-    assert response.success, f"Request failed: {response.error_message}"
+    assert response.success, "Request failed"

tests/e2e/online_serving/test_omnivoice_expansion.py

@@ -105,16 +105,3 @@ class TestOmniVoiceVoiceCloning:
             "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
         }
         openai_client.send_audio_speech_request(request_config)
-
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_voice_clone_invalid_ref_audio_format(self, omni_server, openai_client) -> None:
-        """Test that invalid ref_audio format returns a clear error."""
-        request_config = {
-            "model": omni_server.model,
-            "input": get_prompt("text"),
-            "ref_audio": "not_a_valid_uri",
-            "response_format": "wav",
-            "timeout": 180.0,
-            "status_code": (400, 422),
-        }
-        openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -49,34 +49,17 @@ def get_batch_token_config(default_path):
     )
 
 
-def get_default_config(default_path):
-    """Flip async_chunk on and bump stage 0 thinker output to 2048 tokens.
-
-    Pipeline registry (qwen3_omni/pipeline.py) already wires
-    thinker2talker_async_chunk / talker2code2wav_async_chunk on stages 0/1,
-    so no per-stage processor override is needed. Using only flat-schema
-    writes so _parse_stage_deploy stays in its flat branch (nested
-    ``engine_args:`` would drop other overlay fields).
-    """
-    return modify_stage_config(
-        default_path,
-        updates={
-            "stages": {0: {"default_sampling_params.max_tokens": 2048}},
-        },
-    )
-
-
 # CI deploy YAML (single file; xpu deltas applied via ``platforms:`` section).
 # The overlay explicitly sets ``async_chunk: False``, so ``default`` tests the
 # sync path and ``async_chunk`` tests the streaming path with a longer thinker
 # output — two distinct scenarios, kept as separate parametrizations.
-default_path = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+default_path = get_deploy_config_path("qwen3_omni_moe.yaml")
 
 test_params = [
     pytest.param(
         OmniServerParams(
             model=model,
-            stage_config_path=get_default_config(default_path),
+            stage_config_path=default_path,
             use_stage_cli=True,
             server_args=["--no-async-chunk"],
         ),
@@ -85,7 +68,7 @@ test_params = [
     pytest.param(
         OmniServerParams(
             model=model,
-            stage_config_path=get_default_config(default_path),
+            stage_config_path=default_path,
             use_stage_cli=True,
             server_args=["--async-chunk"],
         ),

tests/e2e/online_serving/test_voxtral_tts.py

@@ -91,20 +91,3 @@ class TestVoxtralTTSFixedVoice:
                     "timeout": 120.0,
                 }
             )
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "H100"}, num_cards=1)
-    def test_speech_invalid_voice_rejected(self, omni_server, openai_client) -> None:
-        """Request with a non-existent voice should be rejected (HTTP 200 with JSON or 4xx, depending on server)."""
-        openai_client.send_audio_speech_request(
-            {
-                "model": omni_server.model,
-                "input": "This voice does not exist.",
-                "voice": "nonexistent_voice_xyz",
-                "language": "English",
-                "response_format": "wav",
-                "timeout": 120.0,
-                "err_message": "Invalid speaker",
-            }
-        )

tests/engine/test_async_omni_engine_stage_init.py

@@ -171,18 +171,6 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
     assert cleanup_clients == [shared, extra]
 
 
-def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode():
-    engine = object.__new__(AsyncOmniEngine)
-    engine.single_stage_mode = True
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})]
-
-    with pytest.raises(
-        ValueError,
-        match="single_stage_mode only supports num_replicas > 1 for diffusion stages",
-    ):
-        engine._validate_single_stage_mode_replica_constraints()
-
-
 def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
@@ -427,6 +415,7 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
+    engine._coordinator_runtime = None
     engine.stage_configs = [types.SimpleNamespace()]
 
     cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
@@ -471,6 +460,7 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
+    engine._coordinator_runtime = None
     engine.stage_configs = [types.SimpleNamespace()]
 
     cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))

tests/engine/test_orchestrator.py

@@ -1052,7 +1052,7 @@ async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) ->
         target_logger.setLevel(prev_level)
 
     assert not stage0.abort_calls
-    assert "abort: no binding for req=missing-req in stage-0" in caplog.text
+    assert "abort: no live binding for req=missing-req in stage-0" in caplog.text
 
 
 @pytest.mark.asyncio

tests/engine/test_single_stage_mode.py

@@ -203,6 +203,32 @@ class TestOmniMasterServerAllocation:
         server = OmniMasterServer(master_address="127.0.0.1", master_port=15004, stage_ids=[3])
         assert server.get_allocation(3) is server._stage_routes[(3, 0)]
 
+    def test_next_free_replica_id_skips_head_local_slot_until_filled(self):
+        # Head pre-allocates slot (0, 0) for its own register_stage_with_omni_master
+        # call. A same-host headless that registers with auto-assign BEFORE the
+        # head's own registration must NOT be handed slot 0 — it should land on
+        # slot 1 instead. Without the head_local_replicas reservation,
+        # _next_free_replica_id would see (0, 0) absent from _stage_configs and
+        # return 0, colliding with the head's own bound sockets.
+        server = OmniMasterServer(
+            master_address="127.0.0.1",
+            master_port=15010,
+            stage_ids=[0],
+            head_local_replicas={0: [0]},
+        )
+        assert server._next_free_replica_id(0) == 1
+
+    def test_next_free_replica_id_uses_remote_slot_when_unowned(self):
+        # When the head pre-allocates a remote-only slot (the head's _initialize_*
+        # path waits on get_stage_config), auto-assign SHOULD fill it so the
+        # head's wait unblocks. This is the original behavior, preserved.
+        server = OmniMasterServer(
+            master_address="127.0.0.1",
+            master_port=15011,
+            stage_ids=[0],
+        )
+        assert server._next_free_replica_id(0) == 0
+
 
 # ---------------------------------------------------------------------------
 # OmniMasterServer registration flow
@@ -233,36 +259,6 @@ class TestOmniMasterServerRegistration:
             ctx.term()
             server.stop()
 
-    def test_server_handles_unknown_stage_id_gracefully(self):
-        import msgspec
-        import zmq
-        from vllm.utils.network_utils import get_open_port
-
-        master_port = get_open_port()
-        server = OmniMasterServer(master_address="127.0.0.1", master_port=master_port, stage_ids=[0])
-        server.start()
-
-        ctx = zmq.Context()
-        bad_sock = None
-        good_sock = None
-        try:
-            bad_sock = ctx.socket(zmq.DEALER)
-            bad_sock.connect(f"tcp://127.0.0.1:{master_port}")
-            bad_sock.send(msgspec.msgpack.encode({"stage_id": 99}))
-            assert not bad_sock.poll(timeout=500)
-
-            good_sock = ctx.socket(zmq.DEALER)
-            good_sock.connect(f"tcp://127.0.0.1:{master_port}")
-            good_sock.send(msgspec.msgpack.encode({"stage_id": 0}))
-            assert good_sock.poll(timeout=2_000)
-            good_sock.recv()
-        finally:
-            for sock in (bad_sock, good_sock):
-                if sock is not None:
-                    sock.close(linger=0)
-            ctx.term()
-            server.stop()
-
     def test_registration_stores_stage_config(self):
         import msgspec
         import zmq
@@ -474,6 +470,8 @@ class TestSingleStageInitialization:
         engine._omni_master_address = "127.0.0.1"
         engine._omni_master_port = 26000
         engine._omni_master_server = None
+        engine._omni_heartbeat_timeout = 30.0
+        engine._coordinator_runtime = None
         engine.async_chunk = False
         engine.diffusion_batch_size = 2
         return engine
@@ -508,10 +506,16 @@ class TestSingleStageInitialization:
 
     def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture):
         import vllm_omni.engine.async_omni_engine as engine_mod
+        from vllm_omni.distributed import omni_coordinator as omni_coord_mod
 
         engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7)
         mock_oms = mocker.Mock(spec=OmniMasterServer)
         mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms)
+        mocker.patch.object(
+            omni_coord_mod,
+            "OmniCoordinatorRuntime",
+            return_value=mocker.Mock(router_address="tcp://127.0.0.1:9999"),
+        )
 
         stage_plans = [
             _make_llm_plan(0, configured_stage_id=7, launch_mode="local"),
@@ -520,12 +524,16 @@ class TestSingleStageInitialization:
 
         engine._start_omni_master_server(stage_plans)
 
-        engine_mod.OmniMasterServer.assert_called_once_with(
-            master_address="127.0.0.1",
-            master_port=26000,
-            stage_ids=[7, 11],
-            stage_replica_counts={7: 1, 11: 1},
-        )
+        call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs
+        assert call_kwargs["master_address"] == "127.0.0.1"
+        assert call_kwargs["master_port"] == 26000
+        assert call_kwargs["stage_ids"] == [7, 11]
+        assert call_kwargs["stage_replica_counts"] == {7: 1, 11: 1}
+        # head_local_replicas reserves slots that the head will register
+        # itself (launch_mode == "local"). Stage 11 is remote, so it must
+        # NOT appear in the head-owned set — that slot is for the headless
+        # to fill via auto-assign.
+        assert call_kwargs["head_local_replicas"] == {7: [0]}
         mock_oms.start.assert_called_once()
 
     def test_start_omni_master_server_duplicate_stage_ids_raise(self):
@@ -545,7 +553,9 @@ class TestSingleStageInitialization:
         with pytest.raises(ValueError, match="requires both"):
             engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")])
 
-    def test_build_logical_stage_init_plans_clears_runtime_cfg_in_single_stage_mode(self, mocker: MockerFixture):
+    def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_single_stage_mode(
+        self, mocker: MockerFixture
+    ):
         import vllm_omni.engine.async_omni_engine as engine_mod
 
         engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7)
@@ -574,7 +584,7 @@ class TestSingleStageInitialization:
         finally:
             monkeypatch.undo()
 
-        assert stage_plans[0].replicas[0].metadata.runtime_cfg is None
+        assert stage_plans[0].replicas[0].metadata.runtime_cfg == {"devices": "0"}
 
     def test_validate_single_stage_mode_allows_diffusion_replicas(self):
         stage_cfg = _make_stage_cfg(0, stage_type="diffusion")
@@ -583,14 +593,6 @@ class TestSingleStageInitialization:
 
         engine._validate_single_stage_mode_replica_constraints()
 
-    def test_validate_single_stage_mode_rejects_llm_replicas(self):
-        stage_cfg = _make_stage_cfg(0, stage_type="llm")
-        stage_cfg.runtime.num_replicas = 2
-        engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0)
-
-        with pytest.raises(ValueError, match="only supports num_replicas > 1 for diffusion"):
-            engine._validate_single_stage_mode_replica_constraints()
-
     def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode(
         self, mocker: MockerFixture
     ):
@@ -809,6 +811,7 @@ class TestSingleStageReplicaInitialization:
         engine.model = "fake-model"
         engine.single_stage_mode = True
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
+        engine._coordinator_runtime = None
         engine.stage_configs = []
 
         fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace())
@@ -896,6 +899,7 @@ class TestSingleStageReplicaInitialization:
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         engine._omni_master_server.address = "127.0.0.1"
         engine._omni_master_server.port = 25000
+        engine._coordinator_runtime = None
 
         plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
         sentinel_client = SimpleNamespace()
@@ -950,6 +954,9 @@ class TestSingleStageReplicaInitialization:
             handshake_address="tcp://hs",
             request_address="tcp://req",
             response_address="tcp://resp",
+            omni_coordinator_address=None,
+            omni_stage_id=5,
+            omni_replica_id=0,
         )
         mock_handshake.assert_called_once_with(proc, "tcp://hs", 60)
         mock_from_addresses.assert_called_once_with(
@@ -972,6 +979,7 @@ class TestSingleStageReplicaInitialization:
         engine._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         engine._omni_master_server.address = "127.0.0.1"
         engine._omni_master_server.port = 25000
+        engine._coordinator_runtime = None
 
         plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0]
         proc = mocker.Mock()

tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py

@@ -5,6 +5,10 @@ E2E tests for speaker_embedding API parameter.
 
 Tests Base models (0.6B-Base and 1.7B-Base) which support voice cloning
 via pre-computed ECAPA-TDNN embeddings, bypassing ref_audio extraction.
+
+Invalid combinations (mutually exclusive with ref_audio, empty embedding,
+wrong embedding dimensions) live in
+``tests/dfx/reliability/http_invalid/test_invalid_audio_speech.py``.
 """
 
 import os
@@ -147,45 +151,6 @@ class TestSpeakerEmbeddingBase:
         assert len(response.content) > MIN_AUDIO_BYTES
         assert_not_silence(response.content)
 
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speaker_embedding_mutually_exclusive_with_ref_audio(self, omni_server) -> None:
-        """Sending both speaker_embedding and ref_audio returns 400."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": DUMMY_EMBEDDING_1024,
-            "ref_audio": "https://example.com/audio.wav",
-            "response_format": "wav",
-        }
-        with httpx.Client(timeout=30.0) as client:
-            response = client.post(url, json=payload)
-
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "mutually exclusive" in response.text
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_speaker_embedding_empty_rejected(self, omni_server) -> None:
-        """Empty speaker_embedding list returns 400."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": [],
-            "response_format": "wav",
-        }
-        with httpx.Client(timeout=30.0) as client:
-            response = client.post(url, json=payload)
-
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "non-empty" in response.text
-
 
 # ── 1.7B-Base model tests (2048-dim embeddings) ──
 
@@ -216,26 +181,3 @@ class TestSpeakerEmbedding1_7B:
         assert response.headers.get("content-type") == "audio/wav"
         assert verify_wav_audio(response.content), "Response is not valid WAV"
         assert len(response.content) > MIN_AUDIO_BYTES, f"Audio too small: {len(response.content)} bytes"
-
-    @pytest.mark.advanced_model
-    @pytest.mark.tts
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
-    def test_1024_dim_on_1_7b_model_rejected(self, omni_server) -> None:
-        """1024-dim speaker_embedding on a 1.7B model (expects 2048) should fail gracefully."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
-        payload = {
-            "model": omni_server.model,
-            "input": SYN_TEXT,
-            "task_type": "Base",
-            "speaker_embedding": DUMMY_EMBEDDING_1024,
-            "x_vector_only_mode": True,
-            "response_format": "wav",
-            "max_new_tokens": MAX_NEW_TOKENS,
-        }
-        with httpx.Client(timeout=120.0) as client:
-            response = client.post(url, json=payload)
-
-        # Wrong dimensions should produce an error, not silent garbage.
-        assert response.status_code == 400, f"Expected 400, got {response.status_code}: {response.text}"
-        assert "speaker_embedding has 1024 dimensions" in response.text
-        assert "expected 2048" in response.text

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -1,13 +1,16 @@
-# tests/entrypoints/openai_api/test_serving_audio_generate.py
 import logging
 from inspect import Signature, signature
-from unittest.mock import MagicMock, patch
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 import torch
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.testclient import TestClient
+from vllm.v1.engine.exceptions import EngineGenerateError
 
+from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
+from vllm_omni.entrypoints.openai import api_server as api_server_module
 from vllm_omni.entrypoints.openai.protocol.audio import (
     CreateAudio,
     OpenAICreateAudioGenerateRequest,
@@ -106,6 +109,29 @@ def client(test_app):
     return TestClient(test_app)
 
 
+def _make_api_server_test_app(handler, *, request_id: str = "audio-gen-req-1"):
+    app = FastAPI()
+    app.state.openai_serving_audio_generate = handler
+    app.state.engine_client = SimpleNamespace(
+        engine=SimpleNamespace(is_alive=lambda: False),
+        errored=True,
+    )
+    app.state.server = SimpleNamespace()
+
+    @app.middleware("http")
+    async def add_request_metadata(request: Request, call_next):
+        request.state.request_metadata = SimpleNamespace(request_id=request_id)
+        return await call_next(request)
+
+    app.add_api_route(
+        "/v1/audio/generate",
+        api_server_module.create_audio_generate,
+        methods=["POST"],
+        response_model=None,
+    )
+    return app
+
+
 # Request Validation (Pydantic model)
 class TestRequestValidation:
     """Validate OpenAICreateAudioGenerateRequest pydantic constraints."""
@@ -472,6 +498,35 @@ class TestErrorHandling:
 
         assert "Audio generation failed" in resp.error.message
 
+    @pytest.mark.parametrize(
+        ("exc", "expected_message", "expected_stage_id"),
+        [
+            (OmniEngineDeadError("engine dead", error_stage_id=2), "engine dead", 2),
+            (EngineGenerateError("engine generate failed"), "engine generate failed", None),
+        ],
+    )
+    def test_api_server_engine_error_response_includes_request_and_stage_id(
+        self,
+        exc,
+        expected_message,
+        expected_stage_id,
+    ):
+        handler = MagicMock()
+        handler.create_audio_generate = AsyncMock(side_effect=exc)
+        app = _make_api_server_test_app(handler)
+
+        with patch.object(api_server_module, "terminate_if_errored") as terminate_mock:
+            with TestClient(app) as client:
+                response = client.post("/v1/audio/generate", json={"input": "Hello"})
+
+        assert response.status_code == 500
+        payload = response.json()
+        assert payload["error"]["message"] == expected_message
+        assert payload["error"]["code"] == 500
+        assert payload["error"]["request_id"] == "audio-gen-req-1"
+        assert payload["error"]["error_stage_id"] == expected_stage_id
+        terminate_mock.assert_called_once()
+
 
 # End-to-End via TestClient
 class TestAudioGenerateAPI:

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,5 +1,6 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import asyncio
+import json
 import logging
 import os
 import struct
@@ -564,7 +565,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_success(self, client):
         """Upload a voice via speaker_embedding JSON."""
-        import json
 
         emb = [0.1] * 1024
         data = {
@@ -585,7 +585,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_appears_in_listing(self, client):
         """Embedding-uploaded voice appears in list with correct source."""
-        import json
 
         emb = [0.2] * 2048
         data = {
@@ -606,7 +605,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_and_audio_mutually_exclusive(self, client):
         """Providing both audio_sample and speaker_embedding returns 400."""
-        import json
 
         emb = [0.1] * 1024
         files = {"audio_sample": ("test.wav", b"fake", "audio/wav")}
@@ -632,7 +630,6 @@ class TestSpeechAPI:
 
     def test_upload_voice_embedding_nan_rejected(self, client):
         """NaN values in speaker_embedding return 400."""
-        import json
 
         data = {
             "speaker_embedding": json.dumps([0.1] * 1023 + [float("nan")]),
@@ -930,7 +927,6 @@ class TestTTSMethods:
 
     def test_upload_voice_embedding_wrong_dims_rejected(self, speech_server):
         """Embedding uploads must match the loaded Qwen3-TTS model before being stored."""
-        import json
 
         speech_server._tts_model_type = "qwen3_tts"
         speech_server.engine_client.model_config = SimpleNamespace(

tests/entrypoints/test_serve.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import argparse
+from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
@@ -31,143 +33,118 @@ def test_serve_parser_accepts_no_async_chunk() -> None:
     assert args.async_chunk is False
 
 
-def _make_headless_args() -> argparse.Namespace:
-    return argparse.Namespace(
+# ---------------------------------------------------------------------------
+# run_headless validation
+# ---------------------------------------------------------------------------
+
+
+def _make_headless_args(**overrides: Any) -> argparse.Namespace:
+    """Build an argparse.Namespace shaped like the headless CLI passes in.
+
+    Defaults pass every validation gate so individual tests can mutate just
+    the field they're exercising.
+    """
+    defaults = dict(
         model="fake-model",
-        stage_id=3,
+        stage_id=0,
         replica_id=0,
         omni_master_address="127.0.0.1",
         omni_master_port=26000,
+        omni_replica_address=None,
+        omni_dp_size_local=1,
         api_server_count=0,
         worker_backend="multi_process",
         stage_configs_path=None,
+        deploy_config=None,
         log_stats=False,
         disable_log_stats=False,
         stage_init_timeout=600,
+        tokenizer=None,
     )
+    defaults.update(overrides)
+    return argparse.Namespace(**defaults)
 
 
-def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
+def test_run_headless_requires_stage_id() -> None:
+    args = _make_headless_args(stage_id=None)
+    with pytest.raises(ValueError, match="--stage-id is required"):
+        run_headless(args)
 
+
+def test_run_headless_requires_master_address() -> None:
+    args = _make_headless_args(omni_master_address=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_requires_master_port() -> None:
+    args = _make_headless_args(omni_master_port=None)
+    with pytest.raises(ValueError, match="--omni-master-address and --omni-master-port"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_multi_api_server_count() -> None:
+    args = _make_headless_args(api_server_count=2)
+    with pytest.raises(ValueError, match="api_server_count can't be set"):
+        run_headless(args)
+
+
+def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
+    args = _make_headless_args(worker_backend="ray")
+    with pytest.raises(ValueError, match="worker_backend=multi_process"):
+        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", stage_cfgs),
+        return_value=("/fake/stages.yaml", [other_stage]),
     )
-    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=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mock_build_vllm_config = mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mock_register = mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
 
-    mock_build_vllm_config.assert_called_once_with(
-        stage_cfg,
-        "fake-model",
-        stage_connector_spec={},
-        engine_args_dict={},
-        headless=True,
-    )
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        coordinator=None,
-        replica_id=0,
-    )
-    mock_manager_cls.assert_called_once()
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["local_engine_count"] == 2
-    assert manager_kwargs["start_index"] == 4
-    assert manager_kwargs["local_start_index"] == 0
-    assert manager_kwargs["local_client"] is False
-    assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
-    assert manager_kwargs["log_stats"] is False
-    engine_manager.monitor_engine_liveness.assert_called_once_with()
-    engine_manager.shutdown.assert_called_once_with()
+    args = _make_headless_args(stage_id=0)
+    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
+        run_headless(args)
 
 
-def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.log_stats = True
-    stage_cfg = mocker.Mock(stage_id=3)
-    stage_cfgs = [stage_cfg]
-    parallel_config = mocker.Mock(
-        data_parallel_size_local=2,
-        data_parallel_rank=4,
-        data_parallel_rank_local=1,
-        node_rank_within_dp=0,
-    )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    executor_class = mocker.Mock()
-    engine_manager = mocker.Mock()
-
-    mocker.patch(
-        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
-    )
-    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=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
-    mocker.patch(
-        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
-        return_value=(None, None, None),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, executor_class),
-    )
-    mocker.patch(
-        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
-    )
-    mock_manager_cls = mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
-    mocker.patch("signal.signal")
-    run_headless(args)
-
-    manager_kwargs = mock_manager_cls.call_args.kwargs
-    assert manager_kwargs["log_stats"] is True
+# ---------------------------------------------------------------------------
+# run_headless happy paths
+# ---------------------------------------------------------------------------
 
 
-def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 2
-    stage_cfg = mocker.Mock(stage_id=3)
-    parallel_config = mocker.Mock(
+def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace:
+    """Build a stage config that satisfies every attribute run_headless reads.
+
+    Notably ``engine_args`` is a real dict (not a Mock) so
+    ``get_stage_devices_per_replica`` can call ``.get("tensor_parallel_size")``
+    and feed the result through ``int()`` without TypeError.
+    """
+    return SimpleNamespace(
+        stage_id=stage_id,
+        stage_type=stage_type,
+        # No "devices" key -> split_devices_for_replicas skipped, each replica
+        # inherits the launcher's CUDA_VISIBLE_DEVICES.
+        runtime=None,
+        engine_args={},
+    )
+
+
+def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None:
+    """LLM headless: each loop iteration registers with auto-assigned
+    replica_id (master picks a free slot) and spawns one
+    ``OmniCoreEngineProcManager`` per local replica."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
         data_parallel_size_local=1,
         data_parallel_rank=0,
         data_parallel_rank_local=0,
         node_rank_within_dp=0,
     )
-    vllm_config = mocker.Mock(parallel_config=parallel_config)
-    vllm_config.needs_dp_coordinator = False
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
     engine_manager = mocker.Mock()
 
     mocker.patch(
@@ -175,82 +152,250 @@ def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> 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=mocker.Mock())
-    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
-    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
     mocker.patch(
         "vllm_omni.engine.stage_init_utils.build_vllm_config",
-        return_value=(vllm_config, mocker.Mock()),
+        return_value=(vllm_config, object),
     )
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value="tcp://127.0.0.1:26001",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address="tcp://127.0.0.1:26100",
+        ),
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        return_value=engine_manager,
     )
-    mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
     mocker.patch("signal.signal")
 
-    run_headless(args)
+    run_headless(_make_headless_args(stage_id=0))
 
-    assert mock_register.call_args.kwargs["replica_id"] == 2
+    # The launcher must request auto-assignment (replica_id=None) and the
+    # full response so it can wire the master-allocated coordinator into the
+    # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM
+    # because the head binds all three sockets (handshake, input, output).
+    assert mock_register.call_count == 1
+    kwargs = mock_register.call_args.kwargs
+    assert kwargs["omni_master_address"] == "127.0.0.1"
+    assert kwargs["omni_master_port"] == 26000
+    assert kwargs["omni_stage_id"] == 0
+    assert kwargs["omni_stage_config"] is stage_cfg
+    assert kwargs["replica_id"] is None
+    assert kwargs["return_full_response"] is True
+    assert kwargs["replica_binds_sockets"] is False
+
+    assert mock_manager_cls.call_count == 1
+    mgr_kwargs = mock_manager_cls.call_args.kwargs
+    assert mgr_kwargs["local_engine_count"] == 1
+    assert mgr_kwargs["local_client"] is False
+    assert mgr_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert mgr_kwargs["omni_stage_id"] == 0
+    assert mgr_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert mgr_kwargs["omni_replica_base_id"] == 0
+
+    engine_manager.monitor_engine_liveness.assert_called_once_with()
+    engine_manager.shutdown.assert_called_once_with()
 
 
-def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None:
-    args = _make_headless_args()
-    args.replica_id = 1
-    stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion")
-    stage_cfg.engine_args = mocker.Mock()
-    stage_cfg.engine_input_source = []
-    stage_cfgs = [stage_cfg]
-    metadata = mocker.Mock(stage_id=3, stage_type="diffusion")
+def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: MockerFixture) -> None:
+    """``--omni-dp-size-local=N`` must spawn N managers, each with its own
+    master-assigned replica_id, and join all of them before returning."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(0, stage_type="llm")
+    parallel_config = SimpleNamespace(
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+        node_rank_within_dp=0,
+    )
+    vllm_config = SimpleNamespace(parallel_config=parallel_config, needs_dp_coordinator=False)
+    manager_a = mocker.Mock()
+    manager_b = mocker.Mock()
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        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)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.build_vllm_config",
+        return_value=(vllm_config, object),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        side_effect=[
+            StageRegistrationResponse(
+                handshake_address=f"tcp://127.0.0.1:2700{idx}",
+                input_address=f"tcp://127.0.0.1:2710{idx}",
+                output_address=f"tcp://127.0.0.1:2720{idx}",
+                replica_id=idx,
+                coordinator_router_address=None,
+            )
+            for idx in (0, 1)
+        ],
+    )
+    mock_manager_cls = mocker.patch(
+        "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager",
+        side_effect=[manager_a, manager_b],
+    )
+    mocker.patch("signal.signal")
+
+    run_headless(_make_headless_args(stage_id=0, omni_dp_size_local=2))
+
+    assert mock_manager_cls.call_count == 2
+    assigned_ids = [call.kwargs["omni_replica_base_id"] for call in mock_manager_cls.call_args_list]
+    assert assigned_ids == [0, 1]
+
+    # Multi-replica path joins the monitor threads instead of calling
+    # ``monitor_engine_liveness`` synchronously on the main thread, but every
+    # manager must still be shut down in the finally block.
+    manager_a.shutdown.assert_called_once_with()
+    manager_b.shutdown.assert_called_once_with()
+
+
+def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) -> None:
+    """Diffusion headless: registers as auto-assign, spawns a single
+    ``StageDiffusionProc`` per local replica, and waits for it via
+    ``multiprocessing.connection.wait``."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
     od_config = mocker.Mock()
-    proc = mocker.Mock()
-    proc.exitcode = 0
+    proc = mocker.Mock(sentinel=object(), exitcode=0)
     proc.is_alive.return_value = False
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", stage_cfgs),
+        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=mocker.Mock())
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
     mocker.patch(
         "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
         return_value=(None, None, None),
     )
-    mocker.patch("vllm_omni.engine.stage_init_utils.extract_stage_metadata", return_value=metadata)
-    mock_inject_stage_info = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mock_inject = mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
     mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=od_config)
     mock_register = mocker.patch(
         "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
-        return_value=("tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address="tcp://127.0.0.1:26100",
+        ),
     )
     mock_spawn = mocker.patch(
         "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
-        return_value=(proc, "tcp://127.0.0.1:26001", "tcp://127.0.0.1:26002", "tcp://127.0.0.1:26003"),
+        return_value=(proc, None, None, None),
     )
     mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
+    # Replace the blocking wait with one that returns the only proc's sentinel
+    # immediately so the test does not hang.
+    mocker.patch(
+        "multiprocessing.connection.wait",
+        side_effect=lambda sentinels: [sentinels[0]],
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
     mocker.patch("signal.signal")
-    run_headless(args)
 
-    mock_inject_stage_info.assert_called_once_with(stage_cfg, 3)
-    mock_register.assert_called_once_with(
-        omni_master_address="127.0.0.1",
-        omni_master_port=26000,
-        omni_stage_id=3,
-        omni_stage_config=stage_cfg,
-        return_addresses=True,
-        replica_id=1,
-    )
-    mock_spawn.assert_called_once_with(
-        "fake-model",
-        od_config,
-        handshake_address="tcp://127.0.0.1:26001",
-        request_address="tcp://127.0.0.1:26002",
-        response_address="tcp://127.0.0.1:26003",
-    )
+    run_headless(_make_headless_args(stage_id=1))
+
+    mock_inject.assert_called_once()
+    assert mock_inject.call_args.args[0] is stage_cfg
+    assert mock_inject.call_args.args[1] == 1
+    assert mock_inject.call_args.args[2] == [stage_cfg]
+
+    reg_kwargs = mock_register.call_args.kwargs
+    assert reg_kwargs["omni_master_address"] == "127.0.0.1"
+    assert reg_kwargs["omni_master_port"] == 26000
+    assert reg_kwargs["omni_stage_id"] == 1
+    assert reg_kwargs["omni_stage_config"] is stage_cfg
+    assert reg_kwargs["replica_id"] is None
+    assert reg_kwargs["return_full_response"] is True
+
+    spawn_kwargs = mock_spawn.call_args.kwargs
+    assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
+    assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002"
+    assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003"
+    assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100"
+    assert spawn_kwargs["omni_stage_id"] == 1
+    assert spawn_kwargs["omni_replica_id"] == 0
+
     mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600)
-    proc.join.assert_called_once_with()
+
+
+def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None:
+    """A diffusion replica that exits with a non-zero code must surface as a
+    RuntimeError from ``run_headless`` (the head needs the signal to roll
+    back its own stage init)."""
+    from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
+
+    stage_cfg = _make_stage_cfg(1, stage_type="diffusion")
+    proc = mocker.Mock(sentinel=object(), exitcode=137, name="proc-stage1-rep0")
+    proc.is_alive.return_value = False
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        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)
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.extract_stage_metadata",
+        return_value=SimpleNamespace(stage_id=1, stage_type="diffusion"),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.inject_kv_stage_info")
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_diffusion_config", return_value=mocker.Mock())
+    mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        return_value=StageRegistrationResponse(
+            handshake_address="tcp://127.0.0.1:26001",
+            input_address="tcp://127.0.0.1:26002",
+            output_address="tcp://127.0.0.1:26003",
+            replica_id=0,
+            coordinator_router_address=None,
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc",
+        return_value=(proc, None, None, None),
+    )
+    mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake")
+    mocker.patch(
+        "multiprocessing.connection.wait",
+        side_effect=lambda sentinels: [sentinels[0]],
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc")
+    mocker.patch("signal.signal")
+
+    with pytest.raises(RuntimeError, match=r"exited with code 137"):
+        run_headless(_make_headless_args(stage_id=1))

tests/helpers/assertions.py

@@ -1,12 +1,13 @@
 """Assertion and response validation helpers for tests."""
 
 import io
+import json
 import tempfile
 import threading
 import wave
 from io import BytesIO
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 if TYPE_CHECKING:
     from tests.helpers.runtime import DiffusionResponse
@@ -495,62 +496,12 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                 )
 
 
-def _speech_combined_error_text_for_assert(response: Any) -> str:
-    """Join ``error_message`` and UTF-8 body (e.g. HTTP 200 + JSON) for substring checks."""
-    parts: list[str] = []
-    em = getattr(response, "error_message", None)
-    if em:
-        parts.append(str(em))
-    ab = getattr(response, "audio_bytes", None)
-    if ab:
-        parts.append(ab.decode("utf-8", errors="replace"))
-    return "\n".join(parts)
-
-
-def _assert_speech_error_substrings(haystack: str, err_expected: Any) -> None:
-    if isinstance(err_expected, (list, tuple)):
-        assert any(str(s) in haystack for s in err_expected), (
-            f"Expected error text to contain one of {err_expected!r}, got: {haystack!r}"
-        )
-    else:
-        assert str(err_expected) in haystack, f"Expected error text to contain {str(err_expected)!r}, got: {haystack!r}"
-
-
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
-    """Validate speech API results.
+    """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.
 
-    If ``status_code`` is set: assert HTTP status, then optional ``err_message`` (4xx: ``error_message`` only;
-    HTTP 200: ``error_message`` and/or body, e.g. JSON with HTTP 200).
-
-    If only ``err_message`` is set: assert substring(s) on combined error text. Otherwise success/min_audio/etc.
+    Success path only. Use :meth:`OpenAIClientHandler.send_audio_speech_http_request` with ``err_code`` /
+    ``err_message`` for expected HTTP errors.
     """
-    # When present, only compare HTTP status and skip all other checks (e.g. expected 4xx).
-    # ``status_code`` may be a single int or a collection of allowed codes (e.g. ``(400, 422)``).
-    expected_status = request_config.get("status_code")
-    err_expected = request_config.get("err_message")
-    if expected_status is not None:
-        actual = getattr(response, "status_code", None)
-        assert actual is not None, (
-            "request_config has status_code but response has no .status_code "
-            "(set OmniResponse.status_code or pass a response object with status_code)"
-        )
-        if isinstance(expected_status, (list, tuple, set, frozenset)):
-            assert actual in expected_status, f"Expected HTTP status in {expected_status!r}, got {actual}"
-        else:
-            assert actual == int(expected_status), f"Expected HTTP status {int(expected_status)}, got {actual}"
-
-        if err_expected is not None:
-            if actual is not None and actual != 200:
-                msg = (getattr(response, "error_message", None) or "") + ""
-                _assert_speech_error_substrings(msg, err_expected)
-            else:
-                _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected)
-        return
-
-    if err_expected is not None:
-        _assert_speech_error_substrings(_speech_combined_error_text_for_assert(response), err_expected)
-        return
-
     assert response.success, "The request failed."
 
     # Optional floor on decoded audio size (models with very short clips may use a lower value).
@@ -599,10 +550,128 @@ def assert_diffusion_response(response: "DiffusionResponse", request_config: dic
         assert_audio_diffusion_response(response=response, request_config=request_config, run_level=run_level)
 
 
+def _http_response_body_materialize(resp: Any) -> tuple[bytes, dict[str, Any] | None]:
+    """Serialize ``HttpResponse``-like body to UTF-8 bytes and parse a JSON object when possible."""
+    jb = getattr(resp, "json_body", None)
+    if jb is not None:
+        raw = json.dumps(jb, ensure_ascii=False).encode("utf-8")
+        if isinstance(jb, dict):
+            return raw, jb
+        try:
+            parsed = json.loads(raw.decode("utf-8", errors="replace"))
+        except json.JSONDecodeError:
+            return raw, None
+        return raw, parsed if isinstance(parsed, dict) else None
+    err = getattr(resp, "error_message", None)
+    raw = (err or "").encode("utf-8", errors="replace")
+    try:
+        parsed = json.loads(raw.decode("utf-8", errors="replace"))
+    except json.JSONDecodeError:
+        return raw, None
+    return raw, parsed if isinstance(parsed, dict) else None
+
+
+def assert_err_message_in_text(
+    haystack: str,
+    err_message: str | tuple[str, ...] | list[str] | set[str] | frozenset[str],
+    *,
+    sequence_match: Literal["all", "any"] = "all",
+) -> None:
+    """Assert ``err_message`` appears in ``haystack`` (case-insensitive).
+
+    For non-string sequences: ``sequence_match='all'`` requires every substring (HTTP-style);
+    ``sequence_match='any'`` requires at least one (WebSocket first-frame JSON helpers).
+    """
+    hl = haystack.lower()
+    if isinstance(err_message, (list, tuple, set, frozenset)):
+        if sequence_match == "all":
+            missing = [s for s in err_message if str(s).lower() not in hl]
+            assert not missing, (
+                f"Expected error text to contain all of {err_message!r}; missing {missing!r}. haystack={haystack!r}"
+            )
+        else:
+            assert any(str(s).lower() in hl for s in err_message), (
+                f"Expected error text to contain one of {err_message!r}. haystack={haystack!r}"
+            )
+    else:
+        assert str(err_message).lower() in hl, f"Expected error text to contain {str(err_message)!r}, got: {haystack!r}"
+
+
+def assert_http_error(
+    resp: Any,
+    *,
+    err_code: int | tuple[int, ...] | list[int] | None = None,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+    websocket_json_message: bool = False,
+) -> dict[str, Any] | None:
+    """Validate a raw-HTTP :class:`~tests.helpers.runtime.HttpResponse`-like object.
+
+    Used by :class:`~tests.helpers.runtime.OpenAIClientHandler` ``send_*_http_request`` helpers when
+    ``request_config`` contains optional ``err_code`` and/or ``err_message``.
+
+    When ``websocket_json_message=True``, only ``json_body`` is checked (first JSON WebSocket text frame).
+    Tuple/list ``err_message`` then uses **any** substring match; HTTP mode still requires **all** pieces.
+
+    - ``err_code``: exact HTTP ``int``, or membership if a non-string sequence (e.g. ``(400, 422)``).
+      When ``err_code`` is set and the actual status is a client error in ``400..499`` other than ``404``,
+      the JSON body must include FastAPI ``detail`` and/or OpenAI-style ``error`` (2xx skips this).
+    - ``err_message``: substring match (case-insensitive) against serialized ``json_body`` and ``error_message``.
+      If ``err_message`` is a non-string sequence (``list`` / ``tuple`` / ``set`` / ``frozenset``), **every**
+      element must appear as a substring; a plain ``str`` still requires that single substring.
+    """
+    if websocket_json_message:
+        if err_code is not None:
+            raise ValueError("assert_http_error: err_code is incompatible with websocket_json_message=True")
+        jb_ws = getattr(resp, "json_body", None)
+        assert jb_ws is not None, resp
+        if err_message is None:
+            return jb_ws if isinstance(jb_ws, dict) else None
+        assert_err_message_in_text(
+            json.dumps(jb_ws, ensure_ascii=False),
+            err_message,
+            sequence_match="any",
+        )
+        return jb_ws if isinstance(jb_ws, dict) else None
+
+    if err_code is None and err_message is None:
+        return None
+
+    actual = getattr(resp, "status_code", None)
+    assert actual is not None, "response missing status_code"
+
+    if err_code is not None:
+        if isinstance(err_code, int):
+            assert actual == err_code, (resp, err_code)
+        else:
+            allowed = tuple(err_code)
+            assert actual in allowed, (resp, allowed)
+
+    body_bytes, payload = _http_response_body_materialize(resp)
+
+    if err_code is not None and actual is not None and 400 <= actual < 500 and actual != 404:
+        assert payload is not None, getattr(resp, "error_message", resp)
+        assert "detail" in payload or "error" in payload, payload
+
+    if err_message is not None:
+        pieces: list[str] = []
+        jb = getattr(resp, "json_body", None)
+        if jb is not None:
+            pieces.append(json.dumps(jb, ensure_ascii=False))
+        em = getattr(resp, "error_message", None)
+        if em:
+            pieces.append(str(em))
+        haystack = "\n".join(pieces) if pieces else body_bytes.decode("utf-8", errors="replace")
+        assert_err_message_in_text(haystack, err_message, sequence_match="all")
+
+    return payload
+
+
 __all__ = [
     "assert_audio_diffusion_response",
     "assert_audio_speech_response",
     "assert_diffusion_response",
+    "assert_err_message_in_text",
+    "assert_http_error",
     "assert_image_diffusion_response",
     "assert_image_valid",
     "assert_omni_response",

tests/helpers/runtime.py

@@ -1,5 +1,6 @@
 """Server/client/runner runtime primitives for tests."""
 
+import asyncio
 import base64
 import concurrent.futures
 import copy
@@ -19,6 +20,7 @@ from dataclasses import asdict, dataclass
 from io import BytesIO
 from pathlib import Path
 from typing import Any, NamedTuple
+from urllib.parse import quote
 
 import psutil
 import requests
@@ -37,6 +39,7 @@ from vllm.logger import init_logger
 from tests.helpers.assertions import (
     assert_audio_speech_response,
     assert_diffusion_response,
+    assert_http_error,
     assert_omni_response,
 )
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
@@ -86,6 +89,16 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
             logger.warning("torch._C._host_emptyCache() only available in Pytorch >=2.5")
 
 
+def _parse_response_json(r: requests.Response) -> dict[str, Any] | list[Any] | None:
+    try:
+        data = r.json()
+        if isinstance(data, (dict, list)):
+            return data
+    except Exception:
+        pass
+    return None
+
+
 def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[str, Any]] | None:
     """If ``extra_body`` has list ``height``/``width``, return one config per index (scalar h/w, ``num_outputs_per_prompt=1``)."""
     eb = cfg.get("extra_body")
@@ -647,6 +660,8 @@ class OmniServerStageCli(OmniServer):
 
 @dataclass
 class OmniResponse:
+    """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
+
     text_content: str | None = None
     audio_data: list[str] | None = None
     audio_content: str | None = None
@@ -656,16 +671,15 @@ class OmniResponse:
     #: OpenAI client call through response parsing and local post-process (e.g. audio decode).
     e2e_latency: float | None = None
     success: bool = False
-    error_message: str | None = None
     prompt_tokens: int | None = None
     cached_tokens: int | None = None
-    # Set for HTTP-level assertions (e.g. ``assert_audio_speech_response`` with ``status_code`` in config).
-    status_code: int | None = None
     logprobs: list | None = None
 
 
 @dataclass
 class DiffusionResponse:
+    """Decoded diffusion output from chat completions or offline runner (not raw ``requests``)."""
+
     text_content: str | None = None
     images: list[Image.Image] | None = None
     audios: list[Any] | None = None
@@ -674,7 +688,67 @@ class DiffusionResponse:
     #: ``chat.completions.create`` through local image / audio decode.
     e2e_latency: float | None = None
     success: bool = False
+
+
+@dataclass
+class HttpResponse:
+    """Normalized view of a ``requests`` response from :class:`OpenAIClientHandler` HTTP helpers."""
+
+    status_code: int
+    success: bool
     error_message: str | None = None
+    json_body: dict[str, Any] | list[Any] | None = None
+
+
+@dataclass
+class WebSocketJsonResponse:
+    """First JSON object delivered as a text WebSocket frame (streaming endpoints)."""
+
+    json_body: dict[str, Any]
+
+
+def _merge_http_expectation_kwargs(
+    base: dict[str, Any] | None,
+    *,
+    err_code: int | tuple[int, ...] | list[int] | None = None,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+) -> dict[str, Any]:
+    cfg = dict(base or {})
+    if err_code is not None:
+        cfg["err_code"] = err_code
+    if err_message is not None:
+        cfg["err_message"] = err_message
+    return cfg
+
+
+def _merge_ws_expectation_kwargs(
+    base: dict[str, Any] | None,
+    *,
+    err_message: str | tuple[str, ...] | list[str] | None = None,
+    ws_json_type: str | None = None,
+    ws_error_code: str | None = None,
+) -> dict[str, Any]:
+    cfg = dict(base or {})
+    if err_message is not None:
+        cfg["err_message"] = err_message
+    if ws_json_type is not None:
+        cfg["ws_json_type"] = ws_json_type
+    if ws_error_code is not None:
+        cfg["ws_error_code"] = ws_error_code
+    return cfg
+
+
+def _run_ws_expectations_from_request_config(cfg: dict[str, Any], resp: WebSocketJsonResponse) -> None:
+    jb = resp.json_body
+    want_type = cfg.get("ws_json_type")
+    if want_type is not None:
+        assert jb.get("type") == want_type, (jb, want_type)
+    want_code = cfg.get("ws_error_code")
+    if want_code is not None:
+        assert jb.get("code") == want_code, (jb, want_code)
+    err_message = cfg.get("err_message")
+    if err_message is not None:
+        assert_http_error(resp, err_message=err_message, websocket_json_message=True)
 
 
 def _merge_diffusion_responses(parts: list[DiffusionResponse]) -> DiffusionResponse:
@@ -739,8 +813,8 @@ class OpenAIClientHandler:
             result.e2e_latency = time.perf_counter() - wall_start
             result.success = True
         except Exception as e:
-            result.error_message = f"Stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Stream processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
     def _process_non_stream_omni_response(self, chat_completion, *, wall_start: float) -> OmniResponse:
@@ -771,8 +845,8 @@ class OpenAIClientHandler:
                 result.logprobs = chat_completion.choices[0].logprobs.content
             result.success = True
         except Exception as e:
-            result.error_message = f"Non-stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Non-stream processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
     def _process_diffusion_response(self, chat_completion, *, wall_start: float) -> DiffusionResponse:
@@ -811,11 +885,592 @@ class OpenAIClientHandler:
             result.e2e_latency = time.perf_counter() - wall_start
             result.success = True
         except Exception as e:
-            result.error_message = f"Diffusion response processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Diffusion response processing error: {str(e)}"
+            print(f"Error: {msg}")
         return result
 
+    def _http_response_from_requests(self, r: requests.Response) -> HttpResponse:
+        payload = _parse_response_json(r)
+        ok = 200 <= r.status_code < 300
+        return HttpResponse(
+            status_code=r.status_code,
+            success=ok,
+            error_message=None if ok else (r.text[:8000] if r.text else None),
+            json_body=payload,
+        )
+
+    def send_health_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/health`` (raw ``requests``).
+
+        ``request_config``: optional ``timeout`` plus optional ``err_code`` / ``err_message`` for
+        :func:`~tests.helpers.assertions.assert_http_error` (also as keyword-only args).
+        """
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(self._build_url("/health"), timeout=float(cfg.get("timeout", 120.0)))
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_models_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/models``. Optional ``timeout`` and HTTP assertions (see :func:`~tests.helpers.assertions.assert_http_error`)."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/models"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_chat_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/chat/completions`` with ``json`` or ``raw_body`` (malformed-body / contract tests)."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/chat/completions", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_omni_sleep_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/omni/sleep`` — ``json`` or ``raw_body``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/omni/sleep", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_omni_wakeup_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/omni/wakeup``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/omni/wakeup", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_list_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/audio/voices``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/audio/voices"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_create_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/voices`` (multipart): ``data`` / ``files`` / ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/audio/voices", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_voices_delete_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """DELETE ``/v1/audio/voices/{name}`` — requires ``name``, optional ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        name = cfg["name"]
+        timeout = float(cfg.get("timeout", 120.0))
+        path = f"/v1/audio/voices/{quote(str(name), safe='')}"
+        r = requests.delete(
+            self._build_url(path),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_speech_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/speech`` with ``json`` or ``raw_body``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/speech", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_speech_batch_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/speech/batch``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/speech/batch", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_audio_generate_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/audio/generate``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/audio/generate", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_images_generations_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/images/generations`` — ``json`` or ``raw_body``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/images/generations", cfg, default_timeout=300.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_images_edits_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/images/edits`` — ``data`` / ``files`` / ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/images/edits", cfg, default_timeout=300.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_create_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/videos`` (async job) — multipart ``data`` / ``files``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/videos", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_sync_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/videos/sync``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_form_endpoint("/v1/videos/sync", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_videos_list_http_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos`` — optional ``params``, ``timeout``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = requests.get(
+            self._build_url("/v1/videos"),
+            params=cfg.get("params"),
+            headers={"Accept": "application/json"},
+            timeout=float(cfg.get("timeout", 120.0)),
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_retrieve_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos/{video_id}``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.get(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_delete_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """DELETE ``/v1/videos/{video_id}``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.delete(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}"),
+            headers={"Accept": "application/json"},
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def send_video_content_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """GET ``/v1/videos/{video_id}/content``."""
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        video_id = cfg["video_id"]
+        timeout = float(cfg.get("timeout", 120.0))
+        r = requests.get(
+            self._build_url(f"/v1/videos/{quote(str(video_id), safe='')}/content"),
+            timeout=timeout,
+        )
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
+    def _build_ws_url(self, path: str) -> str:
+        """Turn HTTP ``base_url`` into ``ws`` / ``wss`` for WebSocket helpers."""
+        base = self.base_url.rstrip("/")
+        suffix = "/" + path.lstrip("/")
+        if base.startswith("http://"):
+            return "ws://" + base.removeprefix("http://") + suffix
+        if base.startswith("https://"):
+            return "wss://" + base.removeprefix("https://") + suffix
+        raise ValueError(f"Unsupported base_url for WebSocket: {base!r}")
+
+    def _send_websocket_first_json_request(
+        self,
+        path: str,
+        cfg: dict[str, Any],
+    ) -> list[WebSocketJsonResponse]:
+        """Connect, optionally send text frames, return first JSON text frame as :class:`WebSocketJsonResponse`.
+
+        ``request_config`` keys:
+
+        - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
+          speaks first, e.g. ``/v1/realtime`` rejection path).
+        - ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
+        - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
+        """
+        send_frames_raw = cfg.get("send_frames")
+        if send_frames_raw is None:
+            frames: list[str] = []
+        elif isinstance(send_frames_raw, str):
+            frames = [send_frames_raw]
+        else:
+            frames = list(send_frames_raw)
+
+        timeout = float(cfg.get("timeout", 120.0))
+        uri = self._build_ws_url(path)
+
+        connect_kw: dict[str, Any] = {}
+        if "ws_max_size" in cfg:
+            connect_kw["max_size"] = cfg["ws_max_size"]
+
+        async def _recv_first_json_object() -> WebSocketJsonResponse:
+            import websockets
+
+            async with websockets.connect(uri, **connect_kw) as ws:
+                for frame in frames:
+                    await ws.send(frame)
+                raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
+                if not isinstance(raw, str):
+                    raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
+                try:
+                    data = json.loads(raw)
+                except json.JSONDecodeError as exc:
+                    raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
+                if not isinstance(data, dict):
+                    raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
+                return WebSocketJsonResponse(json_body=data)
+
+        resp = asyncio.run(_recv_first_json_object())
+        _run_ws_expectations_from_request_config(cfg, resp)
+        return [resp]
+
+    def send_audio_speech_stream_ws_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/audio/speech/stream`` — send ``send_frames`` then read first JSON text frame."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/audio/speech/stream", cfg)
+
+    def send_video_chat_stream_ws_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/video/chat/stream`` — send ``send_frames`` then read first JSON text frame."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/video/chat/stream", cfg)
+
+    def send_realtime_ws_request(
+        self,
+        request_config: dict[str, Any] | None = None,
+        *,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+        ws_json_type: str | None = None,
+        ws_error_code: str | None = None,
+    ) -> list[WebSocketJsonResponse]:
+        """WebSocket ``/v1/realtime`` — optional outbound frames, then first JSON text frame (often server-initiated)."""
+        cfg = _merge_ws_expectation_kwargs(
+            request_config,
+            err_message=err_message,
+            ws_json_type=ws_json_type,
+            ws_error_code=ws_error_code,
+        )
+        return self._send_websocket_first_json_request("/v1/realtime", cfg)
+
     def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+        """Chat completions via the OpenAI Python SDK (not raw HTTP)."""
         responses: list[OmniResponse] = []
         stream = request_config.get("stream", False)
         modalities = request_config.get("modalities", ["text", "audio"])
@@ -939,8 +1594,8 @@ class OpenAIClientHandler:
                 result.audio_format = result.audio_format.headers.get("content-type", "")
 
         except Exception as e:
-            result.error_message = f"Audio speech stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Audio speech stream processing error: {str(e)}"
+            print(f"Error: {msg}")
 
         return result
 
@@ -978,8 +1633,8 @@ class OpenAIClientHandler:
                 result.audio_format = result.audio_format.headers.get("content-type", "")
 
         except Exception as e:
-            result.error_message = f"Audio speech non-stream processing error: {str(e)}"
-            print(f"Error: {result.error_message}")
+            msg = f"Audio speech non-stream processing error: {str(e)}"
+            print(f"Error: {msg}")
 
         return result
 
@@ -1201,6 +1856,7 @@ class OpenAIClientHandler:
         Send OpenAI requests for diffusion models.
         If ``extra_body`` has list ``height``/``width``, sends one chat completion per index in parallel
         (scalar h/w, ``num_outputs_per_prompt=1`` each) and merges images in list order.
+
         Args:
             request_config: A single request configuration dict, or a list of
                 request configuration dicts (one request per element)
@@ -1300,10 +1956,13 @@ class OpenAIClientHandler:
         self, request_config: dict[str, Any], request_num: int = 1
     ) -> list[DiffusionResponse]:
         """
-        Send native /v1/videos requests.
+        Send native /v1/videos requests: multipart ``form_data`` job create, poll until done, download content.
+
+        For raw HTTP to video routes without polling, use ``send_videos_create_http_request``, etc.
         """
         if request_num != 1:
             raise NotImplementedError("Concurrent video diffusion requests are not currently implemented")
+
         form_data = request_config.get("form_data")
         if not isinstance(form_data, dict):
             raise ValueError("Video request_config must contain 'form_data'")
@@ -1346,6 +2005,54 @@ class OpenAIClientHandler:
             self._print_client_stat("[diffusion] request#1 completed")
         return [result]
 
+    def _post_json_endpoint(
+        self,
+        path: str,
+        request_config: dict[str, Any],
+        *,
+        default_timeout: float,
+    ) -> requests.Response:
+        url = self._build_url(path)
+        timeout = float(request_config.get("timeout", default_timeout))
+        if "raw_body" in request_config:
+            raw = request_config["raw_body"]
+            payload = raw.encode("utf-8") if isinstance(raw, str) else raw
+            return requests.post(
+                url,
+                data=payload,
+                headers={"Content-Type": "application/json", "Accept": "application/json"},
+                timeout=timeout,
+            )
+        if "json" not in request_config:
+            raise ValueError(f"{path} request_config must include 'json' or 'raw_body'")
+        return requests.post(
+            url,
+            json=request_config["json"],
+            headers={"Accept": "application/json", "Content-Type": "application/json"},
+            timeout=timeout,
+        )
+
+    def _post_form_endpoint(
+        self,
+        path: str,
+        request_config: dict[str, Any],
+        *,
+        default_timeout: float = 120.0,
+    ) -> requests.Response:
+        url = self._build_url(path)
+        timeout = float(request_config.get("timeout", default_timeout))
+        data = request_config.get("data")
+        files = request_config.get("files")
+        if data is None and not files:
+            data = {}
+        return requests.post(
+            url,
+            data=data,
+            files=files,
+            headers={"Accept": "application/json"} if not files else {"Accept": "application/json"},
+            timeout=timeout,
+        )
+
     def _wait_until_video_completed(
         self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300
     ) -> None:
@@ -1693,9 +2400,9 @@ class OmniRunnerHandler:
             result.text_content = text_content
             result.success = True
         except Exception as e:
-            result.error_message = f"Output processing error: {str(e)}"
+            msg = f"Output processing error: {str(e)}"
             result.success = False
-            print(f"Error: {result.error_message}")
+            print(f"Error: {msg}")
         return result
 
     def _process_diffusion_output(self, outputs: list[OmniRequestOutput]) -> DiffusionResponse:
@@ -1831,15 +2538,11 @@ class OmniRunnerHandler:
                 mm_out = stage_out.request_output.outputs[0].multimodal_output
                 break
         if mm_out is None:
-            result = OmniResponse(success=False, error_message="No audio output from pipeline")
-            assert result.success, result.error_message
-            return result
+            raise AssertionError("No audio output from pipeline")
 
         audio_data = mm_out.get("audio")
         if audio_data is None:
-            result = OmniResponse(success=False, error_message="No audio tensor in multimodal output")
-            assert result.success, result.error_message
-            return result
+            raise AssertionError("No audio tensor in multimodal output")
 
         sr_raw = mm_out.get("sr")
         sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
@@ -2004,6 +2707,8 @@ def iter_omni_runner(
 
 __all__ = [
     "DiffusionResponse",
+    "HttpResponse",
+    "WebSocketJsonResponse",
     "OmniResponse",
     "OmniRunner",
     "OmniRunnerHandler",

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -403,3 +403,57 @@ class TestCodePredictorWrapperConfig:
         assert config.use_parallel_embedding is False
         assert config.return_proj_buf is False
         assert config.sampling_mode == "per_call"
+
+    def test_prefix_graph_config_helpers(self, loaded_target_classes) -> None:
+        """Prefix graph helpers parse deploy config values and keep valid seq lens only."""
+        _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        wrapper_cls = common_mod.CodePredictorWrapper
+
+        assert wrapper_cls._parse_positive_int_set("64; 128,0,-1") == {
+            64,
+            128,
+        }
+        assert wrapper_cls._parse_positive_int_set([2, "4", 0]) == {2, 4}
+        with pytest.raises(ValueError, match="Invalid positive int config value 'bad'"):
+            wrapper_cls._parse_positive_int_set("2,bad")
+
+        wrapper = object.__new__(wrapper_cls)
+        wrapper._prefix_graph_seq_lens = {1, 2, 4, 8, 99}
+        assert wrapper._prefix_seq_lens(6) == [2, 4]
+
+    def test_prefix_graph_env_requires_cuda_graphs(
+        self,
+        mocker: MockerFixture,
+        loaded_target_classes,
+    ) -> None:
+        """Avoid prefix warmup on shared code-predictor users that disable CUDA graphs."""
+        _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+
+        cp_config, _ = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
+        vllm_config.model_config.stage_connector_config = {
+            "extra": {
+                "code_predictor_prefix_graphs": True,
+                "code_predictor_prefix_graph_buckets": [2],
+                "code_predictor_prefix_graph_seq_lens": "2,3",
+            }
+        }
+
+        no_graph_wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+        )
+        assert no_graph_wrapper._prefix_graphs_enabled is False
+        assert no_graph_wrapper._prefix_graph_buckets == {2}
+        assert no_graph_wrapper._prefix_graph_seq_lens == {2, 3}
+
+        graph_wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
+        )
+        assert graph_wrapper._prefix_graphs_enabled is True

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -74,6 +74,14 @@ class SyntheticDecoder(nn.Module):
         return self.out(x).clamp(min=-1, max=1)
 
 
+class ShortOutputDecoder(SyntheticDecoder):
+    """Decoder variant that returns fewer samples than seq_len * total_upsample."""
+
+    def forward(self, codes):
+        out = super().forward(codes)
+        return out[..., :-5] if out.shape[-1] > 5 else out
+
+
 @pytest.fixture(scope="module")
 def decoder():
     """Create a synthetic decoder on CUDA with fixed weights."""
@@ -87,6 +95,7 @@ def wrapper(decoder):
     w = CUDAGraphDecoderWrapper(
         decoder=decoder,
         capture_sizes=[25, 50, 100],
+        capture_batch_sizes=[1, 2],
         num_quantizers=NUM_QUANTIZERS,
         enabled=True,
     )
@@ -94,8 +103,8 @@ def wrapper(decoder):
     return w
 
 
-def _random_codes(seq_len, device=DEVICE):
-    return torch.randint(0, 100, (1, NUM_QUANTIZERS, seq_len), dtype=torch.long, device=device)
+def _random_codes(seq_len, batch_size=1, device=DEVICE):
+    return torch.randint(0, 100, (batch_size, NUM_QUANTIZERS, seq_len), dtype=torch.long, device=device)
 
 
 # ──────────────────────────────────────────────────────────────────
@@ -213,6 +222,73 @@ def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len):
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
+def test_chunked_decode_output_survives_later_replay(wrapper):
+    """Chunked output must not alias graph static buffers overwritten by later replays."""
+    codes = _random_codes(100)
+    overwrite_codes = _random_codes(100)
+
+    with torch.no_grad():
+        graph_out = wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0)
+        expected = graph_out.clone()
+        _ = wrapper.decode(overwrite_codes[..., :50])
+        _ = wrapper.decode(overwrite_codes)
+
+    torch.testing.assert_close(graph_out, expected, atol=0, rtol=0)
+
+
+def test_chunked_decode_preserves_short_chunk_concat_semantics():
+    """Chunked decode must allow shorter-than-nominal chunk outputs.
+
+    Qwen3-Omni non-async startup can hit a decoder path where a chunk returns
+    fewer waveform samples than ``code_len * total_upsample``. The wrapper must
+    preserve the original eager concat behavior instead of copying into exact
+    nominal slices.
+    """
+    torch.manual_seed(123)
+    short_decoder = ShortOutputDecoder().to(DEVICE).eval()
+    short_wrapper = CUDAGraphDecoderWrapper(
+        decoder=short_decoder,
+        capture_sizes=[50],
+        capture_batch_sizes=[1],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    short_wrapper.warmup(DEVICE)
+    codes = _random_codes(100)
+
+    with torch.no_grad():
+        eager_out = _eager_chunked(short_decoder, codes, chunk_size=50, left_context_size=0)
+        graph_out = short_wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0)
+
+    assert graph_out.shape == eager_out.shape
+    assert graph_out.shape[-1] < 100 * TOTAL_UPSAMPLE
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper):
+    """Variable-length chunk batching should match independent chunked decodes."""
+    long_codes = _random_codes(100)
+    short_codes = _random_codes(50)
+    padded_codes = torch.zeros(2, NUM_QUANTIZERS, 100, dtype=torch.long, device=DEVICE)
+    padded_codes[0, :, :] = long_codes[0]
+    padded_codes[1, :, :50] = short_codes[0]
+
+    with torch.no_grad():
+        eager_long = _eager_chunked(decoder, long_codes, chunk_size=50, left_context_size=0)
+        eager_short = _eager_chunked(decoder, short_codes, chunk_size=50, left_context_size=0)
+        graph_out = wrapper.batched_chunked_decode_with_cudagraph(
+            padded_codes,
+            [100, 50],
+            chunk_size=50,
+            left_context_size=0,
+            max_batch_size=2,
+        )
+
+    assert graph_out.shape == (2, 1, 100 * TOTAL_UPSAMPLE)
+    torch.testing.assert_close(graph_out[0:1], eager_long, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(graph_out[1:2, :, : 50 * TOTAL_UPSAMPLE], eager_short, atol=1e-6, rtol=1e-6)
+
+
 def _eager_chunked(decoder, codes, chunk_size, left_context_size):
     """Eager chunked decode matching the real decoder's chunked_decode logic."""
     wavs = []
@@ -254,15 +330,95 @@ def test_disabled_wrapper_matches_eager(decoder, wrapper):
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
-def test_batch_size_gt1_falls_back(decoder, wrapper):
-    """Batch size > 1 should fall back to eager (bit-identical)."""
-    codes = torch.randint(0, 100, (2, NUM_QUANTIZERS, 25), dtype=torch.long, device=DEVICE)
+def test_batch_size_gt1_uses_matching_graph(decoder, wrapper):
+    """Captured batch size > 1 should replay a matching graph."""
+    assert (2, 25) in wrapper.graphs
+    codes = _random_codes(25, batch_size=2)
     with torch.no_grad():
         eager_out = decoder(codes)
         graph_out = wrapper.decode(codes)
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
+def test_uncaptured_batch_size_falls_back(decoder, wrapper):
+    """Uncaptured batch sizes should fall back to eager."""
+    assert (3, 25) not in wrapper.graphs
+    codes = _random_codes(25, batch_size=3)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_extra_capture_shape_uses_sparse_graph(decoder):
+    """Extra capture shapes should not expand to a full batch x size product."""
+    sparse_wrapper = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25],
+        capture_batch_sizes=[1],
+        extra_capture_shapes=[(2, 50)],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    sparse_wrapper.warmup(DEVICE)
+
+    assert (1, 25) in sparse_wrapper.graphs
+    assert (2, 50) in sparse_wrapper.graphs
+    assert (2, 25) not in sparse_wrapper.graphs
+
+    codes = _random_codes(50, batch_size=2)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = sparse_wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch):
+    """Configured torch.compile shapes should replay exact and padded CUDA Graph buckets."""
+
+    compile_kwargs = {}
+
+    def _fake_compile(model, **_kwargs):
+        compile_kwargs.update(_kwargs)
+
+        def _compiled(codes):
+            return model(codes) + 0.125
+
+        return _compiled
+
+    monkeypatch.setattr(torch, "compile", _fake_compile)
+
+    compiled_wrapper = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25, 50],
+        capture_batch_sizes=[1],
+        compile_shapes=[(1, 25), (1, 50)],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    compiled_wrapper.warmup(DEVICE)
+
+    exact_codes = _random_codes(25)
+    padded_codes = _random_codes(30)
+    uncaptured_codes = _random_codes(60)
+    padded_static = torch.zeros(1, NUM_QUANTIZERS, 50, dtype=torch.long, device=DEVICE)
+    padded_static[:, :, :30] = padded_codes
+    with torch.no_grad():
+        exact_eager = decoder(exact_codes)
+        exact_out = compiled_wrapper.decode(exact_codes)
+        padded_graph_expected = decoder(padded_static)[..., : 30 * TOTAL_UPSAMPLE]
+        padded_out = compiled_wrapper.decode(padded_codes)
+        uncaptured_eager = decoder(uncaptured_codes)
+        uncaptured_out = compiled_wrapper.decode(uncaptured_codes)
+
+    torch.testing.assert_close(exact_out, exact_eager + 0.125, atol=0, rtol=0)
+    torch.testing.assert_close(padded_out, padded_graph_expected + 0.125, atol=0, rtol=0)
+    torch.testing.assert_close(uncaptured_out, uncaptured_eager, atol=0, rtol=0)
+    assert compile_kwargs["mode"] == "default"
+    assert compile_kwargs["fullgraph"] is False
+    assert compile_kwargs["dynamic"] is False
+
+
 def test_deterministic_across_calls(decoder, wrapper):
     """Same input should produce identical CUDA graph output across calls."""
     codes = _random_codes(30)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -24,6 +24,7 @@ class _FakeDecoder(nn.Module):
         super().__init__()
         self.total_upsample = total_upsample
         self.decode_calls: list[dict[str, int]] = []
+        self.batched_decode_calls: list[dict[str, int]] = []
         self.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
     def to(self, *args, **kwargs):
@@ -40,12 +41,40 @@ class _FakeDecoder(nn.Module):
             {
                 "chunk_size": chunk_size,
                 "left_context_size": left_context_size,
+                "codes_shape": tuple(codes.shape),
             }
         )
+        batch = codes.shape[0]
         frames = codes.shape[-1]
         wav_len = frames * self.total_upsample + 6
-        wav = torch.arange(wav_len, dtype=torch.float32)
-        return wav.view(1, 1, -1)
+        wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1)
+        offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000
+        return wav.expand(batch, 1, wav_len) + offsets
+
+    def batched_chunked_decode(
+        self,
+        codes: torch.Tensor,
+        lengths: list[int],
+        *,
+        chunk_size: int = 300,
+        left_context_size: int = 25,
+        max_batch_size: int = 0,
+    ) -> torch.Tensor:
+        self.batched_decode_calls.append(
+            {
+                "chunk_size": chunk_size,
+                "left_context_size": left_context_size,
+                "max_batch_size": max_batch_size,
+                "codes_shape": tuple(codes.shape),
+                "lengths": tuple(lengths),
+            }
+        )
+        batch = codes.shape[0]
+        frames = codes.shape[-1]
+        wav_len = frames * self.total_upsample + 6
+        wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1)
+        offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000
+        return wav.expand(batch, 1, wav_len) + offsets
 
     def enable_cudagraph(self, **kwargs):
         self.cudagraph_calls.append(kwargs)
@@ -177,6 +206,7 @@ def test_connector_codec_chunking_does_not_override_decode_chunking():
     assert model.decoder.decode_calls[-1] == {
         "chunk_size": 300,
         "left_context_size": 25,
+        "codes_shape": (1, _NUM_QUANTIZERS, 6),
     }
 
 
@@ -199,6 +229,175 @@ def test_decode_chunking_can_be_overridden_separately():
     assert model._decode_left_context_frames == 17
 
 
+def test_malformed_codec_length_warning_is_rate_limited():
+    model = _make_model()
+
+    with patch("vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.logger.warning") as warning:
+        out1 = model.forward(input_ids=torch.arange(1, dtype=torch.long))
+        out2 = model.forward(input_ids=torch.arange(1, dtype=torch.long))
+
+    assert warning.call_count == 1
+    assert "not divisible by num_quantizers" in warning.call_args[0][0]
+    assert "suppressing repeats" in warning.call_args[0][0]
+    assert out1.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert out2.multimodal_outputs["model_outputs"][0].numel() == 0
+
+
+def test_forward_batches_equal_length_requests_in_one_decoder_call():
+    model = _make_model()
+
+    out = model.forward(
+        input_ids=torch.arange(12, dtype=torch.long),
+        seq_token_counts=[6, 6],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 1}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (2, _NUM_QUANTIZERS, 3),
+        }
+    ]
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(12, dtype=torch.float32))
+    torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
+
+
+def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group():
+    model = _make_model()
+    model._decode_batch_bucket_frames = [4]
+    model._decode_variable_chunk_batch_min_frames = 1
+
+    out = model.forward(
+        input_ids=torch.arange(12, dtype=torch.long),
+        seq_token_counts=[4, 8],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 2}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == []
+    assert model.decoder.batched_decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "max_batch_size": 0,
+            "codes_shape": (2, _NUM_QUANTIZERS, 4),
+            "lengths": (2, 4),
+        }
+    ]
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
+    torch.testing.assert_close(audios[1], torch.arange(1008, 1016, dtype=torch.float32))
+
+
+def test_forward_bucket_batches_different_length_requests_and_trims_rows():
+    model = _make_model()
+    model._decode_batch_bucket_frames = [4]
+
+    out = model.forward(
+        input_ids=torch.arange(18, dtype=torch.long),
+        seq_token_counts=[4, 6, 8],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 1}},
+            {"meta": {"left_context_size": 2}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (3, _NUM_QUANTIZERS, 4),
+        }
+    ]
+    assert model.decoder.batched_decode_calls == []
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
+    torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32))
+    torch.testing.assert_close(audios[2], torch.arange(2008, 2016, dtype=torch.float32))
+
+
+def test_forward_bucket_pads_only_to_group_max_frame_length():
+    model = _make_model()
+    model._decode_batch_bucket_frames = [8]
+
+    out = model.forward(
+        input_ids=torch.arange(10, dtype=torch.long),
+        seq_token_counts=[4, 6],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 0}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (2, _NUM_QUANTIZERS, 3),
+        }
+    ]
+    assert model.decoder.batched_decode_calls == []
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32))
+    torch.testing.assert_close(audios[1], torch.arange(1000, 1012, dtype=torch.float32))
+
+
+def test_forward_splits_bucket_groups_by_configured_max_batch_size():
+    model = _make_model()
+    model._decode_batch_bucket_frames = [4]
+    model._decode_batch_max_size = 2
+
+    model.forward(
+        input_ids=torch.arange(18, dtype=torch.long),
+        seq_token_counts=[4, 6, 8],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 0}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (2, _NUM_QUANTIZERS, 3),
+        },
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (1, _NUM_QUANTIZERS, 4),
+        },
+    ]
+    assert model.decoder.batched_decode_calls == []
+
+
+def test_forward_does_not_pad_singleton_bucket_group():
+    model = _make_model()
+    model._decode_batch_bucket_frames = [4]
+
+    model.forward(
+        input_ids=torch.arange(4, dtype=torch.long),
+        runtime_additional_information=[{"meta": {"left_context_size": 0}}],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (1, _NUM_QUANTIZERS, 2),
+        }
+    ]
+
+
 def test_decode_chunking_override_is_passed_to_cudagraph():
     model = _make_model(
         async_chunk=True,
@@ -216,6 +415,10 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
     _load_weights_noop(model)
 
     assert model.decoder.cudagraph_calls[-1] == {
+        "capture_sizes": None,
+        "capture_batch_sizes": None,
+        "extra_capture_shapes": None,
+        "compile_shapes": None,
         "device": torch.device("cuda"),
         "codec_chunk_frames": 25,
         "codec_left_context_frames": 72,
@@ -224,6 +427,106 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
     }
 
 
+def test_cudagraph_capture_shapes_can_be_configured():
+    model = _make_model(
+        async_chunk=True,
+        device=torch.device("cuda"),
+        stage_connector_config={
+            "extra": {
+                "decode_cudagraph_capture_sizes": "97,325",
+                "decode_cudagraph_batch_sizes": [1, 2, 4, 8],
+                "decode_cudagraph_extra_capture_shapes": ["3:325", [5, 325]],
+            }
+        },
+    )
+
+    _load_weights_noop(model)
+
+    call = model.decoder.cudagraph_calls[-1]
+    assert call["capture_sizes"] == [97, 325]
+    assert call["capture_batch_sizes"] == [1, 2, 4, 8]
+    assert call["extra_capture_shapes"] == [(3, 325), (5, 325)]
+
+
+def test_decode_compile_shapes_can_be_configured():
+    model = _make_model(
+        async_chunk=True,
+        device=torch.device("cuda"),
+        stage_connector_config={
+            "extra": {
+                "decode_compile_shapes": ["1:325", [1, 73]],
+            }
+        },
+    )
+
+    _load_weights_noop(model)
+
+    call = model.decoder.cudagraph_calls[-1]
+    assert call["compile_shapes"] == [(1, 73), (1, 325)]
+
+
+def test_decode_tf32_can_be_configured():
+    old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32
+    old_cudnn_tf32 = torch.backends.cudnn.allow_tf32
+    old_matmul_precision = torch.get_float32_matmul_precision()
+    try:
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        torch.set_float32_matmul_precision("highest")
+        model = _make_model(
+            async_chunk=True,
+            device=torch.device("cuda"),
+            stage_connector_config={
+                "extra": {
+                    "decode_enable_tf32": "true",
+                }
+            },
+        )
+
+        _load_weights_noop(model)
+
+        assert torch.backends.cuda.matmul.allow_tf32 is True
+        assert torch.backends.cudnn.allow_tf32 is True
+        assert torch.get_float32_matmul_precision() == "high"
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32
+        torch.backends.cudnn.allow_tf32 = old_cudnn_tf32
+        torch.set_float32_matmul_precision(old_matmul_precision)
+
+
+def test_decode_batch_bucket_frames_can_be_configured():
+    model = _make_model(
+        async_chunk=True,
+        stage_connector_config={
+            "extra": {
+                "decode_batch_bucket_frames": "73,169",
+                "decode_batch_max_size": 10,
+                "decode_variable_chunk_batch_min_frames": 512,
+            }
+        },
+    )
+
+    _load_weights_noop(model)
+
+    assert model._decode_batch_bucket_frames == [73, 169]
+    assert model._decode_batch_max_size == 10
+    assert model._decode_variable_chunk_batch_min_frames == 512
+
+
+def test_invalid_decode_batch_max_size_is_rejected():
+    model = _make_model(
+        async_chunk=True,
+        stage_connector_config={
+            "extra": {
+                "decode_batch_max_size": -1,
+            }
+        },
+    )
+
+    with pytest.raises(ValueError, match="decode_batch_max_size"):
+        _load_weights_noop(model)
+
+
 def test_invalid_decode_chunking_is_rejected():
     model = _make_model(
         async_chunk=True,

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -0,0 +1,491 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    _NORMALIZED_REF_AUDIO_KEY,
+    _PRECOMPUTED_REF_CODE_KEY,
+    _PRECOMPUTED_REF_IDS_KEY,
+    _PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+
+
+def _make_minimal_talker():
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    return model
+
+
+def test_single_token_prefill_uses_prefill_path():
+    model = _make_minimal_talker()
+    full_prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    trailing_text = torch.ones((2, 4), dtype=torch.float32)
+    tts_pad = torch.full((1, 4), 0.5, dtype=torch.float32)
+    ref_code = torch.arange(32, dtype=torch.long).reshape(2, 16)
+
+    def fake_build_prompt_embeds(*, task_type, info_dict):
+        return full_prompt_embeds, trailing_text, tts_pad, 2, ref_code
+
+    model._build_prompt_embeds = fake_build_prompt_embeds
+
+    input_ids = torch.tensor([123], dtype=torch.long)
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=input_ids,
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        _omni_is_prefill=True,
+        _omni_num_computed_tokens=0,
+        _omni_prompt_len=3,
+    )
+
+    assert out_ids.tolist() == [7]
+    assert torch.equal(out_embeds.cpu(), full_prompt_embeds[:1].to(torch.bfloat16))
+    assert update["meta"]["talker_prefill_offset"] == 1
+    assert update["meta"]["talker_text_offset"] == 0
+    assert update["meta"]["ref_code_len"] == 2
+    assert torch.equal(update["embed"]["prefill"], full_prompt_embeds)
+    assert torch.equal(update["embed"]["tts_pad"], tts_pad)
+    assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text)
+    assert torch.equal(update["codes"]["ref"], ref_code)
+    assert update["codes"]["audio"].shape == (1, 16)
+
+
+def test_single_token_prefill_can_be_inferred_from_token_progress():
+    model = _make_minimal_talker()
+    full_prompt_embeds = torch.arange(8, dtype=torch.float32).reshape(2, 4)
+    trailing_text = torch.ones((1, 4), dtype=torch.float32)
+    tts_pad = torch.zeros((1, 4), dtype=torch.float32)
+
+    def fake_build_prompt_embeds(*, task_type, info_dict):
+        return full_prompt_embeds, trailing_text, tts_pad, None, None
+
+    model._build_prompt_embeds = fake_build_prompt_embeds
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        _omni_num_computed_tokens=0,
+        _omni_prompt_len=2,
+    )
+
+    assert out_ids.tolist() == [7]
+    assert torch.equal(out_embeds.cpu(), full_prompt_embeds[:1].to(torch.bfloat16))
+    assert update["meta"]["talker_prefill_offset"] == 1
+
+
+def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
+        meta={"talker_text_offset": 1},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2,
+        _omni_prompt_len=2,
+    )
+
+    assert out_ids.tolist() == [123]
+    assert torch.equal(out_embeds.cpu(), torch.full((1, 4), 123.0, dtype=torch.bfloat16))
+    assert "hidden_states" not in update
+    assert update["meta"]["talker_text_offset"] == 2
+    past_hidden, text_step = update["mtp_inputs"]
+    assert torch.equal(past_hidden.cpu(), last_hidden.reshape(1, -1).to(torch.bfloat16))
+    assert torch.equal(text_step.cpu(), trailing_text[1:2].to(torch.bfloat16))
+
+
+def test_decode_advances_trailing_text_offset_across_multiple_steps():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(8, dtype=torch.float32).reshape(2, 4)
+    state_tail = trailing_text
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+    meta = {"talker_text_offset": 0}
+    seen_steps = []
+
+    for _ in range(3):
+        _, _, update = model.preprocess(
+            input_ids=torch.tensor([123], dtype=torch.long),
+            input_embeds=None,
+            text=["hello"],
+            task_type=["CustomVoice"],
+            hidden_states={"trailing_text": state_tail, "last": last_hidden},
+            embed={"tts_pad": tts_pad},
+            meta=meta,
+            _omni_is_prefill=False,
+            _omni_num_computed_tokens=2,
+            _omni_prompt_len=2,
+        )
+        seen_steps.append(update["mtp_inputs"][1].cpu())
+        if "hidden_states" in update and "trailing_text" in update["hidden_states"]:
+            state_tail = update["hidden_states"]["trailing_text"]
+        meta = update["meta"]
+
+    assert torch.equal(seen_steps[0], trailing_text[0:1].to(torch.bfloat16))
+    assert torch.equal(seen_steps[1], trailing_text[1:2].to(torch.bfloat16))
+    assert torch.equal(seen_steps[2], tts_pad.to(torch.bfloat16))
+    assert meta["talker_text_offset"] == 0
+    assert state_tail.numel() == 0
+
+
+def test_decode_compacts_long_trailing_text_after_large_offset():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(130 * 4, dtype=torch.float32).reshape(130, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    _, _, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
+        meta={"talker_text_offset": 64},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2,
+        _omni_prompt_len=2,
+    )
+
+    assert torch.equal(update["mtp_inputs"][1].cpu(), trailing_text[64:65].to(torch.bfloat16))
+    assert update["meta"]["talker_text_offset"] == 0
+    assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
+
+
+def test_decode_batch_preprocess_matches_decode_state_updates():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_a = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    trailing_b = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100
+    last_a = torch.full((4,), 2.0, dtype=torch.float32)
+    last_b = torch.full((4,), 3.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
+        input_ids=torch.tensor([101, 202], dtype=torch.long),
+        req_infos=[
+            {
+                "text": ["hello"],
+                "task_type": ["Base"],
+                "hidden_states": {"trailing_text": trailing_a, "last": last_a},
+                "embed": {"tts_pad": tts_pad},
+                "meta": {"talker_text_offset": 1},
+            },
+            {
+                "text": ["world"],
+                "task_type": ["CustomVoice"],
+                "hidden_states": {"trailing_text": trailing_b, "last": last_b},
+                "embed": {"tts_pad": tts_pad},
+                "meta": {"talker_text_offset": 2},
+            },
+        ],
+    )
+
+    assert out_ids.tolist() == [101, 202]
+    assert torch.equal(out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16))
+    assert torch.equal(past_hidden.cpu(), torch.stack([last_a, last_b]).to(torch.bfloat16))
+    assert torch.equal(text_step[0].cpu(), trailing_a[1].to(torch.bfloat16))
+    assert torch.equal(text_step[1].cpu(), tts_pad.reshape(-1).to(torch.bfloat16))
+    assert updates[0]["meta"]["talker_text_offset"] == 2
+    assert updates[0]["meta"]["codec_streaming"] is True
+    assert "hidden_states" not in updates[0]
+    assert updates[1]["meta"]["talker_text_offset"] == 0
+    assert updates[1]["meta"]["codec_streaming"] is False
+    assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
+
+
+def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
+    model = _make_minimal_talker()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+
+    normalize_calls = []
+    ref_audio = np.arange(1024, dtype=np.float32)
+    model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
+
+    ref_audio_ids = []
+    model._encode_ref_audio_to_code = lambda wav, _sr: (
+        ref_audio_ids.append(id(wav)) or torch.ones((2, 2), dtype=torch.long)
+    )
+    model._extract_speaker_embedding = lambda wav, _sr: (
+        ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
+    )
+
+    _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_audio": ["ref.wav"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+        },
+    )
+
+    assert normalize_calls == ["ref.wav"]
+    assert ref_audio_ids == [id(ref_audio), id(ref_audio)]
+    assert ref_code_len == 2
+    assert torch.equal(ref_code, torch.ones((2, 2), dtype=torch.long))
+
+
+def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
+    model = _make_minimal_talker()
+    wav1 = np.arange(2048, dtype=np.float32)
+    wav2 = np.arange(3072, dtype=np.float32)
+    normalize_calls = []
+    model._normalize_ref_audio = lambda raw: (
+        normalize_calls.append(raw)
+        or (
+            wav1 if raw == "a.wav" else wav2,
+            16000,
+        )
+    )
+
+    class FakeSpeechTokenizer:
+        def __init__(self):
+            self.calls = []
+
+        def encode(self, audios, *, sr, return_dict):
+            self.calls.append((audios, sr, return_dict))
+            return SimpleNamespace(
+                audio_codes=[
+                    torch.full((2, 2), 11, dtype=torch.long),
+                    torch.full((3, 2), 22, dtype=torch.long),
+                ]
+            )
+
+    tok = FakeSpeechTokenizer()
+    model._ensure_speech_tokenizer_loaded = lambda: tok
+
+    class FakeTextTokenizer:
+        def __init__(self):
+            self.calls = []
+
+        def __call__(self, texts, *, padding=False):
+            self.calls.append((texts, padding))
+            return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
+
+    text_tok = FakeTextTokenizer()
+    model._get_tokenizer = lambda: text_tok
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+        },
+        "r2": {
+            "task_type": ["Base"],
+            "text": ["two"],
+            "ref_audio": ["b.wav"],
+            "ref_text": ["world"],
+            "x_vector_only_mode": [False],
+        },
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1", "r2"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert normalize_calls == ["a.wav", "b.wav"]
+    assert len(tok.calls) == 1
+    audios, sr, return_dict = tok.calls[0]
+    assert audios[0] is wav1
+    assert audios[1] is wav2
+    assert sr == 16000
+    assert return_dict is True
+    assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
+    assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1
+    assert buf["r2"][_NORMALIZED_REF_AUDIO_KEY][0] is wav2
+    assert len(text_tok.calls) == 2
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+
+
+def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
+    model = _make_minimal_talker()
+    wav = np.arange(2048, dtype=np.float32)
+    model._normalize_ref_audio = lambda raw: (wav, 16000)
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+        AssertionError("singleton should not load speech tokenizer")
+    )
+
+    class FakeTextTokenizer:
+        def __call__(self, texts, *, padding=False):
+            return {"input_ids": [[7, 8, 9] for _ in texts]}
+
+    model._get_tokenizer = lambda: FakeTextTokenizer()
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+        }
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
+
+
+def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exists():
+    model = _make_minimal_talker()
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    model._get_tokenizer = lambda: (_ for _ in ()).throw(AssertionError("tokenizer not expected"))
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+        AssertionError("speech tokenizer not expected")
+    )
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+            "embed": {"prefill": torch.ones((1, 4))},
+        }
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert _PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
+    assert _NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+
+
+def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
+    model = _make_minimal_talker()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+
+    ref_audio = np.arange(2048, dtype=np.float32)
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
+    model._encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("serial encode not expected")
+    )
+    speaker_wav_ids = []
+    model._extract_speaker_embedding = lambda wav, _sr: (
+        speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
+    )
+
+    _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_audio": ["ref.wav"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+            "codes": {_PRECOMPUTED_REF_CODE_KEY: ref_code},
+            _NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
+        },
+    )
+
+    assert speaker_wav_ids == [id(ref_audio)]
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -294,6 +294,36 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
     assert len(payload.codes.audio) == _Q * (35 + 2)
 
 
+def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
+    tm = _tm(chunk_frames=4, left_context=3, initial_chunk_frames=4)
+    rid = "r-ref-bounded"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(8)]
+    ref_code = torch.tensor(
+        [
+            [1, 1, 1, 1],
+            [2, 2, 2, 2],
+            [3, 3, 3, 3],
+            [4, 4, 4, 4],
+            [5, 5, 5, 5],
+        ],
+        dtype=torch.long,
+    )
+    tm.request_payload[rid] = ref_code
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        request=_req(rid, finished=False),
+        is_finished=False,
+    )
+
+    assert payload is not None
+    assert payload.meta.left_context_size == 3 + 3
+    assert len(payload.codes.audio) == _Q * (3 + 3 + 4)
+    frames = payload.codes.audio.reshape(_Q, -1).transpose(0, 1)
+    torch.testing.assert_close(frames[:3], ref_code[-3:])
+
+
 def test_ref_code_context_can_be_buffered_before_first_emit():
     tm = _tm()
     rid = "r-ref-buffered"

tests/test_generate_nightly_perf_excel.py

@@ -35,6 +35,8 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path):
 
     omni_record = json.loads(omni_template_path.read_text(encoding="utf-8"))
     diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8"))
+    diffusion_records[0]["endpoint"] = "/v1/videos"
+    diffusion_records[0]["result"]["endpoint"] = "/v1/videos"
 
     input_dir = tmp_path / "input"
     diffusion_input_dir = tmp_path / "diffusion_input"
@@ -69,3 +71,6 @@ def test_generate_excel_report_with_perf_templates(tmp_path: Path):
     ws_omni_summary = wb["omni_summary"]
     assert _cell_value_by_header(ws_omni_summary, "commit_sha") == "test_commit_sha"
     assert _cell_value_by_header(ws_omni_summary, "build_id") == "test_build_id"
+
+    ws_diff_summary = wb["diffusion_summary"]
+    assert _cell_value_by_header(ws_diff_summary, "endpoint") == "/v1/videos"