[Docs] Redirect the former PD guide URL

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Docs] Document current PD design
2026-08-13 12:48:43 +08:00 · 2026-08-13 12:44:42 +08:00 · 2026-08-13 12:23:56 +08:00 · 2026-08-13 11:23:47 +08:00 · 2026-08-13 09:36:17 +08:00 · 2026-08-13 08:42:37 +08:00
40 changed files with 1069 additions and 441 deletions
@@ -345,7 +345,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -115,7 +115,7 @@ steps:
          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -11,10 +11,15 @@ build:
    post_checkout:
      - git fetch --unshallow || true
    pre_install:
+      # `pre_install` runs before RTD upgrades pip, and the seeded pip 23.1
+      # rejects the `typing_extensions` wheel torch needs (fixed in pip 23.2).
+      - python -m pip install --upgrade pip
      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
-      # 1 GB memory cap and time out the mkdocs build.
-      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
+      # 1 GB memory cap and time out the mkdocs build. Wheel-only: `--index-url`
+      # replaces PyPI, so an sdist fallback cannot reach its build deps. The
+      # quotes keep YAML from reading `:all:` as a mapping key.
+      - "python -m pip install --only-binary=:all: --index-url https://download.pytorch.org/whl/cpu torch"

 mkdocs:
  configuration: mkdocs.yml
@@ -167,9 +167,11 @@ def main() -> None:
    ]
    native_cmd = [
        sys.executable,
-        "examples/offline_inference/text_to_video/text_to_video_lingbot.py",
+        "examples/offline_inference/text_to_video/text_to_video.py",
        "--model",
        args.model,
+        "--model-class-name",
+        "LingBotVideoPipeline",
        "--prompt",
        args.prompt,
        "--output",
@@ -43,6 +43,7 @@ nav:
    - Overview: features/README.md
    - Runtime and Stage Execution:
      - Execution Modes and Streaming: user_guide/diffusion/execution_modes.md
+      - Prefill-Decode Disaggregation (experimental): features/pd_disaggregation.md
      - Sleep Mode: features/sleep_mode.md
    - Quantization:
      - Overview: user_guide/quantization/overview.md
@@ -14,10 +14,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m

 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance

-## Multi-Stage Recipes
-
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
-
 ## Optimization Features

 - **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
@@ -1,171 +0,0 @@
-# Prefill-Decode (PD) Disaggregation
-
-PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
-stages so prompt processing and token generation can run on different workers.
-
-This is documented as a stage-config recipe instead of a bundled YAML because the
-deployment-specific values usually change per environment:
-
- GPU placement
- `tensor_parallel_size`
- connector backend and connector ports
- connector IPs or bootstrap addresses
-
-Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
-and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
-changes below.
-
-## Requirements
-
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
- A KV connector supported by vLLM, such as `MooncakeConnector`
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
-
-## 1. Split the thinker into prefill and decode stages
-
-Replace the original thinker stage with two stages:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_prefill_only: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_producer"
-        kv_rank: 0
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25201
-    final_output: false
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: llm
-    is_decode_only: true
-    runtime:
-      devices: "1"
-    engine_args:
-      max_num_seqs: 64
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_consumer"
-        kv_rank: 1
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25202
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-```
-
-Notes:
-
- `is_prefill_only: true` marks the thinker stage that only saves KV.
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
- `kv_transfer_config` is required on both stages.
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
-  prefill side only processes the prompt and exports KV.
-
-## 2. Shift the downstream stages by one index
-
-After inserting the extra thinker stage, renumber the remaining stages:
-
-```yaml
-  - stage_id: 2
-    runtime:
-      devices: "2"
-    engine_input_source: [1]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
-
-  - stage_id: 3
-    runtime:
-      devices: "2"
-    engine_args:
-      max_num_seqs: 1
-    engine_input_source: [2]
-    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
-```
-
-Compared with the default Qwen3-Omni config:
-
- the talker becomes stage `2` instead of stage `1`
- the code2wav stage becomes stage `3` instead of stage `2`
- the talker now reads from decode stage `1`
-
-## 3. Add runtime edges for the four-stage pipeline
-
-```yaml
-runtime:
-  enabled: true
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2
-    - from: 2
-      to: 3
-```
-
-## 4. Launch with your custom config
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --deploy-config /path/to/qwen3_omni_pd.yaml
-```
-
-## Operational Notes
-
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
-  pair. Keep prefill and decode at the same `tensor_parallel_size`.
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
-  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
- Choose connector ports and addresses that match your deployment. The values
-  shown above are examples only.
@@ -20,13 +20,19 @@ an implementation contract has a user-facing workflow.
 | Choose serial, batched, step-wise, or streaming diffusion execution | [Execution Modes and Streaming](../user_guide/diffusion/execution_modes.md) | [Diffusion Continuous Batching](../design/feature/diffusion_continuous_batching.md), [Async Diffusion Output](../design/feature/async_diffusion_output.md) |
 | Reclaim stage memory without restarting the server | [Sleep Mode](sleep_mode.md) | Runtime lifecycle behavior is documented in the user guide |

+The [Prefill-Decode Disaggregation design reference](pd_disaggregation.md)
+describes the experimental Qwen3-Omni runtime topology and its current
+integration boundaries. It is not yet a supported launch recipe.
+
 Some runtime designs are deliberately not promoted as standalone User Guide
 features yet:

- [Disaggregated Inference](../design/feature/disaggregated_inference.md) and
+- Generic, OmniConnector-based
+  [Disaggregated Inference](../design/feature/disaggregated_inference.md) is not
+  yet a standalone User Guide feature. Its
  [OmniConnector implementations](../design/index.md#communication) describe
-  topology and transport contracts. Practical deployment configuration remains
-  under [Pipeline and deploy configurations](../configuration/stage_configs.md).
+  transport contracts, while practical connector configuration remains under
+  [Pipeline and deploy configurations](../configuration/stage_configs.md).
 - [Async Chunk](../design/feature/async_chunk.md) and
  [Async Omni Output Materialization](../design/feature/omni_async_output_materialization.md)
  are model- and pipeline-dependent. Use the selected model's deploy
@@ -0,0 +1,140 @@
+# Prefill-Decode (PD) Disaggregation (experimental)
+
+!!! warning "Experimental design reference"
+    PD disaggregation is only partially integrated on the current `main` branch.
+    The runtime contains PD detection and routing scaffolding, but there is no
+    supported deploy overlay or end-to-end validated launch recipe yet. Treat this
+    page as a description of the current design, not as production guidance.
+
+Prefill-decode disaggregation splits the Qwen3-Omni Thinker into two logical
+stages: a prefill worker that processes the prompt and produces KV cache, and a
+decode worker that imports that cache and generates tokens. The Talker and
+Code2Wav stages remain downstream; PD does not disaggregate the entire
+Qwen3-Omni pipeline.
+
+## Current design
+
+### Runtime topology
+
+```mermaid
+flowchart LR
+    client["Client request"] --> orchestrator["Omni Orchestrator"]
+    orchestrator -->|"original prompt<br/>max_tokens = 1"| prefill["Thinker prefill<br/>KV producer"]
+    prefill -->|"KV blocks"| connector[("vLLM KV connector")]
+    prefill -->|"KV routing metadata"| orchestrator
+    orchestrator -->|"original prompt<br/>remote KV parameters"| decode["Thinker decode<br/>KV consumer"]
+    connector -->|"remote KV load"| decode
+    decode -->|"generated text"| client
+```
+
+The orchestrator owns the request lifecycle, while the vLLM KV connector owns
+the bulk KV transfer. The original prompt is submitted to both Thinker stages:
+the decode worker still needs the request tokens and multimodal features to
+construct its request, while the imported KV cache prevents it from repeating
+the completed prefill computation.
+
+For the full Qwen3-Omni pipeline, splitting the Thinker adds one logical stage:
+
+| Stage | Role | Output or handoff |
+| --- | --- | --- |
+| 0 | Thinker prefill | Saves prompt KV and prompt-side multimodal output |
+| 1 | Thinker decode | Loads remote KV, generates text, and produces decode-side conditioning |
+| 2 | Talker | Converts combined Thinker conditioning into codec codes |
+| 3 | Code2Wav | Converts codec codes into audio |
+
+Text remains a final output of the decode stage, and audio remains a final
+output of Code2Wav.
+
+### Request lifecycle
+
+1. At startup, PD detection looks for one stage marked `is_prefill_only` and
+   one downstream stage marked `is_decode_only`. The decode stage must list the
+   prefill stage as an input source.
+2. The orchestrator clones the Thinker sampling parameters for prefill, forces
+   `max_tokens=1`, clears stop conditions, and sets producer-side
+   `kv_transfer_params`. This makes prefill finish by length after exporting KV.
+3. When prefill finishes, the orchestrator captures its KV routing metadata and
+   Qwen multimodal output under the original request ID.
+4. The orchestrator submits the original prompt to decode with consumer-side
+   `kv_transfer_params`. Those parameters identify the transfer, prefill engine,
+   bootstrap endpoint, and, when required by the connector, remote request.
+5. The decode engine imports the KV cache and continues generation. The normal
+   Omni stage path then carries Thinker conditioning through Talker and
+   Code2Wav.
+
+### Control and data paths
+
+PD uses two different connector layers. They solve different handoffs and are
+not interchangeable.
+
+| Path | Payload | Owner |
+| --- | --- | --- |
+| PD control path | Request ID, original prompt, `transfer_id`, bootstrap address, engine ID, and connector-specific request metadata | Omni Orchestrator |
+| PD data path | Thinker attention KV blocks | vLLM KV connector |
+| Downstream model path | Thinker hidden states and embeddings, Talker codec codes, and audio tensors | vLLM-Omni stage connectors and processors |
+
+### Detection and validation
+
+The current runtime accepts one logical prefill/decode pair per pipeline. Each
+stage may use tensor parallelism, but both sides of the pair must satisfy these
+invariants:
+
+- both stages define `kv_transfer_config`
+- prefill uses `kv_producer` or `kv_both`
+- decode uses `kv_consumer` or `kv_both`
+- both stages use the same KV connector
+- `kv_buffer_device` and `kv_buffer_size` agree when set on both stages
+- both stages use the same `tensor_parallel_size`
+
+These checks happen before request processing so an inconsistent PD topology
+fails during engine initialization.
+
+## What is implemented today
+
+| Area | Current state |
+| --- | --- |
+| Pair detection and topology validation | Implemented in the shared PD helper |
+| Prefill sampling preparation | Implemented for sync and async entrypoints |
+| Prefill-to-decode routing | Implemented in the current Orchestrator |
+| Deploy-based configuration | Not yet migrated; the supported `stages:` deploy schema cannot currently expand Qwen3-Omni into the four-stage PD topology |
+| Qwen3 Thinker-to-Talker handoff | Bridge state and embedding-merge helpers exist, but the merge is not wired into the live stage processor |
+| Connector compatibility | Incomplete; the Orchestrator currently requires `remote_request_id`, while the pinned Mooncake connector completes prefill without returning that metadata |
+| Validation coverage | PD entrypoint tests and the PD deploy overlay in online end-to-end tests are temporarily disabled during migration |
+| Performance | No non-regression or throughput/latency result is established for the current path |
+
+Because these pieces are still converging, the previous legacy `stage_args` YAML
+and `vllm serve` command have been removed. The serve CLI now accepts the
+deploy-based `stages:` schema, so retaining the legacy recipe would imply a
+launch path that current `main` rejects.
+
+## Configuration contract under development
+
+The intended deploy integration starts from the
+[default Qwen3-Omni deploy config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
+and expands its frozen three-stage pipeline into the four stages shown above.
+A complete integration must:
+
+- split the Thinker definition without duplicating the Talker or Code2Wav model
+  contracts
+- remap downstream stage IDs, input sources, and stage connectors
+- inject producer and consumer `kv_transfer_config` values into the two Thinker
+  stages
+- allocate separate device sets for prefill and decode
+- preserve both final outputs: decode text and Code2Wav audio
+- provide connector-specific routing metadata without hard-coding one connector
+  contract
+
+The required number of GPUs depends on tensor-parallel sizes and whether Talker
+and Code2Wav share a device. A basic TP=1 layout uses one device for prefill, one
+for decode, and one for the two downstream audio stages.
+
+## Implementation map
+
+- PD detection, validation, and sampling-parameter preparation:
+  [`PDDisaggregationMixin`](gh-file:vllm_omni/entrypoints/pd_utils.py)
+- startup detection and connector metadata:
+  [`AsyncOmniEngine`](gh-file:vllm_omni/engine/async_omni_engine.py)
+- request-scoped prefill-to-decode routing:
+  [`Orchestrator`](gh-file:vllm_omni/engine/orchestrator.py)
+- Qwen3 downstream conditioning helpers:
+  [Qwen3-Omni stage input processor](gh-file:vllm_omni/model_executor/stage_input_processors/qwen3_omni.py)
@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
 | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
 | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -35,6 +36,25 @@ python text_to_video.py \
  --output t2v_out.mp4
 ```

+### LingBot-Video
+
+The shared runner recognizes both official LingBot checkpoint IDs and selects
+the native `LingBotVideoPipeline` automatically:
+
+```bash
+python text_to_video.py \
+  --model robbyant/lingbot-video-dense-1.3b \
+  --prompt "a robotic arm picks up a red block" \
+  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
+  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
+  --output lingbot_t2v.mp4
+```
+
+Use `robbyant/lingbot-video-moe-30b-a3b` for the MoE checkpoint. For local
+paths without `lingbot` in the directory name, add
+`--model-class-name LingBotVideoPipeline`. Model-only options use
+`--extra-body`, for example `'{"batch_cfg": true, "output_type": "np"}'`.
+
 ### LTX-2

 ```bash
@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `robbyant/lingbot-video-dense-1.3b` / `robbyant/lingbot-video-moe-30b-a3b` | 192x320 | 9 | 2 | 3.0 | ~68 GiB (MoE smoke) |
 | `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
@@ -57,18 +58,31 @@ python text_to_video.py \

 ### LingBot-Video

-The dedicated LingBot example explicitly selects the video output modality:
+The shared runner recognizes the official dense and MoE checkpoint IDs, selects
+`LingBotVideoPipeline`, and builds the canonical video request envelope:

 ```bash
-python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
+python text_to_video.py \
  --model robbyant/lingbot-video-dense-1.3b \
+  --prompt "a robotic arm picks up a red block" \
  --height 192 --width 320 --num-frames 9 --num-inference-steps 2 \
+  --guidance-scale 3.0 --flow-shift 3.0 --fps 24 \
  --output lingbot_t2v.mp4
 ```

 Use the same command with `robbyant/lingbot-video-moe-30b-a3b` for the MoE
 checkpoint. Requested frame counts are rounded upward to the causal
-VAE's `4n+1` grid.
+VAE's `4n+1` grid. For a local checkpoint path whose name does not contain
+`lingbot`, pass `--model-class-name LingBotVideoPipeline`.
+
+LingBot-only options such as `batch_cfg` and `output_type` use the shared
+model-extra channel, for example:
+
+```bash
+python text_to_video.py \
+  --model robbyant/lingbot-video-dense-1.3b \
+  --extra-body '{"batch_cfg": true, "output_type": "np"}'
+```

 ### LTX-2

@@ -16,7 +16,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
+from vllm_omni.model_extras import get_extra_body_params, get_model_class_name, get_output_tensor_range
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform

@@ -80,6 +80,17 @@ _MODEL_PRESETS = {
        "fps": 16,
        "output": "helios_output.mp4",
    },
+    "lingbot": {
+        "model_class_name": "LingBotVideoPipeline",
+        "height": 192,
+        "width": 320,
+        "num_frames": 9,
+        "num_inference_steps": 2,
+        "guidance_scale": 3.0,
+        "fps": 24,
+        "flow_shift": 3.0,
+        "output": "lingbot_video_output.mp4",
+    },
    "ltx2": {
        "height": 512,
        "width": 768,
@@ -110,27 +121,51 @@ _MODEL_PRESETS = {
 def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
    model_lower = model.lower()
    class_lower = (model_class_name or "").lower()
-    if "distilled" in class_lower or "distilled" in model_lower:
-        return _MODEL_PRESETS["ltx2_distilled"]
-    if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
-        return _MODEL_PRESETS["ltx23"]
-    if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
+    if "lingbot" in model_lower or "lingbotvideo" in class_lower:
+        return _MODEL_PRESETS["lingbot"]
+    if "ltx" in class_lower or "ltx" in model_lower:
+        if "distilled" in class_lower or "distilled" in model_lower:
+            return _MODEL_PRESETS["ltx2_distilled"]
+        if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
+            return _MODEL_PRESETS["ltx23"]
        return _MODEL_PRESETS["ltx2"]
-    if "vace" in model_lower:
+    if "vace" in model_lower or "vace" in class_lower:
        return _MODEL_PRESETS["vace"]
    # Edge must be matched before the generic cosmos branch (its "cosmos" substring would
    # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
    if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
        return _MODEL_PRESETS["cosmos3_edge"]
-    if "cosmos" in model_lower:
+    if "cosmos" in model_lower or "cosmos" in class_lower:
        return _MODEL_PRESETS["cosmos"]
-    if "hunyuan" in model_lower:
+    if "hunyuan" in model_lower or "hunyuan" in class_lower:
        return _MODEL_PRESETS["hunyuan"]
-    if "helios" in model_lower:
+    if "helios" in model_lower or "helios" in class_lower:
        return _MODEL_PRESETS["helios"]
    return _MODEL_PRESETS["wan"]


+def build_text_to_video_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
+    """Build the canonical request envelope for the shared T2V example."""
+    result: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["video"],
+    }
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result
+
+
+def _normalize_float_tensor(tensor: torch.Tensor, source_range: str) -> torch.Tensor:
+    """Normalize decoded frames according to the pipeline's output contract."""
+    if not tensor.is_floating_point():
+        return tensor
+    if source_range == "negative_one_to_one":
+        return tensor.clamp(-1, 1) * 0.5 + 0.5
+    if source_range == "zero_to_one":
+        return tensor.clamp(0, 1)
+    raise ValueError(f"Unsupported floating-point tensor range: {source_range!r}")
+
+
 def parse_profiler_config(value: str) -> dict[str, Any]:
    try:
        config = json.loads(value)
@@ -160,7 +195,7 @@ def parse_extra_body(value: str) -> dict[str, Any]:
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
+        "Supports Wan2.2, HunyuanVideo-1.5, Helios, LingBot-Video, and other text-to-video models."
    )
    parser.add_argument(
        "--model",
@@ -408,6 +443,7 @@ def main():
    for key, default_val in preset.items():
        if getattr(args, key.replace("-", "_"), None) is None:
            setattr(args, key.replace("-", "_"), default_val)
+    model_class_name = args.model_class_name

    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
    # Cache-dit config (Wan2.2 only)
@@ -480,6 +516,7 @@ def main():
    omni = Omni(**omni_kwargs)
    model_class_name = get_model_class_name(omni) or model_class_name
    declared_extra_body_params = get_extra_body_params(model_class_name)
+    output_tensor_range = get_output_tensor_range(model_class_name)

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
@@ -513,12 +550,11 @@ def main():
            lora_path=lora_path,
        )

-    prompt_dict = {"prompt": args.prompt}
-    if args.negative_prompt is not None:
-        prompt_dict["negative_prompt"] = args.negative_prompt
-    elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
+    negative_prompt = args.negative_prompt
+    if negative_prompt is None and all(preset is not _MODEL_PRESETS[name] for name in ("lingbot", "ltx2", "ltx23")):
        # Preserve the historical empty-prompt behavior for non-LTX examples.
-        prompt_dict["negative_prompt"] = ""
+        negative_prompt = ""
+    prompt_dict = build_text_to_video_prompt(args.prompt, negative_prompt)

    extra_args = {}
    if lora_request:
@@ -630,8 +666,7 @@ def main():
                frame_tensor = frame_tensor[0]
            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
                frame_tensor = frame_tensor.permute(1, 2, 0)
-            if frame_tensor.is_floating_point():
-                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            frame_tensor = _normalize_float_tensor(frame_tensor, output_tensor_range)
            return frame_tensor.float().numpy()
        if isinstance(frame, np.ndarray):
            frame_array = frame
@@ -683,8 +718,7 @@ def main():
                video_tensor = video_tensor[0]
        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
            video_tensor = video_tensor.permute(1, 2, 3, 0)
-        if video_tensor.is_floating_point():
-            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_tensor = _normalize_float_tensor(video_tensor, output_tensor_range)
        video_array = video_tensor.float().numpy()
    elif isinstance(frames, np.ndarray):
        video_array = frames
@@ -1,131 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import argparse
-import time
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate text-to-video with a LingBot-Video checkpoint.")
-    parser.add_argument("--model", default="robbyant/lingbot-video-dense-1.3b")
-    parser.add_argument("--prompt", default="a robotic arm picks up a red block")
-    parser.add_argument("--negative-prompt", default=None)
-    parser.add_argument("--output", default="lingbot_video_output.mp4")
-    parser.add_argument("--height", type=int, default=192)
-    parser.add_argument("--width", type=int, default=320)
-    parser.add_argument("--num-frames", type=int, default=9)
-    parser.add_argument("--num-inference-steps", type=int, default=2)
-    parser.add_argument("--guidance-scale", type=float, default=3.0)
-    parser.add_argument("--flow-shift", type=float, default=3.0)
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--output-type", choices=["pt", "np"], default="pt")
-    parser.add_argument("--batch-cfg", action="store_true")
-    return parser.parse_args()
-
-
-def _peak_memory_mb(output: Any) -> float:
-    result = output[0] if isinstance(output, list) and output else output
-    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
-
-
-def _extract_video(output: Any) -> Any:
-    if isinstance(output, list):
-        output = output[0] if output else None
-    if isinstance(output, OmniRequestOutput):
-        if output.final_output_type != "image":
-            raise ValueError(f"Unexpected output type {output.final_output_type!r}.")
-        if not output.images:
-            raise ValueError("No video frames found in OmniRequestOutput.")
-        return output.images[0] if len(output.images) == 1 else output.images
-    return output
-
-
-def _to_frame_array(video: Any) -> np.ndarray:
-    if isinstance(video, torch.Tensor):
-        tensor = video.detach().cpu()
-        if tensor.dim() == 5:
-            tensor = tensor[0]
-        if tensor.dim() == 4 and tensor.shape[-1] not in (3, 4):
-            if tensor.shape[0] in (3, 4):
-                tensor = tensor.permute(1, 2, 3, 0)
-            elif tensor.shape[1] in (3, 4):
-                tensor = tensor.permute(0, 2, 3, 1)
-        if tensor.is_floating_point():
-            if float(tensor.min()) < 0.0:
-                tensor = tensor.clamp(-1, 1) * 0.5 + 0.5
-            else:
-                tensor = tensor.clamp(0, 1)
-        return tensor.float().numpy()
-    if isinstance(video, np.ndarray):
-        array = video[0] if video.ndim == 5 else video
-        if np.issubdtype(array.dtype, np.integer):
-            return array.astype(np.float32) / 255.0
-        return array.astype(np.float32, copy=False)
-    if isinstance(video, list):
-        return np.stack([_to_frame_array(frame) for frame in video], axis=0)
-    return np.asarray(video, dtype=np.float32)
-
-
-def main() -> None:
-    args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    omni = Omni(
-        model=args.model,
-        model_class_name="LingBotVideoPipeline",
-        flow_shift=args.flow_shift,
-        parallel_config=DiffusionParallelConfig(),
-    )
-
-    prompt: dict[str, Any] = {
-        "prompt": args.prompt,
-        "modalities": ["video"],
-        "num_frames": args.num_frames,
-    }
-    if args.negative_prompt is not None:
-        prompt["negative_prompt"] = args.negative_prompt
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_frames=args.num_frames,
-        num_inference_steps=args.num_inference_steps,
-        guidance_scale=args.guidance_scale,
-        generator=generator,
-        output_type=args.output_type,
-        extra_args={"batch_cfg": args.batch_cfg, "shift": args.flow_shift},
-    )
-
-    start = time.perf_counter()
-    output = omni.generate(prompt, sampling_params)
-    elapsed = time.perf_counter() - start
-    print(f"Total generation time: {elapsed:.4f} seconds ({elapsed * 1000:.2f} ms)")
-    peak_mb = _peak_memory_mb(output)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
-
-    video = _to_frame_array(_extract_video(output))
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-    from diffusers.utils import export_to_video
-
-    export_to_video(list(video), str(output_path), fps=args.fps)
-    print(f"Saved generated video to {output_path}")
-    omni.close()
-
-
-if __name__ == "__main__":
-    main()
@@ -77,6 +77,7 @@ plugins:
        user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
        user_guide/diffusion/cpu_offload_diffusion.md: user_guide/diffusion/cpu_offload.md
        user_guide/quantization/index.md: user_guide/quantization/overview.md
+        configuration/pd_disaggregation.md: features/pd_disaggregation.md
        design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
        design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
  - glightbox
@@ -27,7 +27,7 @@ parallel, cache, quantized, or expert-kernel backends.
 - Dense checkpoint: <https://huggingface.co/robbyant/lingbot-video-dense-1.3b>
 - MoE checkpoint: <https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b>
 - Upstream project: <https://github.com/Robbyant/lingbot-video>
- Related offline example: [`examples/offline_inference/text_to_video/text_to_video_lingbot.py`](../../examples/offline_inference/text_to_video/text_to_video_lingbot.py)
+- Related offline example: [`examples/offline_inference/text_to_video/text_to_video.py`](../../examples/offline_inference/text_to_video/text_to_video.py)
 - Related online video API docs: [`docs/serving/videos_api.md`](../../docs/serving/videos_api.md)

 ## Hardware Support
@@ -52,7 +52,7 @@ The MoE path is validated with BF16 expert weights.

 ```bash
 CUDA_VISIBLE_DEVICES=0 \
-python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
+python examples/offline_inference/text_to_video/text_to_video.py \
  --model robbyant/lingbot-video-dense-1.3b \
  --prompt "a robotic arm picks up a red block" \
  --output lingbot_t2v.mp4 \
@@ -70,7 +70,7 @@ python examples/offline_inference/text_to_video/text_to_video_lingbot.py \

 ```bash
 CUDA_VISIBLE_DEVICES=0 \
-python examples/offline_inference/text_to_video/text_to_video_lingbot.py \
+python examples/offline_inference/text_to_video/text_to_video.py \
  --model robbyant/lingbot-video-moe-30b-a3b \
  --prompt "a robotic arm picks up a red block" \
  --output lingbot_moe_t2v.mp4 \
@@ -82,7 +82,7 @@ example and the DiT sections below.
 - Optional environment variables:

 ```bash
-export CACHE_DIT_VERSION=1.3.0
+export CACHE_DIT_VERSION=1.5.0
 ```

 HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
@@ -16,7 +16,7 @@ diffusers==0.38.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
-cache-dit==1.3.0
+cache-dit==1.5.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625
@@ -34,3 +34,5 @@ pydantic-settings>=2.0.0
 # base image does not ship it, so declare it explicitly. The loader also guards
 # this import so a missing gguf no longer breaks all diffusion model loading.
 gguf>=0.10.0
+# Required for running the Cosmos3 model with guardrails enabled (default behavior).
+cosmos-guardrail>=0.3.1
@@ -189,7 +189,7 @@ def _generate_single_stage_video(


 def _generate_bagel_image(
-    quantization_config: str | None = None,
+    diffusion_quantization_config: str | None = None,
    num_inference_steps: int = 15,
 ) -> tuple[Any, float]:
    """Generate an image with BAGEL (multi-stage: LLM + Diffusion).
@@ -202,8 +202,8 @@ def _generate_bagel_image(
        "stage_configs_path": config_path,
        "stage_init_timeout": 300,
    }
-    if quantization_config:
-        omni_kwargs["quantization_config"] = quantization_config
+    if diffusion_quantization_config:
+        omni_kwargs["diffusion_quantization_config"] = diffusion_quantization_config

    model_name = omni_kwargs.pop("model")
    with OmniRunner(model_name, **omni_kwargs) as runner:
@@ -380,21 +380,21 @@ def test_single_stage_ltx2_fp8_uses_less_memory():

@hardware_test(res={"cuda": "H100"})
 def test_bagel_fp8_generates_image():
-    """BAGEL with FP8 quantization_config generates a valid image.
+    """BAGEL with diffusion-stage FP8 generates a valid image.

    FP8 should only apply to the diffusion stage (Stage-1), not the
    LLM stage (Stage-0). We verify this by checking:
      1. Image is generated successfully
      2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
    """
-    image, _ = _generate_bagel_image(quantization_config="fp8")
+    image, _ = _generate_bagel_image(diffusion_quantization_config="fp8")
    image.save("test_bagel_fp8.png")


@hardware_test(res={"cuda": "H100"})
 def test_bagel_bf16_generates_image():
    """BAGEL without quantization generates a valid image (baseline)."""
-    image, _ = _generate_bagel_image(quantization_config=None)
+    image, _ = _generate_bagel_image(diffusion_quantization_config=None)
    image.save("test_bagel_bf16.png")


@@ -680,7 +680,7 @@ class TestSpeechAPI:

    @pytest.mark.asyncio
    async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
-        """Test that extra_params are correctly applied to sampling_params_list in diffusion mode."""
+        """Test public diffusion speech success and extra_params propagation."""
        # Mock the engine client
        mock_engine = mocker.MagicMock()

@@ -704,7 +704,11 @@ class TestSpeechAPI:

        req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"})

-        await server._create_diffusion_speech(req)
+        response = await server.create_speech(req)
+
+        assert response.status_code == 200
+        assert response.media_type == "audio/wav"
+        assert response.body == b"dummy"

        # Verify generate was called
        mock_engine.generate.assert_called_once()
@@ -737,6 +741,19 @@ class TestTTSMethods:
        yield server
        server.shutdown()

+    def test_diffusion_audio_encode_speed_skips_ar_adapter(self, mocker: MockerFixture):
+        """Pure-diffusion speech keeps generic speed handling outside AR adapters."""
+        server = OmniOpenAIServingSpeech.for_diffusion(
+            diffusion_engine=mocker.MagicMock(),
+            model_name="test-model",
+        )
+        resolve_adapter = mocker.patch.object(serving_speech_module, "resolve_adapter")
+
+        speed = server._audio_encode_speed(OpenAICreateSpeechRequest(input="Hello", speed=1.25))
+
+        assert speed == 1.25
+        resolve_adapter.assert_not_called()
+
    def test_is_tts_detection_no_stage(self, speech_server):
        """Test TTS model detection when no TTS stage exists."""
        # Fixture creates server with stage_configs = [] -> _is_tts should be False
@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from examples.offline_inference.text_to_video.text_to_video import (
+    _MODEL_PRESETS,
+    _detect_preset,
+    _normalize_float_tensor,
+    build_text_to_video_prompt,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("negative_prompt", "expected_negative_prompt"),
+    [(None, None), ("", ""), ("blurry", "blurry")],
+)
+def test_text_to_video_builds_canonical_prompt(
+    negative_prompt: str | None,
+    expected_negative_prompt: str | None,
+) -> None:
+    result = build_text_to_video_prompt("a robot waves", negative_prompt)
+
+    assert result["prompt"] == "a robot waves"
+    assert result["modalities"] == ["video"]
+    if expected_negative_prompt is None:
+        assert "negative_prompt" not in result
+    else:
+        assert result["negative_prompt"] == expected_negative_prompt
+
+
+@pytest.mark.parametrize(
+    ("model", "model_class_name", "preset_name"),
+    [
+        ("robbyant/lingbot-video-dense-1.3b", None, "lingbot"),
+        ("/models/custom-checkpoint", "LingBotVideoPipeline", "lingbot"),
+        ("Lightricks/LTX-2", "LTX2DistilledPipeline", "ltx2_distilled"),
+        ("BestWishYsh/Helios-Distilled", None, "helios"),
+        ("/models/cosmos-edge", "Cosmos3EdgeVFMTransformer", "cosmos3_edge"),
+        ("/models/vace", "WanVACEPipeline", "vace"),
+    ],
+)
+def test_detect_preset_is_scoped_to_model_family(
+    model: str,
+    model_class_name: str | None,
+    preset_name: str,
+) -> None:
+    assert _detect_preset(model, model_class_name) is _MODEL_PRESETS[preset_name]
+
+
+def test_lingbot_preset_matches_lingbot_defaults() -> None:
+    assert _MODEL_PRESETS["lingbot"] == {
+        "model_class_name": "LingBotVideoPipeline",
+        "height": 192,
+        "width": 320,
+        "num_frames": 9,
+        "num_inference_steps": 2,
+        "guidance_scale": 3.0,
+        "fps": 24,
+        "flow_shift": 3.0,
+        "output": "lingbot_video_output.mp4",
+    }
+
+
+def test_normalize_float_tensor_preserves_zero_to_one_frames() -> None:
+    frames = torch.tensor([0.0, 0.25, 1.0, 1.5])
+
+    assert torch.equal(_normalize_float_tensor(frames, "zero_to_one"), torch.tensor([0.0, 0.25, 1.0, 1.0]))
+
+
+def test_normalize_float_tensor_maps_negative_one_to_one_frames() -> None:
+    frames = torch.tensor([-1.5, -1.0, 0.0, 1.0, 1.5])
+
+    assert torch.equal(
+        _normalize_float_tensor(frames, "negative_one_to_one"),
+        torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
+    )
+
+
+def test_normalize_float_tensor_uses_one_contract_for_ambiguous_frames() -> None:
+    frames = [torch.tensor([0.0, 0.5, 1.0]), torch.tensor([-1.0, 0.0, 1.0])]
+
+    normalized = [_normalize_float_tensor(frame, "negative_one_to_one") for frame in frames]
+
+    assert torch.equal(normalized[0], torch.tensor([0.5, 0.75, 1.0]))
+    assert torch.equal(normalized[1], torch.tensor([0.0, 0.5, 1.0]))
@@ -1,21 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from functools import lru_cache
 from types import SimpleNamespace
 from unittest.mock import Mock

 import pytest
 import torch
 import torch.nn as nn
+from vllm.platforms import current_platform

 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
    HiFTGenerator,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
+    CFMGraphWrapper,
    HiFTGraphWrapper,
 )

-pytestmark = [pytest.mark.core_model]
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]


 class _F0Predictor(nn.Module):
@@ -142,3 +145,248 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch)
    wrapper._capture.assert_not_called()
    wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
    assert result is wrapper.decode_fn.return_value
+
+
+# ---------------------------------------------------------------------------
+# CFMGraphWrapper tests
+# ---------------------------------------------------------------------------
+
+
+class _MiniDiT(nn.Module):
+    """Minimal DiT-like module with blocks_forward_chunk for CFM graph testing.
+
+    Mimics the upstream cosyvoice2 DiT's cache semantics:
+    - CausalConv1d.forward_chunk: cat([cnn_cache, x], dim=time) when cache is not None
+    - Attention.forward_chunk: cat([k, k_cache], dim=seq) when att_cache is not None
+    """
+
+    def __init__(self, in_dim: int = 16, hidden: int = 8, depth: int = 2) -> None:
+        super().__init__()
+        self.in_proj = nn.Linear(in_dim, hidden)
+        self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)])
+        self.final_layer = nn.Linear(hidden, in_dim)
+
+    def blocks_forward_chunk(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        mask: torch.Tensor | None,
+        cnn_cache: torch.Tensor | None = None,
+        att_cache: torch.Tensor | None = None,
+        cnn_cache_buffer: torch.Tensor | None = None,
+        att_cache_buffer: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        x = self.in_proj(x)
+        for b_idx in range(len(self.blocks)):
+            # Simulate CausalConv1d: cnn_cache contributes to the first frames.
+            # Upstream: if cnn_cache[b_idx] is None, creates zeros internally (same as zeros).
+            cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None
+            if cnn_b is not None:
+                x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2)
+            # Simulate Attention: att_cache contributes a bias when non-empty.
+            # Upstream: if att_cache[b_idx] is None, skips cat entirely (different path).
+            att_b = att_cache[b_idx] if att_cache is not None else None
+            if att_b is not None and att_b.shape[3] > 0:
+                x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1)
+            x = self.blocks[b_idx](x)
+            x = x + t
+            cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous()
+            dt = x.shape[1]
+            att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+def _cfm_inputs(
+    batch_size: int, chunk_size: int, old_att_len: int, *, device: str = "cuda"
+) -> tuple[torch.Tensor, ...]:
+    depth = 2
+    hidden = 8
+    estimator_input = torch.randn(batch_size, 16, chunk_size, device=device)
+    time_emb = torch.randn(batch_size, 1, hidden, device=device)
+    cnn_cache = torch.randn(depth, batch_size, hidden, 2, device=device)
+    att_cache = torch.randn(depth, batch_size, 1, old_att_len, hidden, device=device)
+    cnn_out = torch.empty(depth, batch_size, hidden, 2, device=device)
+    att_out = torch.empty(depth, batch_size, 1, old_att_len + chunk_size, hidden, device=device)
+    return estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    estimator = _MiniDiT().eval().cuda()
+    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
+
+    with torch.inference_mode():
+        for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)):
+            inputs = _cfm_inputs(2, chunk_size, old_att_len)
+
+            eager_inputs = tuple(v.clone() for v in inputs)
+            with torch.no_grad():
+                eager_result = estimator.blocks_forward_chunk(
+                    eager_inputs[0],
+                    eager_inputs[1],
+                    None,
+                    eager_inputs[2],
+                    eager_inputs[3],
+                    eager_inputs[4],
+                    eager_inputs[5],
+                )
+
+            wrapper.replay(*inputs)
+            replay_inputs = tuple(v.clone() for v in inputs)
+            graph_result, graph_cnn, graph_att = wrapper.replay(*replay_inputs)
+
+            torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
+            torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5)
+            torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5)
+
+    wrapper._capture_graph.cache_clear()
+
+
+def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) -> CFMGraphWrapper:
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    wrapper = object.__new__(CFMGraphWrapper)
+    wrapper.max_graphs = max_graphs
+    wrapper.graph_fn = Mock(return_value=torch.tensor([42.0]))
+    wrapper.device = torch.device("cuda")
+
+    @lru_cache(maxsize=max_graphs)
+    def _capture_graph(key: tuple):
+        return wrapper._do_capture(key)
+
+    wrapper._capture_graph = _capture_graph
+    wrapper._do_capture = Mock(return_value=None)
+    return wrapper
+
+
+def test_cfm_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
+    wrapper = _cfm_mock_wrapper(monkeypatch)
+    wrapper._capture_graph = Mock(return_value=None)
+
+    inputs = _cfm_inputs(2, 10, 0)
+    wrapper.replay(*inputs)
+
+    wrapper._capture_graph.assert_called_once()
+    wrapper.graph_fn.assert_called_once()
+
+
+def test_cfm_capture_failure_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
+    wrapper = _cfm_mock_wrapper(monkeypatch)
+    wrapper._capture_graph = lru_cache(maxsize=1)(lambda key: None)
+
+    inputs = _cfm_inputs(2, 10, 0)
+    result = wrapper.replay(*inputs)
+
+    wrapper.graph_fn.assert_called_once()
+    assert result[0] is wrapper.graph_fn.return_value
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_lru_eviction(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify that LRU cache evicts oldest entries when full."""
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    estimator = _MiniDiT().eval().cuda()
+    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=2)
+
+    with torch.inference_mode():
+        inputs_a = _cfm_inputs(2, 10, 0)
+        wrapper.replay(*inputs_a)
+        assert wrapper._capture_graph.cache_info().currsize == 1
+
+        inputs_b = _cfm_inputs(2, 12, 0)
+        wrapper.replay(*inputs_b)
+        assert wrapper._capture_graph.cache_info().currsize == 2
+
+        wrapper.replay(*inputs_a)
+        assert wrapper._capture_graph.cache_info().currsize == 2
+
+        inputs_c = _cfm_inputs(2, 14, 0)
+        wrapper.replay(*inputs_c)
+        assert wrapper._capture_graph.cache_info().currsize == 2
+
+        wrapper.replay(*inputs_b)
+        hits_before = wrapper._capture_graph.cache_info().hits
+        wrapper.replay(*inputs_b)
+        hits_after = wrapper._capture_graph.cache_info().hits
+        assert hits_after == hits_before + 1
+
+    wrapper._capture_graph.cache_clear()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+def test_cfm_none_cache_parity_between_graph_and_eager(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Verify that None→zeros (graph path) matches [None]*depth (eager path).
+
+    setup_batch calls _decode_cfm with cnn_cache=None on every request's
+    prompt pass. The graph path replaces None with zeros; the eager path
+    passes [None]*depth. Both must produce identical outputs.
+    """
+    torch.accelerator.synchronize()
+    pool = torch.cuda.graph_pool_handle()
+    monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
+
+    torch.manual_seed(0)
+    estimator = _MiniDiT().eval().cuda()
+    wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=32)
+
+    depth = len(estimator.blocks)
+    hidden = 8
+    batch_size = 2
+    chunk_size = 20
+
+    estimator_input = torch.randn(batch_size, 16, chunk_size, device="cuda")
+    time_emb = torch.randn(batch_size, 1, hidden, device="cuda")
+    cnn_out = torch.empty(depth, batch_size, hidden, 2, device="cuda")
+    att_out = torch.empty(depth, batch_size, 1, chunk_size, hidden, device="cuda")
+
+    # Eager path: cnn_cache=[None]*depth, att_cache=[None]*depth
+    eager_input = estimator_input.clone()
+    eager_time = time_emb.clone()
+    eager_cnn_out = torch.empty_like(cnn_out)
+    eager_att_out = torch.empty_like(att_out)
+    with torch.no_grad():
+        eager_result = estimator.blocks_forward_chunk(
+            eager_input,
+            eager_time,
+            None,
+            [None] * depth,
+            [None] * depth,
+            eager_cnn_out,
+            eager_att_out,
+        )
+
+    # Graph path: cnn_cache=zeros, att_cache=zero-length (simulating _estimator_step's None→zeros)
+    zero_cnn = torch.zeros_like(cnn_out)
+    zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+    wrapper.replay(estimator_input, time_emb, zero_cnn, zero_att, cnn_out, att_out)
+    graph_input = estimator_input.clone()
+    graph_time = time_emb.clone()
+    graph_cnn_out = torch.empty_like(cnn_out)
+    graph_att_out = torch.empty_like(att_out)
+    graph_zero_cnn = torch.zeros_like(cnn_out)
+    graph_zero_att = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+    graph_result, graph_cnn, graph_att = wrapper.replay(
+        graph_input,
+        graph_time,
+        graph_zero_cnn,
+        graph_zero_att,
+        graph_cnn_out,
+        graph_att_out,
+    )
+
+    torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5)
+    torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5)
+    torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5)
@@ -17,6 +17,7 @@ from vllm_omni.model_extras import (
    build_x_to_text_prompt,
    get_extra_body_params,
    get_extra_output_params,
+    get_output_tensor_range,
    get_x_to_text_model_family,
    should_init_extra_args_for_non_diffusion_stages,
 )
@@ -134,6 +135,8 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
        }
    )
    assert get_extra_output_params("LingBotVideoPipeline") == frozenset()
+    assert get_output_tensor_range("LingBotVideoPipeline") == "zero_to_one"
+    assert get_output_tensor_range("WanPipeline") == "negative_one_to_one"
    assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False


@@ -13,10 +13,13 @@ def _shrink_dit_rope_config(
    target_head_dim: int = 32,
    target_heads: int = 4,
    default_axes_dims_rope: list[int] | None = None,
+    joint_attention_dim: int | None = None,
 ) -> dict:
    config["num_layers"] = num_layers
    if num_single_layers is not None:
        config["num_single_layers"] = num_single_layers
+    if joint_attention_dim is not None:
+        config["joint_attention_dim"] = joint_attention_dim
    axes_dims_rope = config.get("axes_dims_rope", default_axes_dims_rope)
    factor = config["attention_head_dim"] / target_head_dim
    config["attention_head_dim"] = target_head_dim
@@ -50,22 +53,55 @@ def _shrink_flux_t5_text_encoder(config: dict) -> dict:
    return config


-def tiny_qwen_image_builder() -> str:
-    def shrink_text_encoder(config: dict) -> dict:
-        config["num_hidden_layers"] = 2
-        config["intermediate_size"] = 64
-        config["text_config"]["num_hidden_layers"] = 2
-        config["text_config"]["intermediate_size"] = 64
-        config["text_config"]["layer_types"] = config["text_config"]["layer_types"][:2]
-        config["vision_config"]["depth"] = 2
-        config["vision_config"]["intermediate_size"] = 64
-        config["vision_config"]["fullatt_block_indexes"] = [0, 1]
-        return config
+def _shrink_qwen_text_encoder_config(config: dict, hidden_size: int = 64) -> dict:
+    text_config = config["text_config"]
+    old_head_dim = text_config["hidden_size"] / text_config["num_attention_heads"]
+    text_config["num_hidden_layers"] = 2
+    text_config["intermediate_size"] = 64
+    text_config["hidden_size"] = hidden_size
+    text_config["num_attention_heads"] = 2
+    text_config["num_key_value_heads"] = 2
+    text_config["layer_types"] = text_config["layer_types"][:2]
+    factor = old_head_dim / (hidden_size / 2)
+    mrope_section = text_config["rope_scaling"]["mrope_section"]
+    text_config["rope_scaling"]["mrope_section"] = [round(d / factor) for d in mrope_section]
+    config["num_hidden_layers"] = 2
+    config["intermediate_size"] = 64
+    config["hidden_size"] = hidden_size
+    config["num_attention_heads"] = 2
+    config["num_key_value_heads"] = 2
+    config["vision_config"]["depth"] = 2
+    config["vision_config"]["intermediate_size"] = 64
+    config["vision_config"]["fullatt_block_indexes"] = [0, 1]
+    config["vision_config"]["out_hidden_size"] = hidden_size
+    return config

+
+def _shrink_qwen_transformer_config(config: dict, joint_attention_dim: int = 64) -> dict:
+    return _shrink_dit_rope_config(config, joint_attention_dim=joint_attention_dim)
+
+
+def tiny_qwen_image_builder() -> str:
    return build_tiny_from_configs(
        "QwenImagePipeline",
        "Qwen/Qwen-Image",
-        transform={"text_encoder": shrink_text_encoder, "transformer": _shrink_dit_rope_config},
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
+    )
+
+
+def tiny_qwen_image_edit_builder() -> str:
+    return build_tiny_from_configs(
+        "QwenImageEditPipeline",
+        "Qwen/Qwen-Image-Edit",
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
+    )
+
+
+def tiny_qwen_image_edit_plus_builder() -> str:
+    return build_tiny_from_configs(
+        "QwenImageEditPlusPipeline",
+        "Qwen/Qwen-Image-Edit-2511",
+        transform={"text_encoder": _shrink_qwen_text_encoder_config, "transformer": _shrink_qwen_transformer_config},
    )


@@ -74,4 +74,14 @@ DIFFUSION_TEST_SETTINGS = {
        builder=diff_model_builders.tiny_flux2_builder,
        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
    ),
+    "QwenImageEditPipeline": DiffusionModelTestOpts(
+        model="Qwen/Qwen-Image-Edit",
+        builder=diff_model_builders.tiny_qwen_image_edit_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+    ),
+    "QwenImageEditPlusPipeline": DiffusionModelTestOpts(
+        model="Qwen/Qwen-Image-Edit-2511",
+        builder=diff_model_builders.tiny_qwen_image_edit_plus_builder,
+        supported_tasks=[DiffusionTasks.IMAGE_TO_IMAGE],
+    ),
 }
@@ -21,8 +21,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # if you are adding a new model and see the tests below fail, please follow the pattern
 # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
 EXCLUDED_MODELS = [
-    "QwenImageEditPipeline",
-    "QwenImageEditPlusPipeline",
    "QwenImageLayeredPipeline",
    "GlmImagePipeline",
    "ZImagePipeline",
@@ -19,6 +19,8 @@ connectors:
      # A future initial_codec_chunk_frames shape will use lazy capture unless
      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
      enable_hift_graph: true
+      enable_cfm_graph: true
+      cfm_max_graphs: 32

 stages:
  - stage_id: 0
@@ -22,6 +22,8 @@ connectors:
      # A future initial_codec_chunk_frames shape will use lazy capture unless
      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
      enable_hift_graph: true
+      enable_cfm_graph: true
+      cfm_max_graphs: 32

 stages:
  - stage_id: 0
@@ -22,6 +22,8 @@ connectors:
      # A future initial_codec_chunk_frames shape will use lazy capture unless
      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
      enable_hift_graph: true
+      enable_cfm_graph: true
+      cfm_max_graphs: 32

 stages:
  - stage_id: 0
@@ -33,6 +33,8 @@ connectors:
      # A future initial_codec_chunk_frames shape will use lazy capture unless
      # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
      enable_hift_graph: true
+      enable_cfm_graph: true
+      cfm_max_graphs: 32

 stages:
  - stage_id: 0
@@ -65,11 +65,10 @@ except ImportError:
        # ``model_config["guardrails"]`` is False.
        def __init__(self, *args, **kwargs):
            raise ValueError(
-                f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
+                f"The guardrails are enabled but not functional for {self.__class__}. This is most likely due to a"
+                f"missing cosmos-guardrail package. Please install cosmos-guardrail to be compliant with the "
                "[NVIDIA Open Model License Agreement]"
-                "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). "
-                f"Please ensure that you are compliant with the license agreement."
-                f"Please install cosmos-guardrail package to enable safety checks."
+                "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)."
            )


@@ -502,11 +502,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
    def _get_tts_adapter(self):
        """Return the per-model serving adapter for the current ``_tts_model_type``.

+        Pure-diffusion speech uses its dedicated request path and does not
+        resolve adapters for AR-stage TTS models.
+
        Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached
        instance was built) so callers that set ``_tts_model_type`` after
        construction still dispatch to the matching adapter. In production
        ``_tts_model_type`` is fixed at init, so the cached instance is reused.
        """
+        if self._diffusion_mode:
+            return None
+
        adapter_cls = resolve_adapter(self._tts_model_type)
        if adapter_cls is None:
            self._adapter = None
@@ -332,7 +332,7 @@ class DiTAttention(nn.Module):
        self,
        x: torch.Tensor,
        padding_mask: torch.Tensor | None = None,
-        rope=None,
+        rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, seq_len = x.shape[0], x.shape[1]
@@ -406,7 +406,7 @@ class DiTBlock(nn.Module):
        x: torch.Tensor,
        t: torch.Tensor,
        padding_mask: torch.Tensor | None = None,
-        rope=None,
+        rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
        attn_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        # Pre-norm & modulation for attention
@@ -575,6 +575,8 @@ class GLMTTSDiT(nn.Module):
        is_causal: bool = False,
        attn_mask: torch.Tensor | None = None,
        block_pattern: list[int] | None = None,
+        text_embed: torch.Tensor | None = None,
+        rope: tuple[torch.Tensor, torch.Tensor | float] | None = None,
    ) -> torch.Tensor:
        """Forward pass.

@@ -588,6 +590,9 @@ class GLMTTSDiT(nn.Module):
            is_causal: Use block-causal attention
            attn_mask: Custom attention mask (overrides is_causal)
            block_pattern: Token-level block sizes for block-causal mask
+            text_embed: Precomputed text embedding (must already be in model
+                dtype); skips the internal embedding + ConvNeXt pass
+            rope: Precomputed rotary embedding tuple for ``seq_len``

        Returns:
            Predicted mel-spectrogram [B, T, mel_dim]
@@ -610,13 +615,21 @@ class GLMTTSDiT(nn.Module):
            time_emb = torch.cat([time_emb, spkr_emb], dim=-1)

        # Text embedding - ensure cast to model dtype
-        text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype)
+        if text_embed is None:
+            text_embed = self.text_emb_layer(text, seq_len, text_lens=text_lens).to(model_dtype)
+        else:
+            # A precomputed text_embed replaces `text` entirely, so its length must
+            # match middle_point; a mismatch would apply the wrong rotary/mask silently.
+            assert text_embed.shape[1] == seq_len, (
+                f"precomputed text_embed len {text_embed.shape[1]} != seq_len {seq_len}"
+            )

        # Input projection
        x = self.emb_concator(middle_point, condition, text_embed, drop_audio_cond=False)

        # Rotary embeddings
-        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+        if rope is None:
+            rope = self.rotary_embed.forward_from_seq_len(seq_len)

        # Build block-causal mask when is_causal and no explicit attn_mask
        if attn_mask is None and is_causal and block_pattern is not None:
@@ -857,6 +857,40 @@ class GLMTTSDiTForGeneration(nn.Module):
        if self.t_scheduler == "cosine":
            t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)

+        # Hoist loop-invariant DiT inputs (text embedding, RoPE table,
+        # block-causal mask, CFG double-batch conditioning) out of the Euler
+        # loop; they depend only on the conditioning and seq_len, not on the
+        # per-step state x or the timestep.
+        seq_len = int(mel_cond.shape[1])
+        # Must match the dtype forward() casts with internally (both are
+        # proj_out.weight.dtype); keep them in sync if that derivation changes.
+        dit_dtype = self.dit.proj_out.weight.dtype
+        # `device` is speech_tokens.device; in this call path the sampled state x
+        # lives on the same device, so the hoisted mask matches the old in-forward one.
+        attn_mask = (
+            self.dit._create_block_causal_mask(block_pattern, seq_len, device, dit_dtype)
+            if block_pattern is not None
+            else None
+        )
+        rope = self.dit.rotary_embed.forward_from_seq_len(seq_len)
+
+        if self.inference_cfg_rate > 0:
+            condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0)
+            if self.speech_token_cfg:
+                text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0)
+            else:
+                text_b = speech_tokens.expand(2, -1)
+            padding_b = padding_mask.expand(2, -1)
+            if spkr_embedding_dit is not None and self.spkr_emb_adaln:
+                spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0)
+            else:
+                spkr_b = None
+            text_lens_b = text_lens.expand(2) if text_lens is not None else None
+            text_embed_b = self.dit.text_emb_layer(text_b, seq_len, text_lens=text_lens_b).to(dit_dtype)
+        else:
+            spkr_uncond = spkr_embedding_dit if (spkr_embedding_dit is not None and self.spkr_emb_adaln) else None
+            text_embed_u = self.dit.text_emb_layer(speech_tokens, seq_len, text_lens=text_lens).to(dit_dtype)
+
        for step in range(n_timesteps):
            cache_step = step + 1
            if last_step_cache is not None and cache_step in last_step_cache:
@@ -872,24 +906,12 @@ class GLMTTSDiTForGeneration(nn.Module):
            t_current = t_span[step]
            dt = t_span[step + 1] - t_current

-            use_causal = block_pattern is not None
            x_dit = x.to(mel_cond.dtype)

            if self.inference_cfg_rate > 0:
                # Batch conditional + unconditional into a single forward
                middle_b = x_dit.expand(2, -1, -1)
-                condition_b = torch.cat([condition_dit, torch.zeros_like(condition_dit)], dim=0)
-                if self.speech_token_cfg:
-                    text_b = torch.cat([speech_tokens, torch.zeros_like(speech_tokens)], dim=0)
-                else:
-                    text_b = speech_tokens.expand(2, -1)
                time_b = t_current.unsqueeze(0).expand(2)
-                padding_b = padding_mask.expand(2, -1)
-                if spkr_embedding_dit is not None and self.spkr_emb_adaln:
-                    spkr_b = torch.cat([spkr_embedding_dit, torch.zeros_like(spkr_embedding_dit)], dim=0)
-                else:
-                    spkr_b = None
-                text_lens_b = text_lens.expand(2) if text_lens is not None else None
                pred_b = self.dit(
                    middle_point=middle_b,
                    condition=condition_b,
@@ -897,9 +919,10 @@ class GLMTTSDiTForGeneration(nn.Module):
                    time_step=time_b,
                    padding_mask=padding_b,
                    spkr_emb=spkr_b,
-                    is_causal=use_causal,
-                    block_pattern=block_pattern,
+                    attn_mask=attn_mask,
                    text_lens=text_lens_b,
+                    text_embed=text_embed_b,
+                    rope=rope,
                ).to(torch.float32)
                dphi_dt, cfg_dphi_dt = pred_b.chunk(2, dim=0)
                dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt
@@ -910,10 +933,11 @@ class GLMTTSDiTForGeneration(nn.Module):
                    text=speech_tokens,
                    time_step=t_current.unsqueeze(0),
                    padding_mask=padding_mask,
-                    spkr_emb=spkr_embedding_dit if spkr_embedding_dit is not None and self.spkr_emb_adaln else None,
-                    is_causal=use_causal,
-                    block_pattern=block_pattern,
+                    spkr_emb=spkr_uncond,
+                    attn_mask=attn_mask,
                    text_lens=text_lens,
+                    text_embed=text_embed_u,
+                    rope=rope,
                ).to(torch.float32)

            x = x + dt * dphi_dt
@@ -14,7 +14,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger

-from .cuda_graph_wrapper import HiFTGraphWrapper
+from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper

 logger = init_logger(__name__)

@@ -73,6 +73,7 @@ class BatchedToken2Wav(nn.Module):
        *,
        connector_config: Mapping[str, int] | None = None,
        hift_graph_config: Mapping[str, Any] | None = None,
+        cfm_graph_config: Mapping[str, Any] | None = None,
    ):
        super().__init__()
        self._token2wav = token2wav
@@ -135,6 +136,22 @@ class BatchedToken2Wav(nn.Module):
                with torch.inference_mode(), _autocast_disabled(hift_parameter.device):
                    self.hift_graph_wrapper.capture()
                logger.info("HiFT CUDA Graph captured successfully")
+        self._cfm_graph_wrapper: CFMGraphWrapper | None = None
+        cfm_graph_cfg = dict(cfm_graph_config or {})
+        if bool(cfm_graph_cfg.get("enabled", False)):
+            flow_parameter = next(self.flow.parameters(), None)
+            if flow_parameter is not None and flow_parameter.device.type == "cuda":
+                estimator = self.flow.decoder.estimator
+                self._cfm_graph_wrapper = CFMGraphWrapper(
+                    graph_fn=estimator.blocks_forward_chunk,
+                    max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)),
+                )
+                logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32)))
+            else:
+                logger.info(
+                    "CFM CUDA Graph is disabled on device type %s",
+                    flow_parameter.device.type if flow_parameter is not None else "unknown",
+                )
        self._prompt_features: dict[tuple[str, str], PromptFeatures] = {}

    def _hift_inference(
@@ -265,6 +282,19 @@ class BatchedToken2Wav(nn.Module):
        speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width)
        estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1)
        cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache)
+        if self._cfm_graph_wrapper is not None:
+            if cnn_cache is None:
+                cnn_cache = torch.zeros_like(cnn_out)
+            if att_cache is None:
+                att_cache = estimator_input.new_zeros(att_out.shape[:3] + (0,) + att_out.shape[4:])
+            return self._cfm_graph_wrapper.replay(
+                estimator_input,
+                time_embedding,
+                cnn_cache,
+                att_cache,
+                cnn_out,
+                att_out,
+            )
        old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks)
        old_att: Any = att_cache if att_cache is not None else [None] * len(estimator.blocks)
        result = estimator.blocks_forward_chunk(
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from functools import lru_cache
+
 import torch
 from torch.cuda import CUDAGraph
 from vllm.logger import init_logger
@@ -133,3 +135,126 @@ class HiFTGraphWrapper:
        cache_source = static_cache_source_output[:batch_size].clone()
        speech = self.finalize_fn(static_magnitude_output[:batch_size], static_phase_output[:batch_size]).clone()
        return speech, cache_source
+
+
+def _tensor_signature(value: torch.Tensor) -> tuple:
+    return tuple(value.shape), str(value.dtype), str(value.device)
+
+
+_DTYPE_MAP = {
+    "torch.float32": torch.float32,
+    "torch.float16": torch.float16,
+    "torch.bfloat16": torch.bfloat16,
+    "torch.float64": torch.float64,
+}
+
+
+def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]:
+    """Rebuild zero tensors from a cache key (shape, dtype, device tuples)."""
+    tensors = []
+    for shape, dtype_str, device_str in key[1:]:
+        dtype = _DTYPE_MAP.get(dtype_str, torch.float32)
+        device = torch.device(device_str)
+        tensors.append(torch.zeros(shape, dtype=dtype, device=device))
+    return tuple(tensors)
+
+
+class CFMGraphWrapper:
+    """Per-shape CUDA graph capture/replay for the CFM DiT estimator.
+
+    Captures one blocks_forward_chunk call (in_proj -> DiT blocks -> final_layer)
+    as the graph target. The 10-step Euler loop stays in Python, replaying
+    the graph 10 times per decode.
+
+    Uses functools.lru_cache for automatic LRU eviction. Cache misses
+    trigger capture; capture failures fall back to eager. Outputs are
+    cloned after replay to prevent streaming cache corruption.
+    """
+
+    def __init__(self, graph_fn, *, max_graphs: int = 32) -> None:
+        self.graph_fn = graph_fn
+        self.max_graphs = int(max_graphs)
+        self.device = next(graph_fn.__self__.parameters()).device
+
+        @lru_cache(maxsize=self.max_graphs)
+        def _capture_graph(key: tuple):
+            return self._capture(key)
+
+        self._capture_graph = _capture_graph
+
+    def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor:
+        return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5])
+
+    def _capture(self, key: tuple) -> tuple | None:
+        """Capture a CUDA graph for the given key. Returns None on failure."""
+        static_inputs = _tensors_from_key(key)
+
+        current_stream = torch.cuda.current_stream(self.device)
+        warmup_stream = torch.cuda.Stream(device=self.device)
+        warmup_stream.wait_stream(current_stream)
+        with torch.cuda.stream(warmup_stream), torch.no_grad():
+            for _ in range(3):
+                self._call_graph_fn(static_inputs)
+        current_stream.wait_stream(warmup_stream)
+
+        try:
+            graph = CUDAGraph()
+            with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
+                static_output = self._call_graph_fn(static_inputs)
+        except Exception:
+            logger.warning(
+                "CFM graph capture failed for shape=%s; using eager",
+                key,
+                exc_info=True,
+            )
+            return None
+
+        logger.info(
+            "Captured CFM CUDA Graph for shape %s (cache=%d/%d, hits=%d, misses=%d)",
+            key,
+            self._capture_graph.cache_info().currsize,
+            self.max_graphs,
+            self._capture_graph.cache_info().hits,
+            self._capture_graph.cache_info().misses,
+        )
+        return (static_inputs, static_output, graph)
+
+    def replay(
+        self,
+        estimator_input: torch.Tensor,
+        time_emb: torch.Tensor,
+        cnn_cache: torch.Tensor,
+        att_cache: torch.Tensor,
+        cnn_out: torch.Tensor,
+        att_out: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out)
+
+        if torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda":
+            with torch.no_grad():
+                result = self._call_graph_fn(inputs)
+            return result, inputs[4], inputs[5]
+
+        key = ("estimator_step",) + tuple(_tensor_signature(v) for v in inputs)
+        entry = self._capture_graph(key)
+
+        if entry is None:
+            logger.debug(
+                "CFM graph eager fallback for shape=%s (hits=%d, misses=%d)",
+                key,
+                self._capture_graph.cache_info().hits,
+                self._capture_graph.cache_info().misses,
+            )
+            with torch.no_grad():
+                result = self._call_graph_fn(inputs)
+            return result, inputs[4], inputs[5]
+
+        static_inputs, static_output, graph = entry
+        for static, current in zip(static_inputs, inputs, strict=True):
+            static.copy_(current)
+        graph.replay()
+        return (
+            static_output.detach().clone(),
+            static_inputs[4].detach().clone(),
+            static_inputs[5].detach().clone(),
+        )
@@ -162,6 +162,10 @@ class MiniCPMO45Code2Wav(nn.Module):
            "enabled": bool(extra.get("enable_hift_graph", False)),
            "capture_batch_sizes": capture_batch_sizes,
        }
+        self._cfm_graph_config = {
+            "enabled": bool(extra.get("enable_cfm_graph", False)),
+            "max_graphs": int(extra.get("cfm_max_graphs", 32)),
+        }
        self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
        if self._min_batch_size < 1:
            raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
@@ -802,4 +806,5 @@ class MiniCPMO45Code2Wav(nn.Module):
            trt_stepper=trt_stepper,
            connector_config=self._connector_config,
            hift_graph_config=self._hift_graph_config,
+            cfm_graph_config=self._cfm_graph_config,
        )
@@ -242,6 +242,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge
 # ---------------------------------------------------------------------------


+# Subclass keeps .weight name + ctor shape so the state_dict loader stays unchanged.
 class OmniVoiceRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
@@ -315,10 +316,9 @@ class OmniVoiceAttention(nn.Module):
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)

-        # Convert [B, 1, S, S] bool mask to float mask for SDPA
-        # (0.0 where attend, -inf where masked)
-        sdpa_mask = None
-        if attention_mask is not None:
+        # Caller passes a float mask; materialize float form if a bool slips through.
+        sdpa_mask = attention_mask
+        if sdpa_mask is not None and sdpa_mask.dtype == torch.bool:
            sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf"))

        out = F.scaled_dot_product_attention(
@@ -400,24 +400,48 @@ def _precompute_rope(
    theta: float = 1000000.0,
    device: torch.device | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Precompute RoPE cos/sin tensors."""
+    """Precompute RoPE cos/sin pre-doubled to full head_dim so the hot path skips a per-call cat."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    t = torch.arange(max_seq_len, device=device, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)
-    cos = freqs.cos()
-    sin = freqs.sin()
+    cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)
+    sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)
    return cos, sin


 def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-    """Apply rotary position embedding. x shape: (B, S, H, D)."""
+    """Apply rotary position embedding. x shape: (B, S, H, D). cos/sin are pre-doubled to width D."""
    seq_len = x.shape[1]
-    cos = cos[:seq_len].unsqueeze(0).unsqueeze(2)  # (1, S, 1, D/2)
+    cos = cos[:seq_len].unsqueeze(0).unsqueeze(2)  # (1, S, 1, D)
    sin = sin[:seq_len].unsqueeze(0).unsqueeze(2)
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    rotated = torch.cat([-x2, x1], dim=-1)
-    return x * torch.cat([cos, cos], dim=-1) + rotated * torch.cat([sin, sin], dim=-1)
+    return x * cos + rotated * sin
+
+
+# ---------------------------------------------------------------------------
+# TF32 opt-in (process-wide; default off)
+# ---------------------------------------------------------------------------
+
+_TF32_ENABLED = False
+
+
+def _maybe_enable_tf32() -> None:
+    """Enable TF32 matmuls process-wide (idempotent). Not bit-identical; opt-in via config.enable_tf32."""
+    global _TF32_ENABLED
+    if _TF32_ENABLED or not torch.cuda.is_available():
+        return
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.set_float32_matmul_precision("high")
+    _TF32_ENABLED = True
+    logger.info(
+        "OmniVoice TF32 enabled process-wide: matmul.allow_tf32=%s cudnn.allow_tf32=%s float32_matmul_precision=%s",
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.backends.cudnn.allow_tf32,
+        torch.get_float32_matmul_precision(),
+    )


 # ---------------------------------------------------------------------------
@@ -425,6 +449,19 @@ def _apply_rotary_pos_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor)
 # ---------------------------------------------------------------------------


+def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    """Convert a boolean attention mask to its additive float form.
+
+    ``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool
+    mask must never be copied straight into a float buffer: the implicit cast
+    maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so
+    silently *unmasks* them.
+    """
+    if mask.dtype != torch.bool:
+        return mask
+    return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf"))
+
+
 class _OmniVoiceCUDAGraphForward:
    """Pre-captures CUDA graphs for predefined sequence-length buckets.

@@ -476,11 +513,10 @@ class _OmniVoiceCUDAGraphForward:
        mask_padded[:, :S] = audio_mask

        if attention_mask is not None:
-            attn_padded = torch.zeros(
-                two_b,
-                1,
-                bucket,
-                bucket,
+            # Callers normalize to the additive float form first, so pad with -inf.
+            attn_padded = torch.full(
+                (two_b, 1, bucket, bucket),
+                float("-inf"),
                dtype=attention_mask.dtype,
                device=attention_mask.device,
            )
@@ -564,7 +600,8 @@ class _OmniVoiceCUDAGraphForward:
            key = (two_b, bucket)
            dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device)
            dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device)
-            dummy_attn = torch.ones(two_b, 1, bucket, bucket, dtype=torch.bool, device=device)
+            # Capture with a float mask to match what forward() feeds at replay time.
+            dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=torch.float32, device=device)
            self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn)
        logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs))

@@ -586,6 +623,11 @@ class _OmniVoiceCUDAGraphForward:
        two_b = input_ids.shape[0]
        bucket = self._find_bucket(seq_len) if two_b == 2 else None

+        # Graphs are captured with (and their static buffers hold) the additive
+        # float mask, so normalize here, before padding or any copy_ into them.
+        if attention_mask is not None:
+            attention_mask = _additive_float_mask(attention_mask)
+
        if bucket is None:
            # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket).
            # Lock prevents concurrent threads from double-capturing the same key.
@@ -655,6 +697,10 @@ class OmniVoiceGenerator(nn.Module):
        super().__init__()
        self.config = config

+        # Opt-in TF32; must run before any CUDA-graph capture so captured kernels honour it.
+        if getattr(config, "enable_tf32", False):
+            _maybe_enable_tf32()
+
        # Text embedding (shared with LLM)
        self.text_embedding = nn.Embedding(config.llm_vocab_size, config.llm_hidden_size)

@@ -703,6 +749,8 @@ class OmniVoiceGenerator(nn.Module):
        self,
        input_ids: torch.Tensor,
        audio_mask: torch.Tensor,
+        text_embeds: torch.Tensor | None = None,
+        audio_mask_3d: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Prepare mixed text+audio embeddings.

@@ -710,30 +758,38 @@ class OmniVoiceGenerator(nn.Module):
            input_ids: [B, 8, S] - text tokens replicated across codebooks,
                       audio positions have per-codebook token IDs
            audio_mask: [B, S] - True for audio positions, False for text
+            text_embeds: optional cached [B, S, H] text-position embeddings
+            audio_mask_3d: optional cached [B, S, 1] audio_mask.unsqueeze(-1)

        Returns:
            embeddings: [B, S, hidden_size]
        """
-        # Text embeddings from first codebook row (all rows identical for text)
-        text_embeds = self.text_embedding(input_ids[:, 0, :])
+        # Cached across the denoising loop since text ids don't change.
+        if text_embeds is None:
+            text_embeds = self.text_embedding(input_ids[:, 0, :])
+        if audio_mask_3d is None:
+            audio_mask_3d = audio_mask.unsqueeze(-1)

        # Audio embeddings: offset per codebook, then sum across codebooks
        shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1)
        audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1)

        # Merge: audio where audio_mask=True, text elsewhere
-        return torch.where(audio_mask.unsqueeze(-1), audio_embeds, text_embeds)
+        return torch.where(audio_mask_3d, audio_embeds, text_embeds)

    def _transformer_forward(
        self,
        inputs_embeds: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
+        cos: torch.Tensor | None = None,
+        sin: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Run through transformer layers.

        Args:
            inputs_embeds: [B, S, hidden_size]
            attention_mask: [B, 1, S, S] or None
+            cos, sin: optional precomputed RoPE tables in target dtype

        Returns:
            hidden_states: [B, S, hidden_size]
@@ -743,8 +799,15 @@ class OmniVoiceGenerator(nn.Module):
        self._ensure_rope(seq_len, device)

        hidden_states = inputs_embeds
-        cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype)
-        sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype)
+        if cos is None or sin is None:
+            cos = self._rope_cos.to(device=device, dtype=hidden_states.dtype)
+            sin = self._rope_sin.to(device=device, dtype=hidden_states.dtype)
+
+        # Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()).
+        if attention_mask is not None and attention_mask.dtype == torch.bool:
+            attention_mask = torch.zeros_like(attention_mask, dtype=hidden_states.dtype).masked_fill_(
+                ~attention_mask, float("-inf")
+            )

        for layer in self.layers:
            hidden_states = layer(
@@ -859,21 +922,37 @@ class OmniVoiceGenerator(nn.Module):

        layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1)

-        # Compute c_lens for extracting target region from full sequence
-        c_lens = []
-        for i in range(B):
-            # Conditional sequence length = number of non-padding positions
-            c_len = attention_mask[i, 0, 0].sum().item()
-            c_lens.append(int(c_len))
+        # Single D2H pull for all conditional lengths instead of B per-item .item() syncs.
+        c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist()
+
+        # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion.
+        sdpa_attn_mask = torch.zeros_like(attention_mask, dtype=torch.float32).masked_fill_(
+            ~attention_mask, float("-inf")
+        )
+
+        use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda
+        if not use_cuda_graph:
+            # Eager-path-only constants (the cuda-graph captures its own).
+            text_embeds_cached = self.text_embedding(input_ids[:, 0, :])
+            audio_mask_3d = audio_mask.unsqueeze(-1)
+            self._ensure_rope(input_ids.shape[-1], device)
+            target_dtype = text_embeds_cached.dtype
+            cos = self._rope_cos.to(device=device, dtype=target_dtype)
+            sin = self._rope_sin.to(device=device, dtype=target_dtype)

        # Main iterative loop
        for step in range(num_step):
-            if self._cuda_graph_fwd is not None and input_ids.is_cuda:
-                batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, attention_mask).to(torch.float32)
+            if use_cuda_graph:
+                # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below.
+                batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask)
            else:
-                inputs_embeds = self._prepare_embeddings(input_ids, audio_mask)
-                hidden_states = self._transformer_forward(inputs_embeds, attention_mask)
-                batch_logits = self._get_logits(hidden_states).to(torch.float32)
+                # Eager fallback reuses hoisted constants (text embeds, sdpa mask, cos/sin).
+                inputs_embeds = self._prepare_embeddings(
+                    input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d
+                )
+                hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, cos=cos, sin=sin)
+                # fp32 cast deferred to the per-item slices below.
+                batch_logits = self._get_logits(hidden_states)
            # batch_logits: [2*B, 8, S, 1025]

            for i in range(B):
@@ -884,16 +963,17 @@ class OmniVoiceGenerator(nn.Module):
                c_len = c_lens[i]
                t_len = target_lens[i]

-                # Extract logits for target region
-                c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :]  # [1, 8, T, 1025]
-                u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :]  # [1, 8, T, 1025]
+                # Extract logits for target region; upcast only the slices we actually consume.
+                c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32)
+                u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32)

-                # Classifier-free guidance
+                # Classifier-free guidance. Fuse the chain: the two inner
+                # log_softmax normalizers are per-position scalars that the final
+                # shift-invariant log_softmax cancels, so guide on the raw logits
+                # with a single softmax: log_softmax((1+s)*c - s*u). Exact.
                if guidance_scale != 0:
-                    c_log_probs = F.log_softmax(c_logits, dim=-1)
-                    u_log_probs = F.log_softmax(u_logits, dim=-1)
-                    log_probs = torch.log_softmax(
-                        c_log_probs + guidance_scale * (c_log_probs - u_log_probs),
+                    log_probs = F.log_softmax(
+                        (1.0 + guidance_scale) * c_logits - guidance_scale * u_logits,
                        dim=-1,
                    )
                else:
@@ -922,15 +1002,13 @@ class OmniVoiceGenerator(nn.Module):
                sample_tokens = tokens[i : i + 1, :, :t_len]
                scores.masked_fill_(sample_tokens != mask_id, -float("inf"))

-                # Select top-k positions to unmask
+                # Select top-k positions to unmask. .flatten() on this non-contiguous view already copies.
                _, topk_idx = torch.topk(scores.flatten(), k)
-                flat_tokens = sample_tokens.flatten().clone()
+                flat_tokens = sample_tokens.flatten()
                flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx]
                sample_tokens.copy_(flat_tokens.view_as(sample_tokens))

-                # Update tokens and batch inputs for next iteration
-                tokens[i : i + 1, :, :t_len] = sample_tokens
-                input_ids = input_ids.clone()
+                # Mirror update into both cond and uncond input_ids halves for the next step.
                input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0)
                input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0)

@@ -9,6 +9,7 @@ from vllm_omni.model_extras.registry import (
    get_extra_body_params,
    get_extra_output_params,
    get_model_class_name,
+    get_output_tensor_range,
    get_x_to_text_model_family,
    should_init_extra_args_for_non_diffusion_stages,
 )
@@ -21,6 +22,7 @@ __all__ = [
    "get_extra_body_params",
    "get_extra_output_params",
    "get_model_class_name",
+    "get_output_tensor_range",
    "get_x_to_text_model_family",
    "should_init_extra_args_for_non_diffusion_stages",
 ]
@@ -4,7 +4,7 @@
 from __future__ import annotations

 from collections.abc import Callable, Mapping
-from typing import Any
+from typing import Any, Literal

 from PIL import Image

@@ -93,6 +93,7 @@ ImageToVideoPromptBuilder = Callable[
    dict[str, Any],
 ]
 XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]]
+OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"]


 def default_x_to_text_prompt(
@@ -193,6 +194,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
    },
    "LingBotVideoPipeline": {
        "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
+        "output_tensor_range": "zero_to_one",
        # Shared T2I/I2V envelopes select the output modality. LingBot's
        # pipeline owns model-specific validation and normalization.
    },
@@ -262,6 +264,18 @@ def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
    return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()


+def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
+    """Return the declared range for floating-point tensor outputs.
+
+    The default preserves the shared examples' historical handling. Pipelines
+    that already return normalized tensors declare ``zero_to_one`` explicitly.
+    """
+    spec = _get_spec(model_class_name)
+    if spec is None:
+        return "negative_one_to_one"
+    return spec.get("output_tensor_range", "negative_one_to_one")
+
+
 def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool:
    spec = _get_spec(model_class_name)
    return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))
@@ -90,6 +90,8 @@ class OmniVoiceConfig(PretrainedConfig):
            "cuda_graph_capture_sizes",
            [128, 192, 256, 320, 384, 448, 512, 640, 768, 1024],
        )
+        # TF32 matmuls: not bit-identical; opt-in (matches vLLM default-off), set OMNIVOICE_TF32=1 to enable.
+        self.enable_tf32 = getattr(self, "enable_tf32", os.environ.get("OMNIVOICE_TF32", "0") != "0")


 AutoConfig.register("omnivoice", OmniVoiceConfig)
Author	SHA1	Message	Date
Hongsheng Liu	fab29e2873	[Docs] Redirect the former PD guide URL Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 12:48:43 +08:00
Hongsheng Liu	1ade14d147	[Docs] Document current PD design Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 12:44:42 +08:00
Hongsheng Liu	b6313918d6	Merge branch 'main' into codex-pd-navigation	2026-08-13 12:23:56 +08:00
NATUREandAlicia	d77a35a32b	[Bugfix] Scope Bagel FP8 config to diffusion stage (#6085 ) Signed-off-by: natureofnature <wzliu@connect.hku.hk> Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>	2026-08-13 11:23:47 +08:00
liwenerandamy-why-3459	877beeaa49	[Perf][MiniCPM-O-4.5]Add CUDA Graph For CFM DiT estimator (#6082 ) Signed-off-by: liwener <liwener.c@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-08-13 09:36:17 +08:00
Hongsheng Liu	8b8d3f56db	[Misc] Consolidate LingBot text-to-video runner (#6076 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 08:42:37 +08:00
Hongsheng Liu	c57178c0e1	[Dependency] Upgrade Cache-DiT to 1.5.0 (#6065 ) Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 08:41:59 +08:00
Hongsheng Liu	f27e70dce7	[Docs] Move PD guide into Features Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 06:54:34 +08:00
mjZhaoElaineandClaude Opus 5	9a59475015	[CI/Build] Fix ReadTheDocs build under RTD's seeded pip 23.1 (#6129 ) Signed-off-by: Mengjie Zhao <zmj0129@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>	2026-08-13 06:52:01 +08:00
Hongsheng Liu	c46abfef59	[Docs] Mark PD disaggregation as experimental Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-08-13 06:50:41 +08:00
Nick Cao	12f8168179	[Tests] Enable tiny model testing for Qwen-Image Edit and EditPlus (#5656 ) Signed-off-by: Nick Cao <ncao@redhat.com>	2026-08-12 16:03:49 -04:00
Zhenghan Song	72a02b492a	[Bugfix] Fix diffusion TTS adapter lookup in speech serving (#6121 ) Signed-off-by: HaningZS <resossr0v0@gmail.com>	2026-08-12 15:20:12 -04:00
Juan Pablo ZuluagaandYueqian Lin	5358f5db01	[Perf][TTS] omnivoice hot path: D2H batching, mask caching, cached text embeddings (#5174 ) Signed-off-by: JuanPZuluaga <juanz9312@gmail.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-08-12 17:15:35 +00:00
MaciejBalaNV	e677360168	Add cosmos-guardrail dependency and update error message (#6107 ) Signed-off-by: Maciej Bala <mbala@nvidia.com>	2026-08-12 17:04:07 +00:00
Juan Pablo ZuluagaandYueqian Lin	208c6776a8	[Perf][TTS] glm_tts: hoist loop-invariant text embed, RoPE, mask and CFG batch out of Euler loop (#5068 ) Signed-off-by: JuanPZuluaga <juanz9312@gmail.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-08-12 17:02:02 +00:00