Merge branch 'main' into codex-pd-navigation

.buildkite/cuda/test-nightly.yml

@@ -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

.buildkite/npu/test-npu-nightly.yml

@@ -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

.readthedocs.yml

@@ -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

benchmarks/lingbot_video/dense_pipeline_parity.py

@@ -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",

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -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

examples/offline_inference/text_to_video/text_to_video.md

@@ -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
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -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

examples/offline_inference/text_to_video/text_to_video_lingbot.py

@@ -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()

recipes/Robbyant/LingBot-Video.md

@@ -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 \

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -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`

requirements/common.txt

@@ -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

tests/diffusion/quantization/test_quantization_fp8.py

@@ -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")
 
 

tests/examples/offline_inference/test_text_to_video_prompts.py

@@ -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]))

tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py

@@ -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)

tests/model_extras/test_model_extras.py

@@ -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
 
 

vllm_omni/deploy/minicpmo_4_5.yaml

@@ -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

vllm_omni/deploy/minicpmo_4_5_2gpu.yaml

@@ -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

vllm_omni/deploy/minicpmo_4_5_3gpu.yaml

@@ -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

vllm_omni/deploy/minicpmo_4_5_8x4090.yaml

@@ -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

vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py

@@ -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(

vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py

@@ -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(),
+        )

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -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,
         )

vllm_omni/model_extras/__init__.py

@@ -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",
 ]

vllm_omni/model_extras/registry.py

@@ -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))