OmniRequestOutput inherits RequestOutput

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>

.buildkite/pipeline-intel.yaml

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

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

@@ -13,12 +13,33 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(
 )"
 
 cd "${omni_source_dir}"
+
+# The XPU base image is ~37GB; the default gzip layer exporter is single-threaded
+# and takes ~16min to compress it. zstd is multi-threaded (uses all cores) at a
+# similar ratio; set EXPORT_COMPRESSION=uncompressed to skip compression entirely
+# for local-only images. This requires the containerd image store (buildx docker
+# driver), which is the default here.
+EXPORT_COMPRESSION="${EXPORT_COMPRESSION:-zstd}"
+if [ "${EXPORT_COMPRESSION}" = "uncompressed" ]; then
+    export_args=(--output "type=image,name={{IMAGE}},compression=uncompressed")
+else
+    export_args=(--output "type=image,name={{IMAGE}},compression=${EXPORT_COMPRESSION},compression-level=3,force-compression=true")
+fi
+
+docker_build() {
+    # $1 = image name; remaining args passed through to docker build.
+    local image="$1"
+    shift
+    local out=("${export_args[@]/'{{IMAGE}}'/${image}}")
+    docker build "${out[@]}" "$@" -f docker/Dockerfile.xpu .
+}
+
 if [ -z "$(docker images -q "${base_image_name}")" ]; then
-    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+    docker_build "${base_image_name}" --target vllm-base --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 fi
 
 # Try building the docker image
-docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 
 # Setup cleanup
 remove_docker_container() {

.buildkite/test-nightly.yml

@@ -1489,9 +1489,12 @@ steps:
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
             EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
+            EXIT2=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT1
+            if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
+            exit $$EXIT2
         agents:
           queue: "mithril-h100-pool"
         plugins:

.gitignore

@@ -248,6 +248,7 @@ datasets/
 
 # Output files
 outputs/
+!vllm_omni/outputs/
 results/
 generated/
 output_*/

benchmarks/diffusion/backends.py

@@ -28,6 +28,7 @@ class RequestFuncInput:
     slo_ms: float | None = None
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
+    video_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
 
@@ -338,6 +339,11 @@ async def async_request_v1_videos(
         form.add_field(k, str(v))
 
     image_file = None
+    if input.image_paths and input.video_paths:
+        output.error = "Only one of image_paths or video_paths can be provided"
+        output.success = False
+        return output
+
     if input.image_paths and len(input.image_paths) > 0:
         image_path = input.image_paths[0]
         image_file = open(image_path, "rb")
@@ -345,7 +351,16 @@ async def async_request_v1_videos(
             "input_reference",
             image_file,
             filename=os.path.basename(image_path),
-            content_type="application/octet-stream",
+            content_type=_guess_mime_type(image_path),
+        )
+    elif input.video_paths and len(input.video_paths) > 0:
+        video_path = input.video_paths[0]
+        image_file = open(video_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(video_path),
+            content_type=_guess_mime_type(video_path),
         )
 
     job_id = None

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -73,6 +73,10 @@ Usage:
             {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
         ]'
 
+    v2v:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --endpoint /v1/videos --dataset random --task v2v --num-prompts 1 \
+        --height 720 --width 1280 --fps 24 --num-frames 189 --num-inference-steps 35
 
 """
 
@@ -542,11 +546,17 @@ class TraceDataset(BaseDataset):
         if not image_paths:
             single = row.get("image_path")
             image_paths = [single] if single else None
+        video_paths = row.get("video_paths")
+        if not video_paths:
+            single_video = row.get("video_path")
+            video_paths = [single_video] if single_video else None
 
         if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
             raise ValueError(
                 f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
             )
+        if not video_paths and self.args.task == "v2v":
+            raise ValueError("Task v2v requires video input, but no video_path or video_paths found in trace row.")
 
         override_w = self.args.width
         override_h = self.args.height
@@ -570,6 +580,7 @@ class TraceDataset(BaseDataset):
             timestamp=timestamp,
             slo_ms=slo_ms,
             image_paths=image_paths,
+            video_paths=video_paths,
             request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
         )
 
@@ -589,6 +600,8 @@ class CustomDataset(BaseDataset):
     - seed (optional): Random seed
     - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
     - image_urls (optional): List of input image URLs (alternative to image_paths)
+    - video_paths (optional): List of input video paths for v2v tasks
+    - video_urls (optional): List of input video URLs (alternative to video_paths)
 
     Example JSONL for ti2i:
     {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
@@ -660,6 +673,42 @@ class CustomDataset(BaseDataset):
 
         return None
 
+    def _resolve_video_paths(self, item: dict) -> list[str] | None:
+        if "video_paths" in item and item["video_paths"]:
+            paths = item["video_paths"]
+            if isinstance(paths, str):
+                paths = [paths]
+
+            valid_paths = []
+            for path in paths:
+                if os.path.exists(path):
+                    valid_paths.append(path)
+                else:
+                    raise ValueError(f"Video file not found: {path}")
+
+            return valid_paths if valid_paths else None
+
+        if "video_urls" in item and item["video_urls"]:
+            urls = item["video_urls"]
+            if isinstance(urls, str):
+                urls = [urls]
+
+            downloaded_paths = []
+            for url in urls:
+                try:
+                    response = requests.get(url, timeout=30)
+                    response.raise_for_status()
+                    suffix = os.path.splitext(url)[1] or ".mp4"
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                        tmp.write(response.content)
+                        downloaded_paths.append(tmp.name)
+                except Exception as e:
+                    raise ValueError(f"Failed to download video from {url}: {e}")
+
+            return downloaded_paths if downloaded_paths else None
+
+        return None
+
     def __len__(self) -> int:
         return len(self.data)
 
@@ -672,11 +721,22 @@ class CustomDataset(BaseDataset):
         height = item.get("height", self.args.height)
         num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
         seed = item.get("seed", self.args.seed)
-        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
+        reserved_keys = {
+            "prompt",
+            "width",
+            "height",
+            "num_inference_steps",
+            "seed",
+            "image_paths",
+            "image_urls",
+            "video_paths",
+            "video_urls",
+        }
         extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
 
         # Handle image paths/URLs
         image_paths = self._resolve_image_paths(item)
+        video_paths = self._resolve_video_paths(item)
 
         return RequestFuncInput(
             prompt=prompt,
@@ -684,6 +744,7 @@ class CustomDataset(BaseDataset):
             model=self.model,
             seed=seed,
             image_paths=image_paths,
+            video_paths=video_paths,
             extra_body=extra_body,
             width=width,
             height=height,
@@ -737,6 +798,10 @@ class RandomDataset(BaseDataset):
             self._random_image_path = self._generate_random_image_paths()
         else:
             self._random_image_path = None
+        if self.args.task == "v2v":
+            self._random_video_path = self._generate_random_video_paths()
+        else:
+            self._random_video_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
@@ -756,13 +821,15 @@ class RandomDataset(BaseDataset):
         if self._sampled_requests:
             profile = self._sampled_requests[idx]
             params.update(profile)
+        prompt = str(params.pop("prompt", f"Random prompt {idx} for benchmarking diffusion models"))
         return RequestFuncInput(
-            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            prompt=prompt,
             api_url=self.api_url,
             model=self.model,
             seed=self.args.seed,
             extra_body=extra_body,
             image_paths=self._random_image_path,
+            video_paths=self._random_video_path,
             **params,
         )
 
@@ -781,6 +848,35 @@ class RandomDataset(BaseDataset):
             image_paths.append(image_path)
         return image_paths
 
+    def _generate_random_video_paths(self) -> list[str]:
+        video_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_video.mp4")
+        try:
+            import cv2
+
+            profile = self.random_request_config[0] if self.random_request_config else {}
+            width = int(self.args.width or profile.get("width") or 1280)
+            height = int(self.args.height or profile.get("height") or 720)
+            num_frames = int(self.args.num_frames or profile.get("num_frames") or 16)
+            fps = float(self.args.fps or profile.get("fps") or 8)
+            writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+            if not writer.isOpened():
+                raise RuntimeError("cv2.VideoWriter failed to open")
+            for frame_idx in range(num_frames):
+                frame = np.zeros((height, width, 3), dtype=np.uint8)
+                x = int((frame_idx / max(num_frames - 1, 1)) * max(width - 96, 1))
+                y = height // 2
+                frame[:, :, 0] = 24
+                frame[:, :, 1] = 32
+                frame[:, :, 2] = 48
+                frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220)
+                writer.write(frame)
+            writer.release()
+            return [video_path]
+        except Exception as e:
+            raise RuntimeError(
+                "Failed to generate synthetic v2v input video. Install opencv-python or provide a custom dataset."
+            ) from e
+
 
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
@@ -1112,7 +1208,7 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
 
 
 def _default_endpoint_for_task(task: str) -> str:
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"i2i", "ti2i", "it2i"}:
         return "/v1/images/edits"
@@ -1126,7 +1222,7 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
+    VIDEO_TASKS = {"t2v", "i2v", "ti2v", "v2v"}
     IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
 
     if args.task in VIDEO_TASKS:
@@ -1177,6 +1273,10 @@ async def benchmark(args):
     if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
         for req in requests_list:
             req.extra_body.setdefault(_RETURN_STAGE_METRICS_FIELD, True)
+    if args.extra_body:
+        extra_body = json.loads(args.extra_body)
+        for req in requests_list:
+            req.extra_body.update(extra_body)
 
     if args.endpoint == "/v1/images/edits":
         for req in requests_list:
@@ -1328,7 +1428,7 @@ if __name__ == "__main__":
         "--task",
         type=str,
         default="t2v",
-        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
+        choices=["t2v", "i2v", "ti2v", "v2v", "ti2i", "i2i", "it2i", "t2i"],
         help="Task type.",
     )
     parser.add_argument(
@@ -1443,6 +1543,12 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--extra-body",
+        type=str,
+        default=None,
+        help="JSON object merged into every generated request body/form.",
+    )
     parser.add_argument(
         "--num-input-images",
         type=int,

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.24.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

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

docs/.nav.yml

@@ -23,7 +23,7 @@ nav:
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
       - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
-      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
@@ -38,7 +38,7 @@ nav:
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
@@ -80,6 +80,7 @@ nav:
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
       - Int8 W8A8: user_guide/quantization/int8.md
+      - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
       - MXFP8 W8A8: user_guide/quantization/mxfp8.md
       - MXFP4 W4A4: user_guide/quantization/mxfp4.md
       - ModelOpt: user_guide/quantization/modelopt.md

docs/api/README.md

@@ -46,17 +46,17 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
 - [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
-- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
-- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
 - [vllm_omni.engine.orchestrator.Orchestrator][]
 - [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
-- [vllm_omni.engine.output_modality.OutputModality][]
-- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
-- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
-- [vllm_omni.engine.output_processor.OmniRequestState][]
 - [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
 - [vllm_omni.engine.stage_init_utils.StageMetadata][]
 - [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
+- [vllm_omni.outputs.mm_outputs.MultimodalCompletionOutput][]
+- [vllm_omni.outputs.mm_outputs.MultimodalPayload][]
+- [vllm_omni.outputs.output_modality.OutputModality][]
+- [vllm_omni.outputs.output_modality.TensorAccumulationStrategy][]
+- [vllm_omni.outputs.output_processor.MultimodalOutputProcessor][]
+- [vllm_omni.outputs.output_processor.OmniRequestState][]
 
 ## Core
 

docs/cli/README.md

@@ -18,11 +18,17 @@ Specify the port:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file, launch the server with command below
+For a migrated model, load a custom deploy configuration with `--deploy-config`:
+
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
 ```
 
+The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
+```
 
 ## bench
 

docs/cli/bench/serve.md

@@ -2,7 +2,7 @@
 The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
 
 ## Notes
-We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
+vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
 
 ## Basic Parameter Description
 You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
   Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
 
 - `--backend`
-  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
+  Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices.
 
 - `--model`
   The model identifier to load, filled according to the models supported by vLLM-Omni.
@@ -55,7 +55,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
                     'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
                     '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
 
-- `-print-stage`
+- `--print-stage`
 Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
 
 - `--save-result`
@@ -312,7 +312,7 @@ vllm bench serve \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl \
   --random-output-len 2 \
-  --extra_body '{"modalities": ["text"]}'
+  --extra-body '{"modalities": ["text"]}'
 ```
 
 If successful, you will see the following output:
@@ -390,7 +390,7 @@ vllm bench serve --omni \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
   --random-output-len 900 \
-  --extra_body '{"modalities": ["text","audio"]}' \
+  --extra-body '{"modalities": ["text","audio"]}' \
   --print-stage
 ```
 

docs/configuration/composable_parallel.md

@@ -92,12 +92,12 @@ Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSp
 
 From highest to lowest, when the same axis is touched by multiple layers:
 
-1. **CLI overrides** (`--stage_N_tensor_parallel_size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
+1. **CLI overrides** (`--tensor-parallel-size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`).
 2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`.
 3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/<model_type>.yaml`).
 4. **Parser defaults.**
 
-After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `--devices` or `--tensor-parallel-size` that conflicts with the strategy's world size cannot slip past the pre-spawn check.
+After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `devices` value passed through `--stage-overrides` or a `--tensor-parallel-size` value that conflicts with the strategy's world size cannot slip past the pre-spawn check.
 
 ## Worked example
 
@@ -214,7 +214,7 @@ Strategy keys MUST be `model_stage` *names* (strings), the same labels stage con
     Use the `model_stage` name (e.g. `thinker`) instead.
 
 !!! warning "Mixing CLI overrides with a strategy on the same axis"
-    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--stage_0_tensor_parallel_size`, `--stage-overrides '{"0": {...}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
+    A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
 
 !!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
     Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:

docs/configuration/stage_configs.md

@@ -1,11 +1,11 @@
-# Stage configs for vLLM-Omni
+# Deploy and legacy stage configurations
 
-In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
+In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
 
 !!! note
     Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
 
-## New deploy schema reference
+## Deploy configuration schema
 
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
 
@@ -82,14 +82,14 @@ stages:
     input_connectors:  {from_stage_0: shm}
 ```
 
-### CLI flags introduced in this refactor
+### CLI flags
 
 | Flag | Description |
 |------|-------------|
-| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
+| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
 | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
 | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
-| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
+| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
 
 ### Stage-Based CLI Paradigm
 
@@ -146,7 +146,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 
 From highest to lowest:
 
-1. Per-stage flags (`--stage-overrides` JSON, `--stage-<id>-<key>` if registered)
+1. Per-stage overrides (`--stage-overrides` JSON)
 2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.)
 3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:`
 4. Overlay YAML (via `base_config:`) on top of the base YAML
@@ -210,7 +210,7 @@ Effective config per stage after the merge:
 | 1 | `max_model_len` | `16384` | global CLI |
 | 2 | (all defaults) | — | base YAML (no overrides apply) |
 
-Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
+Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment configurations have several main functions:
 
 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.
 - The disaggregated configuration for each stage and the communication topology among them.
@@ -219,16 +219,16 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
 - Default input parameters.
 
 To override specific parameters, explicitly inject the customized configuration schema
-in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
-when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
+in both online and offline instantiation flows. Use the `--deploy-config` flag
+when loading a deploy configuration, reserving the `--stage-configs-path` parameter
 exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
 
 Examples:
 
-For offline (Assume necessary dependencies have been imported):
+For offline inference (assuming the necessary dependencies have been imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
 ```
 
 For online serving:
@@ -242,11 +242,14 @@ Legacy online serving:
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 !!! important
-    We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
+    We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
+
+## Legacy stage configuration reference
+
+The remainder of this page documents the deprecated `stage_args` schema for models that have not migrated to deploy configurations. The following historical Qwen2.5-Omni configuration illustrates that schema; current Qwen2.5-Omni deployments should use `vllm_omni/deploy/qwen2_5_omni.yaml` instead.
 
-Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+# Historical stage configuration for Qwen2.5-Omni.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
@@ -338,7 +341,7 @@ runtime:
 
 ```
 
-## Stage Configuration Arguments
+## Legacy stage configuration arguments
 
 Each stage in the `stage_args` list contains the following configuration options:
 

docs/contributing/DOCS_GUIDE.md

@@ -67,6 +67,21 @@ docs/
 └── stylesheets/         # Custom CSS
 ```
 
+## Naming Model Examples
+
+Offline inference and online serving examples for the same model use the same
+directory name and the shared display name in
+`examples/model_display_names.yml`. Use these title forms:
+
+- `# <Model>: Offline inference`
+- `# <Model>: Online serving`
+
+The documentation generator uses the full title for the page H1 and the shared
+display name alone for navigation, and fails the build if a mapped README uses
+a different H1. Keep checkpoint identifiers in commands and prose rather than
+in the display name, and do not rename an existing example directory solely to
+adjust its title because the directory defines its public documentation URL.
+
 ## Publishing Documentation
 
 ### GitHub Pages (Recommended)

docs/contributing/ci/test_guide.md

@@ -12,12 +12,37 @@ When running test cases, you may need to install the following dependencies:
 
 ```bash
 uv pip install ".[dev]"
-apt-get install -y ffmpeg
+apt-get install -y espeak-ng jq
 ```
 
 ## Running Tests
 Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
 
+### CI job runners (L2–L4): logs, timing, and timeouts
+
+[`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh), [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh), and [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh) share the same run-time behavior (via [`tools/run_jobs_common.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_jobs_common.sh)):
+
+**Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`):
+
+| Path | Description |
+|------|-------------|
+| `jobs/*.sh` | Generated per-step bash wrappers |
+| `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) |
+| `<job_key>.log` | Tee output for each job |
+| `timing_summary.log` | Per-job duration and total wall time after the run completes |
+
+**Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example:
+
+```
+=== Job timing summary ===
+  TTS_VoxCPM2_Test  25m 40s  OK
+  Omni_Qwen3-Omni_Test  8m 05s  TIMED OUT
+Total wall time: 33m 45s (2 jobs)
+Failed jobs: 1/2
+```
+
+**Timeouts:** if a Buildkite step defines `timeout_in_minutes`, the runner wraps that job with `timeout ${N}m` (aligned with CI). Steps without `timeout_in_minutes` rely on any inline `timeout …` already present in the extracted pytest command. On timeout the process exits with code `124` and the summary marks the job as `TIMED OUT`.
+
 === "L1 level"
 
     ```bash
@@ -28,61 +53,109 @@ Our test scripts use the pytest framework. First, please use `git clone https://
 
 === "L2 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh). The script reads [`.buildkite/test-ready.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/ready_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts) for timing and timeout behavior). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
+    ```bash
+    # All L2 and L1 jobs (default: every pytest step in test-ready.yml)
+    bash tools/run_ready_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-ready.yml (group or "Simple ·" labels)
+    bash tools/run_ready_jobs.sh --skip-simple
+    bash tools/run_ready_jobs.sh --skip-simple --model-type diffusion
+
+    # Preview extracted commands without running
+    bash tools/run_ready_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_ready_jobs.sh --model-type omni
+    bash tools/run_ready_jobs.sh --model-type tts,diffusion
+
+    # Match a Buildkite label substring (e.g. one model suite)
+    bash tools/run_ready_jobs.sh --model-type tts --label-substr "VoxCPM2"
+
+    ```
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     cd tests
-    pytest -s -v -m "core_model and not cpu" --run-level=core_model
-    ```
-    If you only want to run a specific test case, you can use:
-    ```bash
     pytest -s -v test_xxxx.py --run-level=core_model
+    pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
-    ```
-    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
 
 === "L3 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh). The script reads [`.buildkite/test-merge.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/merge_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts)). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
     ```bash
-    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    # All L3 and L1 jobs (default: every pytest step in test-merge.yml)
+    bash tools/run_merge_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-merge.yml (group or "Simple ·" labels)
+    bash tools/run_merge_jobs.sh --skip-simple
+    bash tools/run_merge_jobs.sh --skip-simple --model-type omni
+
+    # Preview extracted commands without running
+    bash tools/run_merge_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_merge_jobs.sh --model-type diffusion
+    bash tools/run_merge_jobs.sh --model-type omni,tts
+
+    # Match a Buildkite label substring
+    bash tools/run_merge_jobs.sh --model-type diffusion --label-substr "Wan22"
+
     ```
-    If you only want to run a specific test case, you can use:
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     pytest -s -v test_xxxx.py --run-level=advanced_model
+    pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "advanced_model and distributed_cuda and L4"  --run-level=advanced_model
-    ```
-    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
 
 
 === "L4 level"
 
+    **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
     ```bash
-    cd tests
-    pytest -s -v -m "full_model" --run-level=full_model
+    # All nightly jobs (default: test-type all, model-type all)
+    bash tools/nightly/run_nightly_jobs.sh
+
+    # Preview extracted commands without running
+    bash tools/nightly/run_nightly_jobs.sh --dry-run
+
+    # Test kind (--test-type, OR semantics): perf | acc | function | stability | local | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf,acc
+
+    # Model area (--model-type, OR semantics): omni | tts | diffusion | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf --model-type diffusion
+
+    # Stability scripts (tests/dfx/stability/scripts/) and local marker tests
+    bash tools/nightly/run_nightly_jobs.sh --test-type stability --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type local --model-type tts
     ```
-    If you only want to run a specific test case, you can use:
-    ```bash
-    pytest -s -v test_xxxx.py --run-level=full_model
-    ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "full_model and distributed_cuda and L4"  --run-level=full_model
-    ```
-    Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
-    ```bash
-    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-    ```
-    The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
+
+    Perf steps (label contains ``Perf Test``) run before other jobs; Excel output is written under ``logs/``. Diffusion nightly steps may also appear in [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml) on CI.
+
+    **Ad hoc:** run a single test file or marker expression manually:
 
     ```bash
     cd tests
+    pytest -s -v test_xxxx.py --run-level=full_model
+    pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
     pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
     ```
 
+    To run a standalone perf benchmark (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
+
+    ```bash
+    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+    ```
+
 === "L5 level"
 
     L5 includes stability and reliability testing. Typical commands:

docs/contributing/model/adding_diffusion_model.md

@@ -638,7 +638,7 @@ mkdir -p examples/online_serving/your_model_name
 
 - Script: `examples/offline_inference/your_model_name/end2end.py`
   - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.)
-  - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only)
+  - Use `from vllm_omni.entrypoints.omni import Omni` for both multi-stage and diffusion-only models
   - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`)
 - Doc: `examples/offline_inference/your_model_name/README.md`
   - Include at least one runnable command, e.g.:

docs/design/module/dit_module.md

@@ -470,7 +470,7 @@ class Attention(nn.Module):
                  role="self", role_category=None, ...):
         # Resolve backend for this role from the active OmniDiffusionConfig
         config = get_current_diffusion_config_or_none()
-        attention_config = config.attention_config if config is not None else None
+        attention_config = config.diffusion_attention_config if config is not None else None
         attn_backend_cls, spec = get_attn_backend_for_role(
             role=role,
             head_size=head_size,
@@ -493,7 +493,7 @@ These backends provide the **kernel implementations** for attention computation.
 
 #### Backend Selection Mechanism
 
-Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
+Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
 
 ```python
 def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
@@ -909,7 +909,7 @@ def initialize_model_parallel(
 
 ```
 1. User Request
-   └─> OmniDiffusion.generate(prompt)
+   └─> Omni.generate(prompt)
        └─> Prepare OmniDiffusionRequest
            └─> DiffusionEngine.step(requests)
 

docs/features/custom_pipeline.md

@@ -10,7 +10,7 @@ This guide demonstrates how to use the newly added features for extending vLLM-O
 Three main features enable custom pipeline extension:
 
 1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
-2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
+2. **`diffusion_load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
 3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
 
 ## Features
@@ -27,15 +27,15 @@ Three main features enable custom pipeline extension:
 
 **Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
 
-### load_format Parameter
+### diffusion_load_format parameter
 
-The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
+The `diffusion_load_format` parameter controls the initial diffusion model load. The relevant values are:
 
 - **`"default"`**: Standard model loading using the model registry (default behavior)
-- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
-- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
+- **`"dummy"`**: Skip the initial model load; use this with `custom_pipeline_args["pipeline_class"]` for the custom pipeline workflow below
+- **`"diffusers"`**: Load through the Hugging Face Diffusers adapter
 
-**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
+**Location:** `vllm_omni/diffusion/data.py` (`OmniDiffusionConfig`)
 
 ### CustomPipelineWorkerExtension
 

docs/getting_started/installation/npu/npu.inc.md

@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 
 ```bash
 # vLLM-Ascend image
-# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.24.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.24.0-a3
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it quay.io/atlas-ci/vllm-ascend:0.24.0-a3 bash
+    -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.24.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -84,8 +84,8 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2: quay.io/ascend/vllm-omni:v0.24.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.24.0-a3
+# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3
 docker run --rm \
     --name vllm-omni-a3 \
     --shm-size=4g \
@@ -110,7 +110,7 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v ~/.cache:/root/.cache \
     -p 8091:8091 \
-    -it quay.io/ascend/vllm-omni:v0.24.0-a3 bash
+    -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash
 ```
 
 !!! tip

docs/mkdocs/hooks/generate_examples.py

@@ -16,8 +16,35 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MODEL_DISPLAY_NAMES_FILE = EXAMPLE_DIR / "model_display_names.yml"
 MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
+SERVING_MODE_TITLES = {
+    "offline_inference": "Offline inference",
+    "online_serving": "Online serving",
+}
+
+
+def load_model_display_names() -> dict[str, str]:
+    try:
+        with open(MODEL_DISPLAY_NAMES_FILE, encoding="utf-8") as f:
+            model_display_names = yaml.safe_load(f)
+    except (OSError, yaml.YAMLError) as exc:
+        raise ValueError(f"Failed to load MODEL_DISPLAY_NAMES_FILE at {MODEL_DISPLAY_NAMES_FILE}: {exc}") from exc
+
+    if not isinstance(model_display_names, dict) or not all(
+        isinstance(key, str) and isinstance(value, str) for key, value in model_display_names.items()
+    ):
+        raise ValueError(
+            "MODEL_DISPLAY_NAMES_FILE "
+            f"at {MODEL_DISPLAY_NAMES_FILE} must contain a YAML mapping "
+            "of string model IDs to string display names"
+        )
+    return model_display_names
+
+
+MODEL_DISPLAY_NAMES = load_model_display_names()
+
 
 def fix_case(text: str) -> str:
     subs = {
@@ -54,6 +81,7 @@ class Example:
         main_file (Path): The main file in the directory.
         other_files (list[Path]): list of other files in the directory.
         title (str): The title of the document.
+        nav_title (str): The concise title used in navigation.
 
     Methods:
         __post_init__(): Initializes the main_file, other_files, and title attributes.
@@ -78,6 +106,16 @@ class Example:
     def is_code(self) -> bool:
         return self.main_file.suffix != ".md"
 
+    @property
+    def model_display_name(self) -> str | None:
+        if self.category not in SERVING_MODE_TITLES:
+            return None
+        return MODEL_DISPLAY_NAMES.get(self.path.stem)
+
+    @property
+    def nav_title(self) -> str:
+        return self.model_display_name or self.title
+
     def determine_main_file(self) -> Path:
         """
         Determines the main file in the given path.
@@ -151,13 +189,25 @@ class Example:
         return [file for file in self.path.rglob("*") if is_other_file(file)]
 
     def determine_title(self) -> str:
+        model_display_name = self.model_display_name
         if not self.is_code:
             # Specify encoding for building on Windows
             with open(self.main_file, encoding="utf-8") as f:
                 first_line = f.readline().strip()
             match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            if model_display_name:
+                expected_title = f"{model_display_name}: {SERVING_MODE_TITLES[self.category]}"
+                actual_title = match.group("title") if match else first_line
+                if actual_title != expected_title:
+                    raise ValueError(
+                        f"Model example title mismatch in {self.main_file}: "
+                        f"expected '# {expected_title}', got {first_line!r}"
+                    )
+                return expected_title
             if match:
                 return match.group("title")
+        elif model_display_name:
+            raise ValueError(f"Mapped model example must use a Markdown README: {self.path}")
         return fix_case(self.path.stem.replace("_", " ").title())
 
     def fix_relative_links(self, content: str) -> str:
@@ -300,7 +350,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: rel_path.as_posix()})
+            category_items.append({example.nav_title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")

docs/serving/audio_generate_api.md

@@ -66,7 +66,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 
 #### Diffusion Parameters
@@ -90,7 +90,6 @@ Returns binary audio data with the appropriate `Content-Type` header:
 | `mp3` | `audio/mpeg` |
 | `flac` | `audio/flac` |
 | `pcm` | `audio/pcm` |
-| `aac` | `audio/aac` |
 | `opus` | `audio/opus` |
 
 ## Examples
@@ -285,7 +284,7 @@ Pydantic validation failure (e.g. invalid `response_format`, `speed` out of rang
     "detail": [
         {
             "type": "literal_error",
-            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3' or 'opus'",
             ...
         }
     ]

docs/serving/image_edit_api.md

@@ -252,7 +252,7 @@ Validation errors (missing required fields):
 
 ```bash
 # Check if server is responding
-curl -X http://localhost:8000/v1/images/edit \
+curl http://localhost:8000/v1/images/edits \
   -F "prompt='test'"
 ```
 

docs/serving/image_generation_api.md

@@ -138,7 +138,7 @@ Content-Type: application/json
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `n` | integer | 1 | Number of images to generate (1-10) |
 | `size` | string | model defaults | Image dimensions in WxH format (e.g., "1024x1024", "512x512") |
-| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `response_format` | string | "b64_json" | Response format (`"b64_json"` or `"file"`) |
 | `user` | string | null | User identifier for tracking |
 
 #### vllm-omni Extension Parameters

docs/serving/speech_api.md

@@ -109,7 +109,7 @@ Content-Type: application/json
 | `input` | string | **required** | The text to synthesize into speech |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
 
 #### vLLM-Omni Extension Parameters

docs/user_guide/diffusion/attention_backends.md

@@ -83,7 +83,7 @@ When constructing `OmniDiffusionConfig` directly:
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
 
 config = OmniDiffusionConfig(
-    attention_config=AttentionConfig(
+    diffusion_attention_config=AttentionConfig(
         default=AttentionSpec(backend="FLASH_ATTN"),
         per_role={
             "cross": AttentionSpec(backend="TORCH_SDPA"),

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -33,7 +33,7 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
 
 **CLI:**
 ```bash
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
 ```
 
 ### To Support a Model
@@ -82,6 +82,33 @@ mutual exclusion: when one group runs, the other moves to CPU.
 - Support single GPU only for now
 
 
+### Component offloading for split models (e.g. Cosmos3)
+
+Some models split their transformer into mutually-exclusive *components* that run
+in different phases of a single forward pass rather than as separate pipeline
+components -- e.g. Cosmos3's understanding (reasoner) component runs once per
+generation while the generation (generator) component runs every denoising step.
+Such models have no separate text encoder to swap against, so the transformer
+owns a small model-local offload path and wraps each phase with
+`with self._offload_context(name):`
+
+```python
+class Cosmos3VFMTransformer(nn.Module):
+    def forward(self, ...):
+        with self._offload_context("reasoner"):
+            ...  # understanding pass, runs once
+        with self._offload_context("generator"):
+            ...  # denoising pass, runs every step
+```
+
+Model-level offloading then keeps exactly one component GPU-resident at a time
+(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
+pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
+transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
+works for these models too -- each component declares its own block container via
+`_layerwise_offload_blocks_attrs`.
+
+
 ## Layerwise (Blockwise) Offloading
 
 ### How It Works
@@ -119,10 +146,10 @@ m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True
 **CLI:**
 ```bash
 # Text-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
 
 # Or image-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
 ```
 
 ### To Support a Model
@@ -145,9 +172,9 @@ class Flux2Transformer2DModel(nn.Module):
 ```
 
 ### Limitations
-- Cold start latency increases because of
-    1) components are loaded to CPU first at the very first during initialization,
-    2) weight consolidation and pinning
+- Cold start latency increases because offloaded components must be moved to CPU
+  during setup; layerwise offload may add extra weight consolidation and pinning
+  work.
 - Performance depends on compute cost and H2D bandwidth as well
 - Support single GPU only for now
 
@@ -178,11 +205,18 @@ Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelH
 
 ```
 OffloadBackend (base class)
-├── ModelLevelOffloadBackend → uses SequentialOffloadHook
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
+│                              (delegates to a pipeline's enable_omni_model_cpu_offload
+│                               for split models like Cosmos3)
 └── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 ```
 
-Factory function `get_offload_backend()` selects the appropriate backend based on configuration.
+Factory function `get_offload_backend()` selects the appropriate backend based on
+configuration.
+
+For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
+`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
+reasoner/generator components inside the model forward pass.
 
 
 ## Supported Models

docs/user_guide/diffusion/frame_interpolation.md

@@ -47,7 +47,8 @@ For supported Wan2.2 pipelines, the execution order is:
 1. Diffusion worker finishes denoising and decodes the raw video tensor.
 2. Worker-side model-specific post-processing runs.
 3. If frame interpolation is enabled, RIFE interpolates the decoded video
-   tensor on the worker side and records a FPS multiplier in `custom_output`.
+   tensor on the worker side and records a FPS multiplier in
+   `metadata.video.video_fps_multiplier`.
 4. The API server receives the already-interpolated video and only performs
    MP4 export.
 

docs/user_guide/diffusion_features.md

@@ -105,7 +105,7 @@ The following tables show which models support each feature:
 
 > Notes:
 
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
 > 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
@@ -144,7 +144,7 @@ The following tables show which models support each feature:
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
-> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
+> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 
 ### VideoGen

docs/user_guide/examples/offline_inference/bagel.md

@@ -257,7 +257,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
 | :------- | :--- | :------ | :---------- |
 | `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
 | `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
-| `--worker-backend` | choice | `process` | `process` or `ray` |
+| `--worker-backend` | choice | `multi_process` | `multi_process` or `ray` |
 | `--ray-address` | string | `None` | Ray cluster address |
 | `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
 | `--log-stats` | flag | `False` | Enable statistics logging |

docs/user_guide/examples/offline_inference/glm_image.md

@@ -26,18 +26,22 @@ Stage 0 (AR Model)                Stage 1 (Diffusion)
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
     outputs = omni.generate(
         "A photorealistic mountain landscape at sunset",
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```
@@ -46,6 +50,7 @@ if __name__ == "__main__":
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
@@ -56,13 +61,16 @@ if __name__ == "__main__":
                 "image": "input.png",
             },
         },
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```

docs/user_guide/examples/offline_inference/mimo_audio.md

@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>.
 

docs/user_guide/examples/offline_inference/qwen2_5_omni.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 

docs/user_guide/examples/offline_inference/qwen3_omni.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
 

docs/user_guide/examples/online_serving/aura_omni.md

@@ -1,4 +1,4 @@
-# AURA Omni Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
 

docs/user_guide/examples/online_serving/mimo_audio.md

@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/mimo_audio>.
 

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.
 

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
 

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -139,7 +139,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 | `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
 | `audio_start` | float | 0.0 | Audio start time in seconds |

docs/user_guide/quantization/autoround.md

@@ -3,7 +3,7 @@
 ## Overview
 
 [AutoRound](https://github.com/intel/auto-round) produces pre-quantized
-checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
+checkpoints for LLMs, VLMs, diffusion models, and world models. vLLM-Omni reads the
 checkpoint's `config.json` and auto-detects
 `quantization_config.quant_method = "auto-round"`.
 
@@ -36,6 +36,13 @@ guide. AutoRound is Intel-supported.
 | Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 | Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 
+### World Model (Cosmos3)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
 | Model | Checkpoint | Scope | Scheme | Backend |

docs/user_guide/quantization/bitsandbytes.md

@@ -0,0 +1,120 @@
+# BitsAndBytes W4 Quantization
+
+## Overview
+
+BitsAndBytes 4-bit quantization supports weight-only NF4/FP4 diffusion transformer
+inference on CUDA GPUs. It quantizes BF16/FP16 weights at load time; activations
+stay in BF16/FP16. No pre-quantized checkpoint is required.
+
+This is an online (dynamic) path: load the normal HuggingFace checkpoint and
+quantize during model loading via the `bitsandbytes` CUDA kernels.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA CUDA GPU (SM 75+) | ✅ |
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ❌ |
+| Intel XPU | ❌ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in
+`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`).
+
+Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | CUDA | Mode | Recommendation |
+|-------|-----------|:----:|------|----------------|
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs |
+
+Other diffusion models may work if their transformer uses supported linear
+layers, but they are not validated in this guide.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
+| Qwen3-TTS | TTS language-model stage | Not validated | No BitsAndBytes TTS stage support is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes")
+
+omni_with_skips = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization_config={
+        "method": "bitsandbytes",
+        "ignored_layers": ["to_out"],
+    },
+)
+
+outputs = omni.generate(
+    "A cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes --ignored-layers "to_out"
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"bitsandbytes"`) |
+| `quant_type` | str | `"nf4"` | 4-bit data type: `"nf4"` (recommended) or `"fp4"` |
+| `compress_statistics` | bool | `True` | Double-quantize block scaling statistics for better accuracy |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+
+## Validation and Notes
+
+On Z-Image-Turbo (single GPU, 1024×1024, 50 steps), BitsAndBytes W4 typically
+reduces peak VRAM from roughly 24.5 GiB (BF16) to roughly 17 GiB. Compare output
+quality against a BF16 baseline before enabling on new models.
+
+Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for
+BitsAndBytes in diffusion models; each rank quantizes its own weight shard
+independently.
+
+If quality regresses, use `ignored_layers` to keep sensitive projections in BF16
+(for example `to_out` or `w2`).
+
+```python
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization="bitsandbytes",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Only add a new model to the supported table after comparing the BitsAndBytes
+output against a BF16 baseline and documenting any required `ignored_layers`.

examples/model_display_names.yml

@@ -0,0 +1,12 @@
+# Canonical display names for model examples that have both offline and online pages.
+# The docs generator combines these names with the serving mode for each page H1.
+aura_omni: AURA Omni
+covo_audio: Covo-Audio-Chat
+dynin_omni: Dynin-Omni
+lance: Lance
+mimo_audio: MiMo-Audio
+ming_flash_omni: Ming-flash-omni 2.0
+qwen2_5_omni: Qwen2.5-Omni
+qwen3_omni: Qwen3-Omni
+sensenova_u1: SenseNova-U1
+step_audio2: Step-Audio2

examples/offline_inference/aura_omni/README.md

@@ -1,4 +1,4 @@
-# AURA Omni
+# AURA Omni: Offline inference
 
 This example runs the native AURA Omni pipeline offline:
 

examples/offline_inference/covo_audio/README.md

@@ -1,4 +1,4 @@
-# Covo-Audio-Chat (Offline Inference)
+# Covo-Audio-Chat: Offline inference
 
 ## Setup
 

examples/offline_inference/dynin_omni/README.md

@@ -1,4 +1,4 @@
-# Dynin-Omni Offline End2End Example
+# Dynin-Omni: Offline inference
 
 This folder contains a unified offline inference entrypoint:
 

examples/offline_inference/lance/README.md

@@ -1,4 +1,4 @@
-# Lance (ByteDance) — offline inference
+# Lance: Offline inference
 
 [Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
 autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is

examples/offline_inference/lance/end2end.py

@@ -293,16 +293,19 @@ def main():
 
     if args.modality in ("img2text", "video2text"):
         for i, req_output in enumerate(outputs):
-            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
-            if text is None and hasattr(req_output, "custom_output"):
-                text = (req_output.custom_output or {}).get("text_output")
+            multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            text = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+                text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+            )
+            text = text or getattr(req_output, "output", None) or getattr(req_output, "text", None)
             print(f"[Output {i}] {text}")
         return
 
     if args.modality in ("text2video", "video2video", "image2video"):
         for i, req_output in enumerate(outputs):
-            custom = getattr(req_output, "custom_output", None) or {}
-            frames = custom.get("video_frames")
+            frames = getattr(req_output, "images", None)
             if not frames:
                 print(f"[Output {i}] no video frames returned")
                 continue

examples/offline_inference/lance/gradio_demo.py

@@ -835,12 +835,15 @@ def run_task(
     print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True)
 
     outputs = list(omni.generate(prompts=formatted, sampling_params_list=params))
-    custom = getattr(outputs[0], "custom_output", None) or {}
-    # Lance/Bagel return text at outputs[0].output (string) or
-    # custom_output["text_output"].
+    multimodal_output = getattr(outputs[0], "multimodal_output", {}) or {}
+    metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+    text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
     text_out = None
     if task in {"x2t_image", "x2t_video"}:
-        text_out = custom.get("text_output") or getattr(outputs[0], "output", None)
+        text_out = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+            text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+        )
+        text_out = text_out or getattr(outputs[0], "output", None)
         if not isinstance(text_out, str):
             text_out = None
 
@@ -856,13 +859,11 @@ def run_task(
             gr.update(value=str(text_out or "(empty)"), visible=True),
         )
     if task in {"t2v", "i2v", "video_edit"}:
-        frames = custom.get("video_frames")
+        frames = getattr(outputs[0], "images", None)
         if frames is None:
             raise gr.Error("No video frames in output.")
         return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False))
-    # image tasks — orchestrator strips ``outputs[0].output`` across the
-    # IPC boundary but preserves ``custom_output``, so prefer the latter.
-    img_out = custom.get("image") or getattr(outputs[0], "output", None)
+    img_out = (getattr(outputs[0], "images", None) or [None])[0] or getattr(outputs[0], "output", None)
     if img_out is None:
         raise gr.Error("No image in output.")
     return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False))

examples/offline_inference/mimo_audio/README.md

@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
 

examples/offline_inference/ming_flash_omni/README.md

@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Offline inference
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 

examples/offline_inference/qwen2_5_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

examples/offline_inference/qwen3_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

examples/offline_inference/sensenova_u1/README.md

@@ -1,4 +1,4 @@
-# SenseNova-U1-8B-MoT
+# SenseNova-U1: Offline inference
 
 ## Architecture
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -278,12 +278,17 @@ def main():
     )
 
     for req_output in outputs:
-        custom = getattr(req_output, "_custom_output", {}) or {}
-        if args.print_think and custom.get("think_text"):
-            print(f"[Think]\n{custom['think_text']}\n")
+        multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+        metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+        text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+        think_text = text_metadata.get("think_text") if isinstance(text_metadata, dict) else None
+        if args.print_think and think_text:
+            print(f"[Think]\n{think_text}\n")
 
         if is_text_output:
-            text = custom.get("text_output", "")
+            text = multimodal_output.get("text", "") if isinstance(multimodal_output, dict) else ""
+            if not text:
+                text = text_metadata.get("text_output", "") if isinstance(text_metadata, dict) else ""
             if not text:
                 text = getattr(req_output, "text", "") or ""
             print(f"[Response]\n{text}")

examples/offline_inference/step_audio2/README.md

@@ -1,4 +1,4 @@
-# Step-Audio2 Offline Inference Examples
+# Step-Audio2: Offline inference
 
 This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -202,9 +202,10 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8"],
+        choices=["fp8", "int8", "bitsandbytes"],
         help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), "
+        "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). "
         "Default: None (no quantization, uses BF16).",
     )
     parser.add_argument(

examples/online_serving/aura_omni/README.md

@@ -1,4 +1,4 @@
-# AURA Omni Native Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
 

examples/online_serving/covo_audio/README.md

@@ -1,4 +1,4 @@
-# Covo-Audio-Chat
+# Covo-Audio-Chat: Online serving
 
 ## Installation
 

examples/online_serving/dynin_omni/README.md

@@ -1,4 +1,4 @@
-# Dynin-Omni Online Serving Example
+# Dynin-Omni: Online serving
 
 ## Installation
 

examples/online_serving/lance/README.md

@@ -1,4 +1,4 @@
-# Lance — online serving
+# Lance: Online serving
 
 OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
 

examples/online_serving/mimo_audio/README.md

@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/ming_flash_omni/README.md

@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Online serving
 
 ## Installation
 

examples/online_serving/qwen2_5_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/qwen3_omni/README.md

@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 ## 🛠️ Installation
 

examples/online_serving/sensenova_u1/README.md

@@ -1,4 +1,4 @@
-# SenseNova-U1 Online Serving
+# SenseNova-U1: Online serving
 
 ## Launch the Server
 

examples/online_serving/step_audio2/README.md

@@ -1,4 +1,4 @@
-# Step-Audio2 Online Serving
+# Step-Audio2: Online serving
 
 This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
 

recipes/cosmos3/Cosmos3-Nano.md

@@ -422,3 +422,185 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
   `generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared
   once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
   unknown keys for the model are dropped.
+
+## NPU
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM version: match the repository requirements from your current checkout
+- vLLM-Ascend version: match the repository requirements from your current checkout
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+Requires the `vllm-omni` package (or the `quay.io/atlas-ci/vllm-ascend` A2 / A3 container),
+which provides the `vllm serve … --omni` entrypoint used below.
+
+Safety guardrails are **on by default** (NVIDIA Open Model License). They load
+the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
+
+1. `pip install cosmos-guardrail`
+2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
+3. Export a token with access: `export HF_TOKEN=hf_...`
+
+Then launch the recommended server:
+
+```bash
+vllm serve nvidia/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --init-timeout 1800
+```
+
+To run **without** guardrails (you are responsible for license compliance),
+add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
+add `--tensor-parallel-size 8`. `--quantization fp8` and
+`--enable-layerwise-offload` are not supported on NPU.
+The pipeline auto-resolves from `model_index.json`; pass
+`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
+
+#### Verification
+
+Best quality uses the JSON-upsampled prompts from `assets/` (download with
+`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
+self-contained examples:
+
+```bash
+curl http://localhost:8000/v1/models
+
+# Text-to-image -> /v1/images/generations  (1024x1024, 10 steps; base64 PNG)
+curl -sS -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia/Cosmos3-Nano",
+    "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
+    "negative_prompt": "blurry, distorted, low quality",
+    "size": "1024x1024", "n": 1, "response_format": "b64_json",
+    "num_inference_steps": 10, "guidance_scale": 7.0, "seed": 42
+  }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
+
+# Text-to-video -> /v1/videos/sync  (720p, 49 frames @ 24fps)
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=49" -F "fps=24" \
+  -F "num_inference_steps=20" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "seed=123" \
+  -o cosmos3_t2v.mp4
+
+# Image-to-video -> /v1/videos/sync with an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "size=1280x720" -F "num_frames=25" -F "fps=8" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.jpg;type=image/jpeg" \
+  -o cosmos3_i2v.mp4
+
+# Video-to-video -> /v1/videos/sync with an uploaded reference video
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "size=1280x720" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.mp4;type=video/mp4" \
+  -o cosmos3_v2v.mp4
+```
+
+#### Notes
+
+- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
+  - T2I 1024² — 10 steps → ~8 s
+  - T2V 1280×720 @ 20 steps — 49 frames → ~55 s
+  - I2V 1280×720 @ 10 steps — 25 frames → ~25 s
+  - V2V 480×320 @ 10 steps — 17 frames → ~12 s
+- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
+  full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
+  ~33 GB on disk.
+- **Determinism:** identical seed reproduces identical output on the same
+  hardware; outputs are not bit-identical across different GPU/NPU types.
+- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
+  1280×720, 35 steps, guidance 6.0, `flow_shift=10.0`.
+- **Key flags / params:** `--no-guardrails` (optional, to disable guardrails), `--init-timeout 1800`
+  (for model loading), `--tensor-parallel-size 8` for multi-NPU, and
+  `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
+- **Known limitations:**
+  - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
+    resolution-parsing bug; basic V2V without transfer hints works.
+  - FP8 online quantization and layerwise offload are not supported on NPU.
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+The same offline task examples run on NPU; pass model-specific knobs via
+`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a
+quick local run (install `cosmos-guardrail` + accept the gated repo to enable them).
+
+```bash
+# Text-to-image -> examples/offline_inference/text_to_image
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
+  --negative-prompt "blurry, distorted, low quality" \
+  --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
+  --extra-body '{"flow_shift": 3.0, "guardrails": false}' \
+  --output cosmos3_t2i.png
+
+# Text-to-video -> examples/offline_inference/text_to_video
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A robot arm is cleaning a plate in the kitchen." \
+  --negative-prompt "blurry, distorted, low quality, jittery, deformed" \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
+                 "use_resolution_template": false, "use_duration_template": false}' \
+  --output cosmos3_t2v.mp4
+
+# Image-to-video -> examples/offline_inference/image_to_video
+# (Cosmos3 bundles example frames under assets/; any RGB image works too.)
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
+  --prompt "The scene comes to life with smooth, natural motion." \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
+  --output cosmos3_i2v.mp4
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
+ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4
+```
+
+#### Notes
+
+- Guardrails are on by default on NPU with `cosmos-guardrail` installed. Pass
+  `"guardrails": false` in `--extra-body` to disable them (there is no
+  `--no-guardrails` flag for the offline scripts).
+- Video at the 189-frame default takes ~15 min/clip on 1 NPU; reduce
+  `--num-frames` for faster iteration.

tests/ar_diffusion/test_engine_routing.py

@@ -134,3 +134,17 @@ def test_runner_rejects_multi_sequence_config():
     fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
     with pytest.raises(ValueError, match="max_num_seqs=1"):
         ARDiffusionModelRunner._preallocate_kv_cache(fake)
+
+
+def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
+    """Regression: the worker passes ``kv_prefetch_jobs`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_jobs" in params, cls

tests/ar_diffusion/test_pipeline_bridge.py

@@ -203,7 +203,7 @@ def test_failed_forward_tears_down_session(monkeypatch):
     runner.device = None
     runner._perf_e2e_times = []
 
-    def boom(self, req):
+    def boom(self, req, kv_prefetch_jobs=None):
         raise RuntimeError("layer 17 exploded")
 
     monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -36,7 +36,7 @@ import pytest
 
 from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
 
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model, pytest.mark.local_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -807,7 +807,7 @@ def assert_result(
 
 def _default_benchmark_endpoint_for_task(task: str) -> str:
     """Return the default client-side benchmark endpoint for a diffusion task."""
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
         return "/v1/chat/completions"

tests/dfx/perf/tests/test_cosmos3_vllm_omni.json

@@ -0,0 +1,213 @@
+[
+    {
+        "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/generations",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps4",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 4,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "extra-body": {
+                    "guidance_scale": 7.0,
+                    "flow_shift": 3.0,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 1.107,
+                    "latency_mean": 0.8929,
+                    "peak_memory_mb_mean": 80000.0
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_t2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "t2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "A robot arm is cleaning a plate in the kitchen.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "use_resolution_template": false,
+                    "use_duration_template": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.6824,
+                    "peak_memory_mb_mean": 30155.4
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_i2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "The scene comes to life with smooth, natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.968,
+                    "peak_memory_mb_mean": 30443.6
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_cosmos3_v2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "v2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "Continue the same scene with smooth natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "condition_frame_indexes_vision": [
+                        0,
+                        1
+                    ],
+                    "condition_video_keep": "first"
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 37.2117,
+                    "peak_memory_mb_mean": 30135.6
+                }
+            }
+        ]
+    }
+]

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -7,6 +7,7 @@
         "server_params": {
             "model": "tencent/HunyuanImage-3.0-Instruct",
             "serve_args": {
+                "trust-remote-code": true,
                 "deploy-config-inline": {
                     "type": "yaml",
                     "content": {

tests/diffusion/distributed/test_hsdp.py

@@ -1,19 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
-These tests verify HSDP configuration logic without requiring a distributed environment.
-"""
+import gc
+import os
+import socket
 
 import pytest
+import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.tensor import DeviceMesh, DTensor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    _unshardable_parameters,
+    shard_model,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
+class _PackedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
+        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
+
+
+class _PackedModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block = _PackedBlock()
+        self.root_weight = nn.Parameter(torch.ones(2))
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield
+        return
+
+    master_port = _find_free_port()
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+            os.environ.pop(key, None)
+        gc.collect()
+
+
+def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
+    model = _PackedModel()
+    ignored_params = _unshardable_parameters(model)
+    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
+    assert ignored_params == expected_ignored
+
+    packed_weight = model.block.packed_weight
+    input_global_scale = model.block.input_global_scale
+    shard_model(
+        model,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "block"],
+        ignored_params=ignored_params,
+    )
+
+    assert isinstance(model.block.weight, DTensor)
+    assert isinstance(model.root_weight, DTensor)
+    assert model.block.packed_weight is packed_weight
+    assert model.block.input_global_scale is input_global_scale
+    assert not isinstance(model.block.packed_weight, DTensor)
+    assert not isinstance(model.block.input_global_scale, DTensor)
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -179,6 +179,59 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_hsdp_processes_quantized_weights_before_sharding(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: events.append("prepare"),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "process", "prepare", "shard"]
+
+
 def test_get_all_weights(prefetch_helios_model, mock_tp_group):
     """Ensure that get all weights on a tiny model resolves to nonempty weights."""
     od_config = OmniDiffusionConfig(

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -42,6 +42,53 @@ class _QuantizedPackedModelOptModel(nn.Module):
         )
 
 
+class _RemappedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Linear(2, 2, bias=False)
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return {
+            "transformer.orig.proj.weight": "runtime.proj.weight",
+        }.get(name, name)
+
+
+class _QuantizedRemappedModelOptModel(nn.Module):
+    """FP8 target whose params live under a namespace the raw checkpoint keys
+    never use, reachable only via ``remap_checkpoint_key``.
+
+    Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot
+    resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook
+    the adapter treats every tensor as unresolved and silently drops the scales.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        prefix = "transformer.orig.proj."
+        if name.startswith(prefix):
+            return "runtime.proj." + name[len(prefix) :]
+        return name
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -92,3 +139,62 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
         "transformer.block.to_q.weight_scale",
         "transformer.block.to_q.input_scale",
     ]
+
+
+def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype():
+    model = _RemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight", fp8_weight),
+                    ("transformer.orig.proj.weight_scale", scale),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["runtime.proj.weight"]
+    assert adapted[0][1].dtype == model.runtime.proj.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
+    """Regression: scales were dropped for a quantized (FP8) target that only the
+    model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are
+    emitted and the FP8 weight passes through unchanged (no dequantization)."""
+    model = _QuantizedRemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    input_scale = torch.tensor([1.0], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight_scale", weight_scale),
+                    ("transformer.orig.proj.input_scale", input_scale),
+                    ("transformer.orig.proj.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    emitted = dict(adapted)
+    assert len(adapted) == len(emitted) == 3
+
+    # Both scales are retained (the dropped-scale regression). They keep their
+    # source-checkpoint names; the model's load_weights remaps them downstream.
+    assert torch.equal(emitted["transformer.orig.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["transformer.orig.proj.input_scale"], input_scale)
+
+    # The FP8 weight is passed through unchanged under the remapped target name
+    # (the quantized param is FP8, so it must NOT be dequantized). Compare via
+    # float32 since torch.equal has no FP8 CPU kernel.
+    assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
+    assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -20,6 +20,16 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_pipeline_declares_layerwise_offload_components() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    assert Cosmos3OmniDiffusersPipeline._dit_modules == ["transformer.language_model", "transformer"]
+    assert Cosmos3OmniDiffusersPipeline._encoder_modules == []
+    assert Cosmos3OmniDiffusersPipeline._vae_modules == ["vae"]
+    assert Cosmos3OmniDiffusersPipeline._resident_modules == []
+    assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
+
+
 class StubScheduler:
     def __init__(
         self,
@@ -284,7 +294,6 @@ def test_pipeline_registered_and_exported() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
     from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
     from vllm_omni.diffusion.registry import (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
         _DIFFUSION_IR_OP_PRIORITY_FUNCS,
         _DIFFUSION_MODELS,
         _DIFFUSION_POST_PROCESS_FUNCS,
@@ -301,9 +310,6 @@ def test_pipeline_registered_and_exported() -> None:
     )
     assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
     assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
-    assert (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func"
-    )
     assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func"
     assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
     assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
@@ -708,11 +714,11 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None:
 def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
         RoboLabPolicyInputs,
-        get_cosmos3_action_post_process_func,
+        get_cosmos3_post_process_func,
         make_robolab_action_postprocess_inputs,
     )
 
-    func = get_cosmos3_action_post_process_func(SimpleNamespace())
+    func = get_cosmos3_post_process_func(SimpleNamespace())
     inputs = RoboLabPolicyInputs(
         prompt="Pick the cube.",
         video_tensor=torch.zeros(1, 3, 3, 16, 16),
@@ -736,13 +742,41 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     )
 
     action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
-    custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)}
-    processed = func(action, custom_output=custom_output)
+    processed = func(
+        {
+            "payload": {
+                "actions": action,
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
+                },
+            },
+        }
+    )
 
-    assert processed.shape == (1, 2)
-    assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
-    torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]]))
-    assert "robolab_action_postprocess" not in custom_output
+    processed_action = processed["payload"]["actions"]
+    assert processed_action.shape == (1, 2)
+    assert processed_action.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
+    torch.testing.assert_close(torch.from_numpy(processed_action), torch.tensor([[1.0, 0.25]]))
+    assert processed["metadata"] == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
 
 
 def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -1198,8 +1232,8 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     assert captured["flow_shifts"] == [10.0]
     # Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
     assert captured["scheduler_use_karras_sigmas"] == [False]
-    assert output.custom_output["fps"] == expected_fps
-    assert output.output["video"].device.type == "meta"
+    assert output.output["metadata"]["video"]["fps"] == expected_fps
+    assert output.output["payload"]["video"].device.type == "meta"
 
 
 def test_forward_transfer_runs_multichunk_overlap_path(
@@ -1268,12 +1302,12 @@ def test_forward_transfer_runs_multichunk_overlap_path(
 
     assert captured["conditional_frames"] == [2, 1]
     assert len(captured["decode_calls"]) == 2
-    assert output.output["video"].shape == (1, 3, 8, 16, 16)
+    assert output.output["payload"]["video"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(
-        output.output["video"][0, 0, :, 0, 0],
+        output.output["payload"]["video"][0, 0, :, 0, 0],
         torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]),
     )
-    assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16)
+    assert output.output["metadata"]["transfer"]["controls"]["edge"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0))
@@ -1497,8 +1531,13 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "action_only_output" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert "common" not in output.output["metadata"]
 
     def test_forward_dispatches_robolab_policy_flow(
         self,
@@ -1582,12 +1621,15 @@ class TestForwardRouting:
         assert captured["prepare_action_video"]["kwargs"] == {"image_size": None}
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
         assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1]
-        assert output.output == {}
-        assert output.custom_output["action_only_output"] is True
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "actions" not in output.custom_output
-        assert "robolab_action_postprocess" in output.custom_output
-        assert "robolab_policy_inputs" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert output.output["metadata"]["common"]["action_only_output"] is True
+        assert "robolab_action_postprocess" in output.output["metadata"]["internal"]
+        assert "robolab_policy_inputs" not in output.output["metadata"]
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -10,9 +10,19 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture
+def accelerator_device() -> torch.device:
+    """Provide an accelerator device, skipping when none is available."""
+    if current_omni_platform.get_device_count() == 0:
+        pytest.skip("Accelerator required for this test")
+    return current_omni_platform.get_torch_device(0)
+
+
 def _tiny_cosmos3_config(**overrides):
     config = {
         "hidden_size": 8,
@@ -88,7 +98,7 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) ->
 
 
 def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
 
     model = object.__new__(Cosmos3VFMTransformer)
     nn.Module.__init__(model)
@@ -104,6 +114,7 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
     ]
     assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
     assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
     assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
 
     model.latent_patch_size = 2
@@ -131,6 +142,97 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    class ToyReasonerLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, freqs):
+            del freqs
+            kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1)
+            return hidden + self.weight.to(hidden.dtype), kv, kv
+
+    class ToyGeneratorLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, **kwargs):
+            del kwargs
+            return hidden + self.weight.to(hidden.dtype)
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([ToyReasonerLayer()])
+    model.gen_layers = nn.ModuleList([ToyGeneratorLayer()])
+    model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"}
+    assert model.device == torch.device("cpu")
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+    # The reasoner runs once then the generator component stays resident for GEN.
+    assert model._active_model_cpu_offload_component == "generator"
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+
+
+def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
+    monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.to(accelerator_device)
+
+    reasoner_param = model.language_model.layers[0].weight
+    generator_param = model.gen_layers[0].weight
+
+    model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    # On enable, every group is parked on CPU until a phase activates it.
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("reasoner")
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("generator")
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device == accelerator_device
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+    # Disable restores both components onto the device.
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device == accelerator_device
+
+
 def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as qwen_image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_postprocess_preserves_output_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
+) -> None:
+    model_path = tmp_path / "model"
+    vae_path = model_path / "vae"
+    vae_path.mkdir(parents=True)
+    (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8")
+
+    processed = object()
+    captured: dict[str, object] = {}
+
+    class FakeImageProcessor:
+        def __init__(self, *, vae_scale_factor: int) -> None:
+            captured["vae_scale_factor"] = vae_scale_factor
+
+        def postprocess(self, images: torch.Tensor) -> object:
+            captured["images"] = images
+            return processed
+
+    monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor)
+    postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path)))
+
+    image = torch.zeros(1, 3, 4, 4)
+    output = postprocess(
+        {
+            "payload": {"image": image},
+            "metadata": {"prompt_embeddings": {"present": True}},
+        }
+    )
+
+    assert captured["vae_scale_factor"] == 16
+    assert captured["images"] is image
+    assert output == {
+        "payload": {"image": processed},
+        "metadata": {"prompt_embeddings": {"present": True}},
+    }

tests/diffusion/offloader/test_module_collector.py

@@ -149,6 +149,40 @@ class MultiVaePipeline(nn.Module, SupportsComponentDiscovery):
         self.audio_vae = nn.Linear(10, 10)
 
 
+class SiblingDiTsPipeline(nn.Module, SupportsComponentDiscovery):
+    """Two independent DiTs (e.g. a MoE transformer pair)."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.transformer_2 = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class NestedDiTPipeline(nn.Module, SupportsComponentDiscovery):
+    """Cosmos3-like: a DiT plus one of its own submodules listed as a DiT.
+
+    The nested entry exists so offloading can treat the inner stack as an
+    independent ring; the outer DiT's _hsdp_shard_conditions still own all
+    nested blocks for sharding. Order is nested-first, matching Cosmos3.
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Module()
+        self.transformer.language_model = nn.Linear(10, 10)
+        self.transformer.gen_layers = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -238,3 +272,35 @@ class TestProtocolDiscovery:
         assert len(result.vaes) == 2
         assert result.vaes[0] is pipeline.vae
         assert result.vaes[1] is pipeline.audio_vae
+
+
+class TestOutermostDits:
+    """PipelineModules.outermost_dits() drops DiTs nested inside another DiT.
+
+    HSDP shards each returned DiT via its own _hsdp_shard_conditions, so a DiT
+    nested inside another discovered DiT must be excluded: it is already covered
+    by its ancestor's conditions and may not declare conditions of its own.
+    """
+
+    def test_keeps_independent_sibling_dits(self):
+        result = ModuleDiscovery.discover(SiblingDiTsPipeline())
+
+        names, modules = result.outermost_dits()
+
+        assert names == ["transformer", "transformer_2"]
+        assert [id(m) for m in modules] == [id(m) for m in result.dits]
+
+    def test_drops_nested_dit_keeping_only_ancestor(self):
+        pipeline = NestedDiTPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        # Discovery sees both, nested-first (matches Cosmos3's _dit_modules order).
+        assert result.dit_names == ["transformer.language_model", "transformer"]
+
+        names, modules = result.outermost_dits()
+
+        # Only the ancestor survives; sharding the nested condition-less stack
+        # again would double-wrap its blocks and raise in apply_hsdp_to_model.
+        assert names == ["transformer"]
+        assert len(modules) == 1
+        assert modules[0] is pipeline.transformer

tests/diffusion/offloader/test_sequential_backend.py

@@ -7,7 +7,8 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -41,6 +42,40 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
+def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.enable_args = None
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            self.enable_args = kwargs
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert backend.enabled is True
+    assert pipeline.enable_args == {
+        "device": torch.device("cpu"),
+        "pin_memory": False,
+        "use_hsdp": False,
+    }
+
+    backend.disable()
+
+    assert backend.enabled is False
+    assert pipeline.disable_called is True
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""

tests/diffusion/quantization/test_bitsandbytes_config.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BitsAndBytes quantization config."""
+
+import importlib.util
+import sys
+import types
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from torch.nn import Module, Parameter
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.")
+
+bitsandbytes_available = pytest.mark.skipif(
+    importlib.util.find_spec("bitsandbytes") is None,
+    reason="bitsandbytes package not installed",
+)
+
+
+def _ensure_bitsandbytes_importable(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Inject a stub bitsandbytes package when the optional dep is absent.
+
+    ``mocker.patch("bitsandbytes...")`` imports the target module; without a
+    stub, core-model CI without bitsandbytes raises ModuleNotFoundError.
+    """
+    if importlib.util.find_spec("bitsandbytes") is not None:
+        return
+    bnb = types.ModuleType("bitsandbytes")
+    bnb_functional = types.ModuleType("bitsandbytes.functional")
+    bnb.functional = bnb_functional
+    monkeypatch.setitem(sys.modules, "bitsandbytes", bnb)
+    monkeypatch.setitem(sys.modules, "bitsandbytes.functional", bnb_functional)
+
+
+def test_bitsandbytes_config_creation():
+    """Test that BitsAndBytes config can be created."""
+    config = build_quant_config("bitsandbytes")
+    assert config is not None
+    assert config.get_name() == "bitsandbytes"
+
+
+def test_bitsandbytes_config_with_custom_params():
+    """Test BitsAndBytes config with custom parameters."""
+    config = build_quant_config(
+        "bitsandbytes",
+        quant_type="fp4",
+        compress_statistics=False,
+        ignored_layers=["to_out"],
+    )
+    assert config is not None
+    assert config.quant_type == "fp4"
+    assert config.compress_statistics is False
+    assert "to_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list includes bitsandbytes."""
+    assert "bitsandbytes" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_quantization_integration():
+    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(model="test", quantization_config="bitsandbytes")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "bitsandbytes"
+
+    config2 = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "method": "bitsandbytes",
+            "quant_type": "nf4",
+            "compress_statistics": True,
+        },
+    )
+    assert config2.quantization_config is not None
+    assert config2.quantization_config.get_name() == "bitsandbytes"
+    assert config2.quantization_config.quant_type == "nf4"
+    assert config2.quantization_config.compress_statistics is True
+
+
+def test_quantization_dict_not_mutated():
+    """Test that passing a dict to quantization_config doesn't mutate it."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    original_dict = {"method": "bitsandbytes", "quant_type": "nf4"}
+    dict_copy = original_dict.copy()
+
+    OmniDiffusionConfig(model="test", quantization_config=original_dict)
+
+    assert original_dict == dict_copy
+
+
+def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Test get_quant_method routing for CUDA."""
+    from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+    config = build_quant_config("bitsandbytes")
+
+    def _fake_init(self, quant_config):
+        pass
+
+    layer = mocker.Mock(spec=LinearBase)
+    mocker.patch.object(BnBOnlineLinearMethod, "__init__", _fake_init)
+
+    prefix = "test_layer"
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, BnBOnlineLinearMethod)
+
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+class TestBnBOnlineLinearMethod:
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        config = mocker.Mock()
+        config.quant_type = "nf4"
+        config.compress_statistics = True
+        return config
+
+    @pytest.fixture
+    def mock_deps(self, mocker, monkeypatch: pytest.MonkeyPatch):
+        _ensure_bitsandbytes_importable(monkeypatch)
+        mock_qweight = torch.ones((64, 32), dtype=torch.uint8)
+        mock_quant_state = mocker.Mock()
+        mock_quant = mocker.patch(
+            "bitsandbytes.functional.quantize_4bit",
+            return_value=(mock_qweight, mock_quant_state),
+            create=True,
+        )
+        # x is (2, 16, 32) -> flattened (32, 32); matmul must return (32, 64)
+        # so apply can reshape back to (2, 16, 64).
+        mock_matmul = mocker.patch(
+            "bitsandbytes.matmul_4bit",
+            return_value=torch.randn(32, 64),
+            create=True,
+        )
+        return {
+            "quant": mock_quant,
+            "matmul": mock_matmul,
+            "mock_qweight": mock_qweight,
+            "mock_quant_state": mock_quant_state,
+        }
+
+    def test_process_weights_after_loading(self, mock_deps, mock_quant_config, mocker):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.randn(64, 32))
+        mocker.patch.object(torch.Tensor, "cuda", return_value=layer.weight.data)
+
+        method.process_weights_after_loading(layer)
+
+        mock_deps["quant"].assert_called_once()
+        assert torch.equal(layer.weight.data, mock_deps["mock_qweight"])
+        assert layer.quant_state is mock_deps["mock_quant_state"]
+        assert layer.bnb_shape == (64, 32)
+
+    def test_apply(self, mock_deps, mock_quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.ones((64, 32), dtype=torch.uint8), requires_grad=False)
+        layer.quant_state = mock_deps["mock_quant_state"]
+
+        x = torch.randn(2, 16, 32, dtype=torch.float16)
+        bias = torch.randn(64)
+        output = method.apply(layer, x, bias)
+
+        mock_deps["matmul"].assert_called_once()
+        assert output.shape == (2, 16, 64)
+        assert output.dtype == torch.float16
+
+
+@pytest.fixture
+def quant_config():
+    from vllm_omni.quantization.bitsandbytes_config import DiffusionBitsAndBytesConfig
+
+    return DiffusionBitsAndBytesConfig(
+        quant_type="nf4",
+        compress_statistics=True,
+    )
+
+
+@cuda_available
+@bitsandbytes_available
+class TestCudaBnBSmoke:
+    """Smoke tests using real bitsandbytes CUDA kernels."""
+
+    @pytest.fixture
+    def real_layer(self):
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        layer.logical_widths = [128]
+        layer.input_size_per_partition = 64
+        layer.output_size_per_partition = 128
+        layer.orig_dtype = torch.float16
+        return layer
+
+    def test_real_cuda_quantize_4bit_shape_contract(self, quant_config):
+        import bitsandbytes.functional as bnb_F
+
+        weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
+        qweight, quant_state = bnb_F.quantize_4bit(
+            weight,
+            quant_type=quant_config.quant_type,
+            compress_statistics=quant_config.compress_statistics,
+        )
+
+        assert qweight.dtype == torch.uint8
+        assert quant_state is not None
+
+    def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+        method.process_weights_after_loading(real_layer)
+
+        assert real_layer.weight.dtype == torch.uint8
+        assert hasattr(real_layer, "quant_state")
+        assert real_layer.bnb_shape == (128, 64)
+
+    def test_real_cuda_bnb_apply_forward(self, quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        method.process_weights_after_loading(layer)
+
+        x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda")
+        output = method.apply(layer, x)
+
+        assert output.shape == (2, 16, 128)
+        assert output.dtype == torch.float16

tests/diffusion/test_diffusion_model_runner.py

@@ -509,6 +509,25 @@ def test_split_diffusion_output_by_request_slices_tuple_outputs():
     assert torch.equal(outputs[1].output[1], torch.tensor([20]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
+    """execute_model runs the pipeline forward after the KV receive step."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert runner.pipeline.forward_calls == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):

tests/diffusion/test_diffusion_output_formatter.py

@@ -51,21 +51,23 @@ def _timings() -> DiffusionStepTimings:
     )
 
 
-def test_formatter_preserves_single_video_audio_actions_and_custom_output(
+def test_formatter_preserves_single_video_audio_actions_and_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
-    custom_output = {"from_output": "kept"}
     postprocess_output = normalize_diffusion_postprocess_output(
         {
-            "video": ["frame-0"],
-            "audio": "audio-0",
-            "actions": "action-0",
-            "custom_output": {"from_postprocess": "merged"},
-            "audio_sample_rate": 48000,
-            "fps": 24.0,
-        },
-        custom_output,
+            "payload": {
+                "video": ["frame-0"],
+                "audio": "audio-0",
+                "actions": "action-0",
+            },
+            "metadata": {
+                "audio": {"sample_rate": 48000},
+                "video": {"fps": 24.0},
+                "actions": {"action_mode": "policy"},
+            },
+        }
     )
 
     results = format_diffusion_outputs(
@@ -73,7 +75,6 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
         od_config=_config(),
         diffusion_output=DiffusionOutput(
             output=None,
-            custom_output=custom_output,
             stage_durations={"execute": 1.25},
             peak_memory_mb=321.0,
         ),
@@ -87,12 +88,12 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     assert result.images == ["frame-0"]
     assert result.prompt == "prompt-0"
     assert result.final_output_type == "image"
-    assert result.custom_output == {
-        "from_output": "kept",
-        "from_postprocess": "merged",
-    }
-    assert custom_output == {"from_output": "kept"}
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 48000},
+            "video": {"fps": 24.0},
+            "actions": {"action_mode": "policy"},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 48000,
         "fps": 24.0,
@@ -110,18 +111,123 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     }
 
 
-def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_formatter_normalizes_payload_metadata_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
     postprocess_output = normalize_diffusion_postprocess_output(
-        "caption",
-        {"text_output": "caption"},
+        {
+            "payload": {
+                "video": ["frame-0"],
+                "actions": "action-0",
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": object(),
+                },
+            },
+        }
+    )
+
+    assert postprocess_output.outputs == {
+        "video": ["frame-0"],
+        "actions": "action-0",
+    }
+    assert postprocess_output.primary_key == "video"
+    assert postprocess_output.metadata == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
+    assert "internal" not in postprocess_output.metadata
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.multimodal_output == {
+        "metadata": postprocess_output.metadata,
+        "actions": "action-0",
+    }
+
+
+def test_formatter_maps_trajectory_payload_to_request_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    latents = torch.zeros(2, 1, 4)
+    timesteps = torch.tensor([1.0, 0.5])
+    log_probs = torch.tensor([-0.1, -0.2])
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {
+                "image": "image-0",
+                "trajectory": {
+                    "latents": latents,
+                    "timesteps": timesteps,
+                    "log_probs": log_probs,
+                },
+            },
+            "metadata": {"trajectory": {"type": "denoising"}},
+        }
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == ["image-0"]
+    assert result.latents is latents
+    assert result.trajectory_latents is latents
+    assert result.trajectory_timesteps is timesteps
+    assert result.trajectory_log_probs is log_probs
+    assert result.multimodal_output == {
+        "metadata": {"trajectory": {"type": "denoising"}},
+        "trajectory": {
+            "latents": latents,
+            "timesteps": timesteps,
+            "log_probs": log_probs,
+        },
+    }
+
+
+def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {"text": "caption"},
+            "metadata": {"text": {"text_output": "caption"}},
+        }
     )
 
     [result] = format_diffusion_outputs(
         request=_request("describe this"),
         od_config=_config(),
-        diffusion_output=DiffusionOutput(output="caption"),
-        output_data="caption",
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"payload": {"text": "caption"}},
         postprocess_output=postprocess_output,
         timings=_timings(),
     )
@@ -129,7 +235,10 @@ def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch)
     assert result.images == []
     assert result.prompt == "describe this"
     assert result.final_output_type == "text"
-    assert result.custom_output == {"text_output": "caption"}
+    assert result.multimodal_output == {
+        "metadata": {"text": {"text_output": "caption"}},
+        "text": "caption",
+    }
 
 
 def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
@@ -144,7 +253,7 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
         "_try_load_model_cls",
         lambda _: AudioModel,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak"),
@@ -175,8 +284,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
             "actions": "action-0",
             "audio_sample_rate": 16000,
             "fps": 30.0,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -192,6 +300,10 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
     assert result.prompt == "watch and listen"
     assert result.final_output_type == "image"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 16000},
+            "video": {"fps": 30.0},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 16000,
         "fps": 30.0,
@@ -207,8 +319,7 @@ def test_formatter_preserves_audio_only_postprocess_dict(
         {
             "audio": "waveform",
             "audio_sample_rate": 24000,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -224,6 +335,9 @@ def test_formatter_preserves_audio_only_postprocess_dict(
     assert result.prompt == "speak"
     assert result.final_output_type == "audio"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 24000},
+        },
         "audio": "waveform",
         "audio_sample_rate": 24000,
     }
@@ -238,7 +352,7 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs(
         "_try_load_model_cls",
         lambda _: None,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak", num_outputs_per_prompt=2),
@@ -264,10 +378,8 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
             "video": ["frame-0", "frame-1"],
             "audio": ["audio-0", "audio-1"],
             "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
-            "custom_output": {"shared": True},
             "fps": 12.5,
-        },
-        {},
+        }
     )
 
     results = format_diffusion_outputs(
@@ -282,7 +394,7 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
     assert len(results) == 1
     assert results[0].images == ["frame-0", "frame-1"]
     assert results[0].prompt == "prompt-0"
-    assert results[0].custom_output == {"shared": True}
+    assert results[0].multimodal_output["metadata"] == {"video": {"fps": 12.5}}
     assert results[0].multimodal_output["audio"] == ["audio-0", "audio-1"]
     assert results[0].multimodal_output["fps"] == 12.5
     torch.testing.assert_close(

tests/diffusion/test_diffusion_output_metadata.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.outputs.output_metadata import (
+    strip_internal_metadata,
+    validate_diffusion_metadata,
+    validate_public_diffusion_metadata,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_validate_accepts_canonical_metadata_fields() -> None:
+    validate_diffusion_metadata(
+        {
+            "video": {
+                "fps": 24.0,
+                "video_fps_multiplier": 2,
+                "shape": [1, 3, 16, 64, 64],
+            },
+            "audio": {
+                "sample_rate": 48000,
+                "audiox_task": "tts",
+            },
+            "actions": {
+                "raw_action_dim": 7,
+                "action_mode": "policy",
+                "domain_id": 3,
+            },
+            "transfer": {
+                "controls": {"edge": torch.zeros(1, 3, 4, 8, 8)},
+                "hints": ["edge"],
+            },
+            "text": {
+                "text_output": "caption",
+                "think_text": "reasoning",
+            },
+            "common": {
+                "action_only_output": True,
+            },
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    ("metadata", "error_type"),
+    [
+        ({"video": {"fps": 0}}, ValueError),
+        ({"video": {"fps": "24"}}, TypeError),
+        ({"video": {"video_fps_multiplier": 0}}, ValueError),
+        ({"video": {"shape": [1, 0, 64]}}, ValueError),
+        ({"audio": {"sample_rate": 0}}, ValueError),
+        ({"audio": {"audiox_task": ""}}, ValueError),
+        ({"actions": {"raw_action_dim": 0}}, ValueError),
+        ({"actions": {"action_mode": ""}}, ValueError),
+        ({"actions": {"domain_id": "7"}}, TypeError),
+        ({"transfer": {"controls": {"edge": "not-a-tensor"}}}, TypeError),
+        ({"transfer": {"hints": ["edge", 1]}}, TypeError),
+        ({"text": {"text_output": 1}}, TypeError),
+        ({"text": {"think_text": 1}}, TypeError),
+        ({"common": {"action_only_output": "true"}}, TypeError),
+    ],
+)
+def test_validate_rejects_invalid_canonical_metadata(metadata: dict[str, object], error_type: type[Exception]) -> None:
+    with pytest.raises(error_type):
+        validate_diffusion_metadata(metadata)
+
+
+def test_validate_allows_unknown_model_metadata() -> None:
+    validate_diffusion_metadata(
+        {
+            "model_specific": {"opaque": object()},
+            "video": {"model_key": object()},
+        }
+    )
+
+
+def test_public_metadata_rejects_internal_group() -> None:
+    with pytest.raises(ValueError, match="internal"):
+        validate_public_diffusion_metadata({"internal": {"secret": object()}})
+
+
+def test_strip_internal_metadata_removes_private_group() -> None:
+    metadata = {
+        "actions": {"action_mode": "policy"},
+        "internal": {"robolab_action_postprocess": object()},
+    }
+
+    public_metadata = strip_internal_metadata(metadata)
+
+    assert public_metadata == {"actions": {"action_mode": "policy"}}
+    assert "internal" not in public_metadata

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -16,7 +16,6 @@ from unittest.mock import Mock, patch
 import pytest
 
 from vllm_omni.diffusion.registry import (
-    _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
     _DIFFUSION_IR_OP_PRIORITY_FUNCS,
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
@@ -62,7 +61,6 @@ class TestRegisterDiffusionModel:
         original_models = _DIFFUSION_MODELS.copy()
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
-        original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
         original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
@@ -71,8 +69,6 @@ class TestRegisterDiffusionModel:
         _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
         _DIFFUSION_POST_PROCESS_FUNCS.clear()
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
 
@@ -84,7 +80,6 @@ class TestRegisterDiffusionModel:
             class_name="TestPipeline",
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
-            action_post_process_func_name="test_action_post_process",
             ir_op_priority_func_name="test_ir_op_priority",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
@@ -95,9 +90,21 @@ class TestRegisterDiffusionModel:
         )
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
-        assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
         assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
 
+    def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
+        """Deprecated action hook keyword is accepted but not registered."""
+        register_diffusion_model(
+            model_arch="LegacyActionPipeline",
+            module_name="test_plugin.diffusion.pipeline",
+            class_name="LegacyActionPipeline",
+            post_process_func_name="test_post_process",
+            action_post_process_func_name="test_action_post_process",
+        )
+
+        assert "LegacyActionPipeline" in _DIFFUSION_MODELS
+        assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"
+
 
 class TestWorkerUsesHook:
     """Test that DiffusionWorker resolves model runner via platform hook."""

tests/diffusion/test_diffusion_streaming_output.py

@@ -85,19 +85,20 @@ class _PipelineBackedEngine:
             outputs = self.pipeline.step_outputs(request)
         except Exception as exc:
             outputs = [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
+                _streaming_diffusion_output(chunk=0, finished=False),
                 DiffusionOutput(error=str(exc), finished=True),
             ]
         for output in outputs:
             if output.error is not None:
                 yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
                 continue
-            custom_output = output.custom_output or {}
+            payload, metadata = _streaming_payload_and_metadata(output)
+            stream_metadata = metadata.get("stream", {})
             yield [
                 OmniRequestOutput.from_diffusion(
                     request_id=request.request_id,
-                    images=list(custom_output.get("images") or []),
-                    custom_output=custom_output,
+                    images=list(payload.get("image") or []),
+                    multimodal_output={"metadata": {"stream": stream_metadata}},
                     finished=output.finished,
                 )
             ]
@@ -106,6 +107,30 @@ class _PipelineBackedEngine:
         del request_id
 
 
+def _streaming_diffusion_output(
+    *,
+    chunk: int,
+    finished: bool,
+    images: list | None = None,
+) -> DiffusionOutput:
+    payload = {"image": images or []}
+    metadata = {"stream": {"chunk": chunk}}
+    return DiffusionOutput(
+        output={
+            "payload": payload,
+            "metadata": metadata,
+        },
+        finished=finished,
+    )
+
+
+def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
+    envelope = output.output if isinstance(output.output, dict) else {}
+    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
+    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
+    return payload, metadata
+
+
 class TestPipelineStreamingOutputToStageDiffusionClient:
     """Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
 
@@ -114,8 +139,8 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
 
@@ -123,7 +148,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
 
         assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
         assert [output.finished for output in outputs] == [False, True]
-        assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+        assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
         assert pipeline.requests[0].request_id == "req-stream"
 
     @pytest.mark.asyncio
@@ -136,7 +161,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         assert len(outputs) == 2
         assert outputs[0].request_id == "req-error"
         assert outputs[0].finished is False
-        assert outputs[0].custom_output["chunk"] == 0
+        assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0
         assert outputs[1].request_id == "req-error"
         assert outputs[1].finished is True
         assert outputs[1].error is not None
@@ -227,8 +252,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)
@@ -247,7 +272,7 @@ class TestPipelineStreamingOutputToEntrypoint:
 
             await _wait_for(lambda: len(pipeline.requests) == 1)
             assert pipeline.requests[0].request_id.startswith("req-omni-")
-            assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+            assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
             assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
         finally:
             await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@@ -258,8 +283,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True),
+                _streaming_diffusion_output(chunk=0, images=[frames], finished=False),
+                _streaming_diffusion_output(chunk=1, images=[frames], finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -937,6 +937,76 @@ class TestDistributedReceive:
 # ── TP auto-detect ───────────────────────────────────────────────────
 
 
+class TestBuildCfgRankLocalPayloads:
+    """Guard _build_cfg_rank_local_payloads padding logic for the
+    cfg_size==2 / no branch_roles scenario (hunyuan-image shared-prefix reuse)."""
+
+    def test_cfg_size2_no_branch_roles_copies_main_payload_to_follower(self):
+        """When AR doesn't split KV by branch, cfg follower (rank 1) should
+        receive a copy of the positive KV, not None."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # Rank 0: main payload (always present)
+        assert payloads[0] is not None
+        assert "past_key_values" in payloads[0]
+        # Rank 1: must also be non-None — same positive KV for shared-prefix reuse
+        assert payloads[1] is not None
+        assert payloads[1]["past_key_values"] is req.past_key_values
+        assert payloads[1]["kv_metadata"] == {"seq_len": 10}
+
+    def test_cfg_size3_no_branch_roles_keeps_none_padding(self):
+        """When cfg_size > 2 and no branch_roles, padding stays None (safety guard)."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=3)
+
+        assert len(payloads) == 3
+        assert payloads[0] is not None
+        assert payloads[1] is None
+        assert payloads[2] is None
+
+    def test_empty_main_payload_keeps_none_padding(self):
+        """When main_payload is empty, padding stays None even with cfg_size==2."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # No past_key_values → main_payload is empty dict → falsy → padding stays None
+        assert payloads[1] is None
+
+    def test_with_branch_roles_uses_original_behavior(self):
+        """When branch_roles exist, follower payload comes from branch KV, not padding."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(
+                cfg_text_past_key_values=SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+                cfg_text_kv_metadata={"source": "cfg_text"},
+            ),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        assert payloads[0] is not None
+        assert payloads[1] is not None
+        assert payloads[1]["sp.cfg_active_branch"] == "cfg_text"
+
+
 class TestAutoDetectTP:
     def test_auto_detect_when_config_defaults(self):
         """When config from_tp/to_tp == 1 (default), manager should auto-detect."""

tests/docs/test_generate_examples.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+ROOT_DIR = Path(__file__).parents[2]
+HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py"
+SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH)
+assert SPEC and SPEC.loader
+generate_examples = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(generate_examples)
+
+
+def test_model_example_titles_match_shared_display_names():
+    for slug, display_name in generate_examples.MODEL_DISPLAY_NAMES.items():
+        for category, mode_title in generate_examples.SERVING_MODE_TITLES.items():
+            example_path = ROOT_DIR / "examples" / category / slug
+            example = generate_examples.Example(example_path, category)
+
+            assert example.title == f"{display_name}: {mode_title}"
+            assert example.nav_title == display_name
+
+
+def test_model_example_title_mismatch_is_rejected(tmp_path):
+    example_path = tmp_path / "mimo_audio"
+    example_path.mkdir()
+    (example_path / "README.md").write_text("# A different title\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="Model example title mismatch"):
+        generate_examples.Example(example_path, "offline_inference")
+
+
+def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
+    display_names_path = tmp_path / "model_display_names.yml"
+    display_names_path.write_text("- not a mapping\n", encoding="utf-8")
+    monkeypatch.setattr(generate_examples, "MODEL_DISPLAY_NAMES_FILE", display_names_path)
+
+    with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
+        generate_examples.load_model_display_names()

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -331,7 +331,10 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
         if ro and getattr(ro, "outputs", None):
             cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
         if not cot_text:
-            ar_text = getattr(out, "custom_output", {}).get("ar_generated_text")
+            multimodal_output = getattr(out, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            ar_text = text_metadata.get("ar_generated_text") if isinstance(text_metadata, dict) else None
             if isinstance(ar_text, list):
                 cot_text = "\n".join(text for text in ar_text if text)
             else:

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -288,7 +288,7 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
         else:
             # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
             # Return a minimal dummy output to avoid crashing.
-            return DiffusionOutput(output=None, custom_output={})
+            return DiffusionOutput(output=None)
 
         if isinstance(negative_prompt_ids, list):
             negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
@@ -399,14 +399,20 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
             image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
 
         return DiffusionOutput(
-            output=_maybe_to_cpu(image),
+            output={
+                "payload": {
+                    "image": _maybe_to_cpu(image),
+                },
+                "metadata": {
+                    "prompt_embeddings": {
+                        "prompt_embeds": _maybe_to_cpu(prompt_embeds),
+                        "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
+                        "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
+                        "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
+                    },
+                },
+            },
             trajectory_latents=_maybe_to_cpu(all_latents),
             trajectory_log_probs=_maybe_to_cpu(all_log_probs),
             trajectory_timesteps=_maybe_to_cpu(all_timesteps),
-            custom_output={
-                "prompt_embeds": _maybe_to_cpu(prompt_embeds),
-                "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
-                "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
-                "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
-            },
         )

tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py

@@ -390,20 +390,23 @@ class vLLMOmniHttpServerLocal:
 
         diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0
 
-        mm_output = final_res.custom_output or {}
+        metadata = (
+            final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {}
+        )
+        prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {}
 
         if sampling_params.get("logprobs", False):
-            all_log_probs = mm_output.get("all_log_probs")
+            all_log_probs = final_res.trajectory_log_probs
             log_probs = all_log_probs[0] if all_log_probs is not None else None
         else:
             log_probs = None
 
-        all_latents = mm_output.get("all_latents")
-        all_timesteps = mm_output.get("all_timesteps")
-        prompt_embeds = mm_output.get("prompt_embeds")
-        prompt_embeds_mask = mm_output.get("prompt_embeds_mask")
-        negative_prompt_embeds = mm_output.get("negative_prompt_embeds")
-        negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask")
+        all_latents = final_res.trajectory_latents
+        all_timesteps = final_res.trajectory_timesteps
+        prompt_embeds = prompt_embeddings.get("prompt_embeds")
+        prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask")
+        negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds")
+        negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask")
 
         extra_fields = {
             "all_latents": all_latents[0] if all_latents is not None else None,

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -7,6 +7,7 @@ E2E Online tests for Qwen3-Omni model.
 import os
 
 import pytest
+from openai import BadRequestError
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -610,3 +611,21 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     text = responses[0].text_content if responses else ""
     word_count = len(text.split())
     assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
+def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
+    """Ensure invalid audio format is rejected with 400 before streaming starts."""
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt(),
+    )
+    with pytest.raises(BadRequestError, match="audio format"):
+        openai_client.client.chat.completions.create(
+            model=omni_server.model,
+            messages=messages,
+            modalities=["text", "audio"],
+            audio={"voice": "alloy", "format": "aac"},
+            stream=True,
+        )

tests/engine/test_output_modality.py

@@ -14,7 +14,7 @@ import torch
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-_ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine"
+_OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs"
 
 
 def _load_module(name: str, filepath: Path):
@@ -26,12 +26,16 @@ def _load_module(name: str, filepath: Path):
 
 
 _om_mod = _load_module(
-    "vllm_omni.engine.output_modality",
-    _ENGINE_DIR / "output_modality.py",
+    "vllm_omni.outputs.output_modality",
+    _OUTPUTS_DIR / "output_modality.py",
+)
+_load_module(
+    "vllm_omni.outputs.utils",
+    _OUTPUTS_DIR / "utils.py",
 )
 _mm_mod = _load_module(
-    "vllm_omni.engine.mm_outputs",
-    _ENGINE_DIR / "mm_outputs.py",
+    "vllm_omni.outputs.mm_outputs",
+    _OUTPUTS_DIR / "mm_outputs.py",
 )
 
 OutputModality = _om_mod.OutputModality

tests/engine/test_output_processor.py

@@ -11,8 +11,8 @@ from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.engine import FinishReason
 
-from vllm_omni.engine.output_modality import OutputModalityNames
-from vllm_omni.engine.output_processor import OmniRequestState
+from vllm_omni.outputs.output_modality import OutputModalityNames
+from vllm_omni.outputs.output_processor import OmniRequestState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]