[Rebase] Rebase to vllm v0.25.0 (#5042)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

.buildkite/pipeline-intel.yaml

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

.buildkite/test-amd-merge.yml

@@ -7,17 +7,29 @@ steps:
 
 - group: ":card_index_dividers: Simple Test"
   steps:
-    - label: "Simple · Diffusion & Model Executor Test"
+    - label: "Simple · Diffusion & Model Executor Test (1/2)"
       agent_pool: mi325_1
       depends_on: amd-build
-      timeout_in_minutes: 50
+      timeout_in_minutes: 70
       mirror_hardwares: [amdproduction]
       grade: Blocking
       commands:
         - export VLLM_ROCM_USE_AITER=0
         # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
         # that is not supported on CPU
-        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+
+    - label: "Simple · Diffusion & Model Executor Test (2/2)"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      timeout_in_minutes: 70
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
     - label: "Simple · Other Test"
       agent_pool: mi325_1
@@ -158,7 +170,7 @@ steps:
       grade: Blocking
       commands:
         - |
-          timeout 20m bash -c '
+          timeout 50m bash -c '
             export VLLM_LOGGING_LEVEL=DEBUG
             export VLLM_WORKER_MULTIPROC_METHOD=spawn
             export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
@@ -170,13 +182,14 @@ steps:
       depends_on: amd-build
       mirror_hardwares: [amdproduction]
       grade: Blocking
+      timeout_in_minutes: 50
       commands:
         - |
-          timeout 30m bash -c '
+          timeout 40m bash -c '
             export VLLM_LOGGING_LEVEL=DEBUG
             export VLLM_WORKER_MULTIPROC_METHOD=spawn
             export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
           '
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa

.buildkite/test-merge.yml

@@ -627,11 +627,10 @@ steps:
       - label: "Diffusion · FLUX.2 Klein Test"
         source_file_dependencies:
           - tests/e2e/offline_inference/test_flux2_klein.py
-          - tests/e2e/online_serving/test_flux2_klein.py
           - vllm_omni/diffusion/models/flux2_klein/
         timeout_in_minutes: 30
         commands:
-          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
         agents:
           queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
         plugins:

.buildkite/test-nightly.yml

@@ -598,6 +598,38 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+  - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]"
+    key: nightly-tiny-model-test-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test"
+    steps:
+      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
+        timeout_in_minutes: 30
+        # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
+        # so nightly only adds the multi-GPU parallelism configurations.
+        #
+        # This is a bit confusing, but `full_model` is used to mark the multigpu tests, and
+        # --run-level `core_model` is used to indicate that we should use tiny model weights
+        # for the CI, since these tests should pass irrespective of whether or not we use
+        # the full weights or not.
+        commands:
+          - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -618,6 +650,8 @@ steps:
             tests/e2e/online_serving/test_qwen_image_edit_expansion.py
             tests/e2e/online_serving/test_qwen_image_layered_expansion.py
             tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/online_serving/test_krea2_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and not distributed_cuda"
             --run-level "full_model"
         agents:
@@ -661,6 +695,7 @@ steps:
             pytest -sv
             tests/e2e/online_serving/test_qwen_image_expansion.py
             tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+            tests/e2e/offline_inference/test_krea2_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -1287,6 +1322,44 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2V · Doc Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: ":full_moon: Diffusion X2V · Other Function Test"
         timeout_in_minutes: 90
         commands:

.buildkite/test-ready.yml

@@ -157,6 +157,27 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
 
+  - group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
+    depends_on: upload-ready-pipeline
+    steps:
+      - label: "Tiny Model Tests · Diffusion (base)"
+        commands:
+          - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          # Ready label only runs the cases that don't need parallelism for tiny
+          # model tests, which is usually just a smoke test for TTI and I2I
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Engine&Model Executor Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -663,3 +684,24 @@ steps:
                     hostPath:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
+
+      - label: "Diffusion · GGUF Plugin Test"
+        soft_fail: true
+        source_file_dependencies:
+          - tests/plugin/test_gguf_diffusion.py
+          - vllm_omni/diffusion/model_loader
+        commands:
+          - pip install "vllm-gguf-plugin>=0.0.3"
+          - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"

.github/CODEOWNERS

@@ -22,30 +22,30 @@
 /.buildkite/ @yenuo26 @congw729
 
 # Engine and model_executor
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
+/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
+/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
 /vllm_omni/model_executor/ @tzhouam @gcanlin
 /vllm_omni/worker/ @tzhouam @gcanlin
 
 # Diffusion
-/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
+/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
 /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
 
 # Diffusion CPU offloading
-/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
+/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
 
 # Omni and TTS support
 /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
-/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
+/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
 
 # Benchmarks
 /benchmarks/accuracy/ @david6666666
 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
-/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
+/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
 /benchmarks/distributed/ @princepride
 /benchmarks/glm_image/ @princepride @JaredforReal
-/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
+/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
 
 # Usability and Observability
 /vllm_omni/deploy/ @tzhouam @lishunyang12
@@ -53,13 +53,13 @@
 /vllm_omni/profiler/ @gcanlin
 
 # ComfyUI
-/apps/ @fhfuih
+/apps/ @fhfuih @alex-jw-brooks
 
 # Large Scale Serving
 /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
 
 # Quantization and configuration
-/vllm_omni/config/ @lishunyang12
+/vllm_omni/config/ @lishunyang12 @alex-jw-brooks
 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
 
 # Hardware plugins and accelerators

.gitignore

@@ -208,6 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
+!vllm_omni/experimental/ar_diffusion/kv_cache/
 
 # Logs
 logs/

docker/Dockerfile.ci

@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.24.0
+ARG VLLM_BASE_TAG=v0.25.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
+ARG VLLM_ASCEND_TAG=v0.24.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.23.0-4885d8e-a3
+ARG VLLM_ASCEND_TAG=v0.24.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.23.0-4885d8e
+ARG VLLM_ASCEND_TAG=v0.24.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.23.0-4885d8e-a3
+ARG VLLM_ASCEND_TAG=v0.24.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni

docker/Dockerfile.xpu

@@ -107,7 +107,9 @@ ENV VLLM_TARGET_DEVICE=xpu
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --no-build-isolation .
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.7.1 && \
+    uv pip install --no-build-isolation --no-deps .
 
 CMD ["/bin/bash"]
 
@@ -143,12 +145,11 @@ COPY . .
 
 ENV VLLM_OMNI_TARGET_DEVICE=xpu
 RUN --mount=type=cache,target=/root/.cache/uv \
-    echo "triton @ https://example.com/blocked" > /tmp/vllm_omni_constraints.txt && \
-    PIP_CONSTRAINT=/tmp/vllm_omni_constraints.txt \
     uv pip install --no-build-isolation \
     --extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \
     ".[dev]" && \
-    uv pip uninstall oneccl oneccl-devel
+    uv pip uninstall triton oneccl oneccl-devel && \
+    uv pip install triton-xpu==3.7.1 --reinstall
 
 FROM vllm-omni AS vllm-omni-openai
 

docs/community/governance.md

@@ -33,16 +33,20 @@ Lead maintainers are responsible for the overall direction and strategy of the p
 
 Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
 
+- [@alex-jw-brooks](https://github.com/alex-jw-brooks): Frontend, Configuration, and Entrypoint Support
+- [@amy-why-3459](https://github.com/amy-why-3459): Omni Modality Model and World Model Support
 - [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
-- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin, NPU integration, and TTS Model Support
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
 - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
 - [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
 - [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
+- [@Sy0307](https://github.com/Sy0307): TTS Model and AR Engine Support
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
 - [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
+- [@xuechendi](https://github.com/xuechendi): Diffusion and Intel GPU Support
 - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
 - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
@@ -53,7 +57,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
 
 ### Runtime-related
 
-- Diffusion models: @RuixiangMa, @wtomin
+- Diffusion models: @RuixiangMa, @wtomin, @xuechendi
 
 ### Optimization
 
@@ -61,6 +65,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou
 - Attention backend: @ZJY0516
 - Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
 - Offload: @yuanheng-zhao
+- Intel GPU support: @xuechendi
 - Quantization: @lishunyang12, @david6666666
 - Kernel fusion / communication-computation: @ZJY0516
 

docs/configuration/composable_parallel.md

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

docs/configuration/pd_disaggregation.md

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

docs/contributing/model/adding_diffusion_model.md

@@ -382,7 +382,7 @@ def forward(
         for p in req.prompts
     ]
 
-    # Extract sampling parameters
+    # Extract common sampling parameters
     sampling_params = req.sampling_params
     num_inference_steps = sampling_params.num_inference_steps or 50
     guidance_scale = sampling_params.guidance_scale or 7.5

docs/contributing/model/adding_omni_model.md

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

docs/design/feature/cfg_parallel.md

@@ -126,16 +126,16 @@ Call `self.diffuse` in your pipeline's forward function:
 
 ```python
 import torch.nn as nn
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
 class YourModelPipeline(nn.Module, CFGParallelMixin):
     def forward(
         self,
-        prompt: str,
-        negative_prompt: str | None = None,
-        guidance_scale: float = 3.5,
-        num_inference_steps: int = 50,
-        **kwargs,
-    ):
-        # Encode prompts, Initialize latents, Get timesteps
+        req: DiffusionRequestBatch,
+    ) -> list[DiffusionOutput]:
+        # Read prompts and sampling params from the request
+        ...
+        # Encode prompts, initialize latents, get timesteps
         ...
         # Run diffusion loop (calls the mixin's diffuse method)
         latents = self.diffuse(

docs/features/custom_pipeline.md

@@ -54,20 +54,24 @@ Create a custom pipeline class that extends an existing pipeline. In this exampl
 
 ```python
 # custom_pipeline.py
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 import torch
 
 class CustomPipeline(QwenImageEditPipeline):
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
 
-    def forward(self, req, prompt=None, negative_prompt=None, **kwargs):
+    def forward(self, req: DiffusionRequestBatch):
+        # Optionally customize sampling parameters
+        actual_num_steps = req.sampling_params.num_inference_steps or 50
+        req.sampling_params.num_inference_steps = actual_num_steps
+
         # Call parent's forward to get normal output
-        output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs)
+        output = super().forward(req=req)
 
         # Add custom trajectory data
-        actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50)
         output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32)
         output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, and 0.24.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
 
 ```bash
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 ```
 
-The 0.24.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.24.0
+git checkout v0.25.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.24.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.24.0/vllm-0.24.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.24.0 \
+    vllm/vllm-omni:v0.25.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.24.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
 
 ```bash
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 ```
 
-The pre-built 0.24.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.24.0
+git checkout v0.25.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.24.0 \
+  vllm/vllm-omni-rocm:v0.25.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.24.0
+  vllm/vllm-omni-rocm:v0.25.0
 ```
 
 # --8<-- [end:pre-built-images]

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

@@ -1,19 +1,16 @@
 # --8<-- [start:requirements]
 
-For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
+For detailed hardware and software requirements, please refer to the [vLLM-Ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
 # --8<-- [end:requirements]
 # --8<-- [start:installation-release]
 
-The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
+The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built Docker images:
 
 ```bash
-# Update the vllm-ascend image
-# Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
-# Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+# 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
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -31,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it $IMAGE bash
+    -it quay.io/atlas-ci/vllm-ascend:0.24.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.24.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 .
@@ -87,14 +84,11 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
-# Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3
-export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
+# Atlas A2: quay.io/ascend/vllm-omni:v0.24.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.24.0-a3
 docker run --rm \
-    --name vllm-omni-npu \
-    --shm-size=1g \
+    --name vllm-omni-a3 \
+    --shm-size=4g \
     --device /dev/davinci0 \
     --device /dev/davinci1 \
     --device /dev/davinci2 \
@@ -106,10 +100,17 @@ docker run --rm \
     -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
     -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
     -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
-    -v /root/.cache:/root/.cache \
-    -p 8000:8000 \
-    -it $IMAGE bash
+    -v ~/.cache:/root/.cache \
+    -p 8091:8091 \
+    -it quay.io/ascend/vllm-omni:v0.24.0-a3 bash
 ```
 
 !!! tip

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.24.0 --torch-backend=auto
+uv pip install vllm==0.25.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722
+uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.24.0` and vLLM Omni `0.24.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/models/supported_models.md

@@ -34,6 +34,7 @@ th {
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/user_guide/diffusion/attention_backends.md

@@ -20,6 +20,8 @@ The full set of backends and their platform defaults is in the **Backend Options
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
 | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
 | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
+| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. |
+| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. |
 
 
 ## Configuration
@@ -210,6 +212,36 @@ Notes:
 - `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
 - SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
 
+## HuggingFace Kernels Hub Backends
+
+To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout.
+
+The following backend options are supported:
+- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`)
+- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only)
+
+### Installation
+
+To use these backends, you must install the `kernels` library:
+
+```bash
+pip install kernels==0.14.1
+```
+
+If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
+
+### Usage
+
+Select a Hub backend using the global CLI flag or environment variables:
+
+```bash
+# Environment variable
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
+
+# CLI flag
+vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
+```
+
 ## Usage Examples
 
 ### Default (auto-route)

docs/user_guide/diffusion_features.md

@@ -121,12 +121,13 @@ The following tables show which models support each feature:
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
+| **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
 | **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
@@ -144,6 +145,7 @@ The following tables show which models support each feature:
 > 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.
+> 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/image_to_video.md

@@ -3,7 +3,9 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video>.
 
 
-This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
+This example demonstrates how to generate videos from images using Wan2.2,
+LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
+offline inference API.
 
 ## Local CLI Usage
 
@@ -51,13 +53,41 @@ python image_to_video.py \
   --output i2v_output.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python image_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx23_i2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
+the LTX-2.3 image-to-video pipeline.
+
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
+  LTX-2/LTX-2.3).
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
-- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default 81).
+- `--height/--width`: Output resolution (auto-calculated from image if not set).
+  Wan dimensions should be multiples of 16; LTX dimensions should be multiples
+  of 32.
+- `--num-frames`: Number of frames (model-specific default; LTX-style models
+  work best with `8k + 1`).
 - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative-prompt`: Optional list of artifacts to suppress.
 - `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
@@ -65,6 +95,7 @@ Key arguments:
 - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
 - `--num-inference-steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--audio-sample-rate`: fallback audio sample rate for embedded audio.
 - `--output`: Path to save the generated video.
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.

docs/user_guide/examples/offline_inference/mammothmodal2_preview.md

@@ -22,7 +22,7 @@ Text-to-image now runs through the shared offline image example
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
 | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 
@@ -50,6 +51,30 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python text_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23Pipeline \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --num-inference-steps 20 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --audio-sample-rate 48000 \
+  --output ltx23_t2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
+text-to-video pipeline explicitly.
+
 ### HunyuanVideo-1.5 (480p)
 
 ```bash
@@ -110,6 +135,8 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
+- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
+  for LTX-2.3 text-to-video.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -124,7 +151,8 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the
+  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
 

docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md

@@ -9,7 +9,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly
 
 The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online.
 Currently, the following features are NOT yet supported.
-It is not guaranteed whether they will be supported in the future.
+Community contributions are welcome to add these features when they can be
+implemented with clear Diffusers or vLLM-Omni behavior.
 
 - CFG parallel execution
 - Sequence parallel execution
@@ -66,6 +67,35 @@ This is suitable for sampling parameters not available through the vLLM-Omni int
 When a parameter is available in the vLLM-Omni interface, it will be adapted here.
 But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time).
 
+### Quantization
+
+Use `diffusers_load_kwargs` for Diffusers-native quantization options. If
+`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the
+diffusers backend builds a Diffusers quantization config and lets Diffusers
+validate it before calling `DiffusionPipeline.from_pretrained()`.
+
+The diffusers backend also provides a small compatibility shortcut for
+vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already
+contain `quantization_config`. Currently this maps online/dynamic `fp8` and
+online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer
+components such as `transformer` or `transformer_2`. This path requires
+`torchao` to be installed.
+
+For example, the CLI can request dynamic FP8 through the vLLM-Omni interface:
+
+```bash
+vllm serve "Qwen/Qwen-Image" \
+    --omni \
+    --diffusion-load-format diffusers \
+    --quantization-config '{"method": "fp8"}'
+```
+
+Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
+`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
+such as `ignored_layers`, are intentionally not translated by this compatibility
+shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
+a native vLLM-Omni pipeline for those cases.
+
 ### Attention Backends
 
 The diffusers backend converts

docs/user_guide/quantization/fp8.md

@@ -44,10 +44,16 @@ pip install vllm-omni[quack]
 pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
 ```
 
-It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
-on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
-Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
-not installed, FP8 still works — it just keeps the unfused FlashInfer path.
+It is enabled automatically once installed (no flag needed) and is **datacenter
+Blackwell only** (`sm_100` / `sm_101` / `sm_103`, compute capability `10.x`, e.g.
+B200): quack's CuteDSL GEMM uses the 5th-gen `tcgen05` tensor-core MMA, which exists
+only on those parts. On Hopper/Ada the CUTLASS FP8 kernel already fuses bias, and on
+workstation/consumer Blackwell (`sm_120` / `sm_121`, compute capability `12.x`, e.g.
+RTX PRO 6000 / RTX 50-series) `tcgen05` is absent — so quack is **not** auto-enabled
+there and FlashInfer's native FP8 path is used instead. Set
+`VLLM_OMNI_USE_QUACK_FP8=1` to force quack on, or `VLLM_OMNI_USE_QUACK_FP8=0` to force
+the FlashInfer path. If `quack-kernels` is not installed, FP8 still works — it just
+keeps the unfused FlashInfer path.
 
 #### Compile cache and warmup
 

docs/user_guide/quantization/gguf.md

@@ -9,6 +9,15 @@ for the transformer.
 
 GGUF is static quantization: the quantized weights are produced before serving.
 
+!!! note
+    GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model.
+
+Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin):
+
+```bash
+uv pip install vllm-gguf-plugin>=0.0.3
+```
+
 ## Hardware Support
 
 | Device | Support |

docs/user_guide/quantization/mxfp4.md

@@ -39,7 +39,7 @@ vLLM-Omni provides two quantization methods with different scale structures:
 | NVIDIA Blackwell GPU (SM 100+) | ⭕ |
 | NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
 | NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm | ⭕ |
+| AMD ROCm (gfx950 / MI355X) | ✅ |
 | Intel XPU | ⭕ |
 | Ascend NPU (Atlas 950 A5) | ✅ |
 

docs/user_guide/quantization/overview.md

@@ -20,7 +20,7 @@ type has a different quantization scope.
 | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ |
 | Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
 | Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
 

examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py

@@ -2,14 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import logging
-from typing import Any
 
-import PIL.Image
 import torch
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 logger = logging.getLogger(__name__)
 
@@ -18,58 +16,14 @@ class CustomPipeline(QwenImageEditPipeline):
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
 
-    def forward(
-        self,
-        req: OmniDiffusionRequest,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch):
         """Forward pass for image editing with dummy trajectory data."""
-        # Call parent's forward to get the normal output
-        output = super().forward(
-            req=req,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            true_cfg_scale=true_cfg_scale,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            sigmas=sigmas,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=num_images_per_prompt,
-            generator=generator,
-            latents=latents,
-            prompt_embeds=prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            negative_prompt_embeds=negative_prompt_embeds,
-            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
-            output_type=output_type,
-            attention_kwargs=attention_kwargs,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            max_sequence_length=max_sequence_length,
-        )
+        # Customize num_inference_steps that best suits your model
+        actual_num_steps = req.sampling_params.num_inference_steps or 50
+        req.sampling_params.num_inference_steps = actual_num_steps
 
-        # Get actual num_inference_steps used
-        actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        # Call parent's forward to get the normal output
+        output = super().forward(req=req)
 
         # Create dummy trajectory data
         dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)

examples/offline_inference/image_to_image/image_edit.py

@@ -138,6 +138,15 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.",
     )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help=(
+            "Path to a deploy YAML. Required for multi-stage image-edit pipelines "
+            "whose deploy config is not auto-loaded."
+        ),
+    )
     parser.add_argument(
         "--prompt",
         type=str,
@@ -349,7 +358,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        help="Disable torch.compile and force eager execution.",
+        default=None,
+        help=(
+            "Disable torch.compile and force eager execution. Left unset (None) "
+            "so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML value wins."
+        ),
     )
     parser.add_argument(
         "--vae-use-slicing",
@@ -468,7 +482,7 @@ def main():
         }
 
     # Initialize Omni with appropriate pipeline
-    omni = Omni(
+    omni_kwargs: dict[str, Any] = dict(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
         vae_use_slicing=args.vae_use_slicing,
@@ -476,11 +490,15 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
-        enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
     )
+    if args.enforce_eager is not None:
+        omni_kwargs["enforce_eager"] = args.enforce_eager
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
+    omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
     print("Pipeline loaded")

examples/offline_inference/image_to_video/README.md

@@ -1,16 +1,142 @@
 # Image-To-Video
 
-This shared example generates videos from images with VACE, Wan2.2, LTX2,
-HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
+This shared example generates videos from images with VACE, Wan2.2,
+LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
 
-## Local CLI Usage
+- `image_to_video.py`: command-line script for single video generation with advanced options.
 
-Download the example image:
+## Table of Contents
+
+- [Overview](#overview)
+- [Prerequisites](#prerequisites)
+- [Quick Start](#quick-start)
+- [Key Arguments](#key-arguments)
+- [More CLI Examples](#more-cli-examples)
+- [Advanced Features](#advanced-features)
+- [FAQ](#faq)
+
+## Overview
+
+This folder provides a unified CLI script for image-to-video generation using vLLM-Omni diffusion/video pipelines. The script selects practical defaults for supported model families while still exposing common sampling, memory, and parallelism options.
+
+### Supported Models
+
+| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM Notes |
+| ----- | ------------------ | -------------- | ------------- | -------- | ---------- |
+| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
+| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
+| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
+| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
+
+!!! info
+    Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
+
+Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`.
+
+## Prerequisites
+
+Download the example image used in the snippets below:
 
 ```bash
-wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+curl -L -o cherry_blossom.jpg https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
 ```
 
+## Quick Start
+
+### Python API
+
+Single-prompt generation using TI2V-5B (lightweight):
+
+```python
+import PIL.Image
+import torch
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+if __name__ == "__main__":
+    image = PIL.Image.open("cherry_blossom.jpg").convert("RGB")
+    image = image.resize((576, 320))
+
+    omni = Omni(
+        model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        flow_shift=12.0,
+    )
+    outputs = omni.generate(
+        {
+            "prompt": "Cherry blossoms swaying gently in the breeze, petals falling",
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=320,
+            width=576,
+            num_frames=17,
+            num_inference_steps=20,
+            guidance_scale=4.0,
+            generator=torch.Generator(device="cuda").manual_seed(42),
+        ),
+    )
+    from diffusers.utils import export_to_video
+    frames = outputs[0].request_output.images
+    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
+```
+
+### Local CLI Usage
+
+Quick test using TI2V-5B with a small resolution and few frames:
+
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --height 320 \
+  --width 576 \
+  --num-frames 17 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --flow-shift 12.0 \
+  --fps 16 \
+  --output quick_test_i2v.mp4
+```
+
+## Key Arguments
+
+| Argument | Type | Default | Description |
+| -------- | ---- | ------- | ----------- |
+| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
+| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
+| `--image` | str | (required) | Path to input image |
+| `--prompt` | str | `""` | Text description of desired motion/animation |
+| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
+| `--seed` | int | `42` | Random seed for deterministic sampling |
+| `--guidance-scale` | float | `5.0` | CFG scale |
+| `--guidance-scale-high` | float | `None` | Separate CFG for high-noise stage (MoE only) |
+| `--height` | int | auto | Video height (auto-calculated from image if not set). Multiples of 16 |
+| `--width` | int | auto | Video width (auto-calculated from image if not set). Multiples of 16 |
+| `--num-frames` | int | `81` | Number of frames |
+| `--num-inference-steps` | int | `50` | Number of denoising steps |
+| `--boundary-ratio` | float | `0.875` | Boundary split ratio for two-stage MoE models |
+| `--flow-shift` | float | `5.0` | Scheduler flow shift (5.0 for 720p, 12.0 for 480p) |
+| `--sample-solver` | str | `unipc` | Wan2.2 sampling solver (`unipc` or `euler` for Lightning/Distill) |
+| `--fps` | int | `None` | Frames per second for the saved MP4 |
+| `--frame-rate` | float | `None` | Generation frame rate for pipelines that require it (e.g., LTX2) |
+| `--output` | str | `i2v_output.mp4` | Path to save the generated video |
+| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
+| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
+| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
+| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules |
+| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
+| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) |
+| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree |
+| `--ring-degree` | int | `1` | Ring sequence parallel degree |
+| `--cache-backend` | str | `None` | Cache backend: `cache_dit` or `tea_cache` |
+| `--use-hsdp` | flag | off | Enable Hybrid Sharded Data Parallel |
+| `--hsdp-shard-size` | int | `-1` | GPUs per shard group (-1 auto-calculates) |
+| `--hsdp-replicate-size` | int | `1` | Number of replica groups for HSDP |
+
+## More CLI Examples
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
 
 ```bash
@@ -18,7 +144,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "<optional quality filter>" \
+  --negative-prompt "low quality, blurry" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -28,7 +154,7 @@ python image_to_video.py \
   --boundary-ratio 0.875 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_output.mp4
+  --output i2v_wan_moe.mp4
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
@@ -38,7 +164,7 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative-prompt "<optional quality filter>" \
+  --negative-prompt "low quality, blurry" \
   --height 480 \
   --width 832 \
   --num-frames 48 \
@@ -46,9 +172,91 @@ python image_to_video.py \
   --num-inference-steps 40 \
   --flow-shift 12.0 \
   --fps 16 \
-  --output i2v_output.mp4
+  --output i2v_wan_ti2v.mp4
 ```
 
+### HunyuanVideo-1.5 I2V (480p)
+
+```bash
+python image_to_video.py \
+  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
+  --height 480 \
+  --width 832 \
+  --num-frames 121 \
+  --guidance-scale 6.0 \
+  --flow-shift 5.0 \
+  --num-inference-steps 50 \
+  --fps 24 \
+  --enable-cpu-offload \
+  --vae-use-tiling \
+  --vae-use-slicing \
+  --output hunyuan_i2v.mp4
+```
+
+### LTX2 Image-to-Video
+
+```bash
+python image_to_video.py \
+  --model /path/to/LTX-2 \
+  --model-class-name LTX2ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "A cinematic dolly shot of cherry blossoms" \
+  --height 512 \
+  --width 768 \
+  --num-frames 121 \
+  --num-inference-steps 40 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx2_i2v.mp4
+```
+
+## Advanced Features
+
+### CFG Parallel
+
+Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
+See more examples in the [cfg_parallel user guide](../../../docs/user_guide/diffusion/parallelism/cfg_parallel.md).
+
+### Cache Acceleration
+
+Use `--cache-backend cache_dit` for Cache-DiT acceleration or `--cache-backend tea_cache` for Timestep Embedding Aware Cache:
+
+```bash
+python image_to_video.py \
+  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze" \
+  --cache-backend cache_dit \
+  --output i2v_cached.mp4
+```
+
+### LTX-2.3
+
+```bash
+python image_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 20 \
+  --frame-rate 24 \
+  --fps 24 \
+  --output ltx23_i2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
+the LTX-2.3 image-to-video pipeline.
+
 ### Cosmos3
 
 ```bash
@@ -66,15 +274,19 @@ python image_to_video.py \
 
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
+  LTX-2/LTX-2.3, Cosmos3, or VACE).
 - `--image`: Path to the first-frame or source image.
 - `--last-image`: Optional last-frame condition for models such as VACE.
 - `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
 - `--reference-image`: Optional reference image; repeat it to provide multiple references.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
 - `--prompt`: Text description of desired motion/animation.
-- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
+- `--height/--width`: Output resolution (auto-calculated from image if not set).
+  Wan dimensions should be multiples of 16; LTX dimensions should be multiples
+  of 32.
+- `--num-frames`: Number of frames (model-specific default; LTX-style models
+  work best with `8k + 1`; Cosmos3 defaults to 189).
 - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative-prompt`: Optional list of artifacts to suppress.
 - `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
@@ -82,6 +294,7 @@ Key arguments:
 - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
 - `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
+- `--audio-sample-rate`: fallback audio sample rate for embedded audio.
 - `--output`: Path to save the generated video.
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
@@ -191,3 +404,11 @@ The VACE T2V command is documented in the shared
 
 For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
 assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
+
+## FAQ
+
+**OOM errors**: Try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. For very large models, add `--enable-cpu-offload` or `--enable-layerwise-offload`.
+
+**Auto-calculated resolution**: If `--height` and `--width` are not provided, the script calculates output dimensions from the input image while maintaining aspect ratio and targeting 480 x 832 area (or 512 x 768 for LTX2).
+
+**Wan2.2 MoE vs unified**: `I2V-A14B` is a larger Mixture-of-Experts model with separate low/high-noise DiT stages (use `--guidance-scale-high` and `--boundary-ratio`). `TI2V-5B` is a smaller unified T2V+I2V model that does not need these extra arguments.

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3,
+HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
@@ -27,6 +28,13 @@ Usage:
         --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 
+    # LTX-2.3 image-to-video
+    python image_to_video.py --model dg845/LTX-2.3-Diffusers \
+        --model-class-name LTX23ImageToVideoPipeline \
+        --image input.jpg --prompt "A cinematic dolly shot of a boat" \
+        --height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
+        --frame-rate 24 --fps 24 --output ltx23_i2v.mp4
+
     # HunyuanVideo-1.5 I2V (480p)
     python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
         --image input.jpg --prompt "A cat playing with yarn" \
@@ -79,17 +87,20 @@ parse_profiler_config = functools.partial(parse_json_object, flag_name="--profil
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
+        description=(
+            "Generate a video from one or more images "
+            "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)."
+        )
     )
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
+        help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).",
     )
     parser.add_argument(
         "--model-class-name",
         default=None,
-        help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
+        help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
     )
     parser.add_argument("--image", help="Path to the first-frame or source image.")
     parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
@@ -113,7 +124,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
     parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
     parser.add_argument(
-        "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
+        "--num-inference-steps",
+        type=int,
+        default=None,
+        help="Sampling steps. Default: model-specific.",
     )
     parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
     parser.add_argument(
@@ -206,7 +220,7 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
         help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 
@@ -312,10 +326,8 @@ def main():
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     model_name = str(args.model).lower() if args.model is not None else ""
     model_class_name = args.model_class_name
-    is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
+    is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"}
     is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
-    if model_class_name is None and is_ltx2:
-        model_class_name = "LTX2ImageToVideoPipeline"
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
     last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
@@ -340,7 +352,15 @@ def main():
             16,
         )
     elif is_ltx2:
-        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            24,
+            4.0,
+            121,
+            40,
+            5.0,
+            512 * 768,
+            32,
+        )
     else:  # Wan2.2 / HunyuanVideo-1.5
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
 
@@ -462,7 +482,7 @@ def main():
         f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
         f" pipeline_parallel_size={args.pipeline_parallel_size}"
     )
-    print(f"  Video size: {args.width}x{args.height}")
+    print(f"  Video size: {width}x{height}")
     print(f"{'=' * 60}\n")
 
     prompt_dict = build_image_to_video_prompt(
@@ -509,6 +529,7 @@ def main():
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
     audio = None
+    audio_sample_rate = args.audio_sample_rate
     if isinstance(frames, list):
         frames = frames[0] if frames else None
 
@@ -519,11 +540,13 @@ def main():
             )
         if frames.multimodal_output and "audio" in frames.multimodal_output:
             audio = frames.multimodal_output["audio"]
+            audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]
+                    audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate)
                 frames = inner_output
         if isinstance(frames, OmniRequestOutput):
             if frames.images:
@@ -531,6 +554,7 @@ def main():
                     frames, audio = frames.images[0]
                 elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
                     audio = frames.images[0].get("audio")
+                    audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate)
                     frames = frames.images[0].get("frames") or frames.images[0].get("video")
                 else:
                     frames = frames.images
@@ -543,6 +567,7 @@ def main():
             frames, audio = first_item
         elif isinstance(first_item, dict):
             audio = first_item.get("audio")
+            audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate)
             frames = first_item.get("frames") or first_item.get("video")
         elif isinstance(first_item, list):
             frames = first_item
@@ -551,6 +576,7 @@ def main():
         frames, audio = frames
     elif isinstance(frames, dict):
         audio = frames.get("audio")
+        audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
         frames = frames.get("frames") or frames.get("video")
 
     if frames is None:
@@ -669,7 +695,7 @@ def main():
             frames_u8,
             audio_np,
             fps=float(fps),
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=audio_sample_rate,
         )
         with open(str(output_path), "wb") as f:
             f.write(video_bytes)

examples/offline_inference/mammothmodal2_preview/README.md

@@ -19,7 +19,7 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```bash
 python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
   --model ./MammothModa2-Preview \
-  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \
   --question "Summarize this image." \
   --image ./image.png
 ```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -4,7 +4,7 @@ Offline inference example: MammothModa2 image summarization (single AR stage).
 Example:
   uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
     --model path/to/MammothModa2-Preview \
-    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
     --image /path/to/input.jpg \
     --question "Please summarize the content of this image."
 """
@@ -27,9 +27,7 @@ DEFAULT_QUESTION = "Please summarize the content of this image."
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
     parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
-    parser.add_argument(
-        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
-    )
+    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
     parser.add_argument("--image", type=str, required=True, help="Path to input image.")
     parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
     parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
@@ -73,7 +71,7 @@ def main() -> None:
 
     omni = Omni(
         model=args.model,
-        stage_configs_path=args.stage_config,
+        deploy_config=args.deploy_config,
         trust_remote_code=args.trust_remote_code,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )

examples/offline_inference/ming_flash_omni/README.md

@@ -90,14 +90,40 @@ Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to
 
 ### Image generation (text-to-image / img2img)
 
-The diffusion-stage image-generation path is currently only wired through
-the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
-`"modalities": ["image"]`). For payloads, optional knobs
-(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
-reference-image flow, see the
-[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
-`end2end.py` does not yet exercise the imagegen stage; that is tracked as
-follow-up work.
+Image generation is served through the standard task examples.
+The diffusion-stage knobs are declared centrally in `vllm_omni/model_extras/ming_flash_omni.py` and routed via `--extra-body`. With deploy yaml assigned properly, the model can run through the shared example scripts like any other diffusion model:
+
+Text-to-image (offline):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Jonathan1909/Ming-flash-omni-2.0 \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --prompt "Please draw a cute cat." \
+    --height 1024 \
+    --width 1024 \
+    --extra-body '{"steps": 30, "cfg": 2.0, "seed": 42}' \
+    --output ming_flash_omni_t2i.png
+```
+
+Image-to-image (offline):
+```bash
+# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
+# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
+wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
+
+python examples/offline_inference/image_to_image/image_edit.py \
+    --model Jonathan1909/Ming-flash-omni-2.0 \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --image person_gen_05.png \
+    --prompt "Put a pair of sunglasses on the person." \
+    --extra-args '{"steps": 30, "cfg": 2.0, "seed": 42}' \
+    --output ming_flash_omni_i2i.png
+```
+
+For the online path and the full knob list
+(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`/`height`/`width`), see the
+[image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
 
 ## Online serving
 

examples/offline_inference/qwen2_5_omni/README.md

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

examples/offline_inference/qwen2_5_omni/end2end.py

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

examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml

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

examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml

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

examples/offline_inference/text_to_image/README.md

@@ -36,6 +36,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 | `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
+| `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | 1024 (Raw) / 2048 (Turbo) | — | ~30 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -196,6 +197,37 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output /output.png
 ```
 
+### Krea 2
+
+Krea 2 is published as diffusers-format checkpoints: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`.
+The pipeline class (`Krea2Pipeline`) is auto-detected from `model_index.json`. The distilled Turbo checkpoint sets
+`is_distilled=true` and is best run at 2048x2048 with few steps and guidance disabled:
+
+```bash
+# Few-step distilled (Turbo) checkpoint
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 8 \
+  --guidance-scale 0.0 \
+  --height 2048 --width 2048 \
+  --output krea2.png
+```
+
+For the Raw checkpoint, use more steps with CFG enabled:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Raw \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 28 \
+  --guidance-scale 4.5 \
+  --output krea2.png
+```
+
+`guidance-scale` follows the Krea 2 convention (`velocity = cond + guidance_scale * (cond - uncond)`); CFG is enabled
+whenever `guidance-scale > 0`.
+
 ### Batch Requests (Multiple Prompts)
 
 You can pass multiple prompts in a single `Omni.generate` call. `Omni`

examples/offline_inference/text_to_image/text_to_image.py

@@ -69,7 +69,16 @@ def parse_args() -> argparse.Namespace:
         "--stage-configs-path",
         type=str,
         default=None,
-        help="Path to a YAML file containing stage configurations for Omni.",
+        help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help=(
+            "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage "
+            "text-to-image pipelines whose deploy config is not auto-loaded."
+        ),
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
@@ -155,7 +164,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
-        help="Disable torch.compile and force eager execution.",
+        default=None,
+        help=(
+            "Disable torch.compile and force eager execution. Left unset (None) "
+            "so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML value wins."
+        ),
     )
     parser.add_argument(
         "--enable-cpu-offload",
@@ -188,17 +202,11 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8", "gguf"],
+        choices=["fp8", "int8"],
         help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8). "
         "Default: None (no quantization, uses BF16).",
     )
-    parser.add_argument(
-        "--gguf-model",
-        type=str,
-        default=None,
-        help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
-    )
     parser.add_argument(
         "--ignored-layers",
         type=str,
@@ -221,8 +229,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--tensor-parallel-size",
         type=int,
-        default=1,
-        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+        default=None,
+        help=(
+            "Number of GPUs used for tensor parallelism (TP) inside the DiT. "
+            "Left unset so it is only forwarded when explicitly given; "
+            "otherwise the per-stage deploy YAML (or engine default of 1) wins. "
+            "Passing it always overrides the deploy YAML."
+        ),
     )
     parser.add_argument(
         "--enable-expert-parallel",
@@ -383,14 +396,7 @@ def main():
     # ignored_layers is specified so the list flows through OmniDiffusionConfig
     quant_kwargs: dict[str, Any] = {}
     ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
-    if args.quantization == "gguf":
-        if not args.gguf_model:
-            raise ValueError("--gguf-model is required when --quantization gguf is set.")
-        quant_kwargs["quantization_config"] = {
-            "method": "gguf",
-            "gguf_model": args.gguf_model,
-        }
-    elif args.quantization and ignored_layers:
+    if args.quantization and ignored_layers:
         quant_kwargs["quantization_config"] = {
             "method": args.quantization,
             "ignored_layers": ignored_layers,
@@ -410,10 +416,8 @@ def main():
         "ring_degree": args.ring_degree,
         "ulysses_mode": args.ulysses_mode,
         "cfg_parallel_size": args.cfg_parallel_size,
-        "tensor_parallel_size": args.tensor_parallel_size,
         "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
-        "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
         "mode": "text-to-image",
         "log_stats": args.log_stats,
@@ -425,8 +429,14 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
+    if args.tensor_parallel_size is not None:
+        omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
+    if args.enforce_eager is not None:
+        omni_kwargs["enforce_eager"] = args.enforce_eager
     if args.stage_configs_path:
         omni_kwargs["stage_configs_path"] = args.stage_configs_path
+    if args.deploy_config:
+        omni_kwargs["deploy_config"] = args.deploy_config
     if use_nextstep:
         # NextStep-1.1 requires explicit pipeline class
         omni_kwargs["model_class_name"] = "NextStep11Pipeline"
@@ -451,8 +461,9 @@ def main():
     print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
     if ignored_layers:
         print(f"  Ignored layers: {ignored_layers}")
+    tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
     print(
-        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
+        f"  Parallel configuration: tensor_parallel_size={tp_display}, "
         f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
         f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
@@ -522,8 +533,8 @@ def main():
         diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
 
     if lora_request:
-        diffusion_params.extra_args["lora_request"] = lora_request
-        diffusion_params.extra_args["lora_scale"] = args.lora_scale
+        diffusion_params.lora_request = lora_request
+        diffusion_params.lora_scale = args.lora_scale
 
     # Build per-stage sampling params for multi-stage models (e.g. BAGEL),
     # or wrap single diffusion params for single-stage models.

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -34,7 +34,7 @@ def run_e2e():
         help="Model name or path (HuggingFace or local)",
     )
     parser.add_argument(
-        "--stage-config",
+        "--deploy-config",
         type=str,
         default="vllm_omni/deploy/omnivoice.yaml",
     )
@@ -87,14 +87,14 @@ def run_e2e():
     )
     args = parser.parse_args()
 
-    if not os.path.exists(args.stage_config):
-        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
+    if not os.path.exists(args.deploy_config):
+        raise FileNotFoundError(f"Deploy config not found: {args.deploy_config}")
 
     print(f"Initializing OmniVoice with model={args.model}")
 
     omni = Omni(
         model=args.model,
-        stage_configs_path=args.stage_config,
+        deploy_config=args.deploy_config,
         log_stats=True,
     )
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
+| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
@@ -70,6 +71,30 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
+### LTX-2.3
+
+```bash
+python text_to_video.py \
+  --model dg845/LTX-2.3-Diffusers \
+  --model-class-name LTX23Pipeline \
+  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 384 \
+  --width 512 \
+  --num-frames 25 \
+  --num-inference-steps 20 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --fps 24 \
+  --audio-sample-rate 48000 \
+  --output ltx23_t2v_output.mp4
+```
+
+Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the
+upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
+this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
+text-to-video pipeline explicitly.
+
 ### HunyuanVideo-1.5 (480p)
 
 ```bash
@@ -187,6 +212,8 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
+- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
+  for LTX-2.3 text-to-video.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -202,8 +229,9 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
-- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the
+  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
+- `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 

examples/offline_inference/text_to_video/text_to_video.py

@@ -227,7 +227,7 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
+        choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"],
         help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
     )
 
@@ -455,6 +455,7 @@ def main():
         print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
     audio = None
+    audio_sample_rate = args.audio_sample_rate
     if isinstance(frames, list):
         frames = frames[0] if frames else None
 
@@ -465,11 +466,13 @@ def main():
             )
         if frames.multimodal_output and "audio" in frames.multimodal_output:
             audio = frames.multimodal_output["audio"]
+            audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]
+                    audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate)
                 frames = inner_output
         if isinstance(frames, OmniRequestOutput):
             if frames.images:
@@ -477,6 +480,7 @@ def main():
                     frames, audio = frames.images[0]
                 elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
                     audio = frames.images[0].get("audio")
+                    audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate)
                     frames = frames.images[0].get("frames") or frames.images[0].get("video")
                 else:
                     frames = frames.images
@@ -489,6 +493,7 @@ def main():
             frames, audio = first_item
         elif isinstance(first_item, dict):
             audio = first_item.get("audio")
+            audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate)
             frames = first_item.get("frames") or first_item.get("video")
         elif isinstance(first_item, list):
             frames = first_item
@@ -497,6 +502,7 @@ def main():
         frames, audio = frames
     elif isinstance(frames, dict):
         audio = frames.get("audio")
+        audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
         frames = frames.get("frames") or frames.get("video")
 
     if frames is None:
@@ -611,7 +617,7 @@ def main():
             frames_u8,
             audio_np,
             fps=float(args.fps),
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=audio_sample_rate,
         )
         with open(str(output_path), "wb") as f:
             f.write(video_bytes)

examples/online_serving/diffusers_pipeline_adapter/README.md

@@ -6,7 +6,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly
 
 The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online.
 Currently, the following features are NOT yet supported.
-It is not guaranteed whether they will be supported in the future.
+Community contributions are welcome to add these features when they can be
+implemented with clear Diffusers or vLLM-Omni behavior.
 
 - CFG parallel execution
 - Sequence parallel execution
@@ -68,6 +69,35 @@ But if that parameter is simultaneously set in both the vLLM-Omni interface and
     Consider referring to [`OmniDiffusionSamplingParams`](../../../vllm_omni/inputs/data.py) for the default sampling parameters in the vLLM-Omni interface,
     and the corresponding [diffusers pipeline](https://huggingface.co/docs/diffusers/main/en/api/pipelines/overview)'s `__call__` function documentation.
 
+### Quantization
+
+Use `diffusers_load_kwargs` for Diffusers-native quantization options. If
+`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the
+diffusers backend builds a Diffusers quantization config and lets Diffusers
+validate it before calling `DiffusionPipeline.from_pretrained()`.
+
+The diffusers backend also provides a small compatibility shortcut for
+vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already
+contain `quantization_config`. Currently this maps online/dynamic `fp8` and
+online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer
+components such as `transformer` or `transformer_2`. This path requires
+`torchao` to be installed.
+
+For example, the CLI can request dynamic FP8 through the vLLM-Omni interface:
+
+```bash
+vllm serve "Qwen/Qwen-Image" \
+    --omni \
+    --diffusion-load-format diffusers \
+    --quantization-config '{"method": "fp8"}'
+```
+
+Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`,
+`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists
+such as `ignored_layers`, are intentionally not translated by this compatibility
+shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or
+a native vLLM-Omni pipeline for those cases.
+
 ### Attention Backends
 
 The diffusers backend converts

examples/online_serving/ming_flash_omni/README.md

@@ -49,22 +49,14 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
     --modalities text
 ```
 
-Parameterized curl wrapper in this directory:
 
-```bash
-bash run_curl_multimodal_generation.sh text
-bash run_curl_multimodal_generation.sh use_image
-bash run_curl_multimodal_generation.sh use_audio
-bash run_curl_multimodal_generation.sh use_video
-bash run_curl_multimodal_generation.sh use_mixed_modalities
-bash run_curl_multimodal_generation.sh use_image_gen
-```
+## Image generation (text-to-image / img2img)
 
-## Image generation (text-to-image)
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
 
-Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
-Launch with the image deploy YAML, which adds an image-gen stage behind
-the thinker:
+The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
+
+### Launch
 
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
@@ -74,32 +66,58 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
     --port 8091
 ```
 
-Then request image output by passing `"modalities": ["image"]`:
+
+### Text-to-image
+
+Request image output with `"modalities": ["image"]`:
 
 ```bash
 curl -s http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-      "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
-      "modalities": ["image"]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Please draw a cute cat."
+      }
+    ],
+    "modalities": ["image"]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen.png
 ```
 
-### Optional knobs
+Pass generation knobs under a literal `extra_body` object (the OpenAI Python client's `extra_body=` kwarg produces the same request). Keys are filtered against the declared set and routed into every stage's `extra_args`:
 
-Pass image-gen overrides as flat keys on the diffusion-stage `sampling_params_list[1].extra_args`:
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ],
+    "extra_body": {
+      "steps": 6,
+      "cfg": 1.5,
+      "height": 512,
+      "width": 512,
+      "seed": 123,
+      "byte5_text": ["理解与生成统一"],
+      "negative_prompt": "ugly, blurry, distorted"
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_extra_body.png
+```
 
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
-| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
-| `cfg` | 2.0 | Classifier-free guidance scale. |
-| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
-| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
-| `negative_prompt` | empty | Real CFG negative conditioning (set on **stage-0 thinker** `extra_args` so `expand_cfg_prompts` spawns the companion). |
-
-Example with all the knobs:
+NOTE: `extra_body` does **not** set the thinker's own sampling params. To tune the thinker (stage-0) sampling (`temperature` / `top_p` / `top_k` / `max_tokens`) or place knobs explicitly per stage, use `sampling_params_list` (`[thinker, imagegen]`).
+`negative_prompt` must sit on the **stage-0** entry to trigger the real-CFG companion; the imagegen knobs go on **stage-1**:
 
 ```bash
 curl http://127.0.0.1:8091/v1/chat/completions \
@@ -143,27 +161,48 @@ curl http://127.0.0.1:8091/v1/chat/completions \
 
 ### img2img (reference image)
 
-Add an `image_url` content part to the user message; the parser routes it
-into the diffusion stage as `extra[reference_image]`:
+Add an `image_url` content part (a base64 data URL) to the user message; it is routed into the DiT stage as `extra[reference_image]`. base64-encode a local file and stream it through `jq` via stdin - piping `base64 -> jq -> curl` so that avoids the shell `ARG_MAX` limit that inlining a large base64 string in `-d '…'` hits.
 
-```jsonc
-"messages": [{
-  "role": "user",
-  "content": [
-    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
-    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
-  ]
-}]
+```bash
+# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo
+# Check https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py
+wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png
+
+base64 -w0 person_gen_05.png \
+| jq -Rs --arg prompt "Put a pair of sunglasses on the person." '{
+    model: "Jonathan1909/Ming-flash-omni-2.0",
+    modalities: ["image"],
+    messages: [
+      {
+        role: "user",
+        content: [
+          { type: "text", text: $prompt },
+          { type: "image_url", image_url: { url: ("data:image/png;base64," + (. | rtrimstr("\n"))) } }
+        ]
+      }
+    ]
+  }' \
+| curl http://127.0.0.1:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d @- \
+| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+| base64 -d > ming_img2img.png
 ```
 
-### GPU layout
+The reference image can also be a public URL (`"url": "https://…/photo.jpg"`) or the simplified `{"image": "<base64>"}` content-part form — see the [image-to-image request formats](../image_to_image/README.md#request-format).
 
-The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
-(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
-`devices` per stage to relocate; with fewer GPUs available, drop
-the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
-warmup takes roughly an extra 30–60 s on top of the thinker — set
-`--stage-init-timeout 1800` if the default 300 s is too tight.
+### Knobs (declared `extra_body` params)
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed. |
+| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
+| `negative_prompt` | (empty) | Real CFG negative conditioning (text-to-image only). |
+
+For the offline `text_to_image.py` / `image_edit.py` scripts and the full knob reference, see the [image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img).
 
 ## Modality control
 

examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh

@@ -1,169 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-# Server port
-PORT="${PORT:-8091}"
-# Default query type
-QUERY_TYPE="${1:-text}"
-
-# Validate query type
-if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
-    echo "Error: Invalid query type '$QUERY_TYPE'"
-    echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
-    echo "  text: Text-only query"
-    echo "  use_audio: Audio + Text query"
-    echo "  use_image: Image + Text query"
-    echo "  use_video: Video + Text query"
-    echo "  use_mixed_modalities: Audio + Image + Video + Text query"
-    echo "  use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
-    exit 1
-fi
-
-# Define URLs for assets
-MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
-CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-# Build user content based on query type
-case "$QUERY_TYPE" in
-  text)
-    user_content='[
-      {
-        "type": "text",
-        "text": "请详细介绍鹦鹉的生活习性。"
-      }
-    ]'
-    ;;
-  use_image)
-    user_content='[
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe this image in detail."
-        }
-      ]'
-    ;;
-  use_audio)
-    user_content='[
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Please recognize the language of this speech and transcribe it. Format: oral."
-        }
-      ]'
-    ;;
-  use_video)
-    user_content='[
-        {
-          "type": "video_url",
-          "video_url": {
-            "url": "'"$SAMPLE_VIDEO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe what is happening in this video."
-        }
-      ]'
-    ;;
-  use_mixed_modalities)
-    user_content='[
-        {
-          "type": "image_url",
-          "image_url": {
-            "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'"
-          }
-        },
-        {
-          "type": "audio_url",
-          "audio_url": {
-            "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'"
-          }
-        },
-        {
-          "type": "text",
-          "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral"
-        }
-      ]'
-    ;;
-esac
-
-echo "Running query type: $QUERY_TYPE"
-echo ""
-
-if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
-    # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
-    # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3
-    # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
-    # section in README.md for knobs and img2img.
-    out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
-    request_body=$(cat <<'EOF'
-{
-  "model": "Jonathan1909/Ming-flash-omni-2.0",
-  "modalities": ["image"],
-  "messages": [
-    {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
-  ]
-}
-EOF
-)
-    output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
-        -X POST "http://localhost:${PORT}/v1/chat/completions" \
-        -H "Content-Type: application/json" \
-        -d "$request_body")
-    # Server returns base64-encoded PNG either as a data URL in
-    # ``choices[0].message.content[0].image_url.url`` or as a raw string.
-    b64=$(echo "$output" | jq -r '.choices[0].message.content
-        | if type == "array"
-            then (map(select(.image_url? != null)) | .[0].image_url.url // "")
-            elif type == "string" then . else "" end')
-    if [[ -z "$b64" ]]; then
-        echo "Error: no image returned. Raw response:"
-        echo "$output" | jq '.'
-        exit 1
-    fi
-    if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
-    printf '%s' "$b64" | base64 --decode > "$out_path"
-    echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
-    exit 0
-fi
-
-request_body=$(cat <<EOF
-{
-  "model": "Jonathan1909/Ming-flash-omni-2.0",
-  "modalities": ["text"],
-  "messages": [
-    {
-      "role": "system",
-      "content": [
-        {
-          "type": "text",
-          "text": "你是一个友好的AI助手。\n\ndetailed thinking off"
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": $user_content
-    }
-  ]
-}
-EOF
-)
-
-output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
-    -X POST http://localhost:${PORT}/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d "$request_body")
-
-echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

pyproject.toml

@@ -41,6 +41,7 @@ dev = [
     "pytest-cov>=4.0.0",
     "pytest-mock>=3.10.0",
     "websockets>=15.0.1",
+    "pytest-subtests>=0.11.0",
     "datasets>=2.14.0",
     "mypy==1.11.1",
     "pre-commit==4.0.1",

recipes/LTX/LTX-2.3.md

@@ -6,21 +6,23 @@
 
 - Vendor: Lightricks
 - Model: `dg845/LTX-2.3-Diffusers`
-- Task: Text-to-video with synchronized audio generation
+- Task: Text-to-video and image-to-video with synchronized audio generation
 - Mode: Online serving (pure diffusion)
 - Maintainer: @oglok
 
 ## When to use this recipe
 
-Use this recipe when you want to serve LTX-2.3 for text-to-video generation
-with audio. The model generates videos up to 20+ seconds at 768x512 resolution
-with 48kHz audio, all from a single text prompt. Requires a GPU with at least
-96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text
+Use this recipe when you want to serve LTX-2.3 for text-to-video or
+image-to-video generation with audio. The model generates videos up to 20+
+seconds at 768x512 resolution with 48kHz audio from text prompts, and can
+condition video generation on an initial image. Start with a 96GB-class GPU for
+validation because the model combines a 22B parameter transformer with text
 encoder, VAE, and vocoder components.
 
 ## References
 
-- Model: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
+- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
+- Diffusers-format checkpoint: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
 - Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
 
 ## Serving
@@ -34,7 +36,16 @@ vllm serve dg845/LTX-2.3-Diffusers \
   --stage-init-timeout 600
 ```
 
-### Verification
+For image-to-video:
+
+```bash
+vllm serve dg845/LTX-2.3-Diffusers \
+  --omni \
+  --model-class-name LTX23ImageToVideoPipeline \
+  --stage-init-timeout 600
+```
+
+### T2V Verification
 
 ```bash
 # Health check
@@ -73,19 +84,48 @@ curl -X POST http://localhost:8000/v1/videos \
   -F "guidance_scale=4.0"
 ```
 
+### I2V Verification
+
+Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass
+an image reference. The I2V pipeline requires `input_reference` or
+`image_reference` because the first frame is the conditioning input.
+
+```bash
+curl http://localhost:8000/health
+
+curl -X POST http://localhost:8000/v1/videos/sync \
+  -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
+  -F "input_reference=@/absolute/path/to/reference.png" \
+  -F "model=dg845/LTX-2.3-Diffusers" \
+  -F "num_frames=81" \
+  -F "fps=24" \
+  -F "size=768x512" \
+  -F "num_inference_steps=30" \
+  -F "guidance_scale=4.0" \
+  -F "seed=42" \
+  -o ltx23_i2v.mp4
+```
+
 ### Notes
 
-- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference
+- Memory usage: run a latest-head serving smoke or DFX benchmark on the target
+  GPU before publishing load/peak VRAM numbers.
+- Checkpoint format: use a Diffusers-format checkpoint such as
+  `dg845/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
+  raw safetensors and does not contain the subfolder configs required by this
+  pipeline loader.
 - Key flags:
-  - `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request)
-  - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2)
+  - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints
+  - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2)
+  - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline
+  - `input_reference=@...`: Uploads the I2V reference image directly
 - Audio: 48kHz AAC via BWE vocoder, automatically synced with video
 - CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
 - Supported resolutions: 768x512, 512x384 (must be divisible by 32)
 - Frame rate: 24 fps
 - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
 - Known limitations:
-  - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
+  - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract.
   - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
 
 ## Hardware Support
@@ -99,13 +139,16 @@ curl -X POST http://localhost:8000/v1/videos \
 - OS: Ubuntu 22.04
 - Python: 3.10+
 - Driver / runtime: CUDA 13.0, Driver 580.126.09
-- vLLM version: 0.19.x
-- vLLM-Omni version: 0.19.x
+- vLLM version: Match the PR or release checkout being validated
+- vLLM-Omni version or commit: Use the commit you are deploying from
 
-### Validated configurations
+### Benchmarking
 
-| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM |
-|----------|--------|------------|-------|----------|----------------|-----------|
-| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB |
-| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB |
-| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB |
+Do not copy latency or VRAM numbers between commits or machines. Record them
+only after running a latest-head sweep on the target hardware.
+
+For formal PR benchmarking, reuse
+`tests/dfx/perf/tests/test_ltx2_3_vllm_omni.json` with
+`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. That config captures the
+single-device eager baseline and CFG-parallel=2 case for a small 384x512,
+25-frame, 20-step workload.

recipes/MammothModa2/MammothModa2-Preview.md

@@ -32,12 +32,12 @@ flags. Image size uses the standard `--height` / `--width` flags.
   [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 - Declared parameters:
   [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
-- Stage config:
-  [`vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml`](../../vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml)
+- Deploy config:
+  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
 
 ## Hardware Support
 
-The default stage config runs both the AR and DiT stages on a single GPU
+The default deploy config runs both the AR and DiT stages on a single GPU
 (`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
 stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
 after rebalancing the split so the AR weights (~23 GB) leave room for the KV
@@ -48,7 +48,7 @@ cache — see the note under *1x L40S 48GB*.
 ### 1x L40S 48GB
 
 > **48 GB config adjustment:** the committed
-> `vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml` uses
+> `vllm_omni/deploy/mammoth_moda2.yaml` uses
 > `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
 > set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
 > before running. (On an ~80 GB GPU, leave the defaults unchanged.)
@@ -71,13 +71,13 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P
 ```
 
 Run text-to-image with the shared offline example from the repository root. The
-stage config sets `trust_remote_code`, so no extra flag is needed. Forward the
+deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
 MammothModa2 generation parameters as a JSON object through `--extra-body`:
 
 ```bash
 python examples/offline_inference/text_to_image/text_to_image.py \
   --model ./MammothModa2-Preview \
-  --stage-configs-path vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \
   --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
   --height 1024 \
   --width 1024 \

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -102,8 +102,9 @@ curl http://localhost:8099/v1/chat/completions \
 ```
 
 **Text + speech in one response** (the headline 4.5 feature). The TTS
-path is gated by a Jinja flag on the chat template, passed through
-`extra_body.chat_template_kwargs.use_tts_template=true`:
+path is gated by a Jinja flag on the chat template. Pass
+`use_tts_template=true` via the **top-level** `chat_template_kwargs`
+field (curl does not flatten nested `extra_body`):
 
 ```bash
 curl http://localhost:8099/v1/chat/completions \
@@ -112,12 +113,18 @@ curl http://localhost:8099/v1/chat/completions \
         "model": "openbmb/MiniCPM-o-4_5",
         "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
         "modalities": ["text", "audio"],
-        "extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
+        "chat_template_kwargs": {"use_tts_template": true}
     }'
 ```
 
-Response carries text in `choices[0].message.content` and base64 WAV
-in `choices[0].message.audio.data` (24 kHz mono, see Notes).
+When using the OpenAI Python SDK, the same flag can also be sent as
+`extra_body={"chat_template_kwargs": {"use_tts_template": True}}`
+because the client merges `extra_body` into the request root.
+
+Response carries text in one choice's `message.content` and base64 WAV
+in another choice's `message.audio.data` (24 kHz mono, see Notes). With
+`modalities: ["text", "audio"]` you typically get two `choices` entries
+(one text, one audio).
 
 **Gradio demo (text + image + audio + video UI)**:
 
@@ -209,9 +216,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
   missing dep raises `ImportError` at first request with the same
   install hint instead of silently emitting empty audio.
 
-- **TTS trigger**: speech output is only emitted when the client passes
-  `extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
-  the response is text-only (which is also faster).
+- **TTS trigger**: speech output requires
+  `chat_template_kwargs.use_tts_template=true` so the chat template
+  appends `<|tts_bos|>` before generation. Without it, Stage-1 talker
+  receives no TTS token span and returns silent audio (not text-only).
+  For **curl**, put `chat_template_kwargs` at the request root; nested
+  `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK
+  may use `extra_body` because it flattens those fields into the root.
 
 - **Output audio**: 24 kHz mono WAV inside the OpenAI-style
   `message.audio.data` (base64). The Gradio demo's WAV player decodes

recipes/README.md

@@ -44,6 +44,7 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
+| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |

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

@@ -18,6 +18,7 @@ Use this recipe when you want a known-good starting point for serving
 
 - **Thinker only** — multimodal understanding with text output.
 - **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output.
+- **Thinker + Imagegen (DiT)** — text-to-image / img2img
 - **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint.
 
 ## References
@@ -160,6 +161,129 @@ deltas appear at `choices[0].delta.content`.
 - Memory usage: size depends on output modalities and multimodal input; leave
   headroom for video frames and audio caches.
 
+## Image generation (text-to-image / img2img)
+
+Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker.
+
+The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below).
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
+    --stage-init-timeout 1800 \
+    --init-timeout 1800 \
+    --port 8091
+```
+
+With fewer GPUs, copy the YAML and drop the thinker to TP=2 with the DiT on a
+free card.
+
+### Online (text-to-image)
+
+Request image output with `"modalities": ["image"]`:
+
+```bash
+curl -s http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "messages": [
+      {
+        "role": "user",
+        "content": "Please draw a cute cat."
+      }
+    ],
+    "modalities": ["image"]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen.png
+```
+
+**Quick form — `extra_body`** (keys are filtered against the declared set and routed into every stage's `extra_args`). Convenient, but it does **not** let you set the thinker (stage-0) sampling params. Wrap the knobs in a literal `"extra_body"` object for raw curl; the OpenAI Python client's `extra_body=` kwarg produces the same request:
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ],
+    "extra_body": {
+      "steps": 6,
+      "cfg": 1.5,
+      "height": 512,
+      "width": 512,
+      "seed": 123,
+      "byte5_text": ["理解与生成统一"],
+      "negative_prompt": "ugly, blurry, distorted"
+    }
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_extra_body.png
+```
+
+**Full control — `sampling_params_list`** (one entry per stage: `[thinker, imagegen]`).
+Use this when you need to tune the thinker's own sampling (`temperature` / `top_p` / `top_k` / `max_tokens`), or to place knobs explicitly per stage.
+Note `negative_prompt` must sit on the **stage-0 thinker** `extra_args`; the imagegen-stage knobs (`steps` / `cfg` / `height` / `width` / `seed` / `byte5_text`) go on the **stage-1** entry:
+
+```bash
+curl http://127.0.0.1:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Jonathan1909/Ming-flash-omni-2.0",
+    "modalities": ["image"],
+    "sampling_params_list": [
+      {
+        "temperature": 0.4,
+        "top_p": 0.9,
+        "top_k": 1,
+        "max_tokens": 1,
+        "seed": 42,
+        "extra_args": {
+          "negative_prompt": "ugly, blurry, distorted"
+        }
+      },
+      {
+        "seed": 42,
+        "extra_args": {
+          "steps": 6,
+          "cfg": 1.5,
+          "height": 512,
+          "width": 512,
+          "seed": 123,
+          "byte5_text": ["理解与生成统一"]
+        }
+      }
+    ],
+    "messages": [
+      {
+        "role": "user",
+        "content": "Draw a poster."
+      }
+    ]
+  }' \
+  | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \
+  | base64 -d > ming_imagegen_knobs.png
+```
+
+### Knobs (declared `extra_body` params)
+
+| Key | Default | Description |
+| --- | --- | --- |
+| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
+| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
+| `cfg` | 2.0 | Classifier-free guidance scale. |
+| `seed` | 42 | Per-request RNG seed. |
+| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. |
+| `negative_prompt` | (empty) | Real CFG negative conditioning. Spawns a CFG-text companion via `expand_cfg_prompts`; **online / text-to-image only** (offline uses Ming's default zero-negative). |
+
+For img2img, add an `image_url` content part to the user message (online) or pass `--image` (offline); the reference image is routed into the DiT stage as `extra[reference_image]`.
+
 ### 1x H100 80GB — standalone TTS (talker only)
 
 The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint.

recipes/krea/Krea-2.md

@@ -0,0 +1,143 @@
+# Krea 2 for text-to-image (base + LoRA)
+
+## Summary
+
+- Vendor: Krea
+- Model: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`
+- Task: Text-to-image (T2I)
+- Mode: Offline engine and online OpenAI-compatible serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe for a known-good starting point to run the Krea 2 single-stream
+MMDiT with vLLM-Omni on a **single 80 GB NVIDIA H100** — both the offline
+`text_to_image.py` entrypoint and the online `/v1/images/generations` route,
+including PEFT LoRA adapters. `Krea2Pipeline` is auto-detected from
+`model_index.json`; guidance is checkpoint-aware (distilled Turbo runs no-CFG,
+Raw runs CFG).
+
+## References
+
+- Related example under `examples/`:
+  [`examples/offline_inference/text_to_image/README.md`](../../examples/offline_inference/text_to_image/README.md)
+  (see the "Krea 2" section)
+- Feature support:
+  [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md)
+- Related PR: [#4730](https://github.com/vllm-project/vllm-omni/pull/4730)
+
+## Hardware Support
+
+This recipe documents a **single-GPU** CUDA layout on H100 80 GB. Add more
+platforms (ROCm / NPU) or multi-GPU layouts as community validation lands.
+
+## GPU
+
+### 1× H100 80GB
+
+#### Environment
+
+Versions from a working editable install (activate `vllm-omni/.venv` or your equivalent):
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: NVIDIA driver **595.71.05** (CUDA 13.x)
+- torch: **2.11.0+cu130**
+- vLLM: **0.24.0**
+- vLLM-Omni: editable install from this repo (Git **`960ce8f6`**)
+- Transformers: **5.12.1**
+
+Krea 2 weights are ~30 GB and fit comfortably on a single 80 GB H100.
+
+#### Command
+
+Offline — few-step distilled (Turbo) checkpoint (2048×2048, no CFG):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 8 \
+  --guidance-scale 0.0 \
+  --height 2048 --width 2048 \
+  --output krea2_turbo.png
+```
+
+Offline — Raw checkpoint (1024×1024, CFG enabled):
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Raw \
+  --prompt "a fox in the snow" \
+  --num-inference-steps 28 \
+  --guidance-scale 4.5 \
+  --height 1024 --width 1024 \
+  --output krea2_raw.png
+```
+
+LoRA (offline) — pass a PEFT adapter with `--lora-path` / `--lora-scale`.
+`NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA` is a vLLM-Omni-compatible repackaging
+of `krea/Krea-2-LoRA-darkbrush`; its trigger word is `monochrome ink wash style`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo \
+  --prompt "a fox in the snow, monochrome ink wash style" \
+  --num-inference-steps 8 --guidance-scale 0.0 \
+  --lora-path NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA --lora-scale 1.0 \
+  --output krea2_darkbrush.png
+```
+
+Online — serve and query the OpenAI-compatible image route:
+
+```bash
+vllm serve krea/Krea-2-Turbo --omni --enforce-eager --port 8091
+```
+
+```bash
+curl -s http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "a fox in the snow, monochrome ink wash style",
+    "size": "1024x1024",
+    "num_inference_steps": 8,
+    "guidance_scale": 0.0,
+    "seed": 42,
+    "lora": {"name": "darkbrush", "path": "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA", "scale": 1.0}
+  }' | jq -r '.data[0].b64_json' | base64 -d > krea2_online.png
+```
+
+Cache-DiT acceleration (offline) — add `--cache-backend cache_dit`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model krea/Krea-2-Turbo --prompt "a fox in the snow" \
+  --num-inference-steps 8 --guidance-scale 0.0 \
+  --cache-backend cache_dit --output krea2_cachedit.png
+```
+
+#### Verification
+
+```bash
+# Offline: the script prints "Saved generated image to <path>"; confirm the file exists.
+ls -lh krea2_turbo.png
+```
+
+For the online route, decode the returned `b64_json` (see the `curl` above) and
+confirm a valid 1024×1024 PNG is written.
+
+#### Notes
+
+- Key flags: distilled Turbo → few steps (~8) with `guidance_scale=0.0`; Raw →
+  more steps (~28) with `guidance_scale>0` (CFG). `guidance-scale` follows the
+  Krea 2 convention `velocity = cond + guidance_scale * (cond - uncond)`.
+- LoRA: adapters must be PEFT format (`adapter_config.json` +
+  `adapter_model.safetensors`); the projections are `ReplicatedLinear`
+  (264 modules). Offline uses `--lora-path` (which sets
+  `OmniDiffusionSamplingParams.lora_request` / `lora_scale`); online uses the
+  top-level `lora` object. Style LoRAs like darkbrush need their trigger word in
+  the prompt.
+- Cache-DiT is supported; `has_separate_cfg` is checkpoint-aware (False for Turbo
+  no-CFG, True for Raw CFG).
+- Known limitations: TP / SP / CFG-Parallel are not yet wired for Krea 2. HSDP,
+  layerwise CPU offload, and VAE-patch-parallel (decode) are supported.

recipes/mistralai/Voxtral-TTS.md

@@ -27,6 +27,79 @@ examples in this repository.
 
 ## GPU
 
+### 1x NVIDIA GeForce RTX 5090 32GB
+
+#### Environment
+
+- OS: Ubuntu 24.04.4 LTS
+- Python: 3.12.3
+- PyTorch: 2.11.0+cu130
+- Driver / runtime: NVIDIA driver 580.159.03, CUDA 13.0
+- GPU: NVIDIA GeForce RTX 5090, 32 GB (`32607 MiB`)
+- vLLM version: 0.23.0
+- vLLM-Omni version or commit: 0.23.0rc2.dev140+gbe60d7c7c (`be60d7c7`)
+- Tested with: `mistral-common==1.11.5`
+
+#### Command
+
+```bash
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
+    --omni \
+    --port 8091 \
+    --trust-remote-code
+```
+
+The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`.
+
+#### Verification
+
+Health check and voices:
+
+```bash
+curl -s -o /tmp/voxtral_health.txt -w '%{http_code}' http://127.0.0.1:8091/health
+curl -s http://127.0.0.1:8091/v1/audio/voices
+```
+
+Non-streaming WAV:
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS running on an NVIDIA GeForce RTX 5090 with vLLM-Omni.",
+    "voice": "casual_female",
+    "language": "English",
+    "response_format": "wav"
+  }' \
+  --output /tmp/voxtral_5090.wav
+```
+
+Streaming PCM:
+
+```bash
+curl -X POST http://127.0.0.1:8091/v1/audio/speech \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "mistralai/Voxtral-4B-TTS-2603",
+    "input": "Hello, this is Voxtral TTS streaming PCM on an NVIDIA GeForce RTX 5090.",
+    "voice": "casual_female",
+    "language": "English",
+    "stream": true,
+    "response_format": "pcm"
+  }' \
+  --output /tmp/voxtral_5090_stream.pcm
+
+ffmpeg -f s16le -ar 24000 -ac 1 -i /tmp/voxtral_5090_stream.pcm /tmp/voxtral_5090_stream.wav -y
+ffprobe -hide_banner /tmp/voxtral_5090_stream.wav
+```
+
+#### Notes
+
+- Output: 24 kHz mono WAV/PCM; both non-streaming and streaming requests passed.
+- Memory usage: ~27.3 GiB after warmup (Stage 0 ~25.5 GiB, Stage 1 ~1.7 GiB).
+- Key flags: `--omni` is required; `--trust-remote-code` is recommended.
+
 ### 1x RTX 4090 24GB
 
 #### Environment

requirements/common.txt

@@ -13,7 +13,7 @@ transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
-safetensors>=0.8.0rc0
+safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
 cache-dit==1.3.0

tests/ar_diffusion/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
+
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+)
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
+
+
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
+
+
+def test_ar_diffusion_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
+
+
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)

tests/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,254 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
+
+BLOCK = 16
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    spec = make_spec()
+    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+    assert skip % 16 == 0 and skip == 16
+
+
+def test_sink_chunks_protected():
+    # sink = 1 chunk (16 tokens) is never skipped.
+    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
+    assert skip == 32
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- ARDiffusionKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    k_pools[0][BLOCK + 3].fill_(7)
+    v_pools[0][2 * BLOCK + 5].fill_(11)
+    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        cross_attn_length=L,
+        device=torch.device("cpu"),
+    )
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 2 * (2 + 1) + 2  # above the min-blocks floor, so the deduction is what's tested
+
+
+def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9):
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        local_branches=local_branches,
+        num_frame_per_block=num_frame_per_block,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk;
+    a single-process run (both branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids(False, 0, cap + 1)
+
+
+def test_invalid_local_branches_rejected():
+    with pytest.raises(ValueError, match="local_branches"):
+        _make_kv(local_branches=3)

tests/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        used = {int(s) // kv.block_size for s in slots}
+        assert kv.null_block_id not in used
+        assert used <= set(block_table)
+        # Resident window stays bounded by window + the in-flight chunk.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,394 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    is_negative: bool,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(is_negative)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.pos.completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 1
+    assert st._committed[False] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1)
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            is_negative=False,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.pos.completed_chunks
+    st.commit_paged_context(False)
+    assert st.pos.completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        is_negative=False,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    assert torch.equal(paged, direct)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(False)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(False)
+        return out
+
+    torch._dynamo.reset()
+    from torch._dynamo.testing import CompileCounter
+
+    counter = CompileCounter()
+
+    def fn(inputs, q, k, v):
+        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+    compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+    run_one_forward(commit=True)  # history grows between calls ->
+    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+    run_one_forward(commit=False)
+
+    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"

tests/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        cross_attn_length=cross_attn_length,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(is_negative)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for neg in (False, True):
+        contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
+
+
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.pos.completed_chunks == 0
+    assert kv.window_block_ids(st.pos) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.pos.completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 2
+    assert st._committed[False] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.pos)) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
+    st.commit_paged_context(False)
+
+    assert st.pos.completed_chunks == 0
+    assert st._committed[False] == 0
+    assert kv.window_block_ids(st.pos) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, False, 3)
+    assert len(kv.window_block_ids(st.pos)) == 3
+
+    _prepare_and_commit(st, False, 2)
+    ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
+
+
+def test_branches_are_independent():
+    kv, st = make_state()
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)
+    _prepare_and_commit(st, True, 1)
+    assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {False: 0, True: 0}
+    assert st._paged_pending == {False: None, True: None}
+    assert kv.window_block_ids(st.pos) == []
+    assert kv.window_block_ids(st.neg) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, False, 1)
+    assert len(kv.window_block_ids(st.pos)) == 1
+
+
+def test_get_cross_kv_caches_raises_before_populate():
+    """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
+    _, st = make_state()  # cross_attn_length == 0, nothing populated
+    for neg in (False, True):
+        with pytest.raises(RuntimeError, match="cross-attn read before"):
+            st.get_cross_kv_caches(neg)
+    # Guards the full AND: the populated flag alone is not enough without a pool.
+    st._cross_text_populated[False] = True
+    with pytest.raises(RuntimeError, match="cross-attn read before"):
+        st.get_cross_kv_caches(False)
+
+
+def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for i in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        kv.write_cross_kv(i, False, k, v)
+        written.append((k, v))
+    st._cross_text_populated[False] = True
+
+    out = st.get_cross_kv_caches(False)
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
+
+
+def test_kv_create_is_noop_engine_owns_allocation():
+    from unittest.mock import MagicMock
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._ar_diffusion_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
+
+
+def test_get_kv_cache_requires_frame_aligned_seqlen():
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.neg) == []
+
+
+def test_failed_forward_tears_down_session(monkeypatch):
+    """Review follow-up (hsliuustc0106): a forward that dies partway must not
+    leave allocated-but-uncommitted KV behind — the session is torn down (pool
+    blocks freed, model-local state dropped) and the error propagates."""
+    from types import SimpleNamespace
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, False, 2)  # session owns pool blocks
+    free_before_failure = kv.manager.block_pool.get_num_free_blocks()
+
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.kv_cache = kv
+    runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
+    runner._max_ar_diffusion_states = 4
+    runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
+    runner.device = None
+    runner._perf_e2e_times = []
+
+    def boom(self, req):
+        raise RuntimeError("layer 17 exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    req = SimpleNamespace(
+        request_id="r0",
+        sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
+    )
+
+    with pytest.raises(RuntimeError, match="layer 17 exploded"):
+        ARDiffusionModelRunner.execute_model(runner, req)
+
+    assert "s1" not in runner._ar_diffusion_states
+    assert "s1" not in runner.pipeline._states
+    assert runner.pipeline._ar_diffusion_kv_state is None
+    # close() freed the session's resident blocks.
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure

tests/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size

tests/config/composable_parallel/test_apply.py

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

tests/config/composable_parallel/test_factory_integration.py

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

tests/config/composable_parallel/test_spec_types.py

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

tests/config/composable_parallel/test_strategy_loader.py

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

tests/config/composable_parallel/test_translator.py

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

tests/config/test_endpoint_policy.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for endpoint restrictions logic."""
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.config.endpoint_policy import (
+    EndpointRestriction,
+    OmniServingCapability,
+    shutdown_unsupported_routes,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SERVER_ROUTE = "/v1/completions"
+REJECTION_REASON = "This route is banned!"
+
+
+def _make_app_with_server_route():
+    """Create a FastAPI app with a route to reject."""
+    app = FastAPI()
+
+    @app.post(SERVER_ROUTE)
+    async def existing_handler():
+        """Returns a response; we want to disable this."""
+        return {"ok": True}
+
+    return app
+
+
+def test_restricted_completions_returns_400():
+    """Ensure that simple route rejection works."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_unrestricted_completions_not_blocked():
+    """Ensure that if we don't shutdown any routes, the route stays."""
+    app = _make_app_with_server_route()
+    shutdown_unsupported_routes(app, ())
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE)
+    assert resp.status_code == 200
+    assert resp.json() == {"ok": True}
+
+
+def test_route_rejection_is_idempotent():
+    """Ensure that calling rejection twice doesn't break things."""
+    app = _make_app_with_server_route()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"
+
+
+def test_invalid_route_rejection():
+    """Ensure that a restriction won't crash if a route doesn't exist."""
+    # NOTE: This test is for safety since the route organization is still messy.
+    # Since the restriction needs a valid OmniServingCapability, its implied the route
+    # is a valid vLLM route, but was already removed from the app by earlier processing.
+    app = FastAPI()
+    restrictions = (EndpointRestriction(OmniServingCapability.COMPLETIONS, REJECTION_REASON),)
+    shutdown_unsupported_routes(app, restrictions)
+
+    client = TestClient(app)
+    resp = client.post(SERVER_ROUTE, json={"prompt": "hello", "model": "m"})
+    assert resp.status_code == 400
+    body = resp.json()
+    assert body["error"]["message"] == REJECTION_REASON
+    assert body["error"]["type"] == "BadRequestError"

tests/config/test_omni_config.py

@@ -751,3 +751,14 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "distributed_executor_backend",
     } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS
     assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS
+
+
+def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable():
+    quantization_config = {
+        "method": "example_quant",
+        "weights": "weights.bin",
+    }
+
+    cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
+
+    assert cfg.quantization_config == quantization_config

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod
+import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod
+from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _MockQueue:
+    def __init__(self, items: list | None = None) -> None:
+        self._items = list(items or [])
+
+    def __bool__(self) -> bool:
+        return bool(self._items)
+
+    def __iter__(self):
+        return iter(self._items)
+
+    def __len__(self) -> int:
+        return len(self._items)
+
+    def add_request(self, request) -> None:
+        self._items.append(request)
+
+    def prepend_requests(self, requests) -> None:
+        self._items = list(requests) + self._items
+
+
+class _MockRequest:
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        status: RequestStatus = RequestStatus.RUNNING,
+        num_prompt_tokens: int = 4,
+        num_computed_tokens: int = 4,
+        num_output_placeholders: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self.status = status
+        self.num_prompt_tokens = num_prompt_tokens
+        self.num_computed_tokens = num_computed_tokens
+        self.num_output_placeholders = num_output_placeholders
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+@pytest.fixture(autouse=True)
+def _mock_cuda_graph_platform(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True)
+
+
+def _make_scheduler(
+    *,
+    enable_unified_decode_graph: bool | None = True,
+    deterministic_cfm_noise: bool = False,
+    runtime_config_on_model_config: bool = False,
+) -> VoxCPM2OmniARAsyncScheduler:
+    sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler)
+    hf_config = SimpleNamespace()
+    model_config = SimpleNamespace(hf_config=hf_config)
+    if enable_unified_decode_graph is not None or deterministic_cfm_noise:
+        runtime_config = SimpleNamespace(
+            enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph,
+            deterministic_cfm_noise=deterministic_cfm_noise,
+        )
+        if runtime_config_on_model_config:
+            model_config.voxcpm2_runtime_config = runtime_config
+        else:
+            hf_config.voxcpm2_runtime_config = runtime_config
+    sched.vllm_config = SimpleNamespace(model_config=model_config)
+    return sched
+
+
+def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_decode_ready() -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("prefill-running", num_prompt_tokens=8, num_computed_tokens=4)]
+    scheduler.waiting = _MockQueue([_MockRequest("waiting", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=False)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=None)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None:
+    scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None:
+    scheduler = _make_scheduler(deterministic_cfm_noise=True)
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None:
+    monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False)
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
+
+    assert not scheduler._should_defer_waiting_for_unified_decode_graph()
+
+
+def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
+    scheduler = _make_scheduler()
+    scheduler.running = [_MockRequest("decode")]
+    original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
+    deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
+    original_waiting = _MockQueue([original_waiting_req])
+    scheduler.waiting = original_waiting
+    scheduler.policy = "fcfs"
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler._consume_pending_connector_output = lambda model_mode: None
+
+    monkeypatch.setattr(scheduler_mod, "create_request_queue", lambda _policy: _MockQueue())
+
+    def fake_upstream_schedule(self, throttle_prefills: bool = False):
+        assert self.waiting is not original_waiting
+        assert not self.waiting
+        self.waiting.add_request(deferred_by_upstream)
+        raise RuntimeError("stop before output wrapping")
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
+
+    with pytest.raises(RuntimeError, match="stop before output wrapping"):
+        scheduler.schedule()
+
+    assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]

tests/dfx/perf/scripts/run_benchmark.py

@@ -94,7 +94,7 @@ def omni_server(request):
 benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
 
 
-@pytest.fixture(params=benchmark_indices)
+@pytest.fixture
 def benchmark_params(request, omni_server):
     """Benchmark parameters fixture with proper parametrization"""
     test_name, param_index = request.param

tests/dfx/perf/tests/test_voxcpm2.json

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

tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py

@@ -62,8 +62,8 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
         ),
         pytest.param(
             True,
-            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
-            ("Invalid layers", "3", "1"),
+            {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "1"},
+            ("Invalid layers", "2", "10"),
             400,
             id="invalid_layers",
         ),

tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py

@@ -87,8 +87,8 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
         pytest.param(
             {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
         ),
-        pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
-        pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
+        pytest.param({"layers": 1}, ("layers", "value_error", "2", "10"), id="layers_below_min"),
+        pytest.param({"layers": 11}, ("layers", "value_error", "2", "10"), id="layers_above_max"),
         pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
         pytest.param(
             {"seed": 2**32},

tests/diffusion/attention/test_kernels_hub.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+import types
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.platforms import current_omni_platform
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
+    """Test that when kernels package is not available, the platform fallback logic
+
+    routes to local FLASH_ATTN.
+    """
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    # Temporarily hide/remove kernels module if present
+    monkeypatch.setitem(sys.modules, "kernels", None)
+
+    # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+
+    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
+    kernels_module = types.ModuleType("kernels")
+    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
+    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+
+
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_kernels_hub_execution():
+    """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference."""
+    device = torch.device(current_omni_platform.device_type)
+    dtype = torch.bfloat16
+
+    num_heads = 8
+    head_dim = 64
+    seq_len = 32
+    batch_size = 1
+
+    torch.manual_seed(42)
+    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
+    k = q.clone()
+    v = q.clone()
+
+    # 1. Test PyTorch SDPA reference
+    sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
+    attn_metadata_sdpa = AttentionMetadata(attn_mask=None)
+    output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+
+    # 2. Test FlashAttentionHubBackend (FlashAttention 2)
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl
+
+    fa_hub_impl = FlashAttentionHubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa_hub.shape == q.shape
+    assert not torch.isnan(output_fa_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}"
+
+    # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only)
+    major, _minor = torch.cuda.get_device_capability()
+    if major < 9:
+        pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)")
+
+    from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl
+
+    fa3_hub_impl = FlashAttention3HubImpl(
+        num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False
+    )
+    output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa)
+    assert output_fa3_hub.shape == q.shape
+    assert not torch.isnan(output_fa3_hub).any()
+    max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item()
+    assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}"

tests/diffusion/cache/test_cache_dit.py

@@ -13,7 +13,7 @@ import torch
 from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2
 
 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend
-from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
@@ -193,3 +193,12 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
     # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio)
     assert torch.equal(captured["hidden_states"], video_in)
     assert torch.equal(captured["encoder_hidden_states"], audio_in)
+
+
+def test_summary_with_no_transformer_is_nonfatal():
+    """Regression test for https://github.com/vllm-project/vllm-omni/issues/4325."""
+
+    class FakePipeline:
+        pass
+
+    cache_summary(pipeline=FakePipeline())

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -67,6 +67,50 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(**defaults)
 
 
+def _patch_fake_diffusers_quantization_backends(mocker):
+    class FakePipelineQuantizationConfig:
+        def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+            self.quant_mapping = quant_mapping
+            self.quant_backend = quant_backend
+            self.quant_kwargs = quant_kwargs
+            self.components_to_quantize = components_to_quantize
+
+    class FakeTorchAoConfig:
+        def __init__(self, quant_type, modules_to_not_convert=None):
+            self.quant_type = quant_type
+            self.modules_to_not_convert = modules_to_not_convert
+
+        @classmethod
+        def from_dict(cls, config):
+            return cls(quant_type=config["quant_type"])
+
+    class FakeInt8DynamicActivationInt8WeightConfig:
+        pass
+
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.PipelineQuantizationConfig",
+        FakePipelineQuantizationConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.TorchAoConfig",
+        FakeTorchAoConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_torchao_quant_type_cls",
+        return_value=FakeInt8DynamicActivationInt8WeightConfig,
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_diffusers_quantization_config_cls",
+        return_value=FakeTorchAoConfig,
+    )
+    return FakePipelineQuantizationConfig, FakeTorchAoConfig, FakeInt8DynamicActivationInt8WeightConfig
+
+
+class _TransformerBackedPipeline(DiffusionPipeline):
+    def __init__(self, transformer):
+        pass
+
+
 def _make_batch(**overrides) -> DiffusionRequestBatch:
     """Wrap a single request in a DiffusionRequestBatch, matching the forward contract."""
     return DiffusionRequestBatch(requests=[_make_request(**overrides)])
@@ -98,7 +142,7 @@ class TestPipelineArgumentsHandling:
 
     @pytest.mark.parametrize(
         "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"],
+        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
     )
     def test_adapter_guard_unsupported_feature(self, feature_id):
         if feature_id == "cfg_parallel":
@@ -128,14 +172,188 @@ class TestPipelineArgumentsHandling:
             )
         elif feature_id == "enforce_eager":
             od_config = _make_od_config(enforce_eager=True)
-        elif feature_id == "quantization":
-            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8"))
+        elif feature_id == "unsupported_quantization":
+            od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="gguf"))
         else:
             raise ValueError(f"Unknown feature ID: {feature_id}")
 
         with pytest.raises(NotImplementedError):
             DiffusersAdapterPipeline(od_config=od_config)
 
+    def test_adapter_load_weights_injects_supported_quantization_config(self, mocker):
+        (
+            FakePipelineQuantizationConfig,
+            FakeTorchAoConfig,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        ) = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value={
+                "_class_name": "TransformerBackedPipeline",
+                "scheduler": ["diffusers", "SchedulerMixin"],
+                "transformer": ["diffusers", "Transformer2DModel"],
+                "transformer_2": ["diffusers", "Transformer2DModel"],
+                "vae": ["diffusers", "AutoencoderKL"],
+            },
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_pipeline_cls=_TransformerBackedPipeline,
+            )
+        )
+        adapter.load_weights()
+
+        pipeline_quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        torchao_config = pipeline_quant_config.quant_mapping["transformer"]
+        assert isinstance(pipeline_quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(torchao_config, FakeTorchAoConfig)
+        assert isinstance(torchao_config.quant_type, FakeInt8DynamicActivationInt8WeightConfig)
+        assert isinstance(
+            pipeline_quant_config.quant_mapping["transformer_2"].quant_type,
+            FakeInt8DynamicActivationInt8WeightConfig,
+        )
+
+    @pytest.mark.parametrize(
+        ("pipeline_config", "diffusers_load_kwargs", "diffusers_pipeline_cls"),
+        [
+            (
+                {
+                    "_class_name": "NoTransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "unet": ["diffusers", "UNet2DConditionModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {},
+                _TransformerBackedPipeline,
+            ),
+            (
+                {
+                    "_class_name": "TransformerPipeline",
+                    "scheduler": ["diffusers", "SchedulerMixin"],
+                    "transformer": ["diffusers", "Transformer2DModel"],
+                    "vae": ["diffusers", "AutoencoderKL"],
+                },
+                {"transformer": None},
+                None,
+            ),
+        ],
+    )
+    def test_adapter_load_weights_rejects_converted_quantization_without_transformer_component(
+        self,
+        pipeline_config,
+        diffusers_load_kwargs,
+        diffusers_pipeline_cls,
+        mocker,
+    ):
+        _patch_fake_diffusers_quantization_backends(mocker)
+        mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config",
+            return_value=pipeline_config,
+        )
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained"
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="int8"),
+                diffusers_load_kwargs=diffusers_load_kwargs,
+                diffusers_pipeline_cls=diffusers_pipeline_cls,
+            )
+        )
+
+        with pytest.raises(NotImplementedError, match="transformer"):
+            adapter.load_weights()
+        mock_from_pretrained.assert_not_called()
+
+    def test_adapter_rejects_dduf_file_with_converted_quantization(self):
+        with pytest.raises(NotImplementedError, match="dduf_file"):
+            DiffusersAdapterPipeline(
+                od_config=_make_od_config(
+                    quantization_config=SimpleNamespace(quant_method="int8"),
+                    diffusers_load_kwargs={"dduf_file": "model.dduf"},
+                )
+            )
+
+    def test_adapter_load_weights_preserves_diffusers_native_quantization_config(self, mocker):
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        diffusers_native_quantization_config = object()
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                quantization_config=SimpleNamespace(quant_method="gguf"),
+                diffusers_load_kwargs={
+                    "dduf_file": "model.dduf",
+                    "quantization_config": diffusers_native_quantization_config,
+                },
+            )
+        )
+        adapter.load_weights()
+
+        kwargs = mock_from_pretrained.call_args.kwargs
+        assert kwargs["quantization_config"] is diffusers_native_quantization_config
+        assert kwargs["dduf_file"] == "model.dduf"
+
+    def test_adapter_load_weights_builds_diffusers_native_quantization_config_from_dict(self, mocker):
+        FakePipelineQuantizationConfig, FakeTorchAoConfig, _ = _patch_fake_diffusers_quantization_backends(mocker)
+
+        class MockPipeline:
+            def __call__(self, prompt=None):
+                return None
+
+            def to(self, device):
+                return self
+
+        mock_from_pretrained = mocker.patch(
+            "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained",
+            return_value=MockPipeline(),
+        )
+
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                diffusers_load_kwargs={
+                    "quantization_config": {
+                        "quant_mapping": {
+                            "transformer": {
+                                "quant_method": "torchao",
+                                "quant_type": "fake-int8-config",
+                            },
+                        },
+                    },
+                },
+            )
+        )
+        adapter.load_weights()
+
+        quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"]
+        assert isinstance(quant_config, FakePipelineQuantizationConfig)
+        assert isinstance(quant_config.quant_mapping["transformer"], FakeTorchAoConfig)
+        assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"
+
     def test_adapter_guard_unknown_output_type(self, mocker):
         """Test that the adapter wraps an unknown output type as-is.
         This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass."""

tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Optional real-backend checks for Diffusers quantization config conversion."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _load_quantization_utils():
+    module_path = (
+        Path(__file__).resolve().parents[3] / "vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py"
+    )
+    spec = importlib.util.spec_from_file_location("diffusers_quantization_utils_real_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_convert_diffusers_quantization_config_from_real_torchao_dict():
+    diffusers = pytest.importorskip("diffusers")
+    torchao_quantization = pytest.importorskip("torchao.quantization")
+    quant_type_cls = getattr(torchao_quantization, "Int8DynamicActivationInt8WeightConfig", None)
+    if quant_type_cls is None:
+        pytest.skip("torchao does not expose Int8DynamicActivationInt8WeightConfig")
+
+    quantization_utils = _load_quantization_utils()
+
+    real_torchao_config = diffusers.TorchAoConfig(quant_type=quant_type_cls())
+    serialized_torchao_config = real_torchao_config.to_dict()
+    assert serialized_torchao_config["quant_type"].keys() == {"default"}
+
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": serialized_torchao_config,
+            },
+        },
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    component_config = quant_config.quant_mapping["transformer"]
+    assert isinstance(quant_config, diffusers.PipelineQuantizationConfig)
+    assert isinstance(component_config, diffusers.TorchAoConfig)
+    assert isinstance(
+        component_config.quant_type,
+        quant_type_cls,
+    )

tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused tests for Diffusers backend quantization conversion helpers."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _FakePipelineQuantizationConfig:
+    def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None):
+        self.quant_mapping = quant_mapping
+        self.quant_backend = quant_backend
+        self.quant_kwargs = quant_kwargs
+        self.components_to_quantize = components_to_quantize
+
+
+class _FakeTorchAoConfig:
+    def __init__(self, quant_type):
+        self.quant_type = quant_type
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(quant_type=config["quant_type"])
+
+
+class _FakeFloat8DynamicActivationFloat8WeightConfig:
+    pass
+
+
+class _FakeInt8DynamicActivationInt8WeightConfig:
+    pass
+
+
+@pytest.fixture(autouse=True)
+def patch_quantization_backends(monkeypatch):
+    monkeypatch.setattr(
+        quantization_utils,
+        "PipelineQuantizationConfig",
+        _FakePipelineQuantizationConfig,
+    )
+    monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig)
+
+    def fake_get_torchao_quant_type_cls(class_name: str):
+        return {
+            "Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig,
+            "Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig,
+        }[class_name]
+
+    monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls)
+    monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig)
+
+
+def _quant_config(method: str, **kwargs):
+    return SimpleNamespace(get_name=lambda: method, **kwargs)
+
+
+@pytest.mark.parametrize(
+    ("method", "expected_type"),
+    [
+        ("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig),
+        ("int8", _FakeInt8DynamicActivationInt8WeightConfig),
+    ],
+)
+def test_apply_injects_converted_quantization_config(method, expected_type):
+    od_config = SimpleNamespace(quantization_config=_quant_config(method))
+    load_kwargs = {}
+
+    injected = quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {"transformer": ["diffusers", "Transformer2DModel"]},
+    )
+
+    assert injected is True
+    torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"]
+    assert isinstance(torchao_config.quant_type, expected_type)
+
+
+def test_apply_injects_converted_quantization_config_for_transformer_2():
+    od_config = SimpleNamespace(quantization_config=_quant_config("int8"))
+    load_kwargs = {}
+
+    quantization_utils.apply_diffusers_quantization_config(
+        od_config,
+        load_kwargs,
+        {
+            "transformer": ["diffusers", "Transformer2DModel"],
+            "transformer_2": ["diffusers", "Transformer2DModel"],
+        },
+    )
+
+    assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"]
+
+
+@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"])
+def test_unsupported_methods_fail_explicitly(method):
+    with pytest.raises(NotImplementedError, match=method):
+        quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method))
+
+
+@pytest.mark.parametrize(
+    "quant_config",
+    [
+        _quant_config("fp8", activation_scheme="static"),
+        _quant_config("int8", is_checkpoint_int8_serialized=True),
+        _quant_config("fp8", weight_block_size=[128, 128]),
+        _quant_config("int8", ignored_layers=["transformer.proj_out"]),
+    ],
+)
+def test_ambiguous_mappings_fail_explicitly(quant_config):
+    with pytest.raises(NotImplementedError):
+        quantization_utils.ensure_supported_diffusers_quantization(quant_config)
+
+
+def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker):
+    od_config = SimpleNamespace(quantization_config=_quant_config("fp8"))
+    existing = object()
+    load_kwargs = {"quantization_config": existing}
+    mock_warning = mocker.patch.object(quantization_utils.logger, "warning")
+
+    injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {})
+
+    assert injected is False
+    assert load_kwargs["quantization_config"] is existing
+    assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0]
+
+
+def test_convert_diffusers_quantization_config_from_dict():
+    load_kwargs = {
+        "quantization_config": {
+            "quant_mapping": {
+                "transformer": {
+                    "quant_method": "torchao",
+                    "quant_type": "fake-int8-config",
+                }
+            }
+        }
+    }
+
+    quantization_utils.convert_diffusers_quantization_config(load_kwargs)
+
+    quant_config = load_kwargs["quantization_config"]
+    assert isinstance(quant_config, _FakePipelineQuantizationConfig)
+    assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig)
+    assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config"

tests/diffusion/distributed/test_expert_parallel_layout.py

@@ -0,0 +1,313 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.distributed.parallel_state as parallel_state
+
+
+class _FakeGroup:
+    def __init__(
+        self,
+        group_ranks: list[list[int]],
+        local_rank: int,
+        parallel_mode: str,
+        **kwargs,
+    ) -> None:
+        self.group_ranks = group_ranks
+        self.parallel_mode = parallel_mode
+        self.device_group = object()
+        self.device_communicator = kwargs.get("device_communicator")
+        reduce_scatter = kwargs.get("reduce_scatter")
+        if reduce_scatter is not None:
+            self.reduce_scatter = reduce_scatter
+        self.ulysses_group = kwargs.get("ulysses_group")
+        self.ring_group = kwargs.get("ring_group")
+        self.local_group = next(group for group in group_ranks if local_rank in group)
+        self.world_size = len(self.local_group)
+        self.rank_in_group = self.local_group.index(local_rank)
+
+    def destroy(self) -> None:
+        pass
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
+    """MoE+EP rank layout should map SP->PCP, CFG*DP->DP, and TP*SP*CFG*DP->EP."""
+    local_rank = 0
+    world_size = 32
+    created_groups: list[_FakeGroup] = []
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+        created_groups.append(group)
+        return group
+
+    def fake_init_vllm_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        group_name,
+    ):
+        del backend
+        group = _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            f"vllm_{group_name}",
+            device_communicator=object(),
+            reduce_scatter=lambda tensor, **kwargs: tensor,
+        )
+        created_groups.append(group)
+        return group
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=True))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=2,
+        ulysses_degree=2,
+        ring_degree=1,
+        pipeline_parallel_size=2,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=True,
+        backend="gloo",
+    )
+
+    assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP
+    assert parallel_state.vllm_parallel_state._PCP.world_size == 2
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 4
+    assert parallel_state.vllm_parallel_state._EP.world_size == 16
+    assert parallel_state.vllm_parallel_state._TP.world_size == 2
+    assert parallel_state._PP.world_size == 2
+
+    assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._DP.device_communicator is not None
+    assert parallel_state.vllm_parallel_state._EP.device_communicator is not None
+    assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter")
+    assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter")
+
+    assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2]
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24]
+    assert parallel_state.vllm_parallel_state._EP.local_group == [
+        0,
+        1,
+        2,
+        3,
+        8,
+        9,
+        10,
+        11,
+        16,
+        17,
+        18,
+        19,
+        24,
+        25,
+        26,
+        27,
+    ]
+    assert parallel_state.get_expert_parallel_group_ranks() == [
+        [
+            0,
+            1,
+            2,
+            3,
+            8,
+            9,
+            10,
+            11,
+            16,
+            17,
+            18,
+            19,
+            24,
+            25,
+            26,
+            27,
+        ],
+        [
+            4,
+            5,
+            6,
+            7,
+            12,
+            13,
+            14,
+            15,
+            20,
+            21,
+            22,
+            23,
+            28,
+            29,
+            30,
+            31,
+        ],
+    ]
+
+    vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")]
+    assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"]
+    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"]
+    assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group]
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch):
+    """vLLM DP should keep diffusion DP when expert parallelism is not enabled."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    parallel_state.initialize_model_parallel(
+        tensor_parallel_size=2,
+        sequence_parallel_size=1,
+        ulysses_degree=1,
+        ring_degree=1,
+        pipeline_parallel_size=1,
+        cfg_parallel_size=2,
+        data_parallel_size=2,
+        enable_expert_parallel=False,
+        backend="gloo",
+    )
+
+    assert parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP is parallel_state._DP
+    assert parallel_state.vllm_parallel_state._DP.world_size == 2
+    assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4]
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None
+
+
+@pytest.mark.cpu
+@pytest.mark.core_model
+def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch):
+    """Non-MoE diffusion configs should not create vLLM PCP/DP/EP remap state."""
+    local_rank = 0
+    world_size = 8
+
+    def fake_init_model_parallel_group(
+        group_ranks,
+        local_rank,
+        backend,
+        parallel_mode=None,
+        group_name=None,
+        **kwargs,
+    ):
+        del backend, group_name
+        return _FakeGroup(
+            [list(ranks) for ranks in group_ranks],
+            local_rank,
+            parallel_mode or "",
+            **kwargs,
+        )
+
+    fake_world_group = SimpleNamespace(
+        rank_in_group=local_rank,
+        local_rank=local_rank,
+        device_group=object(),
+    )
+    fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=False))
+    monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size)
+    monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo")
+    monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks))
+    monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group)
+    monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context)
+    monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group)
+    monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None)
+
+    for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"):
+        monkeypatch.setattr(parallel_state, name, None)
+    for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"):
+        monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False)
+
+    with pytest.raises(RuntimeError, match="Expert parallelism enabled for a non-MoE model"):
+        parallel_state.initialize_model_parallel(
+            tensor_parallel_size=2,
+            sequence_parallel_size=2,
+            ulysses_degree=2,
+            ring_degree=1,
+            pipeline_parallel_size=1,
+            cfg_parallel_size=2,
+            data_parallel_size=1,
+            enable_expert_parallel=True,
+            backend="gloo",
+        )
+
+    assert parallel_state.vllm_parallel_state._PCP is None
+    assert parallel_state.vllm_parallel_state._EP is None
+    assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None

tests/diffusion/layers/test_rope_broadcast.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Shape broadcasting regression tests for RotaryEmbedding.forward_native.
+
+Covers the batch-dimension shape mismatch fix: forward_native now strips the
+batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda,
+forward_hip, and apply_rotary_emb_mindiesd paths.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_inputs(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+    cos_batch: int | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Generate x, cos, sin tensors for RoPE testing.
+
+    Args:
+        batch: Batch dimension of x.
+        seq_len: Sequence length.
+        num_heads: Number of attention heads.
+        head_dim: Head dimension (must be even).
+        cos_batch: Batch dimension of cos/sin. Defaults to batch if None.
+        dtype: Data type for the tensors.
+    """
+    if cos_batch is None:
+        cos_batch = batch
+
+    torch.manual_seed(42)
+    x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype)
+    cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype)
+    return x, cos, sin
+
+
+class TestRotaryEmbeddingNativeShapeRegression:
+    """Regression tests for forward_native shape handling (#8e297d81)."""
+
+    @pytest.fixture
+    def rope_neox(self) -> RotaryEmbedding:
+        """NeoX-style (non-interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=True)
+
+    @pytest.fixture
+    def rope_gptj(self) -> RotaryEmbedding:
+        """GPT-J style (interleaved) RoPE."""
+        return RotaryEmbedding(is_neox_style=False)
+
+    # ── 2D cos/sin (no batch dim) ──────────────────────────────────────
+
+    def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None:
+        """2D cos/sin [S, D/2] pass through unchanged (no-op path)."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2)
+        # Manually drop batch dim to create 2D cos/sin
+        cos_2d = cos[0]
+        sin_2d = sin[0]
+        assert cos_2d.dim() == 2
+        output = rope_neox.forward_native(x, cos_2d, sin_2d)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()
+        assert not torch.isinf(output).any()
+
+    # ── 3D cos/sin — batch-dim stripping ──────────────────────────────
+
+    @pytest.mark.parametrize("cos_batch", [1, 2, 4])
+    def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None:
+        """3D cos/sin with batch dim should not raise errors."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+
+    @pytest.mark.parametrize("interleaved", [False, True])
+    def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None:
+        """3D cos/sin with identical batch slices == 2D cos/sin result."""
+        rope = RotaryEmbedding(is_neox_style=not interleaved)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+
+        # 3D variant: expand the same 2D data across batch
+        cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1)
+        sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1)
+
+        out_2d = rope.forward_native(x, cos_2d, sin_2d)
+        out_3d = rope.forward_native(x, cos_3d, sin_3d)
+
+        torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0)
+
+    # ── Batch size mismatch — the regression scenario ─────────────────
+
+    @pytest.mark.parametrize(
+        "x_batch,cos_batch",
+        [
+            (2, 1),  # cos has fewer batch elements
+            (1, 2),  # cos has more batch elements (takes [0])
+            (4, 1),  # large x batch, single cos
+        ],
+    )
+    def test_batch_size_mismatch_regression(
+        self,
+        rope_neox: RotaryEmbedding,
+        x_batch: int,
+        cos_batch: int,
+    ) -> None:
+        """forward_native works when cos batch dim differs from x batch dim.
+
+        This is the exact regression scenario — before the fix, a 3D cos/sin
+        with a different batch size than x would cause a shape mismatch error
+        inside apply_rotary_emb_torch during broadcast.
+        """
+        x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert output.dtype == x.dtype
+
+    # ── Batch consistency — all elements get same rotation ────────────
+
+    def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 2D cos, identical x slices → identical output slices."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)  # all batch elements identical
+        cos = torch.randn(S, D // 2)
+        sin = torch.randn(S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        # All output batch slices must be identical
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """With 3D cos (identical slices), identical x → identical output."""
+        B, S, H, D = 4, 16, 8, 64
+        x_template = torch.randn(1, S, H, D)
+        x = x_template.expand(B, -1, -1, -1)
+        cos = torch.randn(1, S, D // 2).expand(B, -1, -1)
+        sin = torch.randn(1, S, D // 2).expand(B, -1, -1)
+
+        output = rope_neox.forward_native(x, cos, sin)
+
+        for i in range(1, B):
+            torch.testing.assert_close(output[i], output[0], atol=0, rtol=0)
+
+    # ── Output shape / dtype invariants ───────────────────────────────
+
+    def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output shape must match input shape for all cos/sin dimensionalities."""
+        B, S, H, D = 2, 32, 4, 128
+        x = torch.randn(B, S, H, D)
+
+        # 2D cos/sin
+        cos_2d = torch.randn(S, D // 2)
+        sin_2d = torch.randn(S, D // 2)
+        assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D)
+
+        # 3D cos/sin, matching batch
+        cos_3d = torch.randn(B, S, D // 2)
+        sin_3d = torch.randn(B, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D)
+
+        # 3D cos/sin, different batch
+        cos_3d_1 = torch.randn(1, S, D // 2)
+        sin_3d_1 = torch.randn(1, S, D // 2)
+        assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D)
+
+    def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None:
+        """Output dtype must match input dtype."""
+        for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype)
+            output = rope_neox.forward_native(x, cos, sin)
+            assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}"
+
+    def test_output_no_nan_inf(
+        self,
+        rope_neox: RotaryEmbedding,
+        rope_gptj: RotaryEmbedding,
+    ) -> None:
+        """Output must be finite for both RoPE styles."""
+        for rope in (rope_neox, rope_gptj):
+            x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64)
+            output = rope.forward_native(x, cos, sin)
+            assert not torch.isnan(output).any(), f"NaN in {rope}"
+            assert not torch.isinf(output).any(), f"Inf in {rope}"
+
+    def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None:
+        """Output should be contiguous."""
+        x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64)
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.is_contiguous()
+
+    # ── CFG-style doubled batch ───────────────────────────────────────
+
+    def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None:
+        """CFG (classifier-free guidance) doubles batch — cos may have
+        original batch size while x has 2× batch."""
+        B, S, H, D = 2, 16, 8, 64
+        # Simulate CFG: x has 2× batch (cond + uncond concatenated)
+        x = torch.randn(B * 2, S, H, D)
+        # cos/sin only have the original batch
+        cos = torch.randn(B, S, D // 2)
+        sin = torch.randn(B, S, D // 2)
+
+        output = rope_neox.forward_native(x, cos, sin)
+        assert output.shape == (B * 2, S, H, D)
+
+    # ── RotaryEmbeddingWan forward_native is unchanged ────────────────
+
+    def test_wan_forward_native_unchanged(self) -> None:
+        """RotaryEmbeddingWan.forward_native still uses its own implementation
+        (unflatten + stack + flatten) and should not crash with 3D cos/sin."""
+        from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+
+        rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
+        B, S, H, D = 2, 16, 8, 64
+        x = torch.randn(B, S, H, D)
+        cos = torch.randn(B, S, 1, D // 2)
+        sin = torch.randn(B, S, 1, D // 2)
+
+        output = rope.forward_native(x, cos, sin)
+        assert output.shape == x.shape
+        assert not torch.isnan(output).any()

tests/diffusion/lora/test_lora_manager.py

@@ -575,3 +575,45 @@ def test_lora_manager_discovers_bagel_component(monkeypatch):
     assert "bagel.language_model.qkv_proj" in manager._lora_modules
     # Verify the module was actually replaced in the tree (not just recorded)
     assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)
+
+
+def test_lora_manager_discovers_unet_component(monkeypatch):
+    """Verify that _replace_layers_with_lora finds layers under 'unet'."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(
+        manager_mod,
+        "replace_submodule",
+        lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+    )
+
+    # Pipeline with a 'unet' component (no 'transformer')
+    pipeline = torch.nn.Module()
+    pipeline.unet = torch.nn.Module()
+    pipeline.unet.down_block = torch.nn.Module()
+    pipeline.unet.down_block.proj = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert "down_block.proj" in replace_calls
+    assert "unet.down_block.proj" in manager._lora_modules
+    # Verify the module was actually replaced in the tree (not just recorded)
+    assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)

tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py

@@ -1,112 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Flux2-Klein GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import (
-    Flux2KleinGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "Flux2Pipeline",
-    model_type: str = "flux",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_flux2_klein_adapter_selected_for_flux_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, Flux2KleinGGUFAdapter)
-
-
-def test_flux2_klein_adapter_matches_flux_model_type():
-    assert Flux2KleinGGUFAdapter.is_compatible(
-        _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"),
-        object(),
-        _make_source(),
-    )
-
-
-def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight", torch.ones((1, 1))),
-                ("time_in.in_layer.qweight_type", torch.tensor(1)),
-                ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)),
-                ("final_layer.linear.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "x_embedder.qweight" in names
-    assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names
-    assert "transformer_blocks.0.attn.norm_q.weight" in names
-    assert "proj_out.qweight" in names
-
-
-def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module
-
-    monkeypatch.setattr(
-        flux_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                (
-                    "final_layer.adaLN_modulation.1.scale",
-                    torch.tensor([1.0, 2.0, 3.0, 4.0]),
-                )
-            ]
-        ),
-    )
-
-    adapter = Flux2KleinGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert len(weights) == 1
-    assert weights[0][0] == "norm_out.linear.weight"
-    assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0]))

tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py

@@ -1,184 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import (
-    QwenImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "QwenImagePipeline",
-    model_type: str = "qwen_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-class _FakeTransformer:
-    def __init__(self) -> None:
-        self._params = {
-            "img_in.qweight": 1,
-            "img_in.qweight_type": 1,
-            "txt_in.qweight": 1,
-            "txt_in.qweight_type": 1,
-            "proj_out.qweight": 1,
-            "proj_out.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_1.qweight_type": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight": 1,
-            "time_text_embed.timestep_embedder.linear_2.qweight_type": 1,
-            "transformer_blocks.0.img_mod.1.qweight": 1,
-            "transformer_blocks.0.img_mod.1.qweight_type": 1,
-            "transformer_blocks.0.txt_mod.1.qweight": 1,
-            "transformer_blocks.0.txt_mod.1.qweight_type": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight": 1,
-            "transformer_blocks.0.attn.to_qkv.qweight_type": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight": 1,
-            "transformer_blocks.0.attn.add_kv_proj.qweight_type": 1,
-            "transformer_blocks.0.attn.to_out.weight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight": 1,
-            "transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1,
-        }
-
-    def named_parameters(self):
-        return list(self._params.items())
-
-    def named_buffers(self):
-        return []
-
-    def get_submodule(self, _prefix: str):
-        return self
-
-
-def test_qwen_adapter_selected_for_qwen_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-    assert isinstance(adapter, QwenImageGGUFAdapter)
-
-
-def test_qwen_adapter_matches_multiple_pipeline_variants():
-    for model_class_name in (
-        "QwenImagePipeline",
-        "QwenImageEditPipeline",
-        "QwenImageEditPlusPipeline",
-        "QwenImageLayeredPipeline",
-    ):
-        assert QwenImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name=model_class_name),
-            _FakeTransformer(),
-            _make_source(),
-        )
-
-
-def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.attn.to_q.qweight_type", 1),
-                ("transformer_blocks.0.attn.to_q.qweight", 2),
-                ("transformer_blocks.0.attn.to_k.qweight_type", 3),
-                ("transformer_blocks.0.attn.to_k.qweight", 4),
-                ("transformer_blocks.0.attn.to_out.0.qweight_type", 5),
-                ("transformer_blocks.0.attn.to_out.0.qweight", 6),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights
-    assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights
-
-
-def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("img_in.qweight_type", 1),
-                ("img_in.qweight", 2),
-                ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3),
-                ("time_text_embed.timestep_embedder.linear_1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("img_in.qweight_type", 1) in weights
-    assert ("img_in.qweight", 2) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights
-    assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights
-
-
-def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module
-
-    monkeypatch.setattr(
-        qwen_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("transformer_blocks.0.img_mod.1.qweight_type", 1),
-                ("transformer_blocks.0.img_mod.1.qweight", 2),
-                ("transformer_blocks.0.txt_mod.1.qweight_type", 3),
-                ("transformer_blocks.0.txt_mod.1.qweight", 4),
-            ]
-        ),
-    )
-
-    adapter = QwenImageGGUFAdapter(
-        "dummy.gguf",
-        _FakeTransformer(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-
-    assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights
-    assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights
-    assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights

tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py

@@ -1,83 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the Z-Image GGUF adapter."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import (
-    ZImageGGUFAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def _make_od_config(
-    model_class_name: str = "ZImagePipeline",
-    model_type: str = "z_image",
-):
-    return SimpleNamespace(
-        model_class_name=model_class_name,
-        tf_model_config={"model_type": model_type},
-    )
-
-
-def _make_source(prefix: str = "", subfolder: str = "transformer"):
-    return SimpleNamespace(prefix=prefix, subfolder=subfolder)
-
-
-def test_z_image_adapter_selected_for_z_image_family():
-    adapter = get_gguf_adapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    assert isinstance(adapter, ZImageGGUFAdapter)
-
-
-def test_z_image_adapter_matches_model_type_variants():
-    for model_type in ("z_image", "zimage", "z-image"):
-        assert ZImageGGUFAdapter.is_compatible(
-            _make_od_config(model_class_name="OtherPipeline", model_type=model_type),
-            object(),
-            _make_source(),
-        )
-
-
-def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch):
-    import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module
-
-    monkeypatch.setattr(
-        z_image_module,
-        "gguf_quant_weights_iterator",
-        lambda _path: iter(
-            [
-                ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))),
-                ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)),
-                ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)),
-                ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)),
-            ]
-        ),
-    )
-
-    adapter = ZImageGGUFAdapter(
-        "dummy.gguf",
-        object(),
-        _make_source(),
-        _make_od_config(),
-    )
-
-    weights = list(adapter.weights_iterator())
-    names = [name for name, _ in weights]
-
-    assert "all_final_layer.2-1.qweight" in names
-    assert "all_x_embedder.2-1.qweight_type" in names
-    assert "transformer_blocks.0.attention.to_out.0.weight" in names
-    assert "transformer_blocks.0.attention.to_qkv.qweight" in names

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -16,7 +16,6 @@ from vllm.config.load import LoadConfig
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
@@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
         loader.load_weights(model)
 
 
-def test_qwen_model_class_selects_qwen_gguf_adapter():
-    od_config = type(
-        "Config",
-        (),
-        {
-            "model_class_name": "QwenImagePipeline",
-            "tf_model_config": {"model_type": "qwen_image"},
-        },
-    )()
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
-
-    assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
-
-
 class _ConfigAwareModel(nn.Module):
     def __init__(self, *, od_config):
         super().__init__()
@@ -186,7 +164,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -1,259 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module
-from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-class _DummyModel(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.transformer = nn.Linear(2, 2, bias=True)
-        self.vae = nn.Linear(2, 2, bias=False)
-        self.register_buffer("transformer_buffer", torch.ones(1))
-        self.calls: list[list[str]] = []
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
-                fall_back_to_pt=True,
-            ),
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path="dummy",
-                subfolder="vae",
-                revision=None,
-                prefix="vae.",
-                fall_back_to_pt=True,
-            ),
-        ]
-
-    def load_weights(self, weights):
-        loadable = dict(self.named_parameters())
-        loadable.update(dict(self.named_buffers()))
-        seen: list[str] = []
-        loaded: set[str] = set()
-        for name, tensor in weights:
-            seen.append(name)
-            if name in loadable:
-                target = loadable[name]
-                target.data.copy_(tensor.to(dtype=target.dtype))
-                loaded.add(name)
-        self.calls.append(seen)
-        return loaded
-
-
-def _make_loader() -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
-    loader.load_config = SimpleNamespace(
-        download_dir="cache-dir",
-        ignore_patterns=["*.tmp"],
-    )
-    loader.od_config = None
-    loader.quant_config = None
-    loader.counter_before_loading_weights = 0.0
-    loader.counter_after_loading_weights = 0.0
-    return loader
-
-
-def test_is_gguf_quantization_accepts_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "gguf",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_rejects_non_gguf_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {
-        "method": "fp8",
-        "gguf_model": "weights.gguf",
-    }
-
-    assert loader._is_gguf_quantization() is False
-
-
-def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
-    loader = _make_loader()
-    loader.quant_config = {"method": "gguf"}
-
-    with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization()
-
-
-def test_is_gguf_quantization_accepts_object_config():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(
-        get_name=lambda: "gguf",
-        gguf_model="weights.gguf",
-    )
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_is_gguf_quantization_uses_fallback_object_without_get_name():
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-
-    assert loader._is_gguf_quantization() is True
-
-
-def test_get_model_loadable_names_collects_parameters_and_buffers():
-    loader = _make_loader()
-    model = _DummyModel()
-
-    names = loader._get_model_loadable_names(model)
-
-    assert "transformer.weight" in names
-    assert "transformer.bias" in names
-    assert "transformer_buffer" in names
-
-
-def test_resolve_gguf_model_path_returns_local_file(tmp_path):
-    loader = _make_loader()
-    gguf_file = tmp_path / "model.gguf"
-    gguf_file.write_bytes(b"gguf")
-
-    resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None)
-
-    assert resolved == str(gguf_file)
-
-
-def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "hf_hub_download",
-        lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf"
-
-
-def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-
-    monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False)
-    monkeypatch.setattr(
-        loader_module,
-        "download_gguf",
-        lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf",
-    )
-
-    resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main")
-
-    assert resolved == "cache-dir/owner/repo/Q4_0.gguf"
-
-
-def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
-    loader.od_config = SimpleNamespace(revision="main")
-    source = DiffusersPipelineLoader.ComponentSource(
-        model_or_path="dummy",
-        subfolder="transformer",
-        revision=None,
-        prefix="transformer.",
-    )
-
-    class _Adapter:
-        def weights_iterator(self):
-            yield "weight", torch.ones((2, 2))
-            yield "bias", torch.zeros(2)
-
-    monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
-    monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
-
-    weights = list(loader._get_gguf_weights_iterator(source, object()))
-
-    assert weights[0][0] == "transformer.weight"
-    assert torch.equal(weights[0][1], torch.ones((2, 2)))
-    assert weights[1][0] == "transformer.bias"
-    assert torch.equal(weights[1][1], torch.zeros(2))
-
-
-def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
-    )
-
-    def _hf_iter(source):
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert model.calls[0] == ["transformer.weight"]
-    assert model.calls[1] == ["transformer.bias"]
-    assert model.calls[2] == ["vae.weight"]
-
-
-def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
-    loader = _make_loader()
-    model = _DummyModel()
-
-    monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
-    monkeypatch.setattr(
-        loader,
-        "_get_gguf_weights_iterator",
-        lambda source, model: iter(
-            [
-                ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
-                ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
-            ]
-        ),
-    )
-
-    hf_calls: list[str] = []
-
-    def _hf_iter(source):
-        hf_calls.append(source.subfolder or "")
-        if source.subfolder == "transformer":
-            return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))])
-        return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))])
-
-    monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter)
-    monkeypatch.setattr(
-        loader,
-        "_get_expected_parameter_names",
-        lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
-    )
-
-    loaded = loader._load_weights_with_gguf(model)
-
-    assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
-    assert hf_calls == ["vae"]

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -916,6 +916,49 @@ def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeli
     assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
 
 
+@pytest.mark.parametrize("mode", ["policy", "forward_dynamics"])
+def test_prepare_action_video_latents_encodes_only_conditioning_frame(make_cosmos3_pipeline, mode: str) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 24, 32)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        mode,
+        24,
+        32,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+        image_size=torch.tensor([1, 3, 16, 24]),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 1, 24, 32)]
+    assert latents.shape == (1, 2, 3, 2, 3)
+    assert condition_latents.shape == latents.shape
+    torch.testing.assert_close(condition_latents[:, :, 0], torch.ones(1, 2, 2, 3))
+    assert torch.count_nonzero(condition_latents[:, :, 1:]) == 0
+    torch.testing.assert_close(latents[:, :, 0:1], condition_latents[:, :, 0:1])
+    assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]]
+
+
+def test_prepare_inverse_dynamics_latents_encodes_full_video(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    video = torch.zeros(1, 3, 9, 16, 24)
+
+    latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video(
+        video,
+        "inverse_dynamics",
+        16,
+        24,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+    )
+
+    assert pipeline.vae.encode_input_shapes == [(1, 3, 9, 16, 24)]
+    assert condition_latents.shape == (1, 2, 3, 2, 3)
+    torch.testing.assert_close(latents, condition_latents)
+    assert torch.count_nonzero(velocity_mask) == 0
+
+
 def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)

tests/diffusion/models/flux/test_flux_pipeline.py

@@ -21,6 +21,8 @@ def _make_flux_sampling(**overrides):
         "true_cfg_scale": 4.0,
         "num_outputs_per_prompt": 0,
         "latents": None,
+        "output_type": None,
+        "max_sequence_length": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -131,7 +133,7 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                     "negative_prompt_embeds": negative_prompt_embeds_a,
                     "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a,
                 },
-                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a),
+                sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"),
             ),
             SimpleNamespace(
                 request_id="flux-prompt-b",
@@ -145,12 +147,12 @@ def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch):
                         "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b],
                     },
                 },
-                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b),
+                sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"),
             ),
         ]
     )
 
-    outputs = pipeline.forward(batch, output_type="latent")
+    outputs = pipeline.forward(batch)
 
     assert encode_calls[0]["prompt"] is None
     assert encode_calls[0]["prompt_2"] is None

tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py

@@ -50,6 +50,48 @@ class TestSetForwardContextNumTokens:
         mock_get.assert_not_called()
 
 
+class TestSetForwardContextDPMetadata:
+    """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch."""
+
+    def test_sets_dp_metadata_from_vllm_dp_group(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = None
+        fake_group = mocker.MagicMock()
+        fake_group.world_size = 2
+        fake_group.cpu_group = object()
+
+        def _all_gather_object(output, obj, group):
+            assert group is fake_group.cpu_group
+            output[:] = [obj, obj + 1]
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True)
+        mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object)
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is not None
+        assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025]
+
+    def test_keeps_existing_dp_metadata(self, mocker):
+        import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe
+
+        existing_metadata = object()
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.dp_metadata = existing_metadata
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+        mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object")
+
+        hunyuan_moe._set_forward_context_dp_metadata(1024)
+
+        assert mock_ctx.dp_metadata is existing_metadata
+        mock_all_gather.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
     """Test platform dispatch via platform qualname hooks."""